PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Mojtaba Nafez, James Henderson

arXiv:2610.01275v1cs.CL

TL;DR

Current uniform-state diffusion models can revise tokens but do not reliably retain correct ones, while their training objective provides little incentive to do so. The paper adds Correct-Token Retention Regularization to the NELBO without changing the sampler. Across three models, it improves clean-token accuracy and more than halves generative perplexity within five greedy-tail steps while preserving diversity.

  • Problem

    Current USDMs reconstruct clean and corrupted tokens with nearly identical accuracy despite clean tokens being easier targets, and their weak correct-token retention accompanies diversity collapse under greedy-tail sampling.

    Source: §1 · p. 2

  • Method

    CTR-Reg adds an auxiliary loss that trains models to retain tokens left unperturbed by the forward process, without modifying the sampler.

    Source: §1 · p. 2

  • Results

    CTR-Reg improves clean-token accuracy by 26.5 points on average across six benchmarks, while five greedy-tail steps more than halve generative perplexity for all three models with diversity preserved.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Correct-token retention is a key ingredient for self-correcting diffusion language models.

    Source: §5 · p. 9

    Models are small (169M parameters) and evaluated on publicly available research corpora, whose web-derived text may contain biases that trained models can reproduce.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Uniform-state diffusion models can revise any token, but current models fail to reliably retain correct tokens, contributing to diversity collapse under greedy-tail decoding. The paper traces this deficit to weak NELBO incentives and introduces CTR-Reg, which improves retention and generation quality without changing the sampler.

1 INTRODUCTION

  • Under greedy-tail decoding, DUO, UDLM, and SEDD exhibit substantial sample-diversity collapse, despite USDMs’ ability to revise tokens at every step.Large, independent revisions can compound into globally inconsistent sequences; ancestral sampling limits average revisions per step but does not select which positions change.Source: §1 · p. 2
  • A random-token corruption experiment finds nearly identical reconstruction accuracy on clean and corrupted tokens, although clean tokens are easier targets.The experiment replaces a random share of tokens with uniform-vocabulary draws and measures accuracy separately on corrupted and unaltered positions.Source: §1 · p. 2
  • Standard NELBO training barely penalizes incorrect predictions at clean positions: their errors cost 0.09–1.32 nats more, versus 13–100 nats at corrupted positions.Mispredicted clean positions account for 23–31% of positions but only 3.5–5.2% of the NELBO.Source: §1 · p. 2
  • CTR-Reg is an auxiliary loss that trains models to retain tokens left unperturbed by the forward process, without changing the sampler.It supplies a retention signal missing from standard NELBO training.Source: §1 · p. 2
  • NELBO-only models revise 173–270 of 512 positions at every greedy-tail step, whereas CTR-Reg plateaus at 3–11 changed positions.The baseline shows no sign of convergence by step 20; CTR-Reg reaches its plateau within about 10 steps.Source: Figure 2 · p. 3, §1 · p. 3
  • With five greedy-tail steps, generative perplexity falls from 105.86 to 49.57 for DUO, 93.28 to 44.31 for UDLM, and 194.59 to 87.19 for SEDD, with diversity above training-data entropy.The reported gains hold across step budgets and are also reflected in MAUVE, which improves with more greedy-tail steps under CTR-Reg.Source: §1 · p. 3, Figure 3 · p. 3
  • CTR-Reg raises clean-token accuracy by 26.5 points on average across six benchmarks while leaving corrupted-token accuracy virtually unchanged.Across models and both OpenWebText and arXiv, the figure reports clean-token gains of 21.5–29.3 percentage points.Source: Figure 1 · p. 2, §1 · p. 3

2 BACKGROUND AND RELATED WORK

The background introduces discrete diffusion as a learned reverse process for categorical data, then distinguishes uniform-state models by their uniform prior and revisable positions. It also describes the shared NELBO objective used by current USDMs despite differences in parameterization.

2.1 DISCRETE DIFFUSION MODELS

  • Discrete diffusion learns a parameterized reverse process that inverts forward corruption of categorical data such as text.The forward process maps clean data to increasingly noisy latent variables over time.Source: §2.1 · p. 3
  • The denoising network predicts a clean sequence from noisy input, and the reverse transition substitutes that prediction into the forward posterior.The reverse process factorizes across positions, and training minimizes the negative evidence lower bound (NELBO).Source: §2.1 · p. 4

2.2 UNIFORM-STATE DIFFUSION LANGUAGE MODELS

  • Uniform-state diffusion sets the prior to uniform over the vocabulary, so any position can be revised at every step rather than locked after decoding.The transition assigns non-zero probability both to keeping the current token and to moving to another token, providing the basis for self-correction.Source: §2.2 · p. 4
  • DUO, UDLM, and SEDD optimize equivalent continuous-time NELBO bounds up to parameterization, although their formulations and computations differ.DUO rewrites per-token terms and uses Rao-Blackwellization; SEDD predicts ratios of marginals and uses diffusion-weighted denoising score entropy.Source: §2.2 · p. 4

3 METHODOLOGY

The section diagnoses a correct-token retention deficit in current USDMs and traces it to weak NELBO incentives. It introduces CTR-Reg to target retention and adapts the regularization for score-parameterized SEDD.

3.1 CURRENT USDMS LACK CORRECT-TOKEN RETENTION

  • Under compute-matched greedy-tail sampling, DUO entropy falls from 5.54 to 2.14 over 20 steps and stays below the real-text level of 5.40 for every g ≥1.The sampler appends deterministic per-position argmax steps after ancestral sampling.Source: §3.1 · p. 5
  • At step 20, 173–270 of 512 positions still change; simultaneous independent edits compound into globally inconsistent, low-diversity sequences.The number of changed positions starts near half the sequence and shows no sign of convergence.Source: §3.1 · p. 5
  • Ancestral persistence is schedule-driven rather than evidence of model retention; near the terminal region, at most roughly one in five disagreeing positions adopts the prediction.For smooth schedules, the adoption probability is O(1/T); greedy-tail decoding removes this stochastic bound and exposes the deficit.Source: §3.1 · p. 6
  • At 1−αt < 0.4, incorrect predictions cost 13–100 more nats at corrupted positions but only 0.09–1.32 at clean positions.Figure 5 depicts the per-token validation NELBO by noise level, position type, and prediction correctness.Source: §3.1 · p. 6, Figure 5 · p. 6

3.2 CORRECT-TOKEN RETENTION REGULARIZATION (CTR-REG)

  • CTR-Reg rewards the model for assigning probability to the clean token only where forward noising left that token unchanged.Its nonnegative penalty preserves a valid, though looser, upper bound on negative log-likelihood.Source: §3.2 · p. 6
  • Clean-token accuracy rises by 26.5 points, while validation perplexity changes by +1.5% for DUO, +1.2% for UDLM, and −2.4% for SEDD.CTR-Reg adds a nonnegative term to the NELBO while retaining a valid, looser likelihood bound.Source: §3.2 · pp. 6–7
  • CTR-Reg acts only during training, reuses the NELBO forward pass, adds no computational overhead, and leaves the sampler unchanged.The paper uses λ = 0.02 for every USDM.Source: §3.2 · p. 7

3.3 APPLYING CTR-REG TO SCORE-PARAMETERIZED MODELS

  • For SEDD, CTR-Reg penalizes off-diagonal scores above the unperturbed-position target with a softplus hinge that fades below that target.The target normalization places the hinge at SEDD’s training target.Source: §3.3 · p. 7

4 EXPERIMENTS

Across three USDMs, CTR-Reg improves clean-token retention while leaving corrupted-token accuracy nearly unchanged, and supports better generations without the diversity collapse seen under NELBO-only greedy-tail decoding. The experiments also show that stronger regularization is not always better.

4 EXPERIMENTS

  • The study evaluates CTR-Reg on DUO, UDLM, and SEDD with uniform noise, comparing NELBO-only and NELBO + CTR-Reg models on random-token retention and revision, plus unconditional generation quality.Source: §4 · p. 7, §4.1 · p. 7

4.1 CTR-REG IMPROVES RETENTION WITHOUT HURTING REVISION

  • CTR-Reg improves Clean-Token Accuracy by 26.5 points on average across six corpora, while changing Corrupted-Token Accuracy by only 0.7 points per cell on average.The evaluation uses single-pass reconstruction after random-token corruption, measuring accuracy separately on corrupted and clean positions.Source: §4.1 · pp. 7–8

4.2 BETTER RETENTION YIELDS BETTER GENERATIONS

  • From g = 0 to g = 5, CTR-Reg lowers Gen PPL from 105.86 to 49.57 for DUO, from 93.28 to 44.31 for UDLM, and from 194.59 to 87.19 for SEDD while keeping entropy at or above real OpenWebText.The gains hold across sampling budgets from 16 to 1024 steps.Source: §4.2 · p. 8
  • With additional greedy-tail steps, CTR-Reg models improve MAUVE to 0.81 for DUO, 0.83 for UDLM, and 0.58 for SEDD, while NELBO-only models fall as low as 0.13–0.17 for DUO and UDLM and 0.15 for SEDD.At g = 0, the two training approaches score comparably.Source: §4.2 · p. 8

4.3 HOW CTR-REG RESOLVES THE DIVERSITY COLLAPSE

  • CTR-Reg reduces changed positions to a stable plateau of 3–11, whereas NELBO-only models still revise 173–270 of 512 positions at the twentieth greedy-tail step.The paper attributes the diversity gap to bounded, targeted revision preserving sequence structure while indiscriminate revision destroys it.Source: §4.3 · pp. 8–9

4.4 REGULARIZATION STRENGTH ANALYSIS

  • At λ = 0.1, clean-token accuracy rises to 87.2% for DUO and 83.6% for UDLM, but λ = 0.02 yields stronger generation gains from additional greedy-tail steps.The copy rate stays below 0.25%; the authors suggest excessive retention can hinder revision of earlier sampling mistakes.Source: §4.4 · p. 9

5 CONCLUSION

The paper concludes that correct-token retention is a missing ingredient for self-correcting USDMs and reports CTR-Reg as an effective fix. Appendix results indicate that generation gains persist across sampling budgets.

5 CONCLUSION

  • CTR-Reg raises clean-token accuracy by 26.5 points on average without hurting revision and more than halves generative perplexity within five greedy-tail steps while preserving diversity across DUO, UDLM, and SEDD.The conclusion identifies correct-token retention as a key ingredient for self-correcting diffusion language models.Source: §5 · p. 9
  • The models are small (169M parameters) and intended solely for research; web-derived training text may contain biases that models can reproduce.The work uses publicly available research corpora and involves no human subjects or personal data collection.Source: ETHICS STATEMENT · p. 10

A APPENDIX

  • Across budgets from 16 to 1024 steps, CTR-Reg matches or outperforms NELBO-only training under default sampling and lowers Gen PPL while preserving diversity under greedy-tail sampling.The appendix contrasts these gains with NELBO-only greedy-tail models, whose lower PPL comes at the cost of diversity.Source: Figure 7 · p. 15

B STEP-BUDGET ANALYSIS

The appendix reports that CTR-Reg's generation benefits persist across sampling budgets and provides further details on the models' training objectives and ancestral sampling. It distinguishes schedule-limited sampler updates from denoiser predictions.

B STEP-BUDGET ANALYSIS

  • Across 16–1024-step budgets, CTR-Reg matches or outperforms NELBO-only training under default sampling and substantially lowers Gen PPL while preserving diversity under greedy-tail sampling.The apparent NELBO-only PPL gains under greedy-tail sampling come at the cost of diversity.Source: B STEP-BUDGET ANALYSIS · p. 15

C PER-MODEL USDM OBJECTIVES

  • DUO and SEDD objectives are presented as equivalent to the continuous-time bound that UDLM optimizes directly; SEDD predicts reverse-rate ratios rather than clean data.The appendix describes DUO's formulation as using Rao-Blackwellization and a training curriculum.Source: C PER-MODEL USDM OBJECTIVES · pp. 15–16

D ANCESTRAL SAMPLING SUPPRESSES INDISCRIMINATE TOKEN CHANGES

  • The ancestral sampler's noise schedule assigns most probability mass to retaining the currently decoded token, limiting how many positions can adopt the denoiser's prediction in one step.The analysis asks how readily the sampler incorporates a change the model has already decided on.Source: D ANCESTRAL SAMPLING SUPPRESSES… · p. 16, §D.1 · p. 16

D.1 SINGLE-POSITION TRANSITION PROBABILITIES

  • The appendix derives single-position transition probabilities and notes that their exact weights normalize to one for any interpolating noise schedule.It drops O(1/K) terms as negligible for K ≈ 50k when collecting category weights.Source: §D.1 · pp. 16–17

D.2 THE SCHEDULE BOUNDS THE NUMBER OF CHANGES PER STEP

The schedule suppresses model-driven updates through most of ancestral sampling, while greedy-tail decoding removes that suppression and exposes unstable denoiser predictions.

D.2 THE SCHEDULE BOUNDS THE NUMBER OF CHANGES PER STEP

  • Through most of the trajectory, the schedule revises only a small fraction of positions that disagree with the denoiser, and at αt = 0.99 roughly four out of five such positions still fail to update in one step.Early in generation, many positions change, but almost all changes are draws from the uniform prior rather than adoptions of the denoiser's prediction.Source: §D.2 · p. 17

D.3 WHY THE GREEDY TAIL REVEALS WHAT ANCESTRAL SAMPLING HIDES

  • Greedy-tail decoding sets the probability of jumping to the model's prediction to 1, so all m disagreeing positions update at once rather than the O(mh) fraction permitted by ancestral sampling.The appendix says this exposes instability that ancestral sampling can hide.Source: §D.3 · p. 17

E LIKELIHOOD ANALYSIS

The likelihood analysis finds that CTR-Reg changes validation perplexity only modestly while the NELBO assigns far less penalty to incorrect predictions at clean than corrupted positions. This imbalance is consistent across models and corpora, with noisy estimates at the highest noise levels.

E.2 DECOMPOSING THE NELBO BY POSITION TYPE AND PREDICTION

  • The validation decomposition uses matched diffusion times and corruptions for NELBO-only and CTR-Reg models of each USDM, averaging per-token NELBO across position and prediction groups.It evaluates six corpora at sequence length 512; UDLM uses a separate data pipeline, and SEDD's per-token magnitudes are not directly comparable with the other models' because of its geometric schedule.Source: E LIKELIHOOD ANALYSIS · p. 18, §E.2 · p. 18

E.1 VALIDATION PERPLEXITY

  • CTR-Reg changes validation perplexity by +1.5% for DUO and +1.2% for UDLM on average, and decreases it by 2.4% for SEDD on average.The largest increase is 4.4% for DUO on PubMed; the in-distribution OWT change is below 1% for DUO and UDLM.Source: §E.1 · p. 18

E.2 DECOMPOSING THE NELBO BY POSITION TYPE AND PREDICTION

  • At noise below 0.4, incorrect corrupted-position predictions cost 13–100 nats more than correct ones, versus only 0.09–1.32 nats at clean positions.Incorrect clean predictions account for 23–31% of positions but only 3.5–5.2% of NELBO; incorrect corrupted predictions account for 90–92%. At the highest noise levels, few clean tokens make estimates noisier.Source: §E.2 · p. 18

F TRAINING ON A LARGER CORPUS DOES NOT FIX THE PROBLEM

Training DUO on 26× more data improves corrupted-token accuracy but leaves the clean-versus-corrupted accuracy gap essentially unchanged, while CTR-Reg closes that gap with less training data. The appendix also describes how SEDD's scores are read as token predictions.

F TRAINING ON A LARGER CORPUS DOES NOT FIX THE PROBLEM

  • The 524B-token DUO checkpoint improves corrupted-token accuracy from 45.2% to 51.6% on average, but clean-token accuracy rises comparably from 47.2% to 52.7%.The clean-minus-corrupted gap is 2.0 points at 20B tokens and 1.1 points at 524B, described as unchanged within noise.Source: F TRAINING ON A LARGER CORPUS DOES NOT… · pp. 18–19, Table 7 · p. 23
  • At clean positions, incorrect predictions incur a near-zero NELBO penalty except at the highest noise levels, consistent with the pattern reported across the remaining four corpora.Figure 8 compares DUO, UDLM, and SEDD on LM1B, AG News, PubMed, and WikiText-103; magnitudes are comparable within, not across, panels.Source: Figure 8 · p. 19
  • Scaling DUO from 20B to 524B training tokens leaves the clean-minus-corrupted accuracy gap essentially unchanged, while CTR-Reg opens a 31.1-point average gap using the 20B-token budget.CTR-Reg's clean-token accuracy averages 77.4% versus 52.7% for the 524B-token checkpoint, despite using 26× fewer training tokens.Source: F TRAINING ON A LARGER CORPUS DOES NOT… · pp. 19, 23, Table 7 · p. 23

G.1 READING PREDICTIONS FROM SCORES

  • SEDD's positive scores rank tokens the same way as its implied clean-token predictions, so greedy-tail argmax can be taken directly over scores.When the displayed token's score is fixed at 1, it is retained exactly when the model ranks it above every alternative.Source: §G.1 · pp. 23–24

F TRAINING ON A LARGER CORPUS DOES NOT FIX THE PROBLEM

  • For the 524B-token DUO checkpoint, greedy-tail perplexity falls as steps increase, but entropy drops below the sequence-length-matched real-data reference for every setting with at least one step.Table 8 uses the checkpoint's native sequence length of 1024 and reference entropy 5.40; its PPL and entropy magnitudes should not be directly compared with the 512-token models.Source: Table 8 · p. 24

G.2 THE TRAINING TARGET ρt

For SEDD, the training objective regresses per-coordinate scores toward marginal ratios, with different targets for clean and corrupted positions. The CTR-Reg hinge is placed at the clean-position target, and its accuracy gains generalize across evaluation corpora.

G.2 THE TRAINING TARGET ρt

  • SEDD's objective regresses each off-diagonal score toward a marginal ratio, while the displayed token's own score is fixed at 1 and is not trained by the sum.The ratio target is defined from forward marginals as r_t(j) = q_t(j | x)/q_t(i | x).Source: §G.2 · p. 24
  • For all but the noisiest times, the small ρt target favors retaining the displayed token at clean positions, while the score for the original token is pushed above it at corrupted positions.Source: §G.2 · p. 24

G.3 PLACING THE HINGE AT THE TARGET

  • CTR-Reg leaves corrupted-token accuracy essentially unchanged while improving clean-token and overall accuracy across six datasets; placing the hinge at ρt makes its push persist until the clean-position target.The softplus threshold c = ρt fades below the clean-position target, unlike c = 1, whose push fades once a score falls below one.Source: Table 9 · p. 25, §G.3 · p. 25

H RANDOM-TOKEN CORRUPTION ON ADDITIONAL CORPORA

The random-token corruption evaluation finds that CTR-Reg improves clean-token accuracy across six corpora while leaving corrupted-token accuracy nearly unchanged. Additional tables report detailed results across corruption rates and generation-sampling evaluations.

H RANDOM-TOKEN CORRUPTION ON ADDITIONAL CORPORA

  • CTR-Reg raises clean-token accuracy by 26.5 percentage points on average across six corpora, while changing corrupted-token accuracy by less than 1.7 points in every cell.The patterns hold across web, news, scientific, biomedical, and encyclopedic text, and at every individual corruption rate in the reported breakdowns.Source: H RANDOM-TOKEN CORRUPTION ON… · p. 25

I ADDITIONAL EXPERIMENTAL RESULTS

I.1 EXACT VALUES FOR THE GREEDY-TAIL STEP SWEEP

  • Table 10 gives exact perplexity and entropy values for each of three USDMs across greedy-tail step counts.Source: §I.1 · p. 26

I.2 EXACT VALUES ACROSS SAMPLING-STEP BUDGETS

  • Table 11 reports exact perplexity and entropy values for sampling budgets from 16 to 1024 under ancestral and greedy-tail sampling.Source: §I.2 · p. 26

I.3 EXACT VALUES FOR THE MAUVE GREEDY-TAIL SWEEP

  • Table 17 reports exact PPL, entropy, and MAUVE values across greedy-tail step counts for all three USDMs, with and without CTR-Reg.Source: §I.3 · p. 26

I.4 PER-CORRUPTION-RATE BREAKDOWN FOR THE 524B-TOKEN CHECKPOINT

  • For the 524B-token DUO checkpoint, clean- and corrupted-token accuracies remain close at every individual corruption rate, not only in the averages.Source: §I.4 · p. 26

I.5 EXACT VALUES FOR THE CTR-REG COEFFICIENT ABLATION

Table 19 gives exact PPL and generation entropy values for DUO and UDLM across CTR-Reg coefficients and greedy-tail step counts, with the total step budget fixed at 512.

I.5 EXACT VALUES FOR THE CTR-REG COEFFICIENT ABLATION

  • Table 19 reports the exact PPL and entropy values underlying Figure 6 for DUO and UDLM at λ = 0, 0.02, and 0.1 across every greedy-tail step count g.Source: §I.5 · p. 26
  • PPL and generation entropy are reported for each coefficient and step-count combination.PPL is computed with GPT2-Large over 1024 generated samples; the table marks the lowest PPL among the three coefficient settings.Source: §I.5 · p. 26

I.6 PER-CORRUPTION-RATE BREAKDOWN FOR THE CTR-REG COEFFICIENT ABLATION

These tables expand coefficient-ablation and corruption-probe averages into per-corruption-rate views across models and corpora, alongside copy-rate and generation-quality results.

I.6 PER-CORRUPTION-RATE BREAKDOWN FOR THE CTR-REG COEFFICIENT ABLATION

  • <0.35% copy rate holds at every corruption rate; DUO’s rate rises with p on AG News, while UDLM’s stays flat or decreases slightly.For DUO at λ = 0.1 on AG News, the copy rate rises from 0.15% to 0.30%.Source: §I.6 · p. 26, Table 21 · p. 38
  • Across greedy-tail steps and sampling budgets, the tables pair PPL with entropy; CTR-Reg keeps entropy at or above the 4.83 reference while matching or improving PPL.The NELBO-only baseline’s lower PPL under greedy-tail sampling is paired with sub-threshold entropy; Table 17 also reports MAUVE, which it says collapses as baseline entropy falls.Source: Table 10 · p. 27, Table 11 · p. 28, Table 17 · p. 34
  • Per-corruption-rate results compare NELBO alone with NELBO + CTR-Reg on AG News, arXiv, LM1B, PubMed, and WT-103 for three USDMs.Each table reports results at corruption rates of 10, 20, 30, and 40%, with averages corresponding to Table 9.Source: Table 12 · p. 29, Table 13 · p. 30, Table 14 · p. 31
  • The 524B-token DUO checkpoint’s Table 7 results are broken down by corruption rate across all six evaluation corpora.Table 18 reports corrupted-token and clean-token accuracy, with averages matching Table 7.Source: Table 18 · p. 35
  • Table 20 breaks out the λ = 0.1 corruption-probe results for DUO and UDLM across six corpora and four corruption rates.The averages correspond to Table 1, which reports corruption-probe accuracy averaged over those corruption rates and coefficient settings.Source: §I.6 · p. 26, Table 20 · p. 37
Loading 2610.01275v1…