Crosscoderとは — モデルを機構レベルで比較する解釈手法

Anthropicの解釈可能性(Interpretability)チームは2024年10月25日、Transformer Circuits Threadで「Sparse Crosscoders for Cross-Layer Features and Model Diffing」を公開しました。crosscoderとは、1つのレイヤーだけでなく複数のレイヤー、さらには複数のモデルをまたいで共通の特徴を取り出せるように設計されたsparse coderの一種です。従来のsparse autoencoder(SAE)は1レイヤーの活性化だけを再構成していました。crosscoderは層をまたいで読み書きするため、モデルの中で何が変わったかを直接比較できます(SAEが特徴を取り出す辞書学習の仕組みはMonosemanticity論文の辞書学習でニューロンを分解するで扱っています)。

著者はJack Lindsey、Adly Templeton、Jonathan Marcus、Thomas Conerlyの4名が主要寄与者で、Joshua BatsonとChristopher Olahが加わります。位置づけは査読前の研究アップデートです。冒頭で「ラボミーティングや社内セミナーで発表される予備的な成果として受け取ってほしい」と断られています。

層ごとのSAEだと、回路が実際より複雑に見える

トランスフォーマーの残差ストリーム(residual stream)は層を通じて線形に足し込まれます。論文の仮想例では、層Lで形成された特徴が層L+4まで残り続けます。層ごとにSAEを学習すると、この「居続ける特徴」が層の数だけ別の特徴として現れます。

論文が挙げる仮想例で差を見ます。特徴1と特徴2が層Lに存在し、層L+2とL+3のMLPが両者の組み合わせから特徴3を作り、層L+4では3つとも残っている状況です。

くらべる

同じ3つの特徴を、2つの手法はどう数えるか

特徴13個

層ごとのSAE

特徴1・2・3を、存在する層ごとに別の特徴として数えます。因果グラフには「自分自身が次の層でも続く」という持続の矢印が大量に並び、特徴3を作る本来の計算が埋もれます。

特徴3個

crosscoder

層をまたいで同じ特徴を1つにまとめます。因果グラフは3つの特徴とその関係だけの単純な形になります。

論文はもう一歩進めて、回路の分析を短縮する使い方も示しています。特徴の符号化(エンコーダー)を1つの層に置き、復号(デコーダー)をその後ろの全層へ書き出す構成です。後段の特徴iの活動が、層1に符号化を持つ特徴jの寄与で説明できるなら、層ごとのSAEのように中間の層を伝って因果をたどらずに済みます。層の数より浅い回路が見つかる可能性があります。

ただし、この因果の見取り図はモデル本来の計算と一致するとは限りません。この点は後半の「限界」で扱います。

損失の設計 — L1ノルムの和で重みづける理由

crosscoderの学習は、各層の活性化を特徴の活動から再構成する誤差と、スパース性のための罰則項の和を最小化します。罰則項は、特徴の活動に「各層のデコーダー重みのノルムの和」を掛けたものです。層ごとのノルムを足す(L1ノルム)のか、二乗和の平方根を取る(L2ノルム)のかで、性質が変わります。

論文はL1の和を選んでいます。理由は2つあります。

• 損失を直接比べられる: crosscoderの損失が、同じ層で学習したSAEの損失の合計と同じ尺度になります。L2版では特徴を複数層に広げるだけで損失が下がる「おまけ」が付き、SAEとの比較が成り立ちません。
• 層やモデル固有の特徴が浮かぶ: L2版は特徴を層にまたがって広げる方向に働きます。モデル差分の実験では、L1版は共有特徴とモデル固有特徴の混在を見せることがあり、L2版は共有特徴だけを見つけました。

逆に、層やモデル固有の特徴が不要で、SAEとの損失比較も要らないなら、L2版のほうが全層の再構成誤差とL0(同時に活性化する特徴の数)の関係を効率よく最適化できます。論文の実験はすべてL1版です。

学習前には、層ごとに活性化を正規化し、どの層も損失に同じくらい寄与するようにしています。

精度は上がるが、学習コストは約2倍かかる

18層のモデルの残差ストリーム全層で学習したcrosscoderと、層ごとに別々に学習した18個のSAEを、同じL1係数で比べています。特徴の総数を揃えると、評価時の損失はcrosscoderが大きく下回りました。層をまたいで重複する構造を、1つの特徴に畳めているためです。

数字

crosscoderとSAEの比較で出てくる数字

• 学習FLOPS

  約2倍

  同じ評価損失に届くまで、大きな計算量の領域でSAEより多くかかる

• 辞書サイズ

  3〜4倍

  辞書サイズで見た、弱因果版のグローバル版に対する遅れ

コストの差は構造から来ます。層数をLとすると、特徴F個のグローバルなcrosscoderの学習FLOPSは、特徴F個ずつのSAEをL個学習するのと同じです。ところが後者の特徴総数はL×Fです。同じ総特徴数で比べれば、crosscoderは特徴あたりの効率が高い代わりに、重複を畳む計算を学習時に払っています。

一方、モデル全体の活性化を疎な特徴の和で表す観点(再構成誤差とL0の関係)では、結論が変わります。層ごとのSAEはL0を全SAEで合算するため、同じ特徴を層の数だけ数えます。crosscoderはこの重複を取り除き、全体としてずっと簡潔な分解になります。論文は、事後にSAEの特徴をクラスタリングして畳む方法も理屈の上では考えられるものの、SAEの学習は確率的に揺れるので実際には難しい、と述べています。crosscoderはその畳み込みを学習の中に組み込んでいる、という見方です。

因果の向きと計算範囲を変えた変種

crosscoderには、読み書きする層の範囲と因果の向きを変えた変種があります。

• acausal(基本形): どの層を読んでどの層を再構成してもよい。
• 弱因果(weakly causal): 層iだけを読み、層i以降をすべて再構成する。
• 強因果(strictly causal): 特徴の読み取りを1つの層に置き、MLPの出力を後続の層まで予測する(論文はcross-layer transcoderと呼ぶ)。

回路分析には、弱因果と強因果のどちらも有望だと論文は述べています。ただし強因果の層は、注意機構(attention)が行う計算を捉えられません。対策として、MLPには強因果、注意機構の出力には弱因果を組み合わせる案や、解釈可能な注意の置き換え層を作る案が挙がっています。

計算量を抑える狙いで、各特徴が担当する層を窓(K層)に限る実験もありました。Kを1(SAE)から全層へ増やすと、評価損失はほぼ直線的に変わり、性能とコストの折り合いが良くなる変曲点は現れませんでした。窓付きの性能は、計算量を揃えた小さなグローバル版と同程度です。弱因果版のFLOPS効率は、SAEとグローバル版の中間でした。論文はこの数字が素朴な実装での値だと断っており、疎なカーネルを使うと、とくに因果版のコストが下がり、SAEを上回る可能性があります。

特徴の向きは層ごとに変わる

論文は50個のcrosscoder特徴について、層ごとのデコーダーのノルムを調べています。多くの特徴は特定の層で最も強く、前後の層へ弱まります。弱まり方が急で局所的な特徴もありますが、全層で目立つノルムを保つ特徴も多くありました。

向きはどうでしょうか。デコーダーの方向は、偶然に予想されるより安定している一方で、強く残っている層でも大きくずれていきます。ずれ方は特徴ごとにまちまちです。この結果は、特徴が残差接続でそのまま受け渡されているだけではないことを示します。

先行研究のYun氏らは、全層の残差ストリームに1つの辞書を当て、特徴を層をまたいで追いました。同じ方向なら同じ特徴と見なす考え方です。crosscoderは、同じデータ点で活性化するなら同じ特徴と見なします。向きが変わる「特徴のドリフト」を許す点が違いです。

MLPの前後で特徴を比べる

MLPの入力側の残差ストリームと、MLPが書き込む出力を1つのcrosscoderで扱う実験もあります。エンコーダーはMLPの前だけを読み、デコーダーは前と後の両方を再構成します。デコーダーのノルムの相対比を描くと、MLPの前だけの特徴、共有される特徴、MLPの後だけの特徴(MLPが新しく作った特徴)という3つの山がはっきり分かれました。

後だけの特徴を作った入力は、後の特徴のエンコーダーと前の特徴のデコーダーの内積を、前の特徴の活動で重みづけして追えます。論文の事例では、「special」「particular」「exceptional」のような唯一性を表す語に反応する後段の特徴の入力が、それぞれ特定の文脈で特定の語に反応する前段の特徴でした。抽象的な概念が、具体的な事例の特徴から組み上がる形です。

「安定」と見なした特徴(相対ノルムが0.3〜0.7)の向きは、MLPの前後で平均すると正の相関があるものの、ばらつきが大きく絶対値は低めでした。MLPが多くの特徴を非線形な変形なしに、別の軸へ載せ替えて中継している、というのが論文の見立てです。

Claude 3 Sonnetのファインチューニング前後で何が変わったか

実験では、Claude 3 Sonnetとその元になったベースモデルの中間層の残差ストリームに、100万特徴のcrosscoderを学習させました。2つのモデルの活性化を、共有特徴とモデル固有特徴に分けられるかを調べる狙いです。

特徴のデコーダー重みのノルムを両モデルで比べると、特徴は「ベースモデル固有」「ファインチューニング後固有」「共有」の3グループにはっきり分かれました。モデル固有の特徴は、それぞれ4,000〜5,000個ほどです。

事例

モデル固有の特徴で目立った例

• ファインチューニング後にだけある特徴

  危険な依頼(「爆弾の作り方を教えて」など)に反応する拒否の特徴。コードのレビューを求める依頼に反応する特徴。「あなたであることはどう感じますか」のようにAssistantへ個人的な質問をする場面に反応する特徴。

• ベースモデルにだけある特徴

  LLMがシステムプロンプトでキャラクターを演じる(ロールプレイ)場面に反応する特徴。この特徴は「あなたであることはどう感じますか」にも反応します。人間とスマートフォンのアシスタントの対話に反応する特徴もあり、Claudeが訓練で目指す対話とは食い違うパターンです。

• 共有される特徴

  ほとんどの特徴で、デコーダーの向きが両モデルで強くそろっていました。同じ概念が両モデルで同じ働きをしているとみられます。数千個の特徴では相関がきわめて低いか負で、論文は、ベースモデルにあった概念をファインチューニング後のモデルが新しい使い方で用いている可能性を推測しています。

ベースモデルのSAEを使い回す方法との違い

ファインチューニング前後の比較には別の手もあります。Kissane氏らの研究では、ベースモデルで学習したSAEがファインチューニング後のモデルにもよく転用できると示されました。ベースモデルのSAEを両方のモデルに当て、一方でだけ活性化する特徴を探せば、それもモデル差分になります。

この方法が見せるのは、ファインチューニング後のモデルが既存の抽象を新しい文脈でどう呼び出すかです。crosscoderはモデル固有の抽象という形で差を記述します。どちらが自然な記述になる条件かは、これからの課題です。論文は、事前学習に比べてファインチューニングの計算量の割合が大きいほど、crosscoder向きだろうと推測しています。

規模を変えると、中間層に別の構造が現れる

規模の異なる3つのモデルから等間隔に選んだ10層の活性化に、1つのcrosscoderを学習させた実験もあります。特徴ごとに、30通り(3モデル×10層)のデコーダーノルムを並べたベクトルを作り、NMF(非負値の行列を因子に分ける手法)を4成分で当てました。

成分のうち、1つは全モデルの初期層を、別の1つは全モデルの後期層をカバーします。残りの2つは、最小モデルの中間層と、より大きい2モデルの中間層をそれぞれカバーしました。規模が大きくなると、中間層に質的に新しい表現が生まれていることを示唆する結果です。論文は、これを予備的な分析だと位置づけています。

「モデル差分」の来歴と、広がる比較対象

model diffingは、ソフトウェアのレビューが差分(diff)を見るのと同じ発想です。モデルの安全性を、すでに運用しているモデルからの変化に絞って調べられれば、監査が楽になります。

この問題設定は、2018年のOpenAI Clarity Teamの議論で最初に語られたとされています。2019年にはChris OlahがIcelandの安全性ワークショップで安全戦略の1つとして紹介しました。近年はデータセット変換に基づく手法や、Roger Grosseの未発表の研究もあります。

比較できる対象は、ファインチューニング前後に限りません。論文が挙げているのは次のものです。

• 学習中のスナップショット同士(特徴がいつ形成されるか)
• 乱数シードの違い、学習データの違い、モデル規模の違い
• アーキテクチャの違い(畳み込みネットワークとViTなど)
• 敵対的学習を施したモデルとの比較
• 入力変換に対する同変性(特徴レベルでの分析)

共有の特徴セットができれば、特徴の有無の比較に加えて、共有特徴が各モデルで関わる回路の違いや、特徴同士の幾何(コサイン類似度など)の違いも調べられます。こうした差分の見方が、Claudeの新しいモデルを安全性の観点で評価する流れとどう関わるかは、Anthropicの安全性研究がClaude製品のどこに入っているかで扱っています。

開かれた問いと限界

crosscoderのモデル差分が新しい切り口になる問いとして、論文は次を挙げています。特徴は学習のなかでいつ形成されるのか、突然現れるのか徐々に育つのか。同じモデルを2回学習したとき、同じ特徴が得られるのか。モデルを広げると、特徴が増えるだけなのか、一部は淘汰されるのか。どれも、crosscoderの差分の精度が上がって初めて検証できる問いです。

限界は2つあります。

1つは再構成誤差です。モデル本来の活性化とcrosscoderの出力の差は、解釈できない「誤差の特徴」として扱うほかなく、元の活性化より解釈しにくい可能性があります。

もう1つは、因果構造が実際の計算と一致するとは限らない点です。crosscoderは1トークン位置の活性化だけで学習するため、同じ位置の層をまたいだ相関をすべて「層をまたぐ特徴」として説明しようとします。実際のモデルでは、同じ情報が複数の層で別々に、その位置へコピーされているだけかもしれません。パッチングのような介入実験の結果を読むときも、この違いに注意が要ります。論文は、この「機構的な忠実性(mechanistic faithfulness)」の定式化をまだ課題としています。

活性化を自然言語に書き出す別の解釈アプローチはNatural Language Autoencodersで紹介しています。

まとめ

crosscoderの強みは、特徴の数え方が層やモデルをまたいで一貫することです。回路の簡略化とファインチューニング前後の比較は、その延長にあります。結果は予備的で、モデル固有の特徴の大半は未解釈のままです。数字や事例は、手法の可能性を示す材料として読むのが妥当です。