Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Residual Context Diffusion Language Models

Yuezhou Hu, Harman Singh, Monishwaran Maheswaran, Haocheng Xi, Coleman Hooper, Jintao Zhang, Aditya Tomar, Michael W. Mahoney, Sewon Min, Mehrdad Farajtabar, Kurt Keutzer, Amir Gholami, Chenfeng Xu

TL;DR

The paper tackles inefficiency in block-wise diffusion LLMs caused by remasking discarded tokens. It introduces Residual Context Diffusion (RCD), which recycles discarded token representations as contextual residuals using an entropy-weighted mechanism and a decoupled two-stage training pipeline. Through a frozen reference model guiding a residual-aware target model, RCD enables stable training and effective residual injection during inference. Empirical results on SDAR and LLaDA across multiple benchmarks show consistent accuracy gains with modest overhead and strong scalability, including notable improvements on challenging tasks like AIME, making diffusion-based decoding more competitive with autoregressive approaches.

Abstract

Diffusion Large Language Models (dLLMs) have emerged as a promising alternative to purely autoregressive language models because they can decode multiple tokens in parallel. However, state-of-the-art block-wise dLLMs rely on a "remasking" mechanism that decodes only the most confident tokens and discards the rest, effectively wasting computation. We demonstrate that recycling computation from the discarded tokens is beneficial, as these tokens retain contextual information useful for subsequent decoding iterations. In light of this, we propose Residual Context Diffusion (RCD), a module that converts these discarded token representations into contextual residuals and injects them back for the next denoising step. RCD uses a decoupled two-stage training pipeline to bypass the memory bottlenecks associated with backpropagation. We validate our method on both long CoT reasoning (SDAR) and short CoT instruction following (LLaDA) models. We demonstrate that a standard dLLM can be efficiently converted to the RCD paradigm with merely ~1 billion tokens. RCD consistently improves frontier dLLMs by 5-10 points in accuracy with minimal extra computation overhead across a wide range of benchmarks. Notably, on the most challenging AIME tasks, RCD nearly doubles baseline accuracy and attains up to 4-5x fewer denoising steps at equivalent accuracy levels.

Residual Context Diffusion Language Models

TL;DR

The paper tackles inefficiency in block-wise diffusion LLMs caused by remasking discarded tokens. It introduces Residual Context Diffusion (RCD), which recycles discarded token representations as contextual residuals using an entropy-weighted mechanism and a decoupled two-stage training pipeline. Through a frozen reference model guiding a residual-aware target model, RCD enables stable training and effective residual injection during inference. Empirical results on SDAR and LLaDA across multiple benchmarks show consistent accuracy gains with modest overhead and strong scalability, including notable improvements on challenging tasks like AIME, making diffusion-based decoding more competitive with autoregressive approaches.

Abstract

Diffusion Large Language Models (dLLMs) have emerged as a promising alternative to purely autoregressive language models because they can decode multiple tokens in parallel. However, state-of-the-art block-wise dLLMs rely on a "remasking" mechanism that decodes only the most confident tokens and discards the rest, effectively wasting computation. We demonstrate that recycling computation from the discarded tokens is beneficial, as these tokens retain contextual information useful for subsequent decoding iterations. In light of this, we propose Residual Context Diffusion (RCD), a module that converts these discarded token representations into contextual residuals and injects them back for the next denoising step. RCD uses a decoupled two-stage training pipeline to bypass the memory bottlenecks associated with backpropagation. We validate our method on both long CoT reasoning (SDAR) and short CoT instruction following (LLaDA) models. We demonstrate that a standard dLLM can be efficiently converted to the RCD paradigm with merely ~1 billion tokens. RCD consistently improves frontier dLLMs by 5-10 points in accuracy with minimal extra computation overhead across a wide range of benchmarks. Notably, on the most challenging AIME tasks, RCD nearly doubles baseline accuracy and attains up to 4-5x fewer denoising steps at equivalent accuracy levels.
Paper Structure (20 sections, 7 equations, 5 figures, 7 tables, 2 algorithms)

This paper contains 20 sections, 7 equations, 5 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Diffusion Large Language Models
  4. Soft Tokens
  5. The RCD Method
  6. Entropy Weighted Residual
  7. Training
  8. Inference
  9. Experiments
  10. Experimental Setup
  11. Efficiency of RCD
  12. Ablation studies
  13. Discussion
  14. Conclusion
  15. Related Works
  16. ...and 5 more sections

Figures (5)

  • Figure 1: Overview of the residual denoising mechanism in Residual Context Diffusion (RCD). Remasking happens during each denoising step, discarding low-confidence token even though they are already explicitly computed. To recycle computation from previous steps, our proposed method extracts information before it is discarded, and forwards it to the next denoising step.
  • Figure 2: Average token recall across LLaDA GSM8K sampling steps. For each step $k$ (x-axis), recall@5 (y-axis) is the fraction of final decoded tokens that appear in the step-$k$ top-5 predictions at the same position.High early-step recall@5 indicates that intermediate distributions contain semantically informative signals.
  • Figure 3: (Left) In vanilla dLLM inference, tokens not selected are reverted to static mask embeddings for the next step, discarding the predicted probability information. (Center) RCD Inference preserves this information by calculating a weighted summation of the unselected probability distributions (residuals) and injects them into the input embeddings of the subsequent step. This allows the model to aggregate more information and refine ambiguous tokens over time. (Right) RCD Training employs a decoupled strategy. A Frozen Reference Model generates stable probability targets and entropy weights ($\alpha$). These are used to construct the residual vectors, which are then added to the Trainable Target Model's embeddings, helping the model to effectively leverage residual context for prediction.
  • Figure 4: Accuracy vs. Token per Step Pareto frontiers for SDAR models. Curves are generated by sweeping confidence thresholds from 0.5 to 1.0. Token per Step measures the average number of tokens generated in each parallel iteration; higher Token per Step indicates fewer total steps for a given length. RCD (green dashed) consistently shifts the frontier toward the upper-left, providing significant speedups at identical accuracy levels. Note that curves may not fully overlap on the x-axis as RCD enables stable generation at higher Token per Step regimes where the Sequential Denoising baseline fails to maintain coherence.
  • Figure 5: Impact of different interpolation strategies $\alpha$ on the Accuracy vs. Token per Step Pareto frontier. (1) Normalized Entropy (Ours): $\alpha$ is based on predictive uncertainty; (2) Linear: $\alpha$ is a fixed-scalar; (3) Confidence: $\alpha$ is max-probability; and (4) Inverse: reversed mappings of entropy and confidence. Normalized entropy interpolation achieves the optimal trade-off, reaching the highest parallelism (Token per Step) with competitive accuracy.