Table of Contents
Fetching ...

Reversible Diffusion Decoding for Diffusion Language Models

Xinyun Wang, Min Zhang, Sen Cui, Zhikang Chen, Bo Jiang, Kun Kuang, Mingbao Lin

TL;DR

Reversible Diffusion Decoding (RDD), a decoding framework that introduces reversibility into block-wise diffusion generation, allows decoding to recover from early commitment errors while maintaining the parallel efficiency of diffusion-based generation.

Abstract

Diffusion language models enable parallel token generation through block-wise decoding, but their irreversible commitments can lead to stagnation, where the reverse diffusion process fails to make further progress under a suboptimal context.We propose Reversible Diffusion Decoding (RDD), a decoding framework that introduces reversibility into block-wise diffusion generation. RDD detects stagnation as a state-dependent failure of the reverse process and enables efficient backtracking to earlier blocks without recomputation via cached model states. To avoid repeated failure trajectories, RDD applies confidence-guided re-masking to selectively reinitialize uncertain tokens while preserving reliable context.This reversible formulation allows decoding to recover from early commitment errors while maintaining the parallel efficiency of diffusion-based generation. Experiments show that RDD improves generation robustness and quality over baselines with minimal computational overhead.

Reversible Diffusion Decoding for Diffusion Language Models

TL;DR

Reversible Diffusion Decoding (RDD), a decoding framework that introduces reversibility into block-wise diffusion generation, allows decoding to recover from early commitment errors while maintaining the parallel efficiency of diffusion-based generation.

Abstract

Diffusion language models enable parallel token generation through block-wise decoding, but their irreversible commitments can lead to stagnation, where the reverse diffusion process fails to make further progress under a suboptimal context.We propose Reversible Diffusion Decoding (RDD), a decoding framework that introduces reversibility into block-wise diffusion generation. RDD detects stagnation as a state-dependent failure of the reverse process and enables efficient backtracking to earlier blocks without recomputation via cached model states. To avoid repeated failure trajectories, RDD applies confidence-guided re-masking to selectively reinitialize uncertain tokens while preserving reliable context.This reversible formulation allows decoding to recover from early commitment errors while maintaining the parallel efficiency of diffusion-based generation. Experiments show that RDD improves generation robustness and quality over baselines with minimal computational overhead.
Paper Structure (16 sections, 13 equations, 8 figures, 2 tables, 1 algorithm)

This paper contains 16 sections, 13 equations, 8 figures, 2 tables, 1 algorithm.

Figures (8)

  • Figure 1: Visualizing the Stagnation Trap and the Impact of Prefix Quality. We compare generation trajectories by switching decoding strategies at token 64 (red token). Token color indicates prediction confidence, darker is better. (a) Decoding the first 64 tokens with RDD yields a high-confidence prefix (Avg Conf: 0.872). This robust context allows the standard Fast-dLLM to maintain high stability in the suffix (Avg Conf: 0.892). (b) Decoding the prefix with Fast-dLLM results in less certain (Avg Conf: 0.820). Crucially, because these suboptimal tokens are locked in the context, even switching to the robust RDD cannot fully recover the generation (Avg Conf: 0.872), which remains less confident than the method in (a). Further details are provided in Appendix \ref{['app:stagnation-trap']}.
  • Figure 2: Illustration of the RDD Process. The diagram depicts the decoding flow across three blocks. Block 0 demonstrates successful parallel generation. In Block 2, the model falls into the Stagnation Trap at the final denoising steps, where the reverse diffusion process fails to confidently unmask the remaining tokens. Instead of forcing a suboptimal commitment, RDD triggers a Rollback operation. As shown in Block 1, the system efficiently retrieves the previous state and applies confidence-guided Re-masking to the uncertain regions, enabling the model to correct the error and resume valid generation.
  • Figure 3: Ablation studies on key RDD hyperparameters evaluated on the GSM8K dataset. (a) The impact of varying the decoding scaling factor $f$ on generation accuracy and throughput. (b) The impact of varying the rollback budget $R$ on generation accuracy and throughput.
  • Figure 4: Performance using random remask or confidence-based remask. The x-axis represents the ratio of random remasking.
  • Figure 5: Fast-dLLM: The model relies on fallback mechanisms to resolve uncertainty, leading to low-confidence predictions. Avg Conf: 0.849, with 22.9% of tokens below 0.7 confidence.
  • ...and 3 more figures