[Paper Review] Training-Free Self-Correction for Multimodal Masked Diffusion Models

Yidong Ouyang, Panwen Hu+8 more|arXiv (Cornell University)|Feb 2, 2026

Generative Adversarial Networks and Image Synthesis0 citations

TL;DR

The paper introduces a training-free self-correction framework for pre-trained multimodal masked diffusion models, enabling token remasking during inference to revise early mistakes without fine-tuning, improving text-to-image generation and multimodal understanding while enabling faster sampling.

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.

Motivation & Objective

• Investigate error accumulation in parallel, irreversible token updates in masked diffusion models.
• Develop a training-free self-correction mechanism that leverages inductive biases of pre-trained backbones.
• Enable token remasking during inference without modifying model parameters or using external evaluators.
• Assess robustness and generalization across different masked diffusion architectures on multimodal tasks.

Proposed method

• Model-agnostic remasking during inference that re-evaluates token probabilities for already-generated positions.
• Use cumulated predicted probabilities across steps to identify low-confidence tokens for remasking.
• Remask a fixed number of tokens per step based on a remasking schedule  to balance fidelity and speed.
• Optionally employ distributional-uncertainty criteria (KL divergence, Wasserstein distance) to select remasked tokens.
• Algorithm 1 outlines training-free self-correction with options for deterministic or stochastic remasking.

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.

Experimental results

Research questions

• RQ1 Can training-free self-correction identify and revise low-confidence tokens during inference in multimodal masked diffusion models?
• RQ2 Does leveraging the inductive bias of pre-trained backbones enable effective remasking without fine-tuning?
• RQ3 How do remasking strategies (deterministic vs stochastic, cumulated vs current step likelihood) affect generation quality and efficiency?
• RQ4 Is the proposed method robust across different masked diffusion backbones?
• RQ5 What is the impact on sampling efficiency (fewer steps) when applying remasking-based self-correction?

Key findings

• The approach yields consistent improvements on GenEval over vanilla Lumina-DiMOO and prior training-free methods.
• On multimodal understanding benchmarks (MMBend, SEED-Bench, MMMU), the method improves performance compared to baselines.
• Ablation shows cumulated likelihood with deterministic remasking often performs best across metrics.
• The method enables comparable or better GenEval performance with as few as 16 sampling steps vs 64 in the baseline.
• Gives evidence of generalization across backbones (e.g., MMaDA-8B-MixCoT) with consistent gains.

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.