The Query Knows What to Forget: A Second Erase Direction for Linear Attention

昔聞いた物語を思い出そうとしている場面を想像してみてください。あなたの脳には、これらの記憶を保持するための限られたスペースがあります。もしその小さなスペースにあまりにも多くの物語を詰め込もうとすると、細部が混ざり合い始め、ある物語の結末を別の物語の始まりと混ぜてしまうかもしれません。これは、現代のAIモデルが長い会話を処理する方法と似ています。彼らは「線形アテンション（linear attention）」と呼ばれる特別な種類のメモリを使用して、容量不足に陥ることなく、これまでに話されたすべての内容を把握しています。すべての単語を巨大で増え続ける積み重ねとして保存する代わりに、新しい言葉が入ってくるたびに自身を更新する「魔法のメモ帳」のような、固定サイズの「状態（state）」へと情報を圧縮するのです。

しかし、落とし穴があります。会話が長くなるにつれて、そのメモ帳にある古い記憶が互いに干渉し始めます。AIが特定の情報（例えば1万語前に言及された名前など）を探そうとするとき、他のすべての言葉から生じるノイズが、正しい答えを見つけることを困難にします。科学者たちは、AIをより賢く「削除」させることでこれを修正しようとしてきました。彼らは、新しい情報を入れるために古い無関係な情報を「消去」できるシステムを構築してきましたが、これまではその消しゴムは少し不器用でした。それらは、現在処理されている「現在の単語」に基づいて何を削除するかは分かっていても、AIが実際に「何を探しているか」については分かっていなかったのです。

この論文は、この不器用さを解決するための、QED（Query-derived Erase Direction：クエリ由来の消去方向）と呼ばれる巧妙なトリックを紹介しています。研究者たちは、3億4,000万個のパラメータを持つモデルを用いて研究を行い、AIの「検索質問（クエリ）」が何を消去するかを決定するのを助けることで、モデルが重要な情報をより長く保持できることを発見しました。彼らは150億トークンの大規模なデータセットでモデルを訓練し、この新しい手法によって、AIが以前よりもずっとはるか遡った会話の詳細を記憶できるようになったことを明らかにしました。具体的には、S-NIAH-1というテストにおいて、この新手法はモデルが確実に読み取れるテキスト量を倍増させ、有用な限界を約8,000語から16,000語へと押し上げました。論文は、これがAIによる長文読解を向上させるための重要な一歩であると示唆していますが、著者らは、この新しいツールの特定のパーツが厳密に必要かどうかを判断するために、さらなるテストが必要であるとも述べています。

問題点： 「一方通行」の消しゴム

この画期的な発見を理解するために、これらのAIモデルが現在どのように機能しているかを見てみましょう。AIのメモリを、本（情報）が絶えず入れ替わっている忙しい図書館だと想像してください。新しい本が届いたとき、司書（モデル）はどの古い本を捨ててスペースを作るかを決めなければなりません。

この論文以前のモデル（Gated DeltaNet-2など）では、司書は古い本を捨てるかどうかを決めるために、新しい本（「キー」）だけを見ています。それは、「この新しい本は猫についてのものだから、猫についての古い本を捨てよう」と言うようなものです。これはそれなりに機能しますが、盲点があります。

問題は、AIは単に新しい本を見ているだけでなく、答えを出そうとしている特定の質問（「クエリ」）を持っているということです。例えば、質問が「主人公は誰ですか？」だとします。司書はその特定の質問に対する答えを見つける必要があります。しかし、従来の「一方通行」の消しゴムは、新しい本のトピックに基づいて何を削除するかは分かっていても、投げかけられている質問については分かっていないのです。

著者らは、「干渉（AIを混乱させるノイズ）」が、まさに質問が探している場所に存在することに気づきました。もし質問が特定の詳細を探しているのに、司書が新しい本のトピックに基づいてしか削除を行わない場合、混乱を招くノイズが質問の経路のすぐそこに残ってしまう可能性があります。それは、針を探しているのに、司書が「針の色」に似た藁だけを取り除こうとし、実際には針が「ストロー」の山の下に隠れているという事実を無視しているようなものです。

解決策： 「双方向」の消しゴム

論文は、QEDと呼ばれる解決策を提案しています。一つの消しゴムの代わりに、質問自体によって導かれる第二の消しゴムを追加します。

このように考えてみてください。AIには現在、「キー指向の消しゴム（Key-Directed Eraser）」（古い方）と、新しい「クエリ指向の消しゴム（Query-Directed Eraser）」があります。

1. キー指向の消しゴムは引き続きその役割を果たします。新しい情報を見て、それに一致する古いものを排除します。
2. クエリ指向の消しゴムは、AIが現在何を求めているかを見ます。もし質問が「悪役は誰ですか？」であれば、この消しゴムは、その質問への答えを隠してしまう可能性のある、古くて混乱を招く情報を標的にして削除します。

決定的なのは、この新しい消しゴムが古いものに対して「直交（orthogonal）」するように設計されていることです。平たく言えば、これは古い消しゴムが手が届かない部分を掃除することを意味します。それは、ほうきが届かない部屋の隅の埃を吸い取る掃除機を持っているようなものです。

研究者たちは、150億トークンのテキストでモデルを訓練することでこれをテストしました。その結果、「質問」が消しゴムを導くようにした場合、モデルは長いテキストの中から特定の事実を見つけることが非常に上手くなったことが分かりました。

何を発見し、（何を発見しなかったのか）

結果は非常に有望なものでしたが、著者らは過剰な期待を持たせないよう注意深く記述しています。データが実際に示したことは以下の通りです。

• ロングコンテキストの超能力： S-NIAH-1（テキストの中の「干し草の中の針」を見つけるテスト）において、新手法は精度を大幅に向上させました。例えば、長さが16,000トークン（単語）のとき、標準的なモデルの正解率は約21.8%でしたが、新しいQEDモデルは39.8%に達しました。32,000トークンの時点では、9.4%から15.8%へと、さらに劇的な跳ね上がりを見せました。
• 利用可能な長さの倍増： 著者らは、この改善により、利用可能なコンテキスト長が実質的に倍増したと述べています。以前は、モデルは約8,000トークンの前後で混乱し始めていましたが、QEDを用いることで、16,000トークスまで信頼性を維持できました。
• 魔法のようなパープレキシティ（Perplexity）はない： 興味深いことに、この新手法は、AIが文章を書いたり次の単語を予測したりする能力（検証損失/validation loss と呼ばれる指標）を向上させることはありませんでした。損失は全く同じでした。これは、QEDがAIを一般的な意味で「賢く」するのではなく、圧縮された状態から特定の記憶を取り出す能力を高めるものであることを示唆しています。

「アブレーション（要素除去）」の謎：どのパーツが重要か？

QEDがなぜ機能したのかを正確に突き止めるため、研究者たちは、新しいシステムの異なる部分をオフにする実験を行いました。これは、車がなぜ速く走るのかを知るために、エンジンの部品を一つずつ外していく作業のようなものです。

彼らは主に3つのコンポーネントをテストしました：

1. 方向（Direction）： 消しゴムを導くために「質問（クエリ）」を使用すること。
2. ゲート（Gate）： いつ消去するかを決定するスイッチ（AIがまだ必要としているものを削除しないようにするため）。
3. 投影（Projection）： 新しい消しゴムが誤って古い消しゴムの仕事を台無しにしないようにするための数学的なルール。

ここで意外な展開がありました。結果はかなり複雑なものでした。

• 「方向」が鍵である： 消しゴムを導くために「質問」を使用したときは、毎回モデルが改善されました。この部分こそが真のヒーローであるようです。
• 「ゲート」と「投影」は不明瞭： 「ゲート」や「投影」をオフにしたときの実験結果は、一貫性がありませんでした。訓練で使用したランダムシード（開始点）によって、モデルが良くなることもあれば、悪くなることもありました。著者らは、これらのパーツが助けになる可能性はあるものの、実験によってそれらが厳密に必要であるとは証明されなかったと示唆しています。使用した特定のモデルサイズにおいては、モデルが自然に設計上の問題を回避していたため、「投影」は実際にはほとんど機能していなかった可能性があります。

結論

論文は、QEDが、AIモデルが長い会話の中で物事を記憶するのを助ける成功した手法であると結論付けています。「質問」が消去を決定するのを助けることで、モデルは通常メモリをブロックするノイズを取り除くことができます。

しかし、著者らは自らの知見の限界についても正直に述べています。彼らは「ゲート」や「投影」が不可欠であることを証明したのではなく、「クエリ」を用いて消しゴムを導くことが最も一貫した改善をもたらすことを証明したに過ぎません。また、この新しい消しゴムの強度は、タスクに応じて調整する必要があることも指摘しています。テスト時に消しゴムの「強さ」を変更すると、モデルごとに異なる結果が得られたためです。

要約すると、もしあなたがAIに小説を一冊読ませ、最終章で悪役の名前を覚えておいてほしいのであれば、「クエリによって導かれた消しゴム」を与えるべきだ、ということをこの論文は示唆しています。ただし、その新しい機械の他のつまみを正確にどう調整すべきかを判断するには、まださらなるテストが必要かもしれません。