技術的概要：多層トランスフォーマー介入のための残差化スパースオートエンコーダ

問題定義

スパースオートエンコーダ（SAE）は、トランスフォーマーの隠れ状態をスパースで解釈可能な特徴に分解する機械的解釈性の主要なツールである。しかし、標準的な慣行では、各層の生アクティベーションに対して SAE を独立して訓練する。このアプローチは、トランスフォーマーの根本的な構造的特性、すなわち残差ストリームが深度にわたって強く結合しているという事実を無視している。任意の層のアクティベーションの相当部分は、以前の層から前方に運ばれたり、軽微に変換されたりした情報で構成されている。

生アクティベーションに対して独立した SAE を訓練することは、多層介入において 2 つの実用的な問題を生じさせる：

1. 容量の冗長性: 異なる層ごとの辞書が、同じ予測可能な運ばれてくる情報を再学習することで容量を浪費する。
2. 予測不可能な相互作用: 複数の層が同時に置換される場合（例えば、特徴の除去や再構成を通じて）、層間の相互作用は加法的ではない。同じ情報が複数の辞書で表現されている場合、同時置換は特徴の重複、相殺、またはアクティベーションを訓練分布から逸脱させる原因となり、誤差の累積とモデル性能（クロスエントロピー）の低下を招く。

手法：残差化スパースオートエンコーダ（ReSAE）

著者は、層固有の情報を予測可能な層間構造から解離させるために、**残差化スパースオートエンコーダ（ReSAE）**を提案する。この手法は 3 つの主要な段階を含む：

1. 対アフィン残差化

選択された層のサブセット $L = (\ell_1, \ell_2, \dots, \ell_M)$ に対して、この手法は連続する選択された層間のアフィン写像を適合させる。対 $(\ell_m, \ell_{m+1})$ に対して、アフィン変換 $A_m h^{(\ell_m)} + c_m$ を、リッジ正則化付きの最小二乗法を通じて、前の層から後の層のアクティベーションを予測するように学習する。
残差化アクティベーション $r^{(\ell_{m+1})}$ は、このアフィン写像によって説明されない後の層のアクティベーションの成分として定義される：
$$r^{(\ell_{m+1})} = h^{(\ell_{m+1})} - (A_m h^{(\ell_m)} + c_m)$$
最初のアンカー層は、その経験的平均によって中心化される。

2. 独立した層別 SAE 訓練

各 SAE は、生アクティベーションではなく残差に対して訓練される。

• 入力: 残差アクティベーションは（ブロックごとの RMS 正規化を通じて）正規化され、標準的な BatchTopK SAE に供給される。
• 目的: SAE は、残差ターゲットの再構成損失を最小化する。
• 独立性: 勾配は異なるブロックの SAE 間を流れることはない。各 SAE は、特定の層で導入される「新しい」情報のみを学習する。

3. 元空間での再構成とオンライン介入

標準プロトコルを用いて ReSAE を評価するために、再構成された残差は、適合されたアフィン連鎖を通じて元のアクティベーション空間にマッピングされる。

• 再構成: $\hat{h}^{(\ell_{m+1})} = A_m \hat{h}^{(\ell_m)} + c_m + S^{-1}_{m+1} \hat{z}^{(m+1)}$。ここで $\hat{z}$ は残差の SAE 再構成である。
• オンライン介入: フック付きのフォワードパス中、層 $\ell_m$ の残差入力は、現在の入力アクティベーションから再構成された前の層 $\hat{h}^{(\ell_{m-1})}$ の予測寄与分を減算することで計算される。これにより、上流の層が修正された場合でも、層 $\ell_m$ の SAE は訓練ターゲットに類似した入力分布を視認することが保証される。

主要な貢献

1. ReSAE 構造: 前の選択された層から線形的に予測可能なアクティベーションの成分に対して SAE を訓練し、再構成をアフィン連鎖を通じて合成する訓練パラダイムの導入。
2. 改善された多層介入: 生 SAE に比べて少ない生分散を再構成するにもかかわらず、ReSAE が同時多層置換下でトランスフォーマーのクロスエントロピーをより忠実に回復することを示す。
3. 削減された冗長性: 残差化がより直交するデコーダ方向を生成し、容量が重複した層間構造ではなく、層局所的な情報に向けられていることを示す証拠。
4. 包括的評価: 再構成の質、クロスエントロピーの劣化、デコーダの類似性、スパースプロービング、ターゲットプローブ摂動（TPP）、および偽相関除去（SCR）を網羅する Pythia-1.4B および Gemma-2-9B における広範なベンチマーク。

結果

実験は、Pythia-1.4B（層 0, 6, 12, 18）およびGemma-2-9B（各種層間隔）において実施された。

• 再構成対機能性: ReSAE は、線形的に予測可能な構造を意図的に破棄するため、生 SAE に比べて元のアクティベーション空間における説明分散（EV）が低い。しかし、この低い EV は低いデコーダコサイン類似性（より少ない冗長性）と相関している。
• クロスエントロピーの回復: 多層逐次オンライン置換（誤差が伝播する状況）下では、ReSAE は特に高いスパースリティレベル（$k \ge 160$）において、モデルのクロスエントロピーの回復において生 SAE を大幅に上回る。教師あり強制設定（誤差伝播なし）では、ReSAE は一貫して最良のクロスエントロピー回復を達成する。
• 過剰相互作用: 生 SAE は、複数の層が置換された際、大きな負の「過剰相互作用」（非加法的な劣化）を示す。ReSAE はこの指標をゼロに近づけ、層ごとの置換がより予測可能に合成されることを示している。
• 特徴の質:
  • スパースプロービングおよび TPP: ReSAE は、スパースプロービングの精度およびターゲットプローブ摂動において一般的に生 SAE を上回り、学習された特徴が線形読み出しおよび選択的除去にとってより有用であることを示唆する。
  • 偽相関除去（SCR）: ReSAE はここで生 SAE を一様に上回らない。いくつかの設定（特に Pythia における層別結合除去）では、生 SAE のスコアが高い。これは、残差化された特徴が、タスク信号を損なうことなく特定の偽属性を除去する上で、必ずしも選択的ではない可能性を示唆する。
• スケールおよび間隔: ReSAE の恩恵は、異なる層間隔を跨いでより大規模な Gemma-2-9B モデルにおいても持続し、デコーダ冗長性の一貫した削減およびクロスエントロピー回復の改善が見られる。

意義および主張

本論文は、トランスフォーマーの残差ストリームを、孤立した層の集合ではなく結合された対象として扱うことが、効果的な多層介入にとって不可欠であると論じている。

• 容量配分: 線形的に予測可能な層間構造を除去することで、ReSAE はスパース辞書に、その特定の層におけるモデルの下流計算に最も関連するアクティベーションの成分に焦点を当てるよう強制する。
• 介入の安定性: この手法は、標準的な SAE 展開を悩ませる誤差の累積と予測不可能な相互作用を軽減し、多層介入のためのより安定したデフォルトを提供する。
• トレードオフ: 著者は、ReSAE がすべての解釈性タスク（例えば、SCR の結果は混合している）に対して万能な改善ではないことを認めている。しかし、多層置換および特徴ベースの介入という特定の目標に対しては、より高い機能的忠実度と削減された冗長性のために生分散を犠牲にするというトレードオフは有益である。

この研究は、多層設定における SAE にとって、アフィン回帰を通じて予測可能な構造を除去することが有用な前処理ステップであることを示唆しているが、予測子の最適な選択（アフィン対非線形）および層の選択は、依然として開かれた設計領域である。