Remasking Discrete Diffusion Models with Inference-Time Scaling

Part of the success of diffusion models stems from their ability to perform iterative refinement, i.e., repeatedly correcting outputs during generation. However, modern masked discrete diffusion lacks this capability: when a token is generated, it cannot be updated again, even when it introduces an error. Here, we address this limitation by introducing the remasking diffusion model (ReMDM) sampler, a method that can be applied to pretrained masked diffusion models in a principled way and that is derived from a discrete diffusion model with a custom remasking backward process. Most interestingly, ReMDM endows discrete diffusion with a form of inference-time compute scaling. By increasing the number of sampling steps, ReMDM generates natural language outputs that approach the quality of autoregressive models, whereas when the computation budget is limited, ReMDM better maintains quality. ReMDM also improves sample quality of masked diffusion models for discretized images, and in scientific domains such as molecule design, ReMDM facilitates diffusion guidance and pushes the Pareto frontier of controllability relative to classical masking and uniform noise diffusion. We provide the code along with a blog post on the project page: https://guanghanwang.com/remdm

한 줄 요약

ReMDM은 마스킹 기반 이산 확산 모델에 재마스킹 기능을 추가하여 추론 시 계산량 조절과 샘플 품질 향상을 실현한 새로운 샘플러이다.

핵심 기여도

ReMDM 샘플러를 도입하여 기존 마스킹 확산 모델의 재마스킹 부재 문제를 해결.
ReMDM은 기존 MDLM 모델 위에서 즉시 사용 가능하며, 샘플링 단계 수를 조절하여 추론 시 계산량을 유연하게 조정.
자연어, 이미지, 분자 설계 등 다양한 도메인에서 샘플 품질과 제어 가능성을 향상.
ReMDM은 이산 확산 모델에서 예측-보정(predictor-corrector) 기법의 이론적 일반화를 제공.

핵심 아이디어

기존 마스킹 기반 이산 확산 모델은 생성된 토큰이 한 번 결정되면 수정할 수 없어 오류 수정 능력이 결여되어 있다. ReMDM은 이 문제를 해결하기 위해 샘플링 과정에서 토큰을 재마스킹할 수 있는 확장된 역확산 과정(remasking backward process)을 도입한다. 이는 확률 모델 내에서 조건부 분포를 재정의함으로써, 토큰이 여러 단계에서 반복적으로 수정될 수 있도록 한다. ReMDM은 기존 마스킹 확산 모델(MDLM)의 목적 함수를 확장한 형태로 설계되어, 기존 학습된 모델 위에서 바로 사용 가능하다. 이론적으로 ReMDM은 예측-보정 기법보다 더 일반적인 샘플링 프레임워크를 제공하며, 샘플링 단계 수를 조절함으로써 샘플 품질과 추론 속도 간의 균형을 유연하게 조정할 수 있다.

기술적 접근법

**ReMDM 샘플러**: 기존 MDLM 모델에 적용 가능한 샘플러로, 사용자가 각 시간 단계에서 토큰의 재마스킹 확률을 직접 지정할 수 있도록 설계됨.
**remasking backward process**: 샘플링 과정에서 토큰이 반복적으로 수정될 수 있도록 확장된 역확산 과정.
**nucleus sampling 및 remasking schedule**: 샘플러 성능을 향상시키기 위해 다양한 구성 요소를 추가.
**ELBO 유사 목적 함수**: ReMDM은 기존 마스킹 확산 모델과 유사한 ELBO(Evidence Lower Bound) 목적 함수를 사용하며, 이는 기존 학습된 MDLM 위에서 샘플링 가능하게 만든다.
**predictor-corrector 해석**: ReMDM 샘플러는 예측-보정 기법의 이론적 확장으로 해석 가능하며, 기존 보정 샘플러보다 더 일반적.

주요 결과

**자연어 생성**: ReMDM 샘플러는 샘플링 단계 수를 늘릴수록 샘플 품질이 향상되며, 동일 아키텍처의 자동회귀 모델에 근접한 성능을 보임.
**이미지 생성**: MaskGiT 등 기존 마스킹 확산 모델보다 추론 시 계산량 증가에 더 잘 스케일링됨.
**분자 설계**: ReMDM은 마스킹 또는 균일 노이즈 확산 기반 모델 대비 분자 신선도와 목표 특성 간 Pareto 최전선을 향상.
**샘플 품질 향상**: ReMDM은 샘플링 단계 수를 줄여도 샘플 품질 저하가 다른 샘플러보다 적음.

의의 및 한계

ReMDM은 기존 마스킹 확산 모델의 핵심 한계인 재마스킹 부재를 해결함으로써 샘플 품질, 제어 가능성, 추론 속도 간의 균형을 개선하는 데 기여한다. 특히, 샘플링 단계 수를 조절하여 추론 시 계산량을 유연하게 조정할 수 있다는 점에서 실용적 가치가 높다. 또한, 분자 설계와 같은 과학적 도메인에서 확산 가이드라인을 효과적으로 활용할 수 있어 학술적 기여도 크다. 한편, ReMDM은 기존 MDLM 모델에 의존하며, 새로운 학습이 필요하지 않다는 점에서 유연성의 한계가 있을 수 있다. 또한, 재마스킹 확률 조절 방식은 사용자 지정이 필요하므로, 최적의 성능을 얻기 위해서는 추가 실험과 조정이 필요할 수 있다.

실용적 활용

ReMDM은 자연어 생성, 이미지 생성, 분자 설계 등 다양한 이산 생성 작업에 적용 가능하다. 특히, 샘플 품질과 추론 속도 간의 균형이 중요한 산업 현장에서 유용하며, 기존 마스킹 확산 모델을 활용한 연구 및 개발에 즉시 적용할 수 있다.