Residual Context Diffusion Language Models

2026-01-30

AI Summary

This work addresses the inefficiency and contextual information loss in existing blockwise diffusion language models (dLLMs), which discard low-confidence tokens during decoding. To mitigate this, the authors propose the Residual Context Diffusion (RCD) module—the first approach to systematically convert representations of discarded tokens into contextual residuals and inject them into subsequent denoising steps, thereby preserving latent semantic information. RCD is integrated with a two-stage decoupled training strategy that alleviates memory bottlenecks and enables parallel multi-token decoding. Experimental results demonstrate that RCD improves state-of-the-art dLLM accuracy by 5–10 percentage points across multiple benchmarks and nearly doubles baseline performance on AIME tasks, while reducing the required number of denoising steps to one-fourth or one-fifth of the original.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: (Large) Language Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

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.

Problem

• Diffusion Language Models
• remasking
• discarded tokens
• contextual information
• computation waste

Innovation

• Residual Context Diffusion
• diffusion language models
• remasking
• contextual residuals
• parallel decoding