技術要約: Prefilling-dLLM

問題提起

拡散大規模言語モデル（dLLM）は、反復的なデノイジングを通じて並列的なトークン生成を可能にすることで、自己回帰（AR）モデルに代わる有望な選択肢を提供します。しかし、dLLMは長文脈シナリオにおいて重大な非効率性に直面しています。すなわち、デノイジングの各ステップにおいて、自己注意（self-attention）を計算するために入力プレフィックス全体を再エンコードしなければならないという点です。これにより、計算量は全シーケンス長に対して二次関数的にスケールします（$O((L_p + L_d)^2 \cdot T)$、ここで $L_p$ はプレフィックス長、$L_d$ はデコード長、$T$ はデノイジング・ステップ数）。ARモデルでは、プリフィル（prefill）フェーズが計算量制限（compute-bound）であり、デコードフェーズがメモリ制限（memory-bound）であるのに対し、dLLMの推論はすべてのイテレーションを通じて計算量制限のままです。これは、フルシーケンス（プレフィックス + レスポンス）が各ステップで共同処理されるためです。このプレフィックスの繰り返される再計算は、長文脈アプリケーション（例：8K–32K+ トークン）において極めて困難なものとなります。

手法: Prefilling-dLLM

著者らは、プレフィックスの再計算を排除し、スパース性を導入するために設計された、学習不要のプリフィル・デコード分離フレームワークである PREFILLING-DLLM を提案しています。この手法は、以下の3つのステージで動作します。

1. チャンク化されたプリフィルとKVキャッシュ

入力プレフィックスを単一のブロックとして処理する代わりに、プレフィックスを固定サイズ $C$ の $N$ 個の重複しないチャンクに分割します。

• 独立した処理: 各チャンクは、チャンク内の双方向アテンションを用いて独立して処理されます。デリミタ（区切り文字）として機能させるため、各チャンクの先頭に特別なBeginning-of-Sequence (BOS) トークンが付加されます。
• 予測的スコアリング: 正解ラベルなしでどのチャンクが関連しているかを判断するために、本フレームワークは予備的な拡散パスを通じて疑似ラベル（$m$）を生成します。その後、以下の2つの信号を用いて各チャンクをスコアリングします。
  • 自己情報量スコア ($s_I$): チャンクを条件としたクエリウィンドウの負の対数尤度。
  • 疑似ラベルスコア ($s_P$): チャンクおよびクエリを条件とした疑似ラベルの尤度。
• 選択: 統合スコア（$s = s_I + s_P$）に基づいて、最も関連性の高い上位 $K$ 個のチャンクが選択されます。これらの選択されたチャンクのKV表現がキャッシュされます。

2. 2レベルのスパース性

本フレームワークは、複雑さを軽減するために、チャンクレベルとトークンレベルの両方でスパース性を導入します。

• チャンク間スパース性: デコーディング中のアテンション機構には、全 $N$ 個のチャンクではなく、上位 $K$ 個のチャンクのみが参加します。
• チャンク内スパース性: 選択された各チャンク内において、トークンはクエリ・トークンとの累積的な双方向アテンション・ウェイトによってランク付けされます。各チャンク内で最も関連性の高い上位 $B$ 個のトークンのみがKVキャッシュに保持され、無関係なトークンは破棄されます。

3. 静的KV再利用によるデコーディング

反復的なデノイジング・プロセス中：

• 選択された $K$ 個のチャンクのキャッシュされたKV表現は静的であり、すべての $T$ デノイジング・ステップにわたって再利用されます。
• モデルは、クエリおよび進化するレスポンス・トークンのためのKV表現のみを再計算します。
• これにより、1ステップあたりの計算複雑度が $O((L_p + L_d)^2 \cdot T)$ から $O(N \cdot C^2 + (L_d^2 + K \cdot C) \cdot T)$ に削減され、デコーディング・コストをフル・プレフィックス長から事実上切り離すことができます。

アテンション・カーネルの最適化

非連続にキャッシュされたチャンクKVキャッシュのロードによるメモリ・ボトルネックに対処するため、著者らは Split-S FlexAttention カーネルを実装しています。このカーネルは、KV次元を分割し、結果をlog-sum-exp還元によってマージする前にチャンクを独立して処理することで、非連続にキャッシュされたチャンクKVによるデコーディングを並列化し、データを連続したメモリに集約（gather）する必要性を回避します。

主な知見と結果

パフォーマンスと効率

• 高速化: LongBenchおよびInfiniteBenchにおいて、PREFILLING-DLLMは、コンテキスト長8Kから32Kにおいて、標準的なTransformerベースラインと比較して 9.1倍から28.0倍の高速化 を達成しました。
• 品質: 本手法は、dLLM加速手法の中で最高水準の品質を実現しています。LongBenchにおいて、Dream-7Bで平均スコア 34.59、UltraLLaMAで 37.02 を記録し、Fast-dLLM、Sparse-dLLM、dKV-Cacheなどのベースラインを上回りました。
• 長文脈スケーリング: InfiniteBench（コンテキスト >128K）において、本手法はDream-7Bで 43.62% の平均精度を達成し、最強のベースラインであるFast-dLLM v2（30.32%）を大幅に上回りました。特に、Passkey Retrievalタスクで 95.42%、Number Retrievalタスクで 70.00% の精度を達成しています。

「Lost-in-the-Middle」への対処

• 現象: 本論文では、dLLMがコンテキスト外挿下で「lost-in-the-middle（中だるみ）」現象を示すことを観察しています。これは、情報の位置による精度低下が見られる現象であり、ARモデルに見られるU字型の曲線ではなく、レスポンスに向かって単調に減衰するバイアスを示します。
• 緩和策: 位置の近接性ではなく、予測的な関連性に基づいてチャンクを選択することにより、PREFILLING-DLLMはこのバイアスを排除します。本手法は、32Kおよび128Kの長さにおいても、コンテキスト内のあらゆる位置に対して高い完全一致精度を維持します。

メカニズムの洞察

• アテンション・アンカー: 各チャンクに付加されたBOSトークンは、周期的なアテンション・アンカーとして機能します。分析によれば、これらのトークンは、不当に大きなアテンション・マスを吸収する「アテンション・シンク（attention sinks）」へと退化することなく、コンテキスト全体にアテンション・マスを均一に分散させます（安定した低振幅のリッジを形成します）。このメカニタズムにより、位置不変の検索が可能になります。
• チャンクサイズのトレードオフ: チャンクサイズと検索精度の間にはトレードオフが存在します。小さなチャンク（例：256トークン）は128Kのコンテキストでも高い精度を維持しますが、大きなチャンク（例：4096トークン）は32Kを超えると急激に劣化します。長文脈推論においては、マルチチャンク選択が不可欠です。

意義と主張

本論文は、PREFILLING-dLLMが、dLLM推論の計算量制限的な性質と、ARモデルで使用されているプリフィル・デコード分離アーキテクチャの効率性の間の溝を埋めるものであると主張しています。その主な貢献は以下の通りです：

1. 学習不要の分離: プレフィックスのKV表現を一度キャッシュして再利用するフレームワークを導入し、モデルの再学習を必要とせずに、毎ステップのプレフィックスの二次的な再計算コストを排除しました。
2. スパース・プリフィル: スパースなプリフィル（関連するチャンクとトークンの選択）が、密なアテンションよりもdLLMにおいて優れた性能を発揮できることを示し、生成品質を維持または向上させつつ複雑さを削減しました。
3. システムレベルの最適化: プレフィックスのチャンク構造を活用して、デコードフェーズにおけるレイテンシの大幅な削減（バニラなFlexAttentionに対して最大10.2倍）を実現する特化したアテンション・カーネル（Split-S FlexAttention）を提供しました。
4. 長文脈への堅牢性: このアプローチが「lost-in-the-middle」現象を緩和し、高精度な検索を伴う128K+トークンへの効果的なスケーリングを可能にすることを実証しました。これは、これまでdLLMの加速において未開拓であった能力です。

著者らは、制限事項についても述べており、これには、チャンク選択の静的な性質（デコーディング中の動的な再選択が行われないこと）、チャンクサイズと$K$のタスク固有のチューニングの必要性、および現在の評価が特定のモデルアーキテクチャ（Dream-7BおよびUltraLLaMA）を用いた英語ベンチマークに限定されていることが含まれます。