Residual Context Diffusion Language Models

📄 arXiv: 2601.22954v1 📥 PDF

作者: 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

分类: cs.CL, cs.AI

发布日期: 2026-01-30

💡 一句话要点

提出残差上下文扩散（RCD）模块，提升扩散语言模型（dLLM）的推理精度和效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散语言模型 并行解码 残差学习 上下文信息 长文本推理 指令跟随 计算效率 两阶段训练

📋 核心要点

  1. 现有块状扩散语言模型（dLLM）的“重掩码”机制丢弃了大量token，造成计算资源的浪费。
  2. 提出残差上下文扩散（RCD）模块，将被丢弃token的表示转换为上下文残差，并注入回后续的去噪步骤。
  3. 实验表明，RCD能显著提升dLLM在长CoT推理和短CoT指令跟随任务上的准确率，并减少去噪步骤。

📝 摘要（中文）

扩散大型语言模型（dLLM）作为一种并行解码多token的方法，正成为自回归语言模型的有希望的替代方案。然而，目前最先进的块状dLLM依赖于“重掩码”机制，该机制仅解码最自信的token并丢弃其余token，从而有效浪费了计算资源。本文证明了回收被丢弃token的计算是有益的，因为这些token保留了对后续解码迭代有用的上下文信息。为此，我们提出残差上下文扩散（RCD），该模块将这些被丢弃的token表示转换为上下文残差，并将其注入回下一个去噪步骤。RCD使用解耦的两阶段训练流程来绕过与反向传播相关的内存瓶颈。我们在长CoT推理（SDAR）和短CoT指令跟随（LLaDA）模型上验证了我们的方法。结果表明，仅需约10亿个token，标准dLLM即可高效转换为RCD范式。RCD在各种基准测试中，以最小的额外计算开销，始终将前沿dLLM的准确率提高5-10个百分点。值得注意的是，在最具挑战性的AIME任务中，RCD几乎使基线准确率翻倍，并在同等准确率水平下，去噪步骤减少高达4-5倍。

🔬 方法详解

问题定义：现有块状扩散语言模型（dLLM）在解码过程中，通常采用“重掩码”机制，仅保留置信度最高的token进行后续处理，而丢弃其余token。这种做法忽略了被丢弃token中包含的上下文信息，造成了计算资源的浪费，并且可能限制模型的性能提升。

核心思路：论文的核心思路是将被丢弃的token表示视为一种“残差上下文”，这些残差包含了对后续解码迭代有用的信息。通过将这些残差信息重新注入到后续的去噪步骤中，可以有效地利用被浪费的计算资源，并提升模型的性能。

技术框架：RCD模块主要包含两个阶段：残差生成阶段和残差注入阶段。在残差生成阶段，被丢弃的token表示被转换为上下文残差。在残差注入阶段，这些残差被注入回下一个去噪步骤。为了避免反向传播带来的内存瓶颈，RCD采用解耦的两阶段训练流程。

关键创新：RCD的关键创新在于它提出了一种有效利用被丢弃token信息的方法，通过将这些信息转换为上下文残差并重新注入到后续的解码步骤中，从而提升了模型的性能。这种方法避免了直接对所有token进行计算，从而在计算效率和性能之间取得了平衡。

关键设计：RCD模块的具体实现细节包括残差的生成方式、残差的注入方式以及两阶段训练流程的具体步骤。论文中可能涉及的关键参数包括残差的缩放因子、残差注入的位置以及两阶段训练的学习率等。具体的网络结构设计可能包括用于生成残差的神经网络层以及用于注入残差的连接方式。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RCD在长CoT推理（SDAR）和短CoT指令跟随（LLaDA）任务上均取得了显著的性能提升。在AIME任务上，RCD几乎使基线准确率翻倍，并在同等准确率水平下，去噪步骤减少高达4-5倍。在其他基准测试中，RCD也始终将前沿dLLM的准确率提高5-10个百分点，同时保持了较低的计算开销。

🎯 应用场景

该研究成果可广泛应用于需要并行解码的自然语言处理任务，例如机器翻译、文本摘要、对话生成等。通过提升扩散语言模型的推理精度和效率，RCD有望在这些领域实现更好的性能和更快的速度。此外，RCD的设计思想也可以启发其他并行计算框架的设计。

📄 摘要（原文）

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.