# Know When to Hold 'em: Correct-Token Retention in Uniform-State Diffusion Language Models

arXiv: https://arxiv.org/abs/2610.01275 · HTML: https://mlsanity.com/paper/2610.01275/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Discrete and text diffusion, one of 14 directions in Diffusion models.
- Builds on 28 of its 31 references that are in the map.
- First posted on arXiv 1 Oct 2026.

## Placed in

- Diffusion models / [Discrete and text diffusion](https://mlsanity.com/diffusion/discrete-and-text-diffusion/)

## Builds on (most cited first, 28 of 28)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Structured Denoising Diffusion Models in Discrete State-Spaces](https://mlsanity.com/paper/2107.03006/) (2021)
- [Simple and Effective Masked Diffusion Language Models](https://mlsanity.com/paper/2406.07524/) (2024)
- [Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution](https://mlsanity.com/paper/2310.16834/) (2023)
- [Movie Gen: A Cast of Media Foundation Models](https://mlsanity.com/paper/2410.13720/) (2024)
- [Simplified and Generalized Masked Diffusion for Discrete Data](https://mlsanity.com/paper/2406.04329/) (2024)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [A Discourse-Aware Attention Model for Abstractive Summarization of Long Documents](https://mlsanity.com/paper/1804.05685/) (2018)
- [MAUVE: Measuring the Gap Between Neural Text and Human Text using Divergence Frontiers](https://mlsanity.com/paper/2102.01454/) (2021)
- [Non-Autoregressive Neural Machine Translation](https://mlsanity.com/paper/1711.02281/) (2017)
- [A Continuous Time Framework for Discrete Denoising Models](https://mlsanity.com/paper/2205.14987/) (2022)
- [One Billion Word Benchmark for Measuring Progress in Statistical Language Modeling](https://mlsanity.com/paper/1312.3005/) (2013)
- [Mercury: Ultra-Fast Language Models Based on Diffusion](https://mlsanity.com/paper/2506.17298/) (2025)
- [Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps](https://mlsanity.com/paper/2501.09732/) (2025)
- [Seed Diffusion: A Large-Scale Diffusion Language Model with High-Speed Inference](https://mlsanity.com/paper/2508.02193/) (2025)
- [Noise2Music: Text-conditioned Music Generation with Diffusion Models](https://mlsanity.com/paper/2302.03917/) (2023)
- [Simple Guidance Mechanisms for Discrete Diffusion Models](https://mlsanity.com/paper/2412.10193/) (2024)
- [Concrete Score Matching: Generalized Score Matching for Discrete Data](https://mlsanity.com/paper/2211.00802/) (2022)
- [The Diffusion Duality](https://mlsanity.com/paper/2506.10892/) (2025)
- [Informed Correctors for Discrete Diffusion Models](https://mlsanity.com/paper/2407.21243/) (2024)

## Abstract (arXiv)

Uniform-state diffusion models (USDMs) can revise any token at any denoising step, which lets them correct their own mistakes, a key advantage over masked diffusion. Self-correction, however, requires both revising incorrect tokens and retaining correct ones, and we show that current USDMs lack the latter. Even under greedy-tail decoding, state-of-the-art USDMs (DUO, UDLM, and uniform-noise SEDD) keep revising 173--270 of 512 positions at every step, and these large, uncoordinated edits collapse sample diversity. A random-token corruption experiment traces this deficit to the models themselves: they reconstruct clean and corrupted tokens with nearly identical accuracy, even though clean tokens are easier targets. A decomposition of the validation NELBO shows that training barely rewards retention: incorrect predictions are heavily penalized at corrupted positions but almost free at clean ones. We propose Correct-Token Retention Regularization (CTR-Reg), a simple but effective auxiliary loss that trains the model to retain tokens left unperturbed by the forward process and requires no change to the sampler. CTR-Reg improves clean-token accuracy by 26.5 percentage points on average across six benchmarks, while leaving corrupted-token accuracy virtually unchanged, and its per-step revisions converge to only 3--11 positions. With just five greedy-tail steps, generative perplexity more than halves under CTR-Reg for all three models while diversity is preserved, and these gains hold across sampling budgets. Our results identify correct-token retention as a key missing ingredient for self-correcting diffusion language models, and demonstrate an effective fix.