[論文レビュー] Focus-dLLM: Accelerating Long-Context Diffusion LLM Inference via Confidence-Guided Context Focusing

Lingkun Long, Yushi Huang+5 more|arXiv (Cornell University)|Feb 2, 2026

Topic Modeling被引用数 0

ひとこと要約

Focus-dLLM は、過去の信頼度から未マスクのトークン位置を予測し、クロス層の sink 再利用を伴う sink-aware 疎注意を適用することで、長文コンテキスト拡散 LLM 推論を訓練不要で高速化するフレームワーク。

ABSTRACT

Diffusion Large Language Models (dLLMs) deliver strong long-context processing capability in a non-autoregressive decoding paradigm. However, the considerable computational cost of bidirectional full attention limits the inference efficiency. Although sparse attention is promising, existing methods remain ineffective. This stems from the need to estimate attention importance for tokens yet to be decoded, while the unmasked token positions are unknown during diffusion. In this paper, we present Focus-dLLM, a novel training-free attention sparsification framework tailored for accurate and efficient long-context dLLM inference. Based on the finding that token confidence strongly correlates across adjacent steps, we first design a past confidence-guided indicator to predict unmasked regions. Built upon this, we propose a sink-aware pruning strategy to accurately estimate and remove redundant attention computation, while preserving highly influential attention sinks. To further reduce overhead, this strategy reuses identified sink locations across layers, leveraging the observed cross-layer consistency. Experimental results show that our method offers more than $29 imes$ lossless speedup under $32K$ context length. The code is publicly available at: https://github.com/Longxmas/Focus-dLLM

研究の動機と目的

• 拡散 LLM（dLLMs）に対して、双方向全注意と未知の未マスク位置を前提とした効率的な長文推論を動機づける。
• トークン信頼度の一貫性と注意ソースの挙動を調査し、スパース化を導く。
• 訓練不要のフレームワークを開発し、未マスク位置を予測し、注意を剪定して計算量を削減する。
• クリティカルな注意ソースを保持し、層間でソースの位置を再利用してオーバーヘッドを節約する。
• 長文ベンチマークで精度低下を最小限に抑えつつ実質的なスピードアップを実証する。

提案手法

• ステップ t-1 の信頼度を用いて次の未マスク位置を予測し、それを活動領域へ展開する過去信頼度ガイド付き指標を導入する。
• 局所性を活用して distant tokens を剪定するためのウィンドウ拡張済みアクティブクエリ集合を開発する。
• 層を跨ぐ注意ソースを保持し、中間の密集層でソースを同定する sink-aware 疎注意機構を提案する。
• 繰り返しのソース同定を避けるため、層間で同定されたソースを再利用する。
• 関連するプロンプトブロックを選択し、ソースと応答トークンを保持しつつ、キー/バリュ履歴を動的ブロック単位で剪定する。
• これらの要素を組み合わせて、予測された未マスククエリと選択されたキー/バリュ対に対して疎な注意を適用し、デコード品質を損なうことなくスピードアップを達成する。

Figure 1: Confidence dynamics analysis for LLaDA-8B-Instruct Nie et al. ( 2025 ) on GSM $8K$ Cobbe et al. ( 2021 ) ( $L=76$ , $N=128$ , and $T=128$ ). ( Left ) Confidence score correlation between adjacent steps. ( Right ) Step-wise recall rates of predicting the unmasked tokens at $t$ using the rem

実験結果

リサーチクエスチョン

• RQ1 現在の拡散ステップでの未マスク位置を、前段のトークン信頼度から信頼性高く推定できるか。
• RQ2 dLLMs において層を跨ぐ注意ソースの一貫性が見られ、計算を削減する再利用が正確さを損なわずに可能か。
• RQ3 過去信頼度信号と局所的な注意パターンを活用して、長文コンテキスト拡散デコードの効果的なスパース注意を設計できるか。
• RQ4 sink-aware 疎注意とブロック単位 KV剪定が、非常に長いコンテキスト長での精度とスループットに与える影響はどの程度か。

主な発見

• 本研究は、隣接する拡散ステップ間でトークン信頼度が強い正の相関を示すことを発見した。
• 従来不明だった未マスク位置は、前のステップの高信頼トークンと大きく重複し、次にどこを未マスクにするかを正確に予測できる。
• dLLMs の注意には局所性と予測可能な注意ソースが見られ、層を跨ってソースを再利用することを可能にする。
• Focus-dLLM は、32K 文脈で Vanilla を上回る最大 29.6 倍の実質的なスピードアップを達成し、評価モデルに対してベースラインと比較して精度を維持または向上させた。
• 過去信頼度の指針と sink-aware 疎注意、および層間ソース再利用を組み合わせることで、ベースライン以上の性能と顕著なスループット向上を達成した。
• アブレーション研究は、PCGI クエリ選択と SA Sparse Attn の両方が精度と速度に寄与し、ソースが一貫して性能を向上させることを示した。

Figure 2: Attention patterns across decoding steps and layers in LLaDA-8B-Instruct Nie et al. ( 2025 ) ( $L=49$ , $N=128$ , $T=128$ ). More visual results can be found in the Appendix.