🤖 machine learning

Task Structure Reverses Layerwise State Encoding in Sequence Models

本論文は、シーケンスモデルにおける状態エンコーディングの層ごとの分布は固定されたアーキテクチャ上の特性ではなく、タスクの計算構造に基づいて反転するものであることを示しており、可換性のような代数的性質は、接頭辞更新（prefix updates）対スタック操作（stack operations）といった基礎的な計算要件よりも、メカニズム的シグネチャーを予測する上で寄与が低いことを明らかにしている。

原著者： Yuhang Jiang

公開日 2026-06-02

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Yuhang Jiang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、異なる種類の「思考マシン」（AIモデル）が、物語を読み進める際にどのように情報を保持しているのかを理解しようとしていると想像してください。長い間、研究者たちは、マシンの「思考スタイル」は人の性格のように固定されたものだと考えてきました。彼らはこう信じていました：

Transformer（現在の標準的なAI）は、思考を脳全体に均等に広げる。
再帰型モデル（Recurrent models）（古い、ループベースのAI）は、主要な思考を脳の最後の方の部分に大切にしまっておく。

この論文はこう言っています：「ちょっと待ってください。」

著者たちは、これらのマシンには固定された「性格」がないことを発見しました。代わりに、彼らは**「何の仕事をしているか」**に応じて、思考をどのように整理するかを変えるのです。それは、野菜を切る時と魚をさばく時で、使い分けるナイフを変えるシェフのようなものです。道具はシェフの正体ではなく、タスクに基づいて変化するのです。

以下に、シンプルな比喩を用いた彼らの発見の解説をまとめます。

1. 二種類の仕事

研究者たちは、マシンに3つの異なるパズルを与えました：

「合計計算」の仕事（Parity & S3）： コインを何回投げたかを数えるようなものです。現在のカウント（0または1）を覚えておき、進むにつれてそれを更新していきます。これは単純な線形更新です。
「スタック」の仕事（Dyck）： (( )) のように、括弧のバランスが取れているかを確認することを想像してください。単に数を数えるだけでは不十分で、どの開き括弧がどの閉じ括弧に対応しているかを覚えておく必要があります。入れ子構造を追跡するために、精神的な「スタック」（積み上げられた皿の山のようなもの）を構築しなければなりません。

2. 大逆転

論文によると、マシンは仕事の内容によって戦略を反転させることがわかりました。

「合計計算」の仕事において：
- **再帰型モデル（Mambaなど）は、「秘密の守り手」**のように振る舞います。序盤に多くの作業を行いますが、最終的な答えは脳の最後の層で初めて明らかになります。それは、影の中で手品を行い、最後にトリックを明かす手品師のようなものです。
- Transformerは、**「チームの作戦会議」**のように振る舞います。最初から最後まで、層を重ねるごとに少しずつ答えを構築していきます。全員が途中で少しずつ貢献していくのです。
「スタック」の仕事において：
- 役割が逆転します！
- Transformerは、突然**「秘密の守り手」**になります。最初の数層ですぐに答えを導き出し、その後はそれを保持し続けます。
- 再帰型モデルは、**「チームの作戦会議」**になります。答えを正しく出すために、層を重ねるごとにゆっくりと解決策を構築しなければなりません。

教訓： 「Transformerは常に情報を広げる」「Mambaは常に情報を隠す」と言うことはできません。戦略は、それが単純な更新なのか、複雑なスタックなのかによって決まります。

3. 「非可換」テスト（S3パズル）

これが単なる数学的なルール（足し算のように順序が関係するかどうか）によるものかどうかを証明するために、彼らはS3と呼ばれる、よりトリッキーな第3のパズルを追加しました。これは、順序が重要となる（靴を履く前に靴下を履くのと、靴を履いた後に靴下を履くのでは異なるような）「合計計算」の仕事でした。

予測： もし違いが単なる数学のルールによるものなら、このトリッキーな仕事は「スタック」の仕事と同じようになるはずです。
現実： マシンはこのトリッキーな仕事を、全く同じ「合計計算」の仕事として扱いました。つまり、以前と同じ「秘密の守り手」または「チームの作戦会議」の戦略を使用したのです。
結論： マシンは数学のルールに反応しているのではなく、計算構造（単純な更新か、複雑なスタックか）に反応しているのです。

4. 「読み取り可能」か「重要」かという罠

研究者たちは、情報の保存場所についても調査しました。彼らは、特に大規模な学習済みモデルにおいて、驚くべき乖離があることを見つけました。

「読み取り可能な」層： これは、マシンの脳内から答えを簡単に「読み取れる」場所です（答えが書かれた付箋を見つけるようなものです）。
「重要な」層： これは、もしその部分を壊すと、マシンが機能しなくなる部分です。

発見：

小さくて単純なモデルでは、「読み取り可能な」層と「重要な」層は通常一致しています。
大規模な学習済みモデルでは、これらはしばしば一致しません。
- 例：「スタック」の仕事において、大きなモデルは、脳の真ん中（レイヤー7）に答えが明確に書かれた付箋を持っているかもしれませんが、脳の最後（レイヤー11）を壊しても、モデルは正常に動作します。しかし、真ん中を壊すと、モデルはクラッシュします。
- 例：「合計計算」の仕事において、答えは最後の方に明確に見えるかもしれませんが、その特定の場所を壊しても、情報は他のあらゆる場所に分散されているため、モデルの動作を止めることはできません。

教訓： 特定の場所に答えが「書かれている」からといって、その部分がマシンを維持している唯一の要素であるとは限りません。マシンは、答えを一度に多くの場所で保持している可能性があるのです。

まとめ

この論文は、AIのアーキテクチャは硬直したものではないことを教えてくれます。それらは柔軟です。

タスクが重要： マシンの内部戦略は、それが単純な更新を行っているのか、複雑なスタックを行っているのかによって変化します。
構造が重要： 問題の「形」（更新かスタックか）は、数学のルールよりも戦略を決定づけます。
可視性 $\neq$ 必要性： 大きなモデルにおいて、答えがどこに「書かれている」かを見つけることは、必ずしもその部分が最も脆弱である（壊れやすい）ことを意味しません。

著者たちは次のように結論づけています。私たちは単に「このAIはどう機能しているのか？」と問うのではなく、「このAIはこの特定のタスクに対してどう機能しているのか？」と問わなければならないのです。

技術要約：タスク構造によるシーケンスモデルの層別状態エンコーディングの逆転

問題提起

シーケンスモデルのメカニスティックな解釈可能性（Mechanistic Interpretability）において、層ごとの状態エンコーディングは固定されたアーキテクチャ上の特性として扱われることが多い。再帰型モデル（LSTM、GRUなど）や状態空間モデル（SSM：Mambaなど）は、後半の層に読み取り可能な状態を集中させる傾向があり、一方でアテンションベースのモデル（Transformer）は、状態の蓄積を層の深さにわたって分散させるという仮説が一般的である。本論文はこの見解に対し、同じアーキテクチャであっても、計算構造の変化によって層別の状態エンコーディング・プロファイルが逆転することを指摘し、これらがアーキテクチャの不変的な特性ではないことを主張する。さらに、先行研究ではこれらのパターンを駆動する要因として、「代数的構造（特に可換性）」と「計算構造（プレフィックス更新 vs スタック追跡）」の二つを混同している場合が多い。著者らはこれらの要因を切り分け、状態エンコーディングがアーキテクチャ単独の特性なのか、それともアーキテクチャとタスクの相互作用によるものなのかを明らかにすることを目指している。

手法

本研究では、Transformer、Mamba、Mamba-2、LSTM、GRUの5つのアーキテクチャファミリーを用いた制御された実験スイートを採用している。

タスク:
著者らは、以下の3つの形式的な状態追跡タスクと、1つの準現実的なベンチマークを用いている：

Parity（パリティ）: 可換な、単一ビットのプレフィックス更新タスク。
$S_3$ Permutation Composition（ $S_3$ 置換合成）: 非可換な、6クラスのプレフィックス更新タスク。これは決定的な識別子となる。もし代数的可換性がエンコーディング・プロファイルを駆動するのであれば、 $S_3$ はDyckと同様の挙動を示すはずであり、もし計算構造が駆動するのであれば、 $S_3$ はParityと同様の挙動を示すはずである。
Dyck- $k$ : 有界な深さのスタック追跡タスク（括弧のマッチング）で、最大深度は10とする。
Code-Depth（コード深度）: CodeSearchNetから抽出されたPythonコードに基づく、トークンレベルのネスト深度タスク。これは準現実的な検証として機能する。

実験セットアップ:

形式的タスク: 4層モデル（ $d_{model}=128$ ）を使用。アーキテクチャ間でのパラメータ数は、可能な限り一致させている（約83万パラメータ）。学習はシーケンス長 $L=40$ で行い、評価は $L \in \{40, 60, 80, 100, 120\}$ で行う。
学習済みスケール: Mamba-130M（24層）およびPythia-160M/410M（12/24層）に対してファインチューニングを行う。
測定手法:
- 線形プロービング（Linear Probing）: 固定された隠れ状態に対してロジスティック回帰分類器を学習させ、各層における状態の読み取りやすさを測定する。
- 因果的介入（Causal Interventions）: ゼロ・アブレーション（プローブ方向の射影）およびサブスペース・アブレーション（上位のPCAまたは特異方向の除去）を用いて、機能的な必要性をテストする。
- OOD汎化: $L=40$ で学習したアブレーションを、未知の長さ（ $L=80, 100, 120$ ）に適用する。
- コンポーネント属性（Component Attribution）: 純粋なPyTorchバックエンド下でのMambaのConv1Dコンポーネントのアブレーションを行い、タスクごとの寄与を分離する。
- アクティベーション・パッチング（Activation Patching）: 学習済みMambaにおいて、クリーンな入力からの中間位置のアクティベーションを破損した入力へとパッチングし、ロジットの回復量を測定する。

主な貢献

タスク依存の逆転: 本論文は、「後半への集中」対「分散型の構築」というパターンがアーキテクチャの不変量ではないことを示している。このパターンはタスクに応じて反転する：
- プレフィックス更新タスク (Parity, $S_3$ ): 再帰型モデルおよびSSMは最終層に読み取り可能な状態を集中させるが、Transformerは層全体にわたって段階的に構築する。
- 有界スタック追跡 (Dyck): このパターンは逆転する。Transformerは即座に（初期層で）ほぼ完璧な読み取りやすさを達成するが、再帰型モデルやSSMは段階的に状態を構築していく。
代数と計算の分離: 非可換な $S_3$ タスクを導入することで、エンコーディング・プロファイルが代数的性質（可換性）ではなく、計算構造（プレフィックス更新 vs スタック）を追跡していることを示す。 $S_3$ はすべてのアーキテクチャにおいてParityと同じグループに属しており、可換性がエンコーディング・プロファイルの主要な駆動要因であるという仮説を否定している。
読み取りやすさと因果性のデカップリング: 状態が最も線形に読み取りやすい層（最良のプローブ層）が、必ずしもタスクに対して因果的に必要な層と一致しないことを明らかにした。このデカップリングは、表現がより冗長で分散している学習済みモデルにおいて特に顕著である。
コンポーネントレベルのタスク依存性: 単一のアーキテクチャ（Mamba）内であっても、特定のコンポーネントの重要性はタスクによって異なる。例えば、Conv1Dコンポーネントは $S_3$ とParityには極めて重要であるが、Dyckの性能には有意な影響を与えない。これは、アーキテクチャに関する単一の「メカニズム」という概念に異を唱えるものである。

主な結果

行動レベル: 再帰型ベースラインはParityと $S_3$ で完璧に汎化し、DyckにおいてはSSMやTransformerを上回る。 $S_3$ の性能順序（再帰型 > SSM > Transformer）は、DyckではなくParityの挙動を反映している。
プロービング・プロファイル:
- Parity/ $S_3$ : MambaとLSTMは、最終層（4層モデルにおける $L_3$ ）でプローブ精度が急上昇する。Transformerは $L_1$ から $L_3$ にかけて緩やかに上昇する。
- Dyck: Transformerは $L_0$ または $L_1$ で高い精度に達する。MambaとLSTMは緩やかな上昇を示し、後半の層でようやく飽和に達する。
- $S_3$ の確認: $S_3$ はParityのプロファイル（再帰型/SSMでは後半に集中、Transformerでは段階的）に従い、計算構造の仮説を裏付けた。
因果的介入:
- 4層の形式的モデルにおいて、最良のプローブ方向のアブレーションは大幅な精度低下を引き起こし、これらの方向は長いシーケンス長（OOD）においても必要であり続ける。
- 学習済みモデル: 分離が発生する：
  - Pythia (Dyck): 強力な中間層のボトルネックが存在する（160Mモデルにおいて $L_6$ – $L_7$ のアブレーションが精度を約81%低下させる）。最良のプローブ層（ $L_{11}$ ）は、最も因果的に重要な層ではない。
  - Mamba (Parity/Dyck/ $S_3$ ): 最終層は非常に読み取りやすいが、単一方向のアブレーションは無視できる程度の低下（<0.2%）しか引き起こさない。しかし、最終層での中間位置のアクティベーション・パッチングは、ロジットの差の約97–98%を回復させる。これは、状態は十分ではあるが、単一の線形プローブでは捉えきれない多くの方向に分散していることを示唆している。
Conv1D属性: Conv1DをMambaから除去すると、 $S_3$ （-12.64 pp）およびParity（-6.83 pp）で大幅な低下を招くが、Dyck（+4.04 pp、有意ではない）には大きな影響を与えない。これは、コンポーネントの重要性がタスク依存であることを示している。

意義と主張

本論文は、「アーキテクチャXはどのように状態をエンコードするか？」という問いは、タスクの計算構造を指定することなしには不完全であると論じている。著者らは以下の通り主張する：

メカニスティックな署名は結合特性である: 状態エンコーディング戦略は固定されたアーキテクチャの特性ではなく、モデルの帰納バイアスとタスクの計算要件（プレフィックス更新 vs スタック）の相互作用から生じるものである。
方法論的な注意: 一つのタスク領域（例：プレフィックス更新）から得られたメカニスティックな知見を、別の領域（例：有界スタック追跡）へ外挿することは、タスク固有の署名をアーキテクチャの不変量と誤認するリスクがある。
読み取りやすさと因果性: プロービングは情報が線形に利用可能な場所を特定するが、それが必ずしも計算のボトルネックとなる場所と一致するわけではない。特に冗長性が高い大規模な学習済みモデルにおいては顕著である。
代数 vs 計算: これらの形式的タスクにおいて、可換性のような代数的性質よりも、計算構造（逐次的なプレフィックス更新 vs 階層的なスタック）の方が、層別エンコーディング・プロファイルのより優れた予測因子となる。

著者らは、これらの知見が形式的タスク、学習済みスケール（最大410Mパラメータ）、および準現実的なコードベンチマークにわたって保持されることを示しており、シーケンスモデリングにおけるアーキテクチャとタスク構造の間の強固な相互作用を示唆している。