Tensor Cache: Eviction-conditioned Associative Memory for Transformers

技術概要：テンソルキャッシュ

問題定義
自己回帰型トランスフォーマー推論は、接頭辞の再計算を回避するためにキー・バリュー（KV）ペアをキャッシュすることに依存しています。しかし、保持される KV 状態は、コンテキスト長、深さ、並行度に対して線形に増大し、メモリボトルネックを生み出します。スライディングウィンドウキャッシュは、最も最近の $W$ トークンのみを保持することでメモリ使用量を制限しますが、このアプローチは「硬い忘却（hard forgetting）」に悩まされます。つまり、ウィンドウからトークンが追い出されると、その情報は完全に破棄され、ウィンドウ外の関連証拠にアクセスできなくなるのです。既存の境界付き状態手法は、追い出されたトークンを完全に破棄するか、すべてのトークンを走行平均に圧縮しようとしますが、これらはノイズを導入したり、最近のコンテキストと遠くのコンテキストを効果的に区別できなかったりします。

手法：テンソルキャッシュ（TC）
テンソルキャッシュは、KV の追い出しを削除イベントから連想メモリの書き込み操作へと変換する、2 段階の境界付き状態アーキテクチャを導入します。

1. 2 段階アーキテクチャ:

   • L1（第 1 レベルキャッシュ）: 最も最近の $W$ トークンを保持する固定サイズの KV ラインバッファです。これは局所コンテキストに対する正確なソフトマックス注意をサポートします。
   • L2（第 2 レベルキャッシュ）: L1 ウィンドウから追い出されたトークンの情報のみを格納する、固定サイズの外積行列 $A$（高速重みメモリ）です。
   • 書き込みメカニズム: ラインバッファがエントリを上書きする際、移動された $(k, v)$ ペアは破棄されません。代わりに、外積更新を通じて L2 行列 $A$ に書き込まれます：$A \leftarrow \lambda A + \eta (k \otimes v)$。ここで、$\lambda$（減衰）と $\eta$（書き込み率）はヘッドごとに学習されるスカラーです。
   • 読み取りメカニズム: 将来のクエリ $q_t$ は、$r_t = q_t A$ を計算することで圧縮された過去にアクセスします。線形注意の恒等式 $q(k \otimes v) = \langle q, k \rangle v$ を活用することで、この操作は個々に格納することなく、すべての追い出されたトークンから情報を効果的に取得します。
   • 融合: 学習されたスカラーゲート $g$ が、局所注意（L1）の出力とメモリ読み取り（L2）の出力を融合します：$y_t = y_{local} + \sigma(g) m_t$。

2. トレーニング最適化（並列スキャン）:
   本論文は、標準的なチャンク化トレーニングのショートカットに重大な問題があることを特定しました。単純なアプローチは、チャンクをその平均（$\bar{k} \otimes \bar{v}$）で要約しますが、これによりストリーミング推論中には存在しない $C^2 - C$ 個の偽のトークン間外積（$i \neq j$ である $k_i \otimes v_j$）が導入されます。

   • 解決策: 著者は、逐次トークンごとの書き込みと数学的に同等である並列重み付き和スキャンを実装しました。これにより、float32 のイプシロン内（相対誤差 $< 10^{-7}$）でトレーニングと推論のギャップが解消され、モデルがトークン間干渉のノイズなしに正しい連想ダイナミクスを学習できるようになります。

主要な貢献

• 追い出し条件付き連想メモリ: 従来の外積メモリ（Infini-attention や mLSTM など）がすべてのトークンを取り込むのとは異なり、TC は外積行列を、スライディングウィンドウからの追い出しによってのみ供給される L2 キャッシュとして厳密に使用します。これにより、トークンは「L1 に常駐」と「L2 に圧縮」に明確に分割されます。
• トレーニングアーティファクトの修正: 本論文は、偽のトークン間外積を導入する、以前は見落とされていたトレーニングショートカット（チャンク平均）を特定し、解決しました。提案された並列重み付き和スキャンは、同じ計算コストでこの不一致を排除します。
• 境界付き状態の効率性: TC は、個々の追い出されたエントリを保持することなく、古いコンテキストが予測に影響を与えるメカニズムを提供し、完全な KV 保持に対する実用的な代替手段となります。

実験結果
著者は、システムのスケーリング、合成連想想起、および実テキストの長コンテキスト言語モデリング（OpenWebText およびシェイクスピア）において TC を評価しました。

• 合成連想想起: スライディングウィンドウを超えた情報の取得を必要とするタスクにおいて、TC は100% の精度を達成し、完全な KV パフォーマンスに匹敵しました。完全な KV に比べて保持される推論状態を**72–84%**削減しながら、ギャップが増大するにつれて性能が低下する Infini-attention や、確率的なレベルでしか機能しないスライディングウィンドウのベースラインを大幅に上回りました。
• 長コンテキスト言語モデリング: OpenWebText において、TC は 1,024 から 32,768 トークンまでのすべての評価コンテキスト長で最低の平均負対数尤度（NLL）を達成しました。唯一の例外は $L=2,048$ であり、ここでは Infini-attention と統計的に区別できませんでした。$L=32,768$ において、TC は NLL 5.14 に達し、完全な KV（6.00）や Infini-attention（5.42）を上回り、ピーク GPU メモリを完全な KV に比べて2.4 倍少なく使用しました。
• スループット: TC は、L2 の読み取り/更新に起因し、純粋なスライディングウィンドウ手法と比較して約 30% のスループットオーバーヘッドを発生させましたが、テストされたすべてのコンテキスト長において Infini-attention よりも一貫して高速でした。
• 容量診断: 生連想メモリ容量に関する実験では、より多くの無関係な関連付けが固定状態に圧縮されるにつれて、想起性能が低下することが示され、メモリの境界性が確認されました。

意義と主張
本論文は、テンソルキャッシュが、スライディングウィンドウ注意のメモリ効率と完全コンテキスト注意の想起能力の間のギャップを成功裏に埋めたと主張しています。KV の追い出しを破壊的な削除ではなく構造化されたメモリ操作として扱うことで、TC はトランスフォーマーが完全 KV キャッシュの線形メモリコストなしに、直近のウィンドウ外の関連証拠にアクセスすることを可能にします。著者は、彼らの主な貢献が外積メモリそのものの発明ではなく、追い出し条件付き L2 キャッシュとしての具体的な応用、および以前はこのような手法を阻害していたトレーニングアーティファクトの厳密な修正にあることを強調しています。TC は、長コンテキスト推論におけるメモリと品質のフロンティアを向上させる、実用的な境界付き状態の代替手段として提示されています。