Google DeepMindが提案した、圧縮メモリとローカルAttentionを組み合わせることで有限のメモリリソースで事実上無限長のコンテキストを処理する手法。従来のAttentionが保持するKVキャッシュを圧縮メモリに蒸留し、メモリ使用量を一定に保ちながら過去の情報にアクセスする。
Infini-Attention(インフィニアテンション)は、Google DeepMindが2024年4月に発表した長コンテキスト処理手法である。従来のTransformerベースのAttentionメカニズムを拡張し、コンパクティブメモリ(Compactive Memory)と呼ばれる圧縮記憶機構を導入することで、理論上無限長の入力シーケンスを有限のメモリで処理することを目指す。
Ring AttentionやContext Window Scalingは物理的なGPUメモリ・計算量の制約を複数デバイスやエンコーディング調整で緩和するが、根本的にはKVキャッシュの線形成長問題を抱えている。100万トークンのKVキャッシュは、Llama 2-7B(32ヘッド × 128次元 × 32レイヤー)で約50GBに達し、A100 80GBの大半を占有する。
Infini-Attentionは各Attentionレイヤーで以下の3成分を統合する。
各セグメント処理後、そのセグメントのK, Vを圧縮メモリに書き込む。書き込みは線形Attention(σ(Q)σ(K)^T V形式)に類似した連想行列更新で行われる。
新しいセグメントのクエリQで圧縮メモリを検索し、過去の情報を取得する。この検索もσ(Q)Mの行列積で効率的に実行される。
| 成分 | メモリ使用量 | 対象範囲 |
|---|---|---|
| ローカルAttention | O(S²)(Sはセグメント長) | 現セグメント |
| 圧縮メモリ | O(d² × H × L)(固定) | 全過去セグメント |
| ゲート | O(d × H × L)(固定) | - |
圧縮メモリのサイズはコンテキスト長に依存せず一定(d_key × d_value × ヘッド数 × レイヤー数)。これにより、100万トークンでも100トークンでもモデルのメモリフットプリントは変わらない。
1M(100万)トークンのPasskey Retrievalタスクで、8Bパラメータモデルが100%の精度を達成した(論文報告)。ファインチューニングには5Kトークン長のデータのみ使用。
500K(50万)トークンの書籍をInfini-Attentionで処理し、要約品質(ROUGE-L)でベースラインを上回った。
| 手法 | メモリ増加 | 理論上限 | 追加パラメータ |
|---|---|---|---|
| KVキャッシュ(標準) | O(N) | GPUメモリ上限 | なし |
| Ring Attention | O(N/devices) | デバイス数依存 | なし |
| Infini-Attention | O(1) | 無制限 | ゲート + メモリ |
| StreamingLLM | O(Window) | 無制限(近傍のみ) | なし |
A1: 圧縮メモリへの書き込みは非可逆であり、理論的には情報損失が発生する。ただしゲート機構によりローカルAttention(情報損失なし)と圧縮メモリの出力を動的に混合するため、直近の情報は正確に保持される。遠い過去の情報は要約的に保持される。
A2: 原理的には可能だが、圧縮メモリとゲートの追加パラメータをファインチューニングする必要がある。論文では事前学習済みの1Bおよび8Bモデルに数千ステップの継続学習で適用している。
A3: StreamingLLMはAttention Sinkトークン(先頭数トークン)と直近のウィンドウのみを保持し、中間のトークンを完全に破棄する。Infini-Attentionは破棄せず圧縮メモリに蒸留するため、過去の情報への(近似的な)アクセスが可能である。