固定状態キャッシュがBlock Diffusionにもたらす長文脈性能

導入

拡散言語モデルは複数のtokenを並列に生成できる一方、双方向のデノイジングを使うため、自動回帰モデルで一般的なkey-value（KV）キャッシュをそのまま利用できない。Block diffusionは生成をblock単位で進めることでキャッシュを可能にするが、これまでの方式の多くは注意機構に依存していた。そのため、コンテキストが長くなるほどメモリ使用量が増え、学習後に追加したキャッシュはモデル本来の計算を近似するだけになる可能性もある。

arXiv論文「Fixed State, Long Reach」は、別の設計を提案する。確定済みblockを系列ミキサーで再利用可能な状態に要約し、その状態を前提としたblock-causal学習目標を組み合わせることで、キャッシュを後付けの近似ではなく、モデルが学習した計算の一部にする。

主なポイント

• 同一条件で比較。 研究者は、注意機構、Mamba、ハイブリッドの3種類の3B block-diffusionデノイザーを、3000億tokensと共通のsingle-frontier目標で事前学習した。
• 状態空間キャッシュは一定サイズ。 Mambaは完了したblockを固定サイズの状態に圧縮するため、キャッシュ容量と各stepの遅延が系列長に依存しない。注意機構のキャッシュはO(L)で増加する。
• 256k tokensで大きな差。 この条件では注意機構のキャッシュが82GB、1 stepあたり29msを要した。Mambaは遅延を4.3倍低減し、メモリを11分の1に抑え、単一ストリームのスループットを2.6倍にした。
• バッチ処理でも有利。 状態のサイズが一定なので、Mambaはバッチ拡大に対応しやすい。報告された集約スループットは注意機構の14倍で、注意機構は単一ストリームを超えて実行できなかった。
• 学習長を超えた検索。 Mambaとハイブリッドは学習長の8〜16倍でも情報検索を維持した一方、注意機構は2倍で検索性能が崩れた。要旨では、測定された品質低下はなかったとされている。

意義と限界

この研究の意義は、単にキャッシュを小さくした点にない。モデルの構造、学習目標、推論インターフェースを一体として設計した点にある。block-causal目標によって要約状態をモデルの想定計算に組み込み、双方向デノイザーへキャッシュを後付けする際のずれを抑えている。

実運用では、固定サイズの状態は長文脈や多ストリーム配信に適している。入力が長くなっても履歴全体に比例してメモリが増えず、より大きなバッチも扱いやすい。また、状態空間モデルの帰納バイアスが、学習範囲を超えた情報利用に寄与する可能性も示された。

ただし、今回の比較は3Bモデル、単一の学習方式、特定のキャッシュインターフェースに基づく。より大規模なモデル、異なるタスクやblock設定でも同じ効果が得られるかは、今後の検証が必要だ。それでも本研究は、block diffusionを長文脈推論へ展開するための一貫した設計方針を示している。

arXiv