残差文脈拡散言語モデル

Residual Context Diffusion Language Models

January 30, 2026

著者: 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.AI

要旨

拡散型大規模言語モデル（dLLM）は、複数のトークンを並列にデコード可能なことから、純粋に自己回帰的な言語モデルの有望な代替として登場した。しかし、最先端のブロック単位dLLMは、最も確信度の高いトークンのみをデコードし、残りを破棄する「リマスキング」機構に依存しており、実質的に計算リソースを浪費している。我々は、破棄されたトークンからの計算結果を再利用することが有益であることを実証する。これらのトークンには、後続のデコード反復で有用な文脈情報が保持されているためである。この知見に基づき、廃棄されたトークンの表現を文脈的残差に変換し、次のノイズ除去ステップで再投入するモジュール、Residual Context Diffusion（RCD）を提案する。RCDは、逆伝播に伴うメモリボトルネックを回避するため、分離された2段階の訓練パイプラインを採用する。本手法を、長いCoT推論（SDAR）と短いCoT指示追従（LLaDA）の両モデルで検証する。標準的なdLLMは、わずか約10億トークンで効率的にRCDパラダイムに変換可能であることを示す。RCDは、様々なベンチマークにおいて、最小限の追加計算オーバーヘッドで、最先端dLLMの精度を一貫して5～10ポイント向上させる。特に最も挑戦的なAIMEタスクでは、ベースライン精度をほぼ2倍にし、同等の精度レベルでノイズ除去ステップ数を最大4～5分の1に削減する。

English

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.