[논문 리뷰] Training-Free Self-Correction for Multimodal Masked Diffusion Models

Yidong Ouyang, Panwen Hu외 8명|arXiv (Cornell University)|2026. 02. 02.

Generative Adversarial Networks and Image Synthesis인용 수 0

한 줄 요약

이 논문은 훈련 없이도 사전 학습된 다중 모달 마스크 확산 모델에 대해 자체 수정 프레임워크를 도입하여 추론 중 토큰 재마스킹을 가능하게 하여 조기 실수를 수정하고 미세 조정 없이 텍스트-이미지 생성과 다중 모달 이해를 향상시키며 더 빠른 샘플링을 가능하게 한다.

ABSTRACT

Masked diffusion models have emerged as a powerful framework for text and multimodal generation. However, their sampling procedure updates multiple tokens simultaneously and treats generated tokens as immutable, which may lead to error accumulation when early mistakes cannot be revised. In this work, we revisit existing self-correction methods and identify limitations stemming from additional training requirements or reliance on misaligned likelihood estimates. We propose a training-free self-correction framework that exploits the inductive biases of pre-trained masked diffusion models. Without modifying model parameters or introducing auxiliary evaluators, our method significantly improves generation quality on text-to-image generation and multimodal understanding tasks with reduced sampling steps. Moreover, the proposed framework generalizes across different masked diffusion architectures, highlighting its robustness and practical applicability. Code can be found in https://github.com/huge123/FreeCorrection.

연구 동기 및 목표

• 마스크된 확산 모델에서 병렬적이고 되돌릴 수 없는 토큰 업데이트에서의 오차 누적을 조사한다.
• 사전 학습된 백본의 귀납적 바이어스를 활용하는 훈련 없는 자체 수정 메커니즘을 개발한다.
• 모델 파라미터를 수정하거나 外部 평가자를 사용하지 않고 추론 중 토큰 재마스킹을 가능하게 한다.
• 다양한 마스크 확산 아키텍처에서 다중 모달 작업에 대한 강건성과 일반화를 평가한다.

제안 방법

• 이미 생성된 위치의 토큰 확률을 재평가하는 추론 중 모델 비의존적 재마스킹.
• 단계 간 누적 예측 확률을 사용하여 재마스킹 대상인 저신뢰 토큰을 식별한다.
• 재마스킹 일정에 따라 매 단계 고정된 수의 토큰을 재마스 masking한다—충실도와 속도의 균형을 맞추기 위함.
• 선택적으로 분포적 불확실성 기준(KL 발산, Wasserstein 거리)을 사용하여 재마스킹할 토큰을 선택한다.
• Algorithm 1은 결정적 또는 확률적 재마스킹 옵션과 함께 훈련 없는 자체 수정의 개요를 제시한다.

Figure 1: Average predicted probability of flipped tokens and correct tokens over 2000 samples. The x-axis denotes the time steps for generation (64 steps in total for text-to-image generation), while the y-axis denotes the average probability over all flipped positions and the correct position.

실험 결과

연구 질문

• RQ1 훈련 없는 자체 수정이 다중 모달 마스크 확산 모델에서 추론 중 저신뢰 토큰을 식별하고 수정할 수 있는가?
• RQ2 사전 학습된 백본의 귀납적 바이어스를 활용하면 미세 조정 없이 효과적인 재마스킹이 가능한가?
• RQ3 재마스킹 전략(결정적 대 확률적, 누적 대 현재 단계 가능성)이 생성 품질과 효율성에 어떤 영향을 미치는가?
• RQ4 제안된 방법이 다양한 마스크 확산 백본에서 강건한가?
• RQ5 재마스킹 기반 자체 수정 적용 시 샘플링 효율성(더 적은 단계)이 어떤 영향을 받는가?

주요 결과

Method	Single	Two	Count	Color	Pos.	Attr.	Overall
Lumina-DiMOO a	0.99	0.93	0.85	0.84	0.84	0.71	0.86
Lumina-DiMOO (ReMDM)	1.00	0.94	0.86	0.87	0.82	0.74	0.87
Lumina-DiMOO (Ours)	0.99	0.94	0.88	0.93	0.87	0.79	0.90

• 이 방법은 일반적인 Lumina-DiMOO 및 이전의 훈련 없는 방법들보다 GenEval에서 일관된 개선을 보인다.
• 다중 모달 이해 벤치마크(MMBend, SEED-Bench, MMMU)에서 이 방법은 베이스라인에 비해 성능을 향상시킨다.
• 구성요소 제거 실험으로 누적 가능도와 결정적 재마스킹이 대부분의 지표에서 가장 우수한 경향을 보인다.
• 제안된 방법은 베이스라인의 64단계에 비해 16단계 샘플링만으로 GenEval 성능을 비교 가능하거나 더 잘 달성할 수 있음을 보여준다.
• 백본 간 일반화의 근거를 제시하며(예: MMaDA-8B-MixCoT) 일관된 이득이 관측된다.

Figure 2: The effectiveness of using accumulated predicted probability. The x-axis denotes the time steps for generation, while the y-axis denotes the average rank of the predicted probabilities of flipped tokens among correct tokens. The larger the rank is, the smaller the probability is.