Key-Value Means

本論文は、標準的な演算と最小限の追加パラメータを用いながら、柔軟な状態の成長と準二次のプリフィル時間を可能にする効率的でチャンク並列化可能なトレーニングを実現することにより、トランスフォーマーと線形 RNN の利点を統合する新しいブロック再帰型アテンション機構であるキー・バリュー・ミーンズ（KVM）を導入する。

要約

非常に長い本を読もうとしていると想像してください。しかし、読みながら物語を頭の中に保持するための「ワーキングメモリ」の容量は限られています。

現在の AI の問題点
現在の AI モデル（トランスフォーマー）は、本で読んだ「すべての単語」を記憶しようとする生徒のように振る舞います。

• 良い点: 物語全体を目の前に持っているため、驚くほど正確です。
• 悪い点: 本が長くなるにつれて、その「ワーキングメモリ」は巨大化します。100 ページの本を読むのはわずかな努力で済みますが、1,000 ページの本を読むには莫大な時間とエネルギーを要します。まるで、一歩歩くたびに重くなるリュックサックを背負い続けようとしているようなものです。

再帰型（RNN 風）モデルの問題点
RNN 風モデルは、異なるアプローチをとります。これらは読んだ内容の要約を常に小さく固定されたサイズで保持し、読み進めるにつれてそれを更新していきます。

• 良い点: 非常に高速で軽量です。本がどれほど長くても、そのリュックサックは重くなりません。
• 悪い点: 物語の冒頭を忘れてしまいます。10 ページ目のプロットについて尋ねると、最後の数ページしか保持していないため、それを覚えていない可能性があります。

新しい解決策：キー・バリュー・ミーンズ（KVM）
この論文の著者は、**キー・バリュー・ミーンズ（KVM）**と呼ばれる新しい手法を導入しました。KVM は、両者の長所を兼ね備えた賢く魔法のようなノートブックだと考えてください。

以下に、簡単な比喩を用いてその仕組みを説明します。

1. 「スライディングウィンドウ」（即時の文脈）

本を読んでいて、最後の数ページだけを鮮明に見せる虫眼鏡を持っていると想像してください。これが「スライディングウィンドウ」です。KVM は、標準的な AI と同様に、最も直近の言葉に完璧に注意を向けます。これにより、直近の文脈を見逃すことがありません。

2. 「圧縮された要約」（長期的な記憶）

その数ページを読み進めると、古いページは虫眼鏡から外れていきます。現在の AI のように本全体を持ったり、RNN 風モデルのように単純に捨てたりするのではなく、KVM はこうして賢いことをします。

• 外れていったページを見て、
• 「これらのページの中で、最も重要でユニークなのはどれか？」と問いかけ、
• その重要なページに関する短く圧縮された要約を、特別なノートブックに書き込みます。
• もしノートブックにある内容と非常に似た新しいページが現れれば、既存のメモを更新します。もし全く新しく驚くべきものであれば、ノートブックに新しい行を追加します。

3. 「賢い統合」（魔法のトリック）

論文では、**「勝者総取り（Winner-Take-All）」**と呼ばれる情報を統合する特定の手法について説明しています。

• 新しい情報（水）と、ノートブック（スポンジ）のバケツを持っていると想像してください。
• 単に水を流し込むのではなく、KVM はスポンジの中で水と最もよく一致する正確な場所を見つけ、そこに吸収させます。
• また、「ジャスト・イン・タイム」正規化も使用します。これは、ノートブックに書き込む間は、合計値やカウント数を「正規化されていない（生の）」形で蓄積しておき、実際にノートブックを読み取る瞬間（クエリ時）になって初めて、その値を総数で割って平均値を計算するという仕組みです。読み取る直前にだけ割り算を行うことで、新しい情報を追加するたびに毎回再計算する必要がなくなり、効率的に情報を保持できます。

なぜこれが重要なのか

• 柔軟なサイズ: 速度のために小さなノートブック（固定サイズ）に保つように指示することも、本が長くなるにつれてノートブックを成長させる（拡張可能サイズ）こともできます。
• 速度とメモリのトレードオフ: 中間の道を選ぶことができます。「超高速だが忘れっぽい」か「超賢いが遅い」かの二者択一をする必要はありません。リアルタイム使用に十分な速さでありながら、物語全体を記憶するのに十分な賢さを持つように調整できます。
• カスタムハードウェア不要: 実行するために特殊で高価なコンピュータチップを必要とする他の新しい手法とは異なり、KVM は通常のソフトウェア操作を使用して標準的なコンピュータで実行できます。

結果

著者は、この手法を言語モデル（テキストを読み書きする AI）でテストしました。

• 短い文脈: 最高の標準 AI モデルと同様のパフォーマンスを発揮しました。
• 長い文脈: 入力テキストが数千トークンに及ぶ場合、拡張可能な KVM バージョンは、固定メモリを持つ RNN 風モデルよりもはるかに詳細を記憶し、完全なアテンションを持つトランスフォーマーモデルよりもはるかに高速でした。
• 「干し草の山の中の針」: 膨大なテキストの中に隠された特定の事実を見つける必要があるテストにおいて、拡張可能な KVM バージョンは非常に良い結果を示し、過去の深い部分から情報を実際に想起できることを証明しました。

要するに、KVM は、疲れず、冒頭を忘れず、無限に重くなるリュックサックを必要とせずに、長い本を読むための AI の新しい方法です。 これを実現するために、現在の明確な視点を保ちながら、過去の賢く圧縮された要約を維持しています。

技術概要

技術的サマリー：キー・バリュー平均（KVM）

問題定義

トランスフォーマーは現代のハードウェア上での効率的なトレーニングを可能にする一方、コンテキスト長に対して出力トークンあたりのメモリと時間の線形スケーリング（$O(N)$ メモリ、$O(N)$ デコード時間）という欠点があります。一方、現代の線形 RNN（LRNN）はトークンあたりのメモリと時間を一定（$O(1)$）に保ちますが、通常、限られた長コンテキストの想起に苦慮します。このギャップを埋めようとする既存のアーキテクチャは、固定サイズの状態（想起の制限）や、実行時オプティマイザーを用いた複雑なテスト時トレーニング（速度への影響）に依存することが多いです。カスタムカーネルやテスト時トレーニングのための複雑なハイパーパラメータ調整を必要とせず、メモリ効率、速度、長コンテキストの想起を両立するアーキテクチャの必要性があります。

手法：キー・バリュー平均（KVM）

KVM は、ブロックスライディングウィンドアテンション（BSWA）と動的に拡張可能な圧縮状態を統合する、新しいブロック再帰的アテンション機構です。これは単一のソフトマックスアテンション層内で動作し、従来のトランスフォーマー（拡張可能なコンテキスト、チャンク単位の並列性）と線形 RNN の両方の利点を統合します。

中核メカニズム

1. 圧縮状態を伴うブロックスライディングウィンドウ:
   KVM は入力チャンク単位で処理を行います。最近のトークンに対して固定サイズの BSWA ウィンドウを維持し、古いトークンに対しては別途、周期的に更新される状態を維持します。トークンのブロックが BSWA ウィンドウを溢れさせた場合、破棄されるのではなく、状態を更新するために処理されます。

2. 状態の圧縮とマージ:
   溢れたトークンは、「勝者総取り（winner-take-all）」の余弦類似度のようなマージ規則を用いて状態に圧縮されます。

   • 類似度指標: 標準的なソフトマックスの代わりに、KVM はオンラインベクトル量子化に触発された、最大にスパースな更新行列を使用します。ここで、各溢れキーは、最も相関の高い状態キーの単一のものに割り当てられます。
   • ジャストインタイム（JIT）再正規化: 直交または反対向きのベクトルを平均化することにより、時間経過とともに状態ベクトルのノルムが縮小するのを防ぐため、KVM は JIT 正規化を適用します。状態キーはアテンション前にレイヤーノルムを用いて正規化されます。状態値は、スロット作成時に決定された固定の「読み出し半径（$\rho_i$）」に正規化され、方向の変化を許容しつつ値の大きさを保持します。
   • マージゲート: データ依存のスカラーゲートが、状態に吸収される流入する溢れキー/バリューの量を調節します。

3. 状態拡張戦略:
   固定サイズの RNN とは異なり、KVM は拡張可能な状態をサポートします。最も「驚き」のある（最も冗長でない）溢れトークンは直接状態に追加され、残りはマージされます。これにより、初期コンテキストの想起を維持しながら、サブ線形なメモリ成長（例：$O(\sqrt{N})$）が可能になります。

4. 位置エンコーディングの処理:
   BSWA ウィンドウ内でロータリー位置埋め込み（RoPE）との互換性を維持しつつ、広範に異なる位置からトークンを集約する圧縮状態では RoPE を使用しないようにするため、KVM は部分的な RoPE ゼロ化を採用します。状態キーのロータリー部分空間をゼロに置き、BSWA ウィンドウは完全な RoPE を維持します。これにより、モデルは同じアテンションパス内で状態に対して回転していないクエリを、ウィンドウに対して回転したクエリを使用することが可能になります。

5. シンクトークン:
   重要な初期コンテキスト情報の劣化を防ぐため、初期状態行の保護されたセット（シンク）が保持されます。これは、シンクトークンが固有の値の大きさを持つという問題に対処します。

主要な貢献

本論文は以下の具体的な貢献を提示します：

• 新しいブロック再帰的定式化: 勝者総取りマージ規則を用いて溢れトークンを動的に再正規化された状態に圧縮する手法。これにより、独立した圧縮層の必要性が排除されます。
• 状態拡張戦略: 新しい溢れトークンを状態に追加するメカニズム。想起を犠牲にすることなくサブ線形なメモリ成長を可能にします。
• JIT 再正規化: ベクトルノルムを維持し、平均化中の破壊的干渉を防ぐために、状態キーと値をジャストインタイムで正規化する方式。
• 部分的な RoPE 共有: 状態キーの RoPE 次元をゼロにすることで、圧縮領域と非圧縮領域間で位置エンコーディングを共有する手法。追加のパラメータや複雑な再マージメカニズムを必要としません。
• 統合アーキテクチャ: 固定状態 RNN と完全アテンショントランスフォーマーの間を補間する単一のアテンション層。メモリ効率と想起の間の連続的なトレードオフを提供します。

実験結果

著者らは、8k コンテキスト長で Prolong データセットを用いて、1 億 2000 万パラメータと 3 億 5000 万パラメータのモデルをトレーニングしました。

• 長コンテキスト性能:
  • 固定状態 KVM（256 トークン）: シーケンス位置損失および短コンテキストベンチマークにおいて、より大きな状態を持つ OVQ/SWA モデルを上回りました。しかし、極端な長さにおける新規の妨害要因を伴う「藁の中の針（NIAH）」テスト（NIAH-S2/S3）では、状態容量がボトルネックとなり苦戦しました。
  • 拡張可能 KVM（べき乗則/飽和スケジュール）: 「KVM sqrt」変種（状態サイズ $\propto \sqrt{N}$）は、長コンテキストベンチマーク（RULER、LongBench、NIAH）で競争力のある結果を達成し、8k のトレーニングコンテキストを超えた外挿領域において、非ハイブリッドの GPTAlpha モデルと同等かそれ以上の性能を示しました。長距離にわたる新規情報の検索を必要とするタスクにおいて、固定状態 KVM および純粋な LRNN（RWKV-7）を大幅に上回りました。
• 短コンテキスト性能: KVM 変種は、短コンテキストベンチマーク（LAMBADA、ARC、HellaSwag など）において標準的なトランスフォーマーと一貫した性能を示し、BSWA ウィンドウが標準的なアテンション能力を保持していることを確認しました。
• アブレーション研究: 値の長さの正規化を除去すると、最も顕著な性能低下が生じました。シンク保護とマージゲートを除去することも、長コンテキストの検索能力を大幅に弱めました。

意義と主張

本論文は、KVM が固定状態 RNN と完全アテンショントランスフォーマーの間のギャップを成功裏に埋めたと主張します。

• 効率性対想起: 状態サイズの柔軟な選択を提供し、メモリ効率と想起の間のトレードオフをユーザーが調整できるようにします。固定状態では $O(N)$ のチャンク化再帰的動作を提供し、拡張可能状態では、強力な長コンテキスト検索を伴うサブ線形なメモリ成長を達成します。
• 実装の簡素さ: KVM はカスタムカーネルなしで標準的な操作を用いて実装可能であり、チャンク単位の並列化可能なトレーニングとプリフィルをサポートします。
• ハイブリッドの可能性: このアーキテクチャは、LRNN レイヤーと組み合わせたハイブリッドソリューションとして使用でき、改善されたサブ線形なメモリ成長と長コンテキストデコード能力でそれらを補完できます。
• 実行時オプティマイザーの不在: テスト時トレーニング（TTT）アプローチとは異なり、KVM は SGD や Adam などの実行時オプティマイザーではなく、単純な状態更新規則に依存しており、関連するハイパーパラメータの課題を回避します。

著者らは、KVM が単純かつ効果的な方法で固定状態 RNN と完全アテンションの間を滑らかに補間することを可能にし、長コンテキストモデリングのための統合パッケージを提供することを示していると結論付けています。