ニュース

スパース オートエンコーダがどのように機能するか、直感的に説明します。

2024-08-05

한어Русский языкEnglishFrançaisIndonesianSanskrit日本語DeutschPortuguêsΕλληνικάespañolItalianoSuomalainenLatina



マシンハートレポート

編集者:パンダ

簡単に言うと、マトリックス → ReLU アクティベーション → マトリックス

スパース オートエンコーダ (SAE) は、機械学習モデルを解釈するためのツールとしてますます一般的になってきています (ただし、SAE は 1997 年から存在しています)。

機械学習モデルと LLM はますます強力で便利になってきていますが、依然としてブラック ボックスであり、それらがどのようにタスクを達成するのかは理解されていません。それらがどのように機能するかを理解することは大いに役立つはずです。

SAE は、モデルの計算を理解可能なコンポーネントに分解するのに役立ちます。最近、LLM 解釈可能性の研究者である Adam Karvonen が、SAE がどのように機能するかを直感的に説明するブログ投稿を公開しました。

解釈可能性の問題

ニューラル ネットワークの最も自然なコンポーネントは個々のニューロンです。残念ながら、単一のニューロンは、学術論文の引用、英会話、HTTP リクエスト、韓国語のテキストなどの単一の概念に都合よく対応しません。ニューラル ネットワークでは、概念はニューロンの組み合わせによって表現されます。これは重ね合わせと呼ばれます。

その理由は、世界の多くの変数が本来スパースであるためです。

たとえば、有名人の出生地は 10 億トレーニング トークンに 1 未満しか表示されない可能性がありますが、現代の LLM は依然としてこの事実と世界に関するその他の豊富な知識を学習できます。トレーニング データには、モデル内のニューロンの数よりも多くの個々の事実や概念が含まれており、これがおそらく重ね合わせが発生する理由です。

最近、ニューラル ネットワークを理解可能なコンポーネントに分解するために、スパース オートエンコーダー (SAE) テクノロジーの使用が増えています。 SAE の設計は、神経科学におけるスパースコーディング仮説にインスピレーションを得ています。今日、SAE は人工ニューラル ネットワークを解釈するための最も有望なツールの 1 つとなっています。 SAE は標準のオートエンコーダに似ています。

従来のオートエンコーダは、入力データの圧縮と再構築に使用されるニューラル ネットワークです。

たとえば、入力が 100 次元のベクトル (100 個の値を含むリスト) の場合、オートエンコーダーはまず入力をエンコーダー層に渡して 50 次元のベクトルに圧縮し、エンコードされた表現が に供給されます。デコーダを使用して 100 次元の出力ベクトルを取得します。圧縮プロセスにより再構築のタスクが非常に困難になるため、再構築プロセスは通常完全ではありません。



1x4 入力ベクトル、1x2 中間状態ベクトル、および 1x4 出力ベクトルを備えた標準オートエンコーダーの概略図。セルの色は活性化値を表します。出力は入力の不完全な再構築です。

スパースオートエンコーダの説明

スパースオートエンコーダの仕組み

スパース オートエンコーダーは、入力ベクトルを、入力よりも高い、等しい、または低い次元を持つ中間ベクトルに変換します。 LLM で使用される場合、中間ベクトルは通常、入力よりも高い次元になります。この場合、追加の制約がないため、タスクは単純であり、SAE は単位行列を使用して、何の驚きもなく入力を完全に再構築できます。ただし、制約を追加します。その 1 つは、トレーニング損失にスパース ペナルティを追加することです。これにより、SAE はスパースな中間ベクトルを作成します。

たとえば、100 次元の入力を 200 次元のエンコード表現ベクトルに拡張でき、エンコード表現に非ゼロ要素が約 20 個だけになるように SAE をトレーニングできます。



スパース オートエンコーダの概略図。中間のアクティベーションはまばらであり、ゼロ以外の値は 2 つだけであることに注意してください。

多くの層が含まれる可能性があるニューラル ネットワーク内の中間活性化に SAE を使用します。順方向パス中に、各層内および各層間に中間アクティベーションが発生します。

たとえば、GPT-3 には 96 層があります。順方向パス中、入力内の各トークンには 12,288 次元のベクトル (12,288 個の値のリスト) が含まれます。このベクトルは、処理の各層で次のトークンを予測するためにモデルによって使用されるすべての情報を蓄積しますが、不透明であるため、どのような情報が含まれているかを理解することが困難になります。

SAE を使用すると、この中間アクティベーションを理解できます。 SAEは基本的に「マトリックス→ReLU活性化→マトリックス」です。

たとえば、GPT-3 SAE の拡張係数が 4 で、その入力アクティベーションが 12,288 次元を持つ場合、その SAE エンコード表現は 49,512 次元 (12,288 x 4) になります。最初の行列は形状 (12,288, 49,512) のエンコーダー行列で、2 番目の行列は形状 (49,512, 12,288) のデコーダー行列です。 GPT アクティベーションをエンコーダーで乗算し、ReLU を使用すると、SAE 損失関数がスパース性を促進するため、49,512 次元の SAE エンコードされたスパース表現を取得できます。

一般に、私たちの目標は、SAE 表現にゼロ以外の値を 100 未満にすることです。 SAE 表現をデコーダで乗算することにより、12,288 次元の再構成されたモデルのアクティベーションが得られます。スパース性の制約により完全な一致を達成することが困難になるため、この再構成は元の GPT アクティベーションと完全には一致しません。

一般に、SAE はモデル内の 1 つの位置にのみ使用されます。たとえば、レイヤー 26 と 27 の間の中間アクティベーションで SAE をトレーニングできます。 GPT-3 の 96 層すべての出力に含まれる情報を分析するために、96 個の個別の SAE (各層の出力に対して 1 つ) をトレーニングできます。各層内のさまざまな中間アクティベーションも分析したい場合は、数百の SAE が必要になります。これらの SAE のトレーニング データを取得するには、大量のさまざまなテキストをこの GPT モデルに入力する必要があり、選択した位置ごとに中間アクティベーションが収集されます。

SAE の PyTorch リファレンス実装を以下に示します。変数には形状の注釈が付けられています。このアイデアは Noam Shazeer から来ています。https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd を参照してください。パフォーマンスを最大化するために、異なる SAE 実装には異なるバイアス項、正規化スキーム、または初期化スキームが含まれる場合が多いことに注意してください。最も一般的な追加の 1 つは、デコーダ ベクトル ノルムに対するある種の制約です。詳細については、以下の実装をご覧ください。

  • OpenAI:https://github.com/openai/sparse_autoencoder/blob/main/sparse_autoencoder/model.py#L16
  • SAELens:https://github.com/jbloomAus/SAELens/blob/main/sae_lens/sae.py#L97
  • 辞書学習:https://github.com/saprmarks/dictionary_learning/blob/main/dictionary.py#L30

輸入トーチ

torch.nnをnnとしてインポートする

# D = d_model、F = 辞書サイズ

# 例えば、d_model = 12288 かつ dictionary_size = 49152 の場合

# 次に model_activations_D.shape = (12288,) かつ encoding_DF.weight.shape = (12288, 49152)

クラス SparseAutoEncoder (nn.Module):

1 層のオートエンコーダー。

def __init__(self、activation_dim: int、dict_size: int):

スーパー().__init__()

自己.アクティベーション_dim = アクティベーション_dim

自己.dict_size = dict_size

self.encoder_DF = nn.Linear (activation_dim、dict_size、bias=True) です。

self.decoder_FD = nn.Linear (dict_size、activation_dim、bias=True) です。

def encode (self, model_activations_D: torch.Tensor) -> torch.Tensor:

nn.ReLU()(self.encoder_DF(model_activations_D)) を返します。

デコード(self, エンコードされた表現F: torch.Tensor) -> torch.Tensor:

self.decoder_FD (encoded_representation_F) を返します。

def forward_pass (self, model_activations_D: torch.Tensor) -> タプル [torch.Tensor, torch.Tensor]:

エンコードされた表現_F = self.encode (モデルアクティベーション_D)

再構築されたモデルアクティベーション_D = self.decode (エンコードされた表現_F)

reconstructed_model_activations_D、encoded_representation_F を返す

標準オートエンコーダの損失関数は、入力再構成結果の精度に基づいています。スパース性を導入するための最も簡単な方法は、SAE の損失関数にスパース性ペナルティ項を追加することです。このペナルティ項を計算する最も一般的な方法は、この SAE の符号化表現の L1 損失 (SAE 重みではない) を取得し、それに L1 係数を乗算することです。この L1 係数は、スパース性の達成と再構成精度の維持の間のトレードオフを決定するため、SAE トレーニングにおける重要なハイパーパラメーターです。

これは解釈しやすさを考慮して最適化されていないことに注意してください。代わりに、解釈可能な SAE 特徴は、スパース性と再構成を最適化することの副作用です。以下は基準損失関数です。

# B = バッチサイズ、D = d_model、F = dictionary_size

def calculate_loss (autoencoder: SparseAutoEncoder、model_activations_BD: torch.Tensor、l1_coeffient: float) -> torch.Tensor:

再構築されたモデルアクティベーションBD、エンコードされた表現BF = autoencoder.forward_pass (モデルアクティベーションBD)

再構築エラーBD = (再構築モデルアクティベーションBD - モデルアクティベーションBD).pow (2)

再構築エラーB = einops.reduce (再構築エラーBD、'BD -> B'、'合計')

l2_loss = 再構築エラーB.mean()

l1_loss = l1_coefficient * エンコードされた表現BF.sum()

損失 = l2_損失 + l1_損失

リターンロス



スパース オートエンコーダのフォワード パスの概略図。

これは、スパース オートエンコーダーの単一の順方向パスです。 1 つ目は、サイズ 1x4 のモデル ベクトルです。次に、これに 4x8 エンコーダ行列を乗算して 1x8 エンコードされたベクトルを取得し、ReLU を適用して負の値をゼロに変換します。このエンコードされたベクトルはスパースです。次に、これに 8x4 デコーダ行列を乗算して、1x4 の不完全に再構築されたモデルのアクティベーションを取得します。

仮説的な SAE 機能のデモンストレーション

理想的には、SAE 表現の各重要な数値は、何らかの理解可能なコンポーネントに対応します。

ここでは説明のためにあるケースを想定します。 GPT-3 の観点から、12,288 次元のベクトル [1.5、0.2、-1.2、...] が「ゴールデン レトリバー」を表すと仮定します。 SAE は形状 (49,512, 12,288) の行列ですが、それぞれが形状 (1, 12,288) を持つ 49,512 個のベクトルのセットと考えることもできます。 SAE デコーダの 317 ベクトルが GPT-3 と同じ「ゴールデン レトリバー」の概念を学習すると、デコーダ ベクトルは [1.5、0.2、-1.2、...] にほぼ等しくなります。

SAE アクティベーションの 317 要素がゼロ以外の場合は常に、「ゴールデン レトリバー」に対応するベクトル (および 317 要素の大きさに基づく) が再構築されたアクティベーションに追加されます。機械的な解釈可能性の用語では、これは「デコーダ ベクトルは残差フロー空間内の特徴の線形表現に対応する」と簡潔に説明できます。

49,512 次元の符号化表現を持つ SAE には 49,512 個の特徴があるとも言えます。特徴は、対応するエンコーダ ベクトルとデコーダ ベクトルで構成されます。エンコーダ ベクトルの役割は、他の概念の干渉を最小限に抑えながらモデルの内部概念を検出することであり、デコーダ ベクトルの役割は「実際の」特徴の方向を表すことです。研究者の実験では、各特徴のエンコーダとデコーダの特徴が異なり、コサイン類似度の中央値が 0.5 であることがわかりました。下の画像では、3 つの赤いボックスが個々のフィーチャに対応しています。



スパース オートエンコーダの概略図。3 つの赤いボックスは SAE 機能 1 に対応し、緑色のボックスは機能 4 に対応します。各機能には、1x4 エンコーダー ベクトル、1x1 機能アクティベーション、および 1x4 デコーダー ベクトルがあります。再構築されたアクティベーションは、SAE 特徴 1 および 4 のデコーダー ベクトルのみを使用して構築されました。赤いボックスが「赤い色」を表し、緑のボックスが「ボール」を表す場合、モデルは「赤いボール」を表す可能性があります。

では、仮説上の特徴 317 が何を表しているのかをどのようにして知ることができるのでしょうか?現在、機能の有効化を最大化し、その解釈可能性に対して直感的な応答を与える入力を探すことが実践されています。各機能をアクティブにする入力は通常、解釈可能です。

たとえば、Anthropic はクロード ソネットで SAE をトレーニングしたところ、ゴールデン ゲート ブリッジ、神経科学、人気の観光スポットに関連するテキストと画像がさまざまな SAE 機能を活性化することがわかりました。他の機能は、明白ではない概念によってアクティブになります。たとえば、Pythia でトレーニングされた SAE の機能は、「文の主語を変更するために使用される関係節または前置詞句の最後のトークン」という概念によってアクティブになります。 」

SAE デコーダ ベクトルは LLM の中間アクティベーションと同じ形状をしているため、モデルのアクティベーションにデコーダ ベクトルを追加するだけで因果的介入を実行できます。この介入の強度は、このデコーダ ベクトルに拡散率を乗算することで調整できます。 Anthropic 研究者がクロードの起動に「ゴールデン ゲート ブリッジ」SAE デコーダ ベクトルを追加したとき、クロードはすべての応答で「ゴールデン ゲート ブリッジ」について言及することを余儀なくされました。

以下は、仮説機能 317 を使用した因果的介入の参照実装です。 「ゴールデン ゲート ブリッジ」クロードと同様に、この非常に単純な介入により、GPT-3 モデルはすべての応答で「ゴールデン レトリバー」について言及するようになります。

def perform_intervention (model_activations_D: torch.Tensor、decoder_FD: torch.Tensor、scale: float) -> torch.Tensor:

介入ベクトルD = デコーダーFD [317, :]

スケール介入ベクトルD = 介入ベクトルD * スケール

修正モデル活性化D = モデル活性化D + スケール介入ベクトルD

modified_model_activations_D を返す

スパースオートエンコーダの評価ジレンマ

SAE を使用する際の大きな課題の 1 つは評価です。言語モデルを解釈するためにスパース オートエンコーダーをトレーニングすることはできますが、自然言語表現の基礎となる測定可能なグラウンド トゥルースはありません。現在、評価は非常に主観的であり、基本的には「一連の機能の活性化入力を研究し、これらの機能の解釈可能性を直感的に説明します。」これが解釈可能性分野の主な制限です。

研究者は、特徴の解釈可能性に対応すると思われるいくつかの一般的なプロキシを発見しました。最も一般的に使用されるのは、L0 と Loss Recovered です。 L0 は、SAE のエンコードされた中間表現内の非ゼロ要素の平均数です。 Loss Recovered は、GPT の元のアクティベーションを再構築されたアクティベーションで置き換え、不完全な再構築結果の追加の損失を測定します。 SAE はスパース性を改善するために再構成精度の低下につながるソリューションを選択する可能性があるため、通常、これら 2 つのメトリクスの間にはトレードオフがあります。

SAE を比較する場合、一般的なアプローチは 2 つの変数をプロットし、それらの間のトレードオフを調べることです。より良いトレードオフを達成するために、多くの新しい SAE メソッド (DeepMind の Gated SAE や OpenAI の TopK SAE など) はスパース ペナルティを変更しました。下の写真は、DeepMind の Gated SAE 論文からのものです。 Gated SAE はグラフの左上にある赤い線で表されており、このトレードオフでより優れたパフォーマンスを発揮していることを示しています。



ゲート SAE L0 と損失の回復

SAE の測定にはいくつかの難易度があります。 L0 と損失回復は 2 つの代理指標です。ただし、L0 は微分可能ではなく、SAE トレーニング中に回復された損失を計算するのは非常に計算コストがかかるため、トレーニング中には使用しません。代わりに、トレーニング損失は、下流損失への影響ではなく、L1 ペナルティ項と内部アクティベーションの再構築の精度によって決まります。

トレーニング損失関数はサロゲート メトリックに直接対応するものではなく、サロゲート メトリックは特徴の解釈可能性の主観的な評価の代用にすぎません。私たちの本当の目標は「モデルがどのように機能するかを理解する」ことであり、主観的な解釈可能性の評価は単なる代用にすぎないため、別の層の不一致が生じることになります。 LLM の一部の重要な概念は解釈するのが容易ではないため、解釈しやすくするために盲目的に最適化している間にこれらの概念を見落とす可能性があります。

要約する

説明可能性の分野にはまだ長い道のりがありますが、SAE は真の進歩です。 SAE により、ゴールデン ゲート ブリッジのステアリング ベクトルのようなステアリング ベクトルを見つけるための教師なし手法など、興味深い新しいアプリケーションが可能になります。 SAE は、言語モデル内のループをより簡単に見つけるのにも役立ちます。これは、モデル内の不要なバイアスを除去するために使用できます。

SAE が解釈可能な特徴を見つけられるという事実は (たとえその目的が単に活性化のパターンを特定することであっても)、SAE が意味のある何かを明らかにできることを示唆しています。また、LLM は単に表面的な統計パターンを記憶するのではなく、実際に何か意味のあることを学習できるという証拠もあります。

SAEは、Anthropicのような企業が目指していた初期のマイルストーン、つまり「機械学習モデル用のMRI(磁気共鳴画像法)」を表す可能性もある。 SAE はまだ完全には理解できませんが、不正な動作を検出するために使用できます。 SAE および SAE 評価の主要な課題は克服できないわけではなく、多くの研究者がすでにこのテーマに取り組んでいます。

スパース オートエンコーダの詳細については、Callum McDougal の Colab ノートブックを参照してください: https://www.lesswrong.com/posts/LnHowHgmrMbWtpkxx/intro-to-superposition-and-sparse-autoencoders-colab

https://www.reddit.com/r/MachineLearning/comments/1eeihdl/d_an_intuitive_explanation_of_sparse_autoencoders/

https://adamkarvonen.github.io/machine_learning/2024/06/11/sae-intuitions.html