Feed-Forward Steering in Transformer Residual Dynamics

コンピューターの脳内で行われている、巨大で目に見えないダンスパーティーを眺めていると想像してみてください。これは音楽やフラッシングライトのあるパーティーではなく、言葉と数字によるパーティーです。人工知能の世界、特にチャットボットや検索エンジンを動かしている「Transformer（トランスフォーマー）」には、「残差ストリーム（residual stream）」と呼ばれる情報の絶え間ない流れがあります。このストリームを、言葉を表す小さなボート（トークン）を運ぶ川だと考えてください。長い間、科学者たちは、これらのボートを操っているのは「Attention（アテンション／注意）」と呼ばれるメカニズムだけだと信じてきました。Attentionは、似たようなことを話しているボート同士を引き寄せる、巨大で目に見えない磁石のようなものだとイメージできます。もしあるボートが「王」と言えば、磁石は「女王」と言っているボートを近くに引き寄せるかもしれません。この引き寄せる力は「集約（aggregation）」と呼ばれ、これによって、意味の通じる言葉たちが同じ方向へと集まり始めるのです。

しかし、ダンスフロアにはもう一つの要素があります。それが「Feed-Forward Network（FFN／フィードフォワード・ネットワーク）」です。もしAttentionがボートを引き寄せる磁石であるなら、FFNは各ボートのすぐ隣に立って、そのボートが話している内容に基づいて特定の方向へと押し進める、ローカルなDJやパーソナルコーチのようなものです。長年、研究者たちはFFNを単なるサイドキック（脇役）であり、おそらくボートの速度を調整したり音量を修正したりする程度のものだと考えてきました。しかし、ある論文は大きな問いを投げかけています。FFNは単なる音量つまみなのか、それともボートがどこへ行くかを決定する「ステアリングホイール（操舵輪）」なのだろうか？ これを理解することは重要です。なぜなら、これらのボートがどのように動くのかを正確に把握できれば、AIをより速く、より賢く、そして間違いを犯したときに修正しやすくできるからです。

論文の大きな発見：DJこそがキャプテンである

「Feed-Forward Steering in Transformer Residual Dynamics」と題されたこの論文は、従来の視点ではパズルの重要なピースが欠けていたことを示唆しています。著者らは、ダンスフロアの新しい見方を提案しています。すなわち、Attentionは全員を共通のグループへと引き寄せる「集約」であるが、FFNは個々のボートを新しい方向へと積極的に押し出す「ステアリング・フィールド（操舵場）」であるという見方です。彼らは、これらの言葉のボートの動きを、球体（地球の表面のようなもの）の上を移動する粒子の物理学の問題として扱っています。

研究者たちは、FFNが単にボートを前後に押す（これは単に速度や大きさを変えるだけになります）のではなく、最も重要な役割は、ボートを横方向、つまり「接線方向（tangentially）」へと押し出すことであることを見出しました。地球儀の上のボートを想像してみてください。中心に向かって押すとサイズが変わりますが、横に押すと地図上の実際の方向が変わります。論文は、この横方向への押しが、言葉の意味と経路を実際に変えるものであることを示しています。

彼らが何を行い、何を見出したか

これを証明するために、チームはGPT-2、Pythia、Mistral、Llama-3を含むいくつかの有名なAIモデルを用いて一連の実験を行いました。彼らはAIを、パーツをオン・オフしたり、仕組みを変更したりできる科学実験室のように扱いました。

まず、Attentionだけでボートがどこへ向かうかを説明できるかどうかを検証しました。その結果、巨大な「角度整合性の欠損（angular-alignment deficit）」が見つかりました。簡単に言えば、もしAttentionの磁石だけを作動させた場合、ボートは実際のAIが送る方向とは異なる方向を向いてしまうということです。この欠損は、モデルが大きく新しくなるにつれて拡大し、GPT-2の約0.41からLlama-3-8Bの0.63へと増加しました。これは、FFNによるステアリングがなければ、AIは道に迷ってしまうことを示唆しています。

次に、彼らはFFNに対して「手術」を行いました。FFNの押しを、「径方向（radial）」（前後への押し）と「接線方向（tangential）」（横への押し）の2つの部分に分割しました。

• 径方向の部分のみを残した場合、AIのパフォーマンスは完全に崩壊しました。これはFFNを完全に削除したときと同じ状態でした。
• 接線方向の部分のみを残した場合、AIは品質のわずかな低下（極めて小さなドロップダウン）のみで、ほぼ完璧に機能し続けました。

これが決定的な証拠（smoking gun）となりました。これは、FFNの主な仕事が言葉の方向を操ること（ステアリング）であり、単にサイズを調整することではないことを証明しました。

また、Attentionの磁石が強く引きすぎて、すべてのボートを一つの小さな塊に押し込もうとする場合（「ランク崩壊（rank collapse）」と呼ばれる問題）、何が起こるかも調査しました。彼らは、FFNがカウンターフォース（抗力）として機能していることを見つけました。Attentionがボートを一つに押し潰そうとする一方で、FFNはボートを押し広げ、グループの多様性を維持し、すべてが同じ言葉になってしまうのを防いでいるのです。それは、まるでDJが「みんなが中央に集まりすぎている」ことに気づき、人々が再び散らばるようなビートを流し始めるようなものです。

実践的なテクニック：ダンスを加速させる

この論文で最もエキサイティングな部分は、彼らが発見した実践的なテクニックです。標準的なAIでは、ボートはまずAttentionの磁石によって引かれ、その後にFFNコーチが押しを進めます。これは逐次的（sequential）に行われます。研究者たちはこう問いかけました。「もし、磁石とコーチを同時に（並列で）働かせたらどうなるだろうか？」

通常、これはできません。なぜなら、コーチは磁石がボートをどこに引いたかを、その後に確認する必要があるからです。しかし、チームは各レイヤーにおいて、この順番がどれほど重要かを測るための「欠損スコア（defect score）」を測定しました。その結果、いくつかのレイヤーにおいては、順番はそれほど重要ではないことがわかりました。これらの「低欠損（low-defect）」レイヤーを特定することで、彼らはAttentionとFFNの部分を同時に実行することができました。

• GPT-2-largeでは、品質の損失をほとんどなくしつつ（わずか+0.02の損失）、プロセスを約1.24倍高速化することに成功しました。
• Mistralでは、1.10倍の高速化を実現し、損失はわずか+0.009でした。

しかし、もし彼らが、順番が重要となる（高欠損の）レイヤーでこれを試みれば、AIのパフォーマンスは崩壊してしまいます。これは、AI全体を即座に並列化して走らせることはできないものの、ステアリングがどのように機能するかを知ることで、特定のパーツを高速化できることを示唆しています。

これが意味すること

この論文は、AIのすべてを解決したと主張しているわけではありませんが、エンジンの見方を変えるものです。それは、FFNが単なるヘルパーではなく、AIの思考の幾何学を形作る「方向的なステアリング・フィールド」であることを示唆しています。FFNは、AIが単一の方向に固執するのを防ぎ、複雑な概念をナビゲートすることを可能にします。FFNが「アクセル」ではなく「ステアリングホイール」であることを理解することで、研究者は、道を見失うことなく、より高速で効率的なAIモデルを構築できるようになるかもしれません。これらの知見は、複数のモデルにわたる注意深い測定とシミュレーションに基づいたものであり、この「集約・ステアリング（aggregation-steering）」という見方が、デジタル脳の動きを理解するための堅牢な方法であることを示しています。