PaperlayerView exact revision on arXiv

Source-linked AI summary

Remasking Discrete Diffusion Models with Inference-Time Scaling

Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, Volodymyr Kuleshov

arXiv:2503.00307v4cs.LGstat.ML

TL;DR

Masked discrete diffusion cannot revise decoded tokens, limiting iterative error correction and generation quality. ReMDM introduces a principled remasking sampler compatible with pretrained masked diffusion models. It improves quality and controlled generation across multiple domains while enabling inference-time compute scaling.

  • Problem

    Masked discrete diffusion locks decoded tokens, preventing error correction and limiting controllable generation, sampling speed, and sample quality.

  • Method

    ReMDM defines a remasking backward process that augments masked diffusion with remasking during generation and can be applied to pretrained models.

  • Results

    ReMDM improves sample quality, controlled generation, and downstream performance across language, discretized images, and molecule string representations, with quality scaling with inference-time computation.

  • Takeaways & Limitations

    Increasing sampling steps can bring ReMDM language generation close to autoregressive quality, while reduced-step sampling better preserves quality than other samplers.

  • Takeaways & Limitations

    Generative perplexity can be misleading for corrector samplers because tuning schedules can lower it while sacrificing sentence entropy and diversity.

Abstract

from arXiv · show

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

1 Introduction

Masked discrete diffusion models cannot revise decoded tokens, limiting error correction and controllable generation. ReMDM introduces remasking-based iterative refinement and improves quality across language, images, and molecule design while scaling with inference-time computation.

  • Masked diffusion locks decoded tokens, so errors cannot be corrected and controllable generation, sampling speed, and sample quality are limited.
  • ReMDM adds remasking during generation, allowing users to specify remasking probabilities and combine them with schedules and nucleus sampling.
  • ReMDM is derived as ancestral sampling in a probabilistic model with a remasking backward process and an objective similar to MDLM.
  • ReMDM provides inference-time scaling: more language-modeling steps improve quality toward autoregressive models, while fewer steps degrade quality less than other samplers.
  • Across natural language, discretized images, and molecule strings, ReMDM improves sample quality, controlled generation, and downstream performance.

2 Background

Discrete diffusion reverses a predefined noising process over discrete token states. Masked diffusion uses a special absorbing mask state, but its posterior makes every decoded token remain fixed throughout denoising.

  • Discrete diffusion models learn parameterized reverse processes for latent variables produced by a fixed noising process over discrete signals.
  • The forward process uses a user-specified prior and a monotonically decreasing schedule α_t across finite time steps.
  • Masked diffusion is an absorbing-state special case whose limiting distribution is a one-hot mask token.
  • In masked diffusion, any unmasked token must remain unchanged at earlier denoising times, creating the failure-to-remask property.
  • A denoising network parameterizes the reverse transition and is trained with a variational objective.

3 Remasking Diffusion Models

ReMDM resolves masked diffusion’s failure to remask by defining a probabilistic model whose posterior can remask decoded tokens. Its marginals remain those of classical masked diffusion, supporting pretrained-model reuse.

  • ReMDM generalizes the MDLM denoising posterior so decoded tokens can be remasked during generation.
  • ReMDM preserves the classical masked-diffusion marginals despite introducing remasking.
  • The remasking construction is non-Markovian because transitions from the mask state back to token values require conditioning on the original data.
  • The remasking parameter σ_t controls the probability of remasking decoded tokens, with σ_t = 0 recovering the MDLM posterior.
  • ReMDM’s NELBO contains a diffusion-loss term that is a reweighted MDLM objective, enabling reuse of pretrained MDLM weights and different σ_t values at training and inference.

4 ReMDM Samplers

ReMDM samplers expose remasking schedules that selectively revisit predictions, including confidence-based, capped, switch, and loop strategies. The framework also subsumes predictor-corrector samplers and is more general than DFM correctors.

  • 4.1 Design Strategies for the Remasking Schedule σ_t: ReMDM sampling can use pretrained denoising networks with a chosen timestep count, noise schedule, and remasking schedule.
  • 4.1 Design Strategies for the Remasking Schedule σ_t: The max-capped schedule limits remasking through σ_t = min{η_cap, 1−α_s/α_t}, while rescaling uses η_rescale to temper remasking probabilities.
  • 4.1 Design Strategies for the Remasking Schedule σ_t: The confidence-based schedule remasks unmasked tokens with probabilities inversely proportional to their stored decoding confidence.
  • 4.2 Design Strategies for ‘Turning On’ ReMDM: ReMDM-switch activates remasking only over a selected interval, whereas ReMDM-loop first decodes normally, then repeatedly remasks and predicts at constant α, and finally completes decoding.
  • 4.2 Design Strategies for ‘Turning On’ ReMDM: ReMDM-loop is designed to replace bad tokens while retaining good ones during its constant-noise phase.
  • 4.3 Comparison with Discrete Predictor-Corrector Samplers: Forward-backward and discrete-flow-matching correctors are ReMDM special cases, while ReMDM additionally supports constant-noise schedules unavailable to DFM.

5 Experiments

Experiments evaluate ReMDM across text, images, and controlled molecule generation. ReMDM improves quality scaling with inference compute, maintains quality better under faster sampling, and strengthens controllability.

  • Text Generation: 15.62× MAUVE versus masked diffusion models and 2.23× versus MDLM with corrector samplers demonstrate favorable inference-time scaling on OWT.Masked diffusion models scale poorly with T, while corrector samplers saturate at large T.
  • Text Generation: ReMDM maintains sample quality better than baselines when fewer sampling steps provide a more limited computational budget.The OWT experiments report both inference-time scaling and faster sampling settings.
  • Text Generation: The largest MAUVE improvement in the ablation comes from ReMDM’s remasking ability, while each proposed sampling improvement increases MAUVE over MDLM.The ablation sequentially removes ReMDM building blocks on OWT generation.
  • Image Generation: ReMDM produces the highest-quality ImageNet images of the three methods at T = 64 and has the best scaling.The comparison uses pretrained MaskGiT outputs with ReMDM, original MaskGiT sampling, and MDLM.
  • Molecule Generation: For QM9 ring-count maximization, ReMDM pushes the novelty-property frontier beyond AR and diffusion baselines under both CFG and CBG.Larger inference budgets also dominate smaller ones, and the strongest strategy combines ReMDM-rescale with ReMDM-conf schedules.
  • Molecule Generation: Using ReMDM-rescale with the confidence-based scheduler improves controlled molecule-generation results.The parameter study identifies this combination as beneficial.

6 Discussion and Related Work

The discussion places ReMDM alongside discrete predictor-corrector samplers and reports downstream-task gains for LLaDA with ReMDM. These results connect remasking to improved performance in diffusion language models.

  • Discussion: LLaDA with ReMDM consistently performs best on downstream tasks and exceeds original LLaDA with statistical significance.The results report means and 95% confidence intervals across multiple random seeds.
  • Related Work: Discrete predictor-corrector samplers use CTMC theory to correct errors in reverse processes, including Forward-Backward, DFM, Stein, and DPC methods.The discussion distinguishes these correctors from plug-and-play methods and situates ReMDM among discrete correction approaches.

7 Conclusion

ReMDM extends masked discrete diffusion with remasking, enabling iterative correction while retaining a close relationship to classical masked diffusion objectives and existing corrector samplers.

  • 7 Conclusion: ReMDM enables pretrained masked diffusion models to remask predicted tokens during sampling, supporting more flexible generation strategies.The method is derived through a custom remasking backward process and can reuse pretrained weights.
  • 7 Conclusion: ReMDM preserves the classical masked-diffusion marginals while producing an objective that is a reweighted version of the MDLM diffusion loss.When σt = 0, the MDLM objective is recovered, and pretrained MDLM weights can be reused at inference.
  • 7 Conclusion: The sampler supports user-specified remasking schedules, including confidence-based strategies and corrector schedules that control remasking probability.These schedules provide the mechanism for adjusting sampling behavior across inference settings.

A.8 Proof of Proposition 4.4

The proposition establishes that ReMDM’s formulation encompasses discrete DDIM-style processes for absorbing-state diffusion through a reparameterization of the remasking parameter.

  • A.8 Proof of Proposition 4.4: The DDIM adaptation replaces the uniform categorical limiting distribution with the absorbing mask state m.This extends the discrete DDIM construction from uniform limiting noise to absorbing-state diffusion.
  • A.8 Proof of Proposition 4.4: ReMDM is equivalent to the adapted absorbing-state DDIM process after reparameterizing σt.The equivalence holds for both the unmasked and masked cases of the posterior.
  • A.8 Proof of Proposition 4.4: The ReMDM parameterization is easier to analyze, including deriving bounds on σt and exploring sampler design choices.The paper presents this as an advantage over directly reparameterizing with σDDIM_t.

C ReMDM Sampler Algorithms

The appendix provides switch and loop implementations of ReMDM that apply remasking schedules on top of a pretrained denoising network.

  • C ReMDM Sampler Algorithms: ReMDM-switch activates remasking after a user-specified switching time and otherwise uses no remasking.Its inputs include timestep count, noise schedule, remasking caps, rescaling, and an optional confidence strategy.
  • C ReMDM Sampler Algorithms: The switch algorithm computes an approximate posterior from the pretrained denoising network and samples each sequence position from it.Confidence-based remasking can optionally modify the per-token remasking values.

D Additional Experimental Details

The experiments reuse pretrained models across language, image, molecule, and LLaDA settings, varying sampling schedules, guidance, temperatures, and evaluation procedures.

  • D Additional Experimental Details: Language experiments reuse 169M-parameter AR, SEDD, and MDLM checkpoints with a log-linear noise schedule and evaluate MAUVE, perplexity, and entropy.Each model or sampler generates 5,000 samples for evaluation.
  • D Additional Experimental Details: Image experiments compare MaskGiT, MDLM, and ReMDM samplers while sweeping remasking schedules and softmax temperatures.ReMDM uses the same pretrained denoising outputs as MDLM but applies the ReMDM posterior.
  • D Additional Experimental Details: QM9 guidance experiments vary D-CFG and D-CBG, guidance strength γ ∈{1, 2, 3, 4, 5}, and diffusion timesteps T ∈{32, 64, 128}.Generated strings are parsed with RDKit, deduplicated, and evaluated for novelty and molecular properties.
  • D Additional Experimental Details: QM9 ReMDM experiments grid-search rescaling, confidence schedules, and switch or loop strategies for QED and ring-count guidance.The loop varies its start and end timing through predefined parameter tuples.
  • D Additional Experimental Details: LLaDA evaluation uses 32-token blocks on Countdown and TruthfulQA, with repeated sampling and reported confidence intervals.Countdown uses pass@1, while TruthfulQA reports the ROUGE difference between correct and plausibly wrong answers.

E.1 Training with MDLM v.s. ReMDM NELBO

Training with ReMDM and MDLM objectives produces comparable validation perplexities on QM9, while inference quality depends on the remasking schedule, step budget, and metric choice.

  • Training objectives: Comparable validation perplexities are obtained when training QM9 models with MDLM or ReMDM NELBO objectives.ReMDM training uses σt = min{ηcap, (1 −αs)/αt}.
  • Metric choice: MAUVE is preferred over generative perplexity because tuning remasking can lower perplexity while sacrificing entropy and diversity.An extreme ReMDM-cap setting generates repetitive text with Gen PPL. 1.5.
  • Sampling schedules: For T ≥ L, ReMDM-loop performs best, whereas ReMDM-cap performs better for faster sampling with T < L.ReMDM-rescale performs less well than the other two strategies in both scenarios.
  • Sampling schedules: For ReMDM-cap, increasing ηcap improves MAUVE at T = 1024 but can reduce it at T = 4096 because diversity loss outweighs perplexity gains.The reported inference-time-scaling choice is ηcap = 0.008; faster-sampling experiments use η = 0.040.
  • Image generation: Softmax temperature τ = 0.8 benefits MDLM and ReMDM, while MaskGiT performs best without temperature scaling at τ = 1.FID determines the best setting for each sampler in the discretized ImageNet comparison.

E.4.1 Drug-likeness Property Maximization

For drug-likeness maximization, ReMDM improves the novelty–property frontier over MDLM, with benefits shaped by guidance type, remasking strength, confidence scheduling, and activation timing.

  • QED maximization: ReMDM improves the novelty-property maximization frontier relative to MDLM for QED under D-CFG and D-CBG, though gains are less pronounced than for ring count.The comparison uses the best ReMDM setting and varies inference compute and guidance strength.
  • Remasking strength: No single ηrescale dominates when confidence-based scheduling is used for QED, while ηrescale = 0.1 improves results without confidence scheduling.The QED hyperparameter pattern is therefore less conclusive than the ring-count pattern.
  • Confidence scheduling: Confidence-based schedules consistently improve QED maximization results.For ring count, the same trend holds, especially at larger ηrescale, but is less pronounced overall.
  • Activation timing: For D-CFG, later ReMDM activation generally helps QED and ring-count maximization, while D-CBG ring-count experiments favor earlier activation.The exception reverses the general timing trend observed for the other listed settings.
  • Predictor-corrector comparison: ReMDM better trades off novel sample generation and ring-count maximization than the compared predictor-corrector samplers.The comparison is reported for conditional generation with γ = 1 using D-CFG.

F Generated Samples

Generated-sample analyses show both the benefits and failure modes of ReMDM: it can produce fluent, semantically coherent text, but poorly chosen remasking can exploit perplexity or reduce diversity.

  • Text samples: ReMDM-cap can generate repetitive, semantically empty text with generative perplexity as low as 1.5.The example uses OWT, ηcap = 0.2, and T = 1024.
  • Molecule generation: The molecule experiments search ring count and QED across D-CFG and D-CBG guidance, remasking strategies, confidence schedules, and γ values.The appendix includes hyperparameter searches for both properties and guidance mechanisms.
  • Text samples: At T = 4096, ReMDM samples are fluent and topic-coherent, unlike MDLM’s fragmented grammar and MDLM+DFM’s topic-focused but internally nonfluent text.The three samplers are compared using generated OpenWebText samples.
  • Molecule generation: Confidence-based and switch or loop configurations are evaluated through marker sizes for γ and remasking-schedule variants.The figures cover D-CFG and D-CBG searches for ring count and QED.
Loading 2503.00307v4…