Adaptive inference and function vectors in deep transformers

問題提起
Transformerは、多様な領域における複雑な相関関係を学習するための汎用的な基盤として台頭しているが、そのインコンテキスト学習（ICL）における内部メカニズムは依然として謎に包まれている。経験的な研究により、ICLの2つの主要なモチーフ——特定のアルゴリズム（例：コピー、勾配降下法）を実装するアテンション回路と、「関数ベクトル（function vectors）」（タスクの内部表現）の形成——が特定されているが、後者の理論的理解は限られている。具体的には、ディープTransformerがどのように関数ベクトルを利用して推論を行うのか、そして深さとフィードフォワード（MLP）ブロックがどのような具体的な計算的役割を果たしているのかが不明である。既存の理論はアテンションベースの回路に焦点を当てることが多く、アテンションとMLPの間の多層的な相互作用によって可能となるメカニズムについては、十分に探求されていない。

手法
著者らは、ディープTransformerを、分散型ベイズ推論を実装する平均場相互作用系としてモデル化する理論的枠組みを提案する。このモデルは、通信（チャネル次元 $d$）、局所性、および深さ（$L$）に関するアーキテクチャ上の制約を受ける。

1. 理論モデル: この枠組みは、潜在的なコンテキスト変数 $\theta$ によって結合されたトークン集合を処理するTransformerを想定している。推論は、圧縮された状態変数として機能する「関数ベクトル」($\Phi_\ell$) を通じて媒介される。

   • 各レイヤー $\ell$ において、局所的な非線形性（MLP）がトークンと現在の状態 $\Phi_\ell$ を処理し、埋め込み $\lambda_{\ell+1}$ を生成する。
   • アテンションのような操作が、これらの埋め込みをトークン間で集約し、新しい統計量 $\phi_{\ell+1}$ を形成する。
   • 状態は直和によって更新される：$\Phi_{\ell+1} = \Phi_\ell \oplus \phi_{\ell+1}$。
   • MLPブロックは、「ルーター」（どの情報を伝達するかを選択する）およびデコーダーとして機能すると仮定される。一方、アテンションは情報を集約する。

2. 最適化戦略: 著者らは、埋め込みの選択を動的計画法問題として定式化する。レイヤー $\ell$ における最適な戦略は、レイヤー $L$ における期待損失を最小化するものであり、これは埋め込みの選択を適応的推論のための実験計画戦略として扱うことを意味する。

3. 解析ケース: 理論は、潜在変数 $\theta$ に関する2つの事前分布の下で分析される：

   • ガウス事前分布: 潜在変数が標準正規分布に従う場合。
   • ツリー事前分布: $\theta$ がバランスの取れた二分木のパスに基づいて構築される、階層構造を持つ事前分布。この構造は、層を追うごとに、より微細なスケールで潜在変数を解決していくための「深さ」の優位性をテストするものである。

4. 実証的検証: 提案された理論は、3つの構成を持つ制約付き線形アテンショントランスフォーマー・アーキテクチャを用いてテストされる：

   • モデルA: 深い（$L=M$）モデルで、チャネル次元 $d=1$、すべてのレイヤーにMLPブロックを備える。
   • モデルB: 単一層のモデルで、チャネル次元 $M$、同一のMLPを備える。
   • モデルC: 単一層のモデルで、チャネル次元 $M$、MLPなし（線形アテンションのみ）。
   • 因果的摂動: 「ドナー（提供者）」コンテキストのキーを「ターゲット」コンテキストに特定のレイヤーで挿入する「キー・パッチング（key-patching）」実験を行い、ルーティング・メカニズムをテストする。

主な貢献

• 関数ベクトルの理論: 本論文は、関数ベクトルをコンテキスト情報を予測するために圧縮する内部状態表現として定義する形式的な理論を提供する。これらのベクトルが、層を追うごとに、より微細なスケールで潜在的なコンテキスト変数を推論できることを示している。
• 深さとMLPの役割: 理論は、非ガウス的かつ階層的な事前分布に対して、深さとMLPブロックが、非適応的な単一層のアプローチを凌駕する適応的推論戦略を実装するために不可欠であることを予測する。逆に、ガウス事前分布の場合、深さは同等の総通信予算を持つ単一層に対して優位性を持たない。
• メカニズムのマッピング: 本研究は、理論的構成要素をトランスフォーマー・アーキテクチャにマッピングし、MLPを「統計量を測定する対象を選択する」ルーターとし、アテンションを「これらの統計量を集約する」役割として割り当てている。
• 実験的確認: 数値実験により、MLPを持つ深いトランスフォーマー（モデルA）が、階層的なツリー事前分布において適応的推論戦略を学習できることが確認された。これは、最適化された単層の非適応的戦略（モデルB）や線形アテンションモデル（モデルC）を大幅に上回る性能である。

結果

• ガウス事前分布: ガウス事前分布の場合、平均二乗誤差（MSE）損失は関数ベクトルの具体的な実現とは無関係であり、測定戦略のみに依存する。理論と実験の両方が、深いモデル（モデルA）は、同等のチャネル容量を持つ単層モデル（モデルB）よりも優れた性能を示さないことを示している。
• ツリー事前分布: 階層的なツリー事前分布の場合、理論は、適応的戦略（前の結果に基づいて一度に一つの座標を測定する）を用いると、MSE損失が深さに対して指数関数的（$\alpha^M$）にスケールする一方で、非適応的戦略は $M$ の累乗としてスケールすることを予測する。
  • 実験では、モデルAがモデルBよりも大幅に低いMSE損失を達成しており、モデルが適応的推論戦略を学習していることが確認された。
  • モデルC（MLPなし）は、非適応的な理論的ベースラインと同様の性能であり、適応的なルーティング・メカニズムにMLPが必要であることが確認された。
• 因果的摂動: キー・パッチング実験により、ドナー・コンテキスト（リーフ $\sigma'$）のキーをターゲット・コンテキスト（リーフ $\sigma$）に挿入すると、パスが早期に分岐し、かつ介入がネットワークの初期段階で行われた場合に、損失が大幅に増大することが明らかになった。これは、モデルの内部状態（関数ベクトル）がコンテキスト・パスに基づいて適応的に進化していることを裏付けている。

意義
本論文は、ICLに関する2つの視点――ICLを潜在変数に対するベイズ推論と見る視点と、ICLをアルゴリズム実行（例：勾配降下法）と見るメカニズム的視点――を統合することを主張している。深いトランスフォーマーが、関数ベクトルを介した分散型かつ適応的な層ごとの手続きとしてベイズ推論を実装できることを示すことで、著者らは、深さとMLPブロックが、従来の記述（前条件付き勾配降下法など）よりも豊かなクラスのICLアルゴリズムを可能にすることを実証している。本研究は、関数ベクトルを回路レベルのフレームワークの中に位置づけており、数十億パラメータを持つモデルの組織化が、この構造を支持していることを示唆している。これは、通信制約下での効率的かつ適応的な推論を可能にするためである。著者らは、本理論が線形集約を用いているものの、交換可能性を尊重しており、ソフトマックス・アテンションによる特定のトークン間リトリーバルがICLの主要な駆動源ではないシナリオにおいても十分である可能性があると述べている。