PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent Refinement Decoding: Enhancing Diffusion-Based Language Models by Refining Belief States

Qinglin Zhu, Yizhen Yao, Runcong Zhao, Yanzheng Xiang, Amrutha Saseendran, Chen Jin, Philip Teare, Bin Liang, Yulan He, Lin Gui

arXiv:2510.11052v2cs.CL

TL;DR

Parallel diffusion decoding avoids strictly sequential generation, but hard masking discards uncertain-token beliefs and can lead to premature commitments. LRD refines soft beliefs before progressively committing confident tokens, and reports accuracy gains across coding and reasoning benchmarks with speedups up to 10.6×.

  • Problem

    Diffusion language models discard predictive distributions for unfinalized positions, while early commitment or fixed iteration counts can impair coordination or waste computation.

    Source: §1 · pp. 1–2

  • Method

    LRD first refines masked positions as mixtures of mask and predicted-token embeddings, then progressively commits confident tokens while retaining uncertain ones for feedback.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    LRD consistently improves accuracy across coding and reasoning benchmarks and achieves speedups up to 10.6×.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The authors conclude LRD can serve as a flexible drop-in decoding module for future diffusion-based language models.

    Source: §6 · p. 10

    Soft embeddings cannot fully collapse distributions to one-hot states, slowing convergence in later stages when sharper updates are needed.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

Autoregressive (AR) models remain the standard for natural language generation but still suffer from high latency due to strictly sequential decoding. Recent diffusion-inspired approaches, such as LlaDA and Dream, mitigate this by generating in parallel, yet they suffer from two core limitations: information loss, as predictive distributions for non-finalized tokens are discarded at each step, and premature commitment, where local decisions are made without sufficient global coordination. We introduce Latent Refinement Decoding (LRD), a two-stage framework with Latent Refinement and a Predictive Feedback Loop. The first stage maintains masked positions as distributional mixtures of predicted tokens and the mask embedding, allowing the model to establish more globally consistent beliefs. The second stage progressively finalizes confident tokens while retaining uncertain ones for iterative feedback. KL-divergence dynamics provide a principled and reliable criterion for convergence and early stopping. Experiments across coding (HumanEval +6.3, MBPP +2.6) and reasoning (GSM8K +2.9, MATH500 +3.8) show that LRD improves accuracy while delivering speedups of up to 10.6x, making it a strong and versatile alternative for parallel sequence generation.

1 INTRODUCTION

Diffusion language models decode in parallel but discard predictive beliefs for unfinalized tokens and can commit before positions are globally coordinated. LRD refines beliefs in embedding space before progressively finalizing confident tokens, and reports gains across coding and reasoning benchmarks.

1 INTRODUCTION

  • Hard masking discards the predictive distributions of uncertain positions, while aggressive commitment can lock in errors and conservative masking slows decoding.Fixed iteration counts can also waste computation on simple cases or underserve complex ones.Source: §1 · pp. 1–2
  • LRD mixes masked embeddings with predicted-token representations, then progressively commits confident tokens while retaining uncertain positions in soft form.The framework uses KL dynamics to monitor convergence and enable adaptive early stopping.Source: §1 · p. 2
  • LRD combines soft refinement with adaptive two-phase sampling, using KL monitoring for phase transitions and early stopping.Its reported gains span HumanEval (+6.3), MBPP (+2.6), GSM8K (+2.9), and MATH500 (+3.8), with speedups up to 10.6×.Source: §1 · p. 2

2 PRELIMINARY

The section describes categorical corruption in the forward process, iterative denoising in reverse, and hard-assignment sampling that discards information at positions left masked.

2 PRELIMINARY

  • The forward process corrupts clean token sequences into noisier versions, with tokens remaining unchanged or transitioning to [MASK].The cumulative transition probability includes the probability that a token remains unmasked from time 0 to t.Source: §2 · p. 2
  • The reverse process iteratively denoises a fully masked sequence toward clean text, predicting a token distribution at each position.The model predicts distributions over tokens for individual positions at each denoising step.Source: §2 · p. 2
  • Traditional sampling selects positions by confidence and assigns them token embeddings, while unselected positions return to a fixed [MASK] embedding.This binary assignment discards predictive distributions for positions below the confidence threshold.Source: §2 · p. 3

3 METHODOLOGY

LRD replaces hard remasking with continuous embedding mixtures, then shifts from belief refinement to discrete commitments using convergence monitoring and predictive feedback.

3 METHODOLOGY

  • LRD interpolates between [MASK] and token embeddings according to prediction uncertainty, preserving a gradual denoising trajectory.High-entropy positions remain more mask-like, while low-entropy positions incorporate more token information.Source: §3 · p. 3

3.1 SOFT DIFFUSION

  • Entropy controls the interpolation strength: uncertain positions stay mask-like, while confident positions move toward token embeddings.The entropy is normalized by log |V|, and rf sets the maximum interpolation strength.Source: §3.1 · p. 4
  • Soft posterior mixtures retain nonzero probability for every token and propagate uncertainty between positions through self-attention.The mixture can be viewed as the expected embedding under the model’s belief at a position.Source: §3.1 · p. 4

3.2 ADAPTIVE SAMPLING WITH SOFT-TO-HARD SCHEDULING

  • Soft mixtures cannot fully collapse to one-hot distributions, so convergence slows in later stages when sharper updates are needed.KL monitoring can stop decoding when whole-sentence predictive distributions stabilize, avoiding redundant iterations.Source: §3.2 · pp. 4–5
  • Phase 1 iteratively refines distributions using predictions conditioned on prior soft embeddings, then transitions when consecutive predictions stabilize or the refinement limit is reached.The schedule moves to hard assignment after stabilization, or after T_refine steps for computational efficiency.Source: §3.2 · pp. 4–5
  • In Phase 2, confident positions become discrete tokens while unselected positions retain soft embeddings for subsequent decoding context.This preserves Phase 1 distributional information in positions that remain uncommitted.Source: §3.2 · p. 5

4 EXPERIMENTS

LRD experiments evaluate two diffusion-based language models on reasoning and code benchmarks, showing accuracy gains and faster decoding across settings. Convergence analysis and ablations examine when refinement helps and how mixing and early stopping affect performance and efficiency.

4.1 IMPLEMENTATION DETAILS

  • Experiments use LLaDA 8B and Dream 7B Base and Instruct variants, with temperature 0 and minimum-entropy token selection.All experiments ran on 8 NVIDIA A100 80GB GPUs.Source: §4.1 · p. 5

4.2 BENCHMARKS AND METRICS

  • Evaluation covers GSM8K and MATH500 for mathematical reasoning and MBPP and HumanEval for code generation, reporting accuracy and pass@1, respectively.Instruct models use zero-shot evaluation; Base models use benchmark-specific few-shot settings.Source: §4.2 · p. 5

4.3 MAIN RESULTS

  • LRD raises HumanEval pass@1 by up to +6.3 points (Dream-Base-7B, 256 tokens) and provides at least 1.2× speedup in all cases.Dream-Ins-7B reaches up to 9.3× faster decoding at length 1024; LLaDA models reach up to 4.3× at that length.Source: Table 1 · p. 6, §4.3 · p. 6

4.4 CONVERGENCE ANALYSIS

  • KL divergence drops rapidly during latent refinement and then declines during decoding, approaching zero within about 300 steps on most benchmarks, but persisting through 512 on MATH500.Dream-Ins approaches zero around 140 steps; Figure 2 marks decoding onset after 20 refinement steps.Source: §4.4 · p. 6, Figure 2 · p. 7
  • On HumanEval with Dream-Ins, 68.9% of samples converge by refinement step 2 and over 85% by step 3; across other benchmarks, over 70% converge within three to four steps.Figure 3 plots convergence ratios starting at step 2 because calculating KL-difference requires three consecutive steps.Source: §4.4 · p. 7, Figure 3 · p. 7

4.5 ABLATION STUDY

  • Repeated latent refinement (LF×2–5) generally degrades accuracy, and adaptive refinement also underperforms the two-stage strategy.The comparison is at length 512; the authors attribute the result to redundant computation after stabilization.Source: §4.5 · p. 7, Table 2 · p. 7
  • Removing mixed embeddings causes accuracy drops up to −2.9 on HumanEval and −2.5 on MATH500, while early stopping incurs almost no accuracy loss and is the primary speedup driver.Removing latent refinement speeds decoding but removes some stability; the full model reduces effective token usage.Source: §4.5 · pp. 7–8, Table 3 · p. 8, Table 4 · p. 8
  • Intermediate rf values give the best accuracy; top-p mixing improves performance quickly, with little benefit beyond p≈0.2 despite mixing only a small fraction of tokens at p≤0.2.The rf=0 baseline and the rf=1 extreme are both suboptimal; aggressive mixing can destabilize refinement.Source: §4.5 · pp. 8–9, Figure 4 · p. 8, Figure 5 · p. 8

5 RELATED WORK

The related work traces diffusion language models from masked-token methods to current efficiency efforts, then frames improving generation quality alongside inference speed as an open priority.

Diffusion LLMs (dLLMs)

Efficient dLLMs

  • The paper positions improving dLLM generation quality while using KL divergence for reliable early stopping as a research priority.The section says existing dLLMs face inference-efficiency constraints, and that reducing inference steps can severely compromise performance.Source: §5 · p. 9
  • dLLM efficiency work targets KV-cache integration or computational load, including adaptive caching, block-wise cache reuse, sparse attention, early commitment, and dynamic length expansion.Examples include dLLM-Cache (Liu et al., 2025), Fast-dLLM (Wu et al., 2025), Sparse-dLLM (Song et al., 2025), Prophet (Li et al., 2025b), and DAEDAL (Li et al., 2025a).Source: §5 · p. 9

6 CONCLUSION

The conclusion presents LRD as a two-stage decoding framework for diffusion language models that refines global beliefs before progressively finalizing tokens. Experiments report gains in output quality and inference efficiency over standard diffusion decoding baselines, and the authors propose LRD as a flexible module for future systems.

6 CONCLUSION

  • LRD addresses information loss from hard masking and suboptimal convergence speed through embedding-space belief refinement followed by a predictive feedback loop with KL-based convergence monitoring.The loop progressively finalizes confident tokens while supporting adaptive early stopping.Source: §6 · p. 9
  • Experiments across code-generation and mathematical-reasoning benchmarks show consistent gains in output quality and inference efficiency over standard diffusion decoding baselines.The conclusion says these gains are particularly apparent as sequence length and complexity increase.Source: §6 · p. 10
  • The authors propose LRD as a flexible drop-in decoding module and identify KV caching, speculative decoding, and hardware-aware techniques as potential efficiency integrations.These are presented as future opportunities rather than evaluated results.Source: §6 · p. 10

THE USE OF LLMS

The authors report limited LLM use for preliminary literature surveying and language polishing, while stating that they independently conducted the core research.

THE USE OF LLMS

  • LLMs were used for a preliminary literature survey to identify research directions and keywords, and for limited language polishing.The authors say the polishing aimed to improve clarity and grammatical correctness.Source: THE USE OF LLMS · p. 13
  • The authors state that they independently conceived and conducted the core research ideas, theoretical contributions, experimental design, implementation, and analysis.They say LLM-generated suggestions were reviewed, verified, and substantially modified before incorporation.Source: THE USE OF LLMS · p. 13
  • The authors take responsibility for all manuscript content, including text refined with LLM assistance.Source: THE USE OF LLMS · p. 13

A EXPERIMENT DETAILS

The appendix specifies benchmark-specific few-shot settings and evaluation metrics, along with sampling and stopping parameters. For LLaDA models, LRD is integrated within the official semi-AR framework, with blocks decoded autoregressively and token distributions refined before discrete commitment.

A EXPERIMENT DETAILS

  • Base-model evaluations use zero-shot HumanEval, 3-shot MBPP, 4-shot MATH500, and 8-shot GSM8K; reporting uses pass@1 for code and accuracy for mathematical reasoning.Source: A EXPERIMENT DETAILS · p. 13
  • The experiments set top-p to 0.9, vary rf from 0.1 to 0.2, set τrefine and τdecode to 0.1, and cap latent refinement at 20 steps.Source: A EXPERIMENT DETAILS · p. 13
  • For LLaDA-Instruct and LLaDA-1.5, LRD is integrated into the official semi-AR framework, with autoregressive block-level decoding and within-block refinement before discrete commitment.The appendix refers to Appendix C for detailed integration steps.Source: A EXPERIMENT DETAILS · p. 13

B DERIVATION OF THE TRUE POSTERIOR IN THE MASKING PROCESS

The appendix derives the true posterior for the absorbing masking forward process, in which tokens can transition to [MASK] but cannot unmask.

B DERIVATION OF THE TRUE POSTERIOR IN THE MASKING PROCESS

C INTEGRATION WITH SEMI-AR FRAMEWORK

The semi-AR integration applies soft refinement within each block before progressively selecting tokens, conditioning each block on those generated earlier.

C INTEGRATION WITH SEMI-AR FRAMEWORK

D STABILITY ANALYSIS OF MIXED EMBEDDING UPDATES

The analysis links stable mixed-embedding updates to small refinement strength and a neighborhood near [MASK], while reporting empirical convergence and acknowledging that global guarantees remain out of reach.

D STABILITY ANALYSIS OF MIXED EMBEDDING UPDATES

Related papers

Loading 2510.11052v2…