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

Source-linked AI summary

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

arXiv:2601.22954v2Jan 30, 2026cs.CLcs.AI

TL;DR

Block-wise diffusion language models waste computation by discarding low-confidence token representations during remasking. RCD recycles those signals as entropy-weighted contextual residuals, consistently improving baseline performance with comparable inference throughput.

Problem
Block-wise dLLMs can trail autoregressive models in accuracy while remasking discards intermediate computation from undecoded tokens.
Method
RCD converts remasked token representations into entropy-weighted contextual residuals and trains the framework with a memory-efficient two-stage pipeline.
Results
RCD consistently outperforms standard baselines with comparable inference throughput, gaining 4–11 percentage points when scaling from 4B to 8B models.
Takeaways & Limitations
Discarded signals in block-wise dLLM decoding can function as structured contextual priors for higher-quality parallel text generation.
Takeaways & Limitations
GSM8K comparisons are subject to potential contamination in the original Chat models, which inflates that baseline.

Abstract

from arXiv · show

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 ~300 million tokens. RCD consistently improves frontier dLLMs by 4-11 percentage 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 baseline's peak accuracy.

1. Introduction

dLLMs offer parallel decoding but still trail autoregressive models in accuracy despite higher denoising compute, partly because remasking discards useful intermediate information. Residual Context Diffusion (RCD) recycles low-confidence token representations as contextual residuals to progressively refine decoding.

Motivation: dLLMs enable parallel decoding and show promising results across instruction following, code generation, long-context understanding, and complex reasoning.They are presented as an alternative to purely autoregressive models, although autoregressive models still dominate large-scale industrial deployment.
Motivation: Despite significantly more computation per token through sequential denoising iterations, dLLMs still trail autoregressive models in accuracy.The introduction characterizes this gap as unsatisfactory because the added computation does not reliably translate into accuracy gains.
Motivation: Remasking commits only the most confident tokens and discards the rest, wasting intermediate computation performed on low-confidence tokens.This inference-time strategy is identified as a factor exacerbating the accuracy gap in state-of-the-art block-wise dLLMs.
Residual Context Diffusion: RCD treats low-confidence token representations as residual updates to the input context instead of discarding them.The mechanism propagates both discrete tokens and continuous embedding vectors enriched with contextual information, enabling progressive decoding refinement.
Challenges: Residual feedback is challenging because contextual signals must be selected and aggregated without destabilizing the denoising input distribution.The self-referential feedback loop also creates a long unrolled computation graph, making training intrinsically difficult.
Residual Context Diffusion: RCD uses entropy-based embedding aggregation and constructs residual context from the model’s own embedding codebooks.This design addresses limitations of reusing hidden states, including failure to account for masked-versus-unmasked structure and potential embedding-magnitude mismatch.

2. Preliminaries

Diffusion-based LLMs generate text through masked latent-space denoising, training to reconstruct corrupted tokens and iteratively committing the most confident predictions. Soft tokens represent probability mixtures through weighted vocabulary embeddings, but naive injection can disrupt masking and create unstable recursive dependencies.

dLLM denoising: dLLMs initialize tokens as masks and progressively denoise them, selecting the highest-confidence token at each forward pass until the output is fully decoded.This frames text generation as progressive denoising within a masked latent space.
dLLM training: During training, dLLMs recover an original b-token sequence from independently masked noisy inputs by minimizing expected cross-entropy over masked positions.The mask indicator is sampled independently for each position, and the corruption level varies through t ∼ U(0, 1).
dLLM inference: During inference, dLLMs predict distributions, commit the top-m positions by confidence, and remask less-confident positions for the next denoising iteration.The reverse process starts from an entirely masked block and repeats these steps for K iterations.
dLLM inference: Standard dLLMs discard the predicted distributions for unselected positions after remasking, causing information loss that RCD seeks to mitigate.Unselected positions are reset to the mask embedding rather than retaining their probability information.
Soft tokens: Soft tokens encode a mixture of vocabulary tokens as a probability-weighted sum of embedding vectors, preserving fine-grained context before discrete determination.For distribution p and embedding codebook E, the soft token is e = E⊤p.
Soft tokens: Naively adding soft tokens to dLLM input embeddings disrupts the discrete masking scheme and creates unstable recursive dependencies during training.These limitations motivate a more careful mechanism for injecting contextual residual information.

3. The RCD Method

Residual Context Diffusion recycles remasked-token information as entropy-weighted residual context, injecting it into mask-token representations during denoising. A decoupled two-stage training procedure and warm-started, temperature-calibrated inference address the recursive mechanism’s optimization and distribution-alignment challenges.

Entropy-Weighted Residual: Entropy-weighted residuals assign greater contribution to discarded tokens with higher normalized Shannon entropy, preserving richer context across denoising steps.Normalized entropy is computed relative to the maximum entropy log V and serves as the residual weight.
Residual Context Construction: RCD converts predicted token distributions into residual information vectors and injects them only into mask-token inputs through a highway connection.The residual contribution is controlled by α_i ∈ [0, 1], avoiding the magnitude misalignment caused by plain summation.
Training: RCD decouples residual generation from residual utilization through two-stage training with a lightweight frozen Reference Model and single-step Target Model supervision.The Reference Model supplies stable proxy probability distributions and residual signals, preventing recursive graph unrolling and self-reinforcing instability.
Training: During Target Model training, the Reference Model provides probabilities and entropy weights, while residual construction uses the Target Model’s embedding codebook before standard cross-entropy optimization.This places the residual vector in the Target Model’s latent space while keeping the Reference Model frozen as a stationary signal source.
Inference: Inference uses a recursive residual loop initialized by either a Reference Model warm start or zero-valued cold start, with temperature-scaled entropy aligning self-generated residual weights.Temperature scaling calibrates entropy when the Target Model’s predictions are sharper or flatter than the Reference Model’s.

4. Experiments

RCD is evaluated on LLaDA and SDAR across mathematical reasoning benchmarks, model scales, and decoding block sizes. It consistently improves accuracy and efficiency over Sequential Denoising, with gains attributed to recovering discarded contextual signals.

Main Results: RCD more than doubles SDAR-8B-b64 accuracy on AIME24, from 7.08% to 18.75%, and improves AIME25 from 9.79% to 16.04%.LLaDA also gains nearly 6% absolute accuracy on MinervaMath.
Pareto Frontier Analysis: 4–5× computation savings are achieved on the accuracy–Token per Step Pareto frontier across model scales and tasks.RCD introduces minimal computational overhead per step, making Token per Step a proxy for generation parallelism.
Throughput-Matched Accuracy: 2–9% accuracy improvements occur across throughput-matched models and benchmarks while RCD remains almost comparable to Sequential Denoising in Tokens per Second.Throughput matching uses Fastdllm for LLaDA and D2F for SDAR.
Ablations: 77.4% accuracy is achieved with Dirichlet residual initialization, while normalized-entropy weighting attains the optimal Pareto frontier and standard residual or top-1 weighting degrades performance.The initialization ablation finds all single-model variants superior to the reference-based approach.
Scalability across Model and Block Sizes: 4–11 percentage-point gains emerge when scaling RCD from 4B to 8B, with larger margins at larger block sizes.The passage attributes this expansion to larger blocks containing more abundant, stabilized contextual prior.
Training Efficiency and Cost: 300M tokens suffice for RCD to reach near-optimal reasoning performance in a single epoch, outperforming Loopholing under the constrained SDAR-4B-b64 training budget.The comparison uses one epoch on OpenR1-Math-220k; RCD reaches near-optimal performance on GSM8K and MATH500.

5. Conclusion

Residual Context Diffusion repurposes discarded dLLM signals as a structured contextual prior, using entropy to dynamically weight injected residuals. Across benchmarks, it improves denoising accuracy over standard baselines while maintaining comparable inference throughput.

RCD repurposes discarded signals in dLLMs as a structured contextual prior.The framework converts otherwise discarded decoding information into contextual residuals.
Entropy dynamically weights injected residuals to improve denoising accuracy.
RCD consistently outperforms standard baselines across benchmarks with comparable inference throughput.The conclusion characterizes the method as scalable, practical, and robust.

A. Related Works

The related work spans latent reasoning, looped transformers, and masked diffusion language models. These approaches target more efficient reasoning or generation through representation compression, iterative computation, or progressive denoising.

Latent Reasoning: Latent reasoning compresses multiple tokens or candidate answers into a single embedding vector to improve reasoning efficiency and accuracy.Early methods optimize continuous latent spaces using supervised finetuning and self-distillation.
Looped Transformers: Looped transformers reuse cyclic transformer layers and shared weights to embed iterative reasoning into the architecture while avoiding sequential CoT bottlenecks.Reusing weights across iterative steps provides test-time efficiency and enables deeper effective computation.
Diffusion Large Language Models: Masked diffusion language models generate text through progressive denoising steps and include bidirectional and block-wise dLLM architectures.This modeling approach differs from conventional sequential models by adopting masked language modeling.

B. Formulation of Residual Context Diffusion

RCD uses a decoupled training framework with proxy residual signals from a frozen reference model, and recursive denoising with temperature-adjusted entropy alignment during inference.

RCD Training: A frozen reference model provides proxy residual signals to the target model in RCD’s decoupled training framework.Algorithm 1 describes this training procedure.
RCD Inference: RCD inference follows a recursive denoising process that incorporates temperature-adjusted entropy alignment to bridge the distribution gap.Algorithm 2 describes the inference procedure.

C. Detailed Training Configurations

This section specifies the hyperparameters used to fine-tune the SDAR and LLaDA model families, with full details provided in Table 7.

C. Detailed Training Configurations: The fine-tuning hyperparameter settings for the SDAR and LLaDA model families are detailed in Table 7.The section presents the specific configurations used for both model families.

C.1. SDAR Family Configuration

SDAR models at 1.7B, 4B, and 8B scales were fine-tuned on filtered OpenR1-Math-220k data retaining reasoning chains of at most 8K tokens. Training uses a constant learning rate to support stability across model scales.

SDAR Family Configuration: SDAR models at 1.7B, 4B, and 8B were fine-tuned on OpenR1-Math-220k.The dataset was filtered to retain samples with reasoning chains ≤8K tokens.
SDAR Family Configuration: Reasoning-chain filtering retained OpenR1-Math-220k samples with chains ≤8K tokens.
SDAR Family Configuration: A constant learning rate was used to ensure stability across different model scales.

C.2. LLaDA Configuration

LLaDA-8B-Base was fine-tuned on a 1M-example OpenMathInstruct-2 subset using distributed FSDP training and standard SFT. Its 2048-token context window supports optimization with global bidirectional context.

C.2. LLaDA Configuration: LLaDA-8B-Base was fine-tuned on a 1M subset of OpenMathInstruct-2 using FSDP for distributed training.Unlike block-wise SDAR, LLaDA uses standard SFT with a larger 2048-token context window to optimize global bidirectional context.

D. Potential Data Contamination on GSM8K

SDAR-Chat models can outperform reasoning-specialized SDAR and RCD models on GSM8K, but cross-benchmark results suggest this advantage reflects potential data contamination rather than robust mathematical reasoning. Performance drops on newly curated or perturbed GSM variants, indicating reliance on memorized training-distribution patterns.

Cross-benchmark evaluation: SDAR-Chat models occasionally outperform reasoning-specialized Sequential Denoising and RCD models on GSM8K.The authors investigate this paradox through evaluation on GSM1K and GSM-Plus.
Cross-benchmark evaluation: Up to 88.40% accuracy on original GSM8K falls by 2–7% on GSM1K for SDAR-Chat models.GSM1K contains newly curated problems designed to match GSM8K difficulty while testing generalization beyond familiar examples.
Interpretation: The GSM8K advantage therefore indicates fragile memorization that fails when problem surface features are slightly perturbed.The authors interpret the cross-benchmark degradation as evidence that models rely partially on seen patterns from the training distribution.
Cross-benchmark evaluation: Over 10% performance plunges across all configurations on GSM-Plus, showing sensitivity to modest problem-surface perturbations.The results indicate that high GSM8K scores are not representative of robust mathematical reasoning.