Research map / Image generation / Fast diffusion sampling
Remasking Discrete Diffusion Models with Inference-Time Scaling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ReMDM, a sampler for pretrained masked diffusion language models that allows remasking of already-decoded tokens, giving inference-time compute scaling for text, image, and molecule generation.
Problem
Masked discrete diffusion models cannot revise tokens once unmasked (the "failure to remask" property), limiting sample quality, speed-quality tradeoffs, and controllable generation.
Approach
A non-Markovian discrete diffusion posterior with a tunable remasking probability sigma_t for decoded tokens; it reuses pretrained MDLM weights (its ELBO is a reweighted MDLM objective) and samples with a remasking schedule, which can be interpreted as a predictor-corrector and generalizes FB and DFM correctors.
Key results
- On OpenWebText with T=4096, ReMDM reaches MAUVE 0.656 versus 0.243 for MDLM+FB and 0.035 for MDLM with nucleus sampling, a large inference-time scaling gain.
- On ImageNet 256 with T=64, ReMDM attains FID 4.45 versus 4.85 for MaskGiT and 4.69 for MDLM.
- On LLaDA 8B Instruct, ReMDM improves Countdown pass@1 from 45.2 to 46.1 and TruthfulQA delta ROUGE-1 from 27.1 to 29.5.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ReMDM instead of MDLM | MDLM text generation · OpenWebText T=4096 | MAUVE: 0.035 -> 0.656 | Helped |
| Use ReMDM instead of MDLM+DFM | MDLM text generation · OpenWebText T=4096 | MAUVE: 0.269 -> 0.656 | Helped |
| Use ReMDM instead of MDLM for faster sampling | MDLM text generation · OpenWebText T=512 | MAUVE: 0.031 -> 0.350 | Helped |
| Increase inference steps with ReMDM | ReMDM text generation · OpenWebText | MAUVE: 0.403 -> 0.610 -> 0.656 | Helped |
| Use ReMDM instead of MaskGiT | MaskGiT image generation · ImageNet 256×256 T=64 | FID: 4.85 -> 4.45 | Helped |
| Use ReMDM instead of MDLM | Image generation · ImageNet 256×256 T=16 | FID: 7.88 -> 7.40 | Hurt |
| Use ReMDM instead of MaskGiT | Image generation · ImageNet 256×256 T=32 | FID: 4.92 -> 4.92 | No clear effect |
| Add remasking ability and the proposed sampling improvements (ablation) | MDLM text generation · OpenWebText | MAUVE | Helped |
| Use ReMDM-rescale rather than ReMDM-cap (ablation) | ReMDM image generation · ImageNet | FID | Helped |
| Increase the schedule η parameter (ablation) | ReMDM image generation · ImageNet | FID | Helped |
| Add confidence-based scheduling to ReMDM-rescale (ablation) | ReMDM molecule generation · QM9 ring-count guidance | novelty-property frontier | Helped |
| Use ReMDM instead of original LLaDA | LLaDA 8B Instruct · Countdown | pass@1%: 45.2 -> 46.1 | Helped |
Surprising findings
- Gen PPL can be made arbitrarily low by degenerate repetitive text, so the authors use MAUVE instead.
- Sampling with 32-bit floating point significantly curbs diversity, so the authors use 64-bit for text sampling.
Limitations
- Gains at small step budgets are less pronounced; for some settings (e.g., T=128 MAUVE) ReMDM is comparable to or below MDLM+FB/DFM.
- Gen PPL is unreliable for corrector samplers, so evaluation relies on MAUVE and entropy.
- Inference cost grows with number of sampling steps.
Experimental setup
Datasets: OpenWebText ImageNet QM9 Countdown TruthfulQA
Benchmarks: OpenWebText MAUVE ImageNet FID/IS QM9 novelty and ring count Countdown pass@1 TruthfulQA ROUGE difference
Metrics: MAUVE generative perplexity entropy FID IS novelty pass@1 delta ROUGE
Baselines: AR SEDD MDLM MDLM+FB corrector MDLM+DFM corrector MaskGiT UDLM FUDGE LLaDA 8B Instruct + DFM
Builds on: MDLM DDIM forward-backward corrector discrete flow matching
Closest papers by meaning
- Learn from Your Mistakes: Self-Correcting Masked Diffusion Models (NeurIPS 2026 · similarity 0.92)
- Simplified and Generalized Masked Diffusion for Discrete Data (NeurIPS 2024 · similarity 0.91)
- Simple and Effective Masked Diffusion Language Models (NeurIPS 2024 · similarity 0.91)
- Encoder-Decoder Diffusion Language Models for Efficient Training and Inference (NeurIPS 2025 · similarity 0.91)
- Diffusion Beats Autoregressive in Data-Constrained Settings (NeurIPS 2025 · similarity 0.91)
- Re-evaluating Confidence Remasking in Masked Diffusion Language Models (NeurIPS 2026 · similarity 0.91)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.