Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Training-Free Self-Correction for Multimodal Masked Diffusion Models

Yidong Ouyang, Panwen Hu, Zhengyan Wan, Zhe Wang, Liyan Xie, Dmitriy Bespalov, Ying Nian Wu, Guang Cheng, Hongyuan Zha, Qiang Sun

TL;DR

This work identifies error accumulation in masked diffusion models caused by parallel, absorbing token updates and proposes a training-free self-correction framework that leverages the inductive biases of pre-trained backbones to reassess and revise uncertain updates at inference time. By designing remasking criteria that rely on accumulated token probabilities and distributional uncertainty, the approach avoids model fine tuning or external evaluators while achieving meaningful gains in text-to-image generation and multimodal understanding tasks. The method demonstrates faster sampling (e.g., 16 steps rivaling 64 steps) and robust performance across different masked diffusion backbones, suggesting practical applicability and broad generalization for multimodal generation regimes.

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.

Training-Free Self-Correction for Multimodal Masked Diffusion Models

TL;DR

This work identifies error accumulation in masked diffusion models caused by parallel, absorbing token updates and proposes a training-free self-correction framework that leverages the inductive biases of pre-trained backbones to reassess and revise uncertain updates at inference time. By designing remasking criteria that rely on accumulated token probabilities and distributional uncertainty, the approach avoids model fine tuning or external evaluators while achieving meaningful gains in text-to-image generation and multimodal understanding tasks. The method demonstrates faster sampling (e.g., 16 steps rivaling 64 steps) and robust performance across different masked diffusion backbones, suggesting practical applicability and broad generalization for multimodal generation regimes.

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.
Paper Structure (22 sections, 11 equations, 5 figures, 5 tables, 1 algorithm)

This paper contains 22 sections, 11 equations, 5 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Masked Diffusion Models.
  4. Self-Correction of Masked Diffusion Models.
  5. Preliminaries
  6. Formulation of Masked Diffusion Model
  7. Forward Process.
  8. Reverse Process.
  9. Training Objective.
  10. Sampling Process.
  11. The Limitation of Masked Diffusion Model
  12. Self-Correction of Masked Diffusion Models
  13. Proposed Method
  14. Rationality of Training-free Self-Correction
  15. In-depth Design of Self-Correction Criteria
  16. ...and 7 more sections

Figures (5)

  • 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.
  • 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.
  • Figure 3: Comparison of generated images of Lumina-DiMOO, ReMDM, and our methods. Our method achieves better quality with less sampling steps.
  • Figure 4: GenEval performance under different inference steps. The x-axis denotes the number of steps for generation, while the y-axis denotes the overall performance on the GenEval benchmark.
  • Figure 5: Illustration of the performance on multimodal understanding.