最近、自己回帰型TransformerのKVキャッシュがコンテキスト長に比例して線形に増大する問題に対処するため、Kabir Swain氏ら5名の研究者による論文「Tensor Cache: Eviction-conditioned Associative Memory for Transformers」がarXivで公開されました。スライディングウィンドウキャッシュはメモリを制限できますが、排除されたトークンは完全に破棄されるため、ウィンドウ外の重要な情報にアクセスできなくなります。
Tensor Cacheは2レベルキャッシュアーキテクチャを採用し、正確なローカルアテンションと圧縮された連想メモリを組み合わせます。第1レベル（L1）はスライディングウィンドウソフトマックスアテンションを使用し、最近のトークンに対して正確なアテンションを保持します。第2レベル（L2）は固定サイズの外積高速重みメモリであり、スライディングウィンドウから排除されたKVペアのみで満たされます。排除されたペアは層ごとの行列Aに圧縮され、将来のクエリは線形アテンション恒等式q_t(k_i ⊗ v_i) = ⟨q_t, k_i⟩v_iを利用した単一の行列乗算で読み出されます。
著者らは、学習可能なスカラーゲートを導入してL1とL2の出力を融合し、ヘッドごとの減衰率と書き込み率パラメータをエンドツーエンドで訓練します。さらに、一般的なチャンク平均学習ショートカットA←λA+η(ķ⊗v)がチャンクごとにC²-C個の偽のクロストークン外積を静かに導入することを特定し、並列加重和スキャン（float32イプシロン内でトークンごとの書き込みと等価）を用いてこのギャップを埋めました。
システムスケーリング、制御された連想想起、長文言語モデリング、メモリ容量診断の各実験において、Tensor Cacheは有状態ベースラインよりもメモリと品質のフロンティアを改善しました。これにより、Transformerのメモリ効率と長コンテキスト性能のトレードオフを改善する新しいアプローチが提供されます。この手法は、ドキュメント分析、対話システム、コード生成などの長いコンテキストを必要とするアプリケーションに影響を与える可能性があります。