Fine-Tuning Masked Diffusion for Provable Self-Correction

• URL: http://arxiv.org/abs/2510.01384v2
• Date: Fri, 07 Nov 2025 04:01:45 GMT
• Title: Fine-Tuning Masked Diffusion for Provable Self-Correction
• Authors: Jaeyeon Kim, Seunggeun Kim, Taekyun Lee, David Z. Pan, Hyeji Kim, Sham Kakade, Sitan Chen,
• Abstract summary: Masked Diffusion Models (MDMs) have emerged as a promising approach for generative modeling in discrete spaces.<n>We introduce PRISM--Plug-in Remasking for Inference-time Self-correction of Masked Diffusions.
• Score: 28.338622227684453
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: A natural desideratum for generative models is self-correction--detecting and revising low-quality tokens at inference. While Masked Diffusion Models (MDMs) have emerged as a promising approach for generative modeling in discrete spaces, their capacity for self-correction remains poorly understood. Prior attempts to incorporate self-correction into MDMs either require overhauling MDM architectures/training or rely on imprecise proxies for token quality, limiting their applicability. Motivated by this, we introduce PRISM--Plug-in Remasking for Inference-time Self-correction of Masked Diffusions--a lightweight, model-agnostic approach that applies to any pretrained MDM. Theoretically, PRISM defines a self-correction loss that provably learns per-token quality scores, without RL or a verifier. These quality scores are computed in the same forward pass with MDM and used to detect low-quality tokens. Empirically, PRISM advances MDM inference across domains and scales: Sudoku; unconditional text (170M); and code with LLaDA (8B).

Related papers

• Learn from Your Mistakes: Self-Correcting Masked Diffusion Models [31.536464269884103]
  Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models.<n>We propose a framework that trains a model to perform both unmasking and correction.<n>We name our training and sampling method Progressive Self-Correction (ProSeCo) for its unique ability to iteratively refine an entire sequence.
  arXiv  Detail & Related papers  (2026-02-12T05:17:31Z)
• Training-Free Self-Correction for Multimodal Masked Diffusion Models [61.84305395626145]
  We propose a training-free self-correction framework that exploits the inductive biases of pre-trained masked diffusion models.<n>Our method significantly improves generation quality on text-to-image generation and multimodal understanding tasks with reduced sampling steps.
  arXiv  Detail & Related papers  (2026-02-02T23:58:15Z)
• MDiff4STR: Mask Diffusion Model for Scene Text Recognition [59.79818820650126]
  Mask Diffusion Models (MDMs) have emerged as a promising alternative to auto-regressive models (ARMs) for vision-language tasks.<n>We show that vanilla MDM lags behind ARMs in terms of accuracy, although it improves recognition efficiency.<n>We propose MDiff4STR, a Mask Diffusion model enhanced with two key improvement strategies tailored for Scene Text Recognition.
  arXiv  Detail & Related papers  (2025-12-01T08:57:51Z)
• Masked Diffusion Models are Secretly Learned-Order Autoregressive Models [21.17429712617749]
  We show that Masked Diffusion Models can identify and optimize for a decoding order during training.<n>We prove that the MDM objective decomposes precisely into a weighted auto-regressive losses over these orders.
  arXiv  Detail & Related papers  (2025-11-24T14:17:56Z)
• Enabling Autoregressive Models to Fill In Masked Tokens [50.9948753314669]
  This work introduces MARIA (Masked and Autoregressive Infilling Architecture), a novel approach that achieves state-of-the-art masked infilling performance.<n>MARIA combines a pre-trained and AR model by training a linear decoder that takes their hidden states as input.<n>Our results demonstrate that MARIA significantly outperforms existing methods, namely discrete diffusion models, on masked infilling tasks.
  arXiv  Detail & Related papers  (2025-02-09T20:02:05Z)
• Steering Masked Discrete Diffusion Models via Discrete Denoising Posterior Prediction [88.65168366064061]
  We introduce Discrete Denoising Posterior Prediction (DDPP), a novel framework that casts the task of steering pre-trained MDMs as a problem of probabilistic inference. Our framework leads to a family of three novel objectives that are all simulation-free, and thus scalable. We substantiate our designs via wet-lab validation, where we observe transient expression of reward-optimized protein sequences.
  arXiv  Detail & Related papers  (2024-10-10T17:18:30Z)
• Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling [47.82616476928464]
  Masked diffusion models (MDMs) have emerged as a popular research topic for generative modeling of discrete data.<n>We show that both training and sampling of MDMs are theoretically free from the time variable.<n>We identify, for the first time, an underlying numerical issue, even with the commonly used 32-bit floating-point precision.
  arXiv  Detail & Related papers  (2024-09-04T17:48:19Z)
• Model Inversion Attacks Through Target-Specific Conditional Diffusion Models [54.69008212790426]
  Model inversion attacks (MIAs) aim to reconstruct private images from a target classifier's training set, thereby raising privacy concerns in AI applications. Previous GAN-based MIAs tend to suffer from inferior generative fidelity due to GAN's inherent flaws and biased optimization within latent space. We propose Diffusion-based Model Inversion (Diff-MI) attacks to alleviate these issues.
  arXiv  Detail & Related papers  (2024-07-16T06:38:49Z)
• Continual-MAE: Adaptive Distribution Masked Autoencoders for Continual Test-Time Adaptation [49.827306773992376]
  Continual Test-Time Adaptation (CTTA) is proposed to migrate a source pre-trained model to continually changing target distributions. Our proposed method attains state-of-the-art performance in both classification and segmentation CTTA tasks.
  arXiv  Detail & Related papers  (2023-12-19T15:34:52Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.