Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

技術要約：トランスフォーマーを用いたカウンタ言語のモデリングにおけるスタック表現の因果的証拠

問題提起
メカニスティック・インタープリタビリティ（機械論的解釈可能性）は、形式言語を学習した言語モデル（LM）における潜在的な特徴を特定することに成功してきたが、これらの特徴の「因果的」な役割を確立するという点において、決定的な空白が残っている。Tiwariら（2025）などの先行研究は、カウンタ言語（例：Dyck-1、Shuffle-k）を学習したトランスフォーマーが、基礎となるスタック構造と一致する表現を学習することを示した。しかし、これらの研究は、これらの表現が単にタスクと相関しているだけなのか、それともモデルの計算において因果的に必要なものなのかを特定できていない。本論文はこの問いに取り組む。すなわち、「学習されたスタック表現は、トランスフォーマーの次トークン予測能力に因果的に関与しているのか？」という問いである。

手法
本研究では、学習、プロービング、および因果的介入の3段階の実験パイプラインを採用している。

1. モデルの学習: エンコーダーのみのトランスフォーマーを、Shuffle-k 言語（$k \in \{2, 4, 6, 8\}$、ここでShuffle-kは$k$個の独立したDyck-1言語のインターリーブとして定義される）の次トークン予測に対して学習させる。Shuffle-kは、$k$個の異なるブラケット・スタックの深さを同時に追跡することをモデルに要求する。モデルのアーキテクチャは、1層の注意ヘッド4つ、埋め込み次元32、隠れ次元64で構成される。モデルは25エポックまでに完全な検証精度を達成し、以降のすべての解析は完全に収束したモデルに対して行われる。
2. スタック深さのプロービング: モデルの隠れ状態（$h^{(l)}_t$）から特定のスタックの深さ（$c_i(t)$）を予測するために、線形プローブを学習させる。プローブが、単なる偽の相関ではなく真の構造的情報を捉えていることを保証するため、摂動を加えたラベルを用いて**制御プローブ（control probe）を学習させる。指標である選択性（Selectivity）**は、実験用プローベと制御プローベの精度の差として定義される。高い選択性は、スタックの深さが隠れ空間の中に真正にエンコードされていることを示す。
3. 因果的介入（アブレーション）:
   • 方向の抽出: プローベの重み行列 $W$ を特異値分解（SVD）によって分解する。右特異ベクトル $w$ を、スタック深さ表現に対応する主成分の方向として特定する。
   • アブレーション: 推論中、隠れ状態から $w$ に整列した成分を、操作 $h' = h - \alpha(w \cdot h)w$ を用いて除去する。ここで $\alpha \in [0, 1]$ はアブレーションの強度を表す。
   • コントロール: ランダムな方向も同様にアブレーションを行い、ベースラインとして使用する。
   • 評価: アブレーションされた状態をデコーダー・ヘッドに直接渡し（トランスフォーマーの内部層を再実行せずに）、**位置精度（Positional Accuracy, $A_{pos}$）および逐次精度（Sequential Accuracy, $A_{seq}$）**への影響を測定する。

主な結果

• 表現の存在: 線形プローベは、高い選択性（例：$k=6$、スタック2において85.8%）をもってスタック深さの情報を抽出することに成功し、モデルがスタック構造と一致する表現を学習していることを確認した。
• 因果的必要性: 主なスタック方向（$w$）を標的としたアブレーションは、モデルの性能を劇的に崩壊させる。アブレーション強度 $\alpha$ が0から1へと増加するにつれて、逐次精度はほぼ0%まで低下する。
• 特異性: ランダムな方向のアブレーションは精度に有意な影響を与えず、この性能崩壊がベクトル操作による一般的なアーティファクトではなく、スタック表現の除去に特有であることを裏付けている。
• 指標の感度: 本研究では、位置精度は低下するものの、逐次精度の方がより急激に崩壊することが観察された。これは、エラーが累積する性質に起因しており、一つの誤ったトークンが文字列全体を無効にしてしまうためである。

意義と主張
本論文は、カウンタ言語を学習したトランスフォーマーにおけるスタック表現が、単なる創発的なアーティファクトではなく、因果的に必要なものであるという強力な経験的証拠を提供すると主張している。

特定の学習された表現の方向を除去することで、モデルがほぼ完全に失敗することを示すことで、著者らは、トランスフォーマーが形式言語に求められる論理的推論タスクを解決するために、これらのスタックのような構造を明示的に利用していると論じている。これは、分野を（プロービングによる）相関分析から、因果的なメカニズムの理解へと進化させるものであり、モデルの「ブラックボックス」的な振る舞いが、特定の識別可能な計算プリミティブに依存していることを検証している。

今後の展望
著者らは、今後の研究として以下の事項を提案している。

• アブレーションまたは編集された状態をフィードフォワードネットワークに通過させ、後続の出力への変化を観察することで、ネットワークの中間層で介入を行うこと。
• 異なるアーキテクチャが、これらの表現の創発、冗長性、および因果性にどのように関与しているかを調査すること。
• トランスフォーマーが、カウンタのような他の創発的な計算プリミティブの操作をどのように扱うかを理解すること。