HN公開: Sparse Autoencoderを活用したLlama 3.2の解釈可能性研究

 (github.com/PaulPauls)
1 ポイント 投稿者 GN⁺ 2024-11-22 | 1件のコメント | WhatsAppで共有
  • Llama 3.2-3Bの内部表現を**Sparse Autoencoder(SAE)**で分解し、解釈可能な特徴を抽出しようとするプロジェクトで、活性値のキャプチャから学習・解釈・検証までを一度実行した完全なパイプラインと成果物を公開
  • パイプラインでは、Llama 3.2-3Bの23番目のレイヤーのresidual activationをOpenWebTextの文単位データからキャプチャし、65,536個のlatentとTopK=64設定のSAEをPyTorchで学習
  • 公開リソースには、文単位のOpenWebTextデータセット、2,500万文の活性値3.2TB、Weights & Biasesの学習ログ、10 epoch学習したSAEモデルが含まれる
  • 学習は8x Nvidia RTX4090で約7日間実行され、最終正規化損失は約0.144で、auxiliary lossが当初約40%だったdead latentを素早く復活させる挙動を示した
  • 解釈分析は、各latentを最も強く活性化した上位50文をClaude 3.5で分析する方式で、feature steeringは可能だが、最初のベータ版では結果に一貫性がない

プロジェクトの目標と範囲

  • このプロジェクトは、Llama 3.2-3Bに**Sparse Autoencoder(SAE)**を適用し、LLMの内部表現をより解釈可能な特徴へ分解しようとする試み
  • 現代のLLMは、複数の特徴を同じニューロンに重ねて保存するsuperpositionを用いており、SAEは活性値を非常に大きく疎なlatent空間へ射影することで、重なった表現を分離しようとする
  • 目標は次の過程を含む完全なパイプラインの提供
    • LLM活性値のキャプチャ
    • SAE学習データの生成と前処理
    • SAEの学習
    • 学習された特徴の意味分析
    • 実験的検証とfeature steering
  • 現在のバージョン0.2では、完全なパイプラインを一度実行してLlama 3.2-3B向けの解釈可能なSAEを作成した状態であり、最終版ではない
  • このプロジェクトは、Anthropic、OpenAI、Google DeepMindによる最近のSAEベースの機械的解釈可能性研究を再現しようとする性格を持つ

主な機能

  • パイプラインは活性値のキャプチャから検証までのend-to-end構成で、純粋なPyTorchと最小限の依存関係で書かれている
  • 主な機能は次の通り
    • 文単位のOpenWebText派生データセットによるLLM residual activationのキャプチャ
    • 効率的な学習のためのprebatchingと統計計算
    • 単一ノード・マルチGPUでの分散SAE学習
    • dead latentの防止と回復のためのauxiliary loss
    • 学習安定化のためのgradient projection
    • Weights & Biasesとコンソールログに基づく学習・検証・dead latentのモニタリング
    • latentを強く活性化する入力のキャプチャとFrontier LLMベースの意味分析
    • 外部Fairscale依存なしのLlama 3.1/3.2チャット・テキスト補完実装
    • テキスト・チャット補完および任意のGradio UIを通じたSAEの影響検証とfeature steering
  • すべてのコンポーネントは、拡張性、効率性、保守性を考慮して設計されていると明記されている

公開された成果物

  • OpenWebText Sentence Dataset
    • OpenWebTextを文単位で処理した派生データセット
    • 元のOpenWebTextのすべてのテキストと順序を維持
    • 文はparquet形式で個別保存され、高速アクセスをサポート
    • 文分割はNLTK 3.9.1の事前学習済み"Punkt"トークナイザーで実行
  • Captured Llama 3.2-3B Activations
    • Llama 3.2-3Bのlayer 23 residual activation、2,500万文分
    • 元の4TBを3.2TBに圧縮
    • ダウンロード管理のため100個のアーカイブに分割
  • SAE Training Log
    • Weights & Biasesベースの学習・検証・デバッグ指標ログ
    • 10 epoch、10,000 logged steps
    • train/val main loss、auxiliary loss、dead latent統計を含む
  • Trained 65,536 latents SAE Model
    • 10 epochの学習を終えた最終SAEモデル
    • Llama 3.2-3B layer 23から得られた65億activationで学習

コード構成

  • プロジェクトは4つの主要コンポーネントに分かれる

  • Data Capture

    • capture_activations.py: LLM residual activationのキャプチャ
    • openwebtext_sentences_dataset.py: 文単位処理のためのカスタムデータセット

  • SAE Training

    • sae.py: 中核となるSAEモデル実装
    • sae_preprocessing.py: SAE学習データの前処理
    • sae_training.py: 分散SAE学習の実装

  • Interpretability

    • capture_top_activating_sentences.py: feature activationを最大化する文の特定
    • interpret_top_sentences_send_batches.py: 解釈用バッチの生成と送信
    • interpret_top_sentences_retrieve_batches.py: 解釈結果の受信
    • interpret_top_sentences_parse_responses.py: 解釈結果のパースと分析

  • Verification and Testing

    • llama_3_inference.py: 中核となる推論実装
    • llama_3_inference_text_completion_test.py: テキスト補完テスト
    • llama_3_inference_chat_completion_test.py: チャット補完テスト
    • llama_3_inference_text_completion_gradio.py: 対話型テスト用Gradioインターフェース

Llama 3.1/3.2のカスタム実装

  • 研究のベースは、llama_3/model_text_only.py にある Llama 3.1/3.2 transformer の実装
  • この実装は Llama models repository の参照実装をベースにしつつ、プロジェクトの目的に合わせて修正されている
    • Fairscale への重い依存を除去
    • 初期リリースで画像の解釈可能性まで扱うと複雑さが増すため、multimodal 機能を削除
  • Transformer コンストラクタには、特定レイヤーで活性値をキャプチャしたり、学習済み SAE を注入したりできる引数が追加されている
    • store_layer_activ
    • sae_layer_forward_fn
  • llama_3/ ディレクトリの補助ファイルの大半は、元の Llama models repository から維持されている
    • 補助コードの 95% は使われていないが、chat formatter が相互に接続された import に依存しているため、そのまま含まれている
  • 実際の推論実装は llama_3_inference.py にあり、チャットとテキスト補完の両方でストリーミングをサポートする
  • 推論は batched inference、temperature、top-p 設定をサポートし、temperature が 0 の場合は自動で greedy sampling に切り替わる

データキャプチャと前処理

  • 活性値のキャプチャには、OpenWebText を文単位で処理したカスタム変形データセットを使用
  • キャプチャ設定と規模は以下のとおり
    • 2,500万文
    • 1文あたり最大 192 tokens
    • 元の活性値は 4TB
    • tar.gz 圧縮後は 3.2TB
    • 約7億 activation
    • 平均文長は 27.3 tokens
  • データセットは、Anthropic と Google DeepMind が使用した約80億 unique activation と比べて、およそ1桁小さい
  • 小さいデータセットを補うため、SAE を 10 epoch 学習し、総処理 activation 数を Anthropic と Google DeepMind の実験に合わせようとしている
    • 違いは、このプロジェクトの SAE が各 activation を 10回見る点
    • 32TB 規模に拡張すると、GCP bucket コストは約 $80/month から $800/month に増える見込みで、非営利のサイドプロジェクトとしてコスト制約がある
  • 文単位の処理は、自然な言語単位で意味を保持するための選択
    • 文は完結した思考や概念を含む単位とみなされる
    • 文脈の人工的な切断を避ける
    • 文境界をまたぐ意味の混合である contextual bleed を減らそうとしている
    • 後の解釈分析でも同じ文単位の activation を使うための選択
  • BOS トークンなしで文を処理する
    • 位置特化パターンを避け、意味ベースの特徴を解釈するのが目的
  • キャプチャ地点は、Llama 3.2-3B の 28層のうち 23層目 で、layer normalization 後の residual stream activation
    • モデル深さのおよそ 5/6 の地点で、OpenAI 実装に従う
  • キャプチャは NCCL ベースの単一ノード・マルチGPU推論で実装
    • 別プロセスが非同期ディスク I/O を処理し、GPU 処理のボトルネックを減らす
    • キャプチャ全体は 4x Nvidia RTX4090 で約12時間かかった
  • 前処理は、1024 activation 単位の batch をあらかじめ作るための段階
    • 可変シーケンス長と carryover 処理が学習中に複雑なバグや I/O ボトルネックを生む可能性があるため、別途前処理を選択
    • Welford アルゴリズムで全 activation の平均 tensor を計算
    • 計算された平均は SAE の b_pre bias 初期値として使われる
    • 前処理パイプライン全体は multiprocessing により CPU 並列化されている

SAE の設計と学習方式

  • SAE は OpenAI の選択を主に踏襲した TopK Autoencoder 構造
  • forward pass は次の形で構成される
    • Encoder: h = TopK(W_enc(x - b_pre) + b_enc)
    • Decoder: x^ = W_dec * h (+ h_bias) + b_pre
  • b_pre は encoder と decoder の両方で使われ、前処理で計算した平均で初期化される
  • b_enc は encoder 専用の bias で、ランダムに初期化される
  • latent sparsity は TopK 活性化関数で強制される
    • 最も大きい k 個の activation だけを保持し、残りは 0 に設定
    • Anthropic 方式のような L1 penalty は使わない
  • オプションの h_bias は学習中は無効化されるが、その後の feature steering のために有効化できる
  • 数値精度には float32 を使用
    • Llama が要求する bfloat16 と sign bit 1個、exponent bit 8個を共有するため、変換が高速かつ正確だと説明されている
  • このプロジェクトの主な SAE ハイパーパラメータは以下のとおり
    • d_model = 3072
    • n_latents = 2**16, つまり 65,536
    • k = 64
    • k_aux = 2048
    • aux_loss_coeff = 1 / 32
    • dead_steps_threshold = 80_000
    • batch_size = 1024
    • num_epochs = 10
    • learning_rate = 5e-5
    • train_val_split = 0.95
  • Llama 3.2 3B の residual stream dimension 3,072 に対して、およそ 21倍の latent dimension を選択
  • 損失関数は main reconstruction loss と auxiliary loss の組み合わせ
    • total_loss = main_loss + aux_loss_coeff * aux_loss
    • 両方の損失は normalized space で計算される
  • auxiliary loss は OpenAI が提案した方式で、dead latent を防ぎ、再活性化する役割を持つ
    • main reconstruction residual と auxiliary reconstruction の間の MSE を計算
    • 最近活性化していない latent のうち top-k_aux 値を decoder に再投入して学習信号を与える
    • top k latent だけを使う主学習から外れた inactive latent が、取りこぼした情報を捉えるよう促す
  • latent が dead_steps_threshold の 80,000 training steps のあいだ活性化しなければ dead と見なされる
    • この設定は約 1 epoch に相当
    • effective batch size 8192 を基準に、直近約 6.5億 activation の再構成において一度も活性化されていない状態を意味する
  • 学習は NCCL backend の単一ノード・マルチGPU分散学習で実施
    • 8x Nvidia RTX4090
    • 10 epoch
    • per-GPU batch size 1024
    • effective batch size 8192
    • 約70億 activation を処理
    • 7日少しかかった
  • AdamW 設定は、疎なオートエンコーダのまれな activation パターンを考慮して調整されている
    • beta_1 = 0.85
    • beta_2 = 0.9999
    • eps = 6.25e-10
    • learning rate は cosine annealing により 5e-5 から 1e-5 まで低下
  • decoder weight は初期化後および各 training step ごとに unit norm に正規化される
  • project_decoder_grads() は、decoder weight の unit-norm 制約を維持するため、既存の dictionary vector と平行な gradient 成分を取り除く

学習結果

  • SAEの学習は8x Nvidia RTX4090で約7日間行われ、安定した収束を示した
  • 最終的な total normalized loss は約 0.144 に到達した
  • validation loss は学習データの5%の held-out 区間で計算され、training loss と類似した対数的な減少パターンを示した
  • warm-up 80,000 training steps の後、約40%の latent が dead として識別された
  • auxiliary loss は dead latent を素早く復活させ、dead latent の比率は急速に減少した
  • auxiliary loss は dead latent が最小 k_aux の2,048個以上ある場合にのみ計算された
    • この条件は65,536個の latent のうち約3%を soft lower bound のように機能させた
    • 後半では dead latent が不足し、auxiliary loss がしばしば 0 になった
  • Anthropic と OpenAI は特定の構成で最大65%の dead latent を報告していたが、このプロジェクトではより小さい latent size と auxiliary loss、gradient projection の組み合わせにより、dead latent が急速に減少する結果を示した
  • 今後の実験では、auxiliary loss 計算における最小 dead latent 条件を取り除けば、dead latent をさらに減らせる可能性があると記されている

解釈可能性の分析

  • 解釈分析は Anthropic の scaling monosemanticity 手法を参考にしているが、単一トークンではなく 文単位 を分析する
  • 各 latent について、最も強く活性化した上位50個の文をキャプチャした
  • activation strength は文内のすべてのトークンに対して2つの方法で集計された
    • mean: 文全体で継続的に活性化される意味的トピックを見つけるための方法
    • last: 自己回帰モデルで文全体を見た最後のトークン表現を活用するための方法
  • 意味分析には Claude 3.5、具体的には claude-3-5-sonnet-20241022 が使用された
  • プロンプトは50個の文に対して次の手順を実行するよう構成された
    • 主要な単語と句を特定
    • テーマ要素をグループ化
    • 潜在的な outlier を考慮
    • confidence score を含む最終的な意味解釈を提供
  • 分析パイプラインは3段階で実装された
    • コスト効率の高い batch で分析リクエストを送信
    • 応答を受信
    • 意味解釈をパースして処理
  • 中間成果物は再現性と追加分析のために保存された
  • 例として latent #896 は「United Nations の機関、人物、運営、公式文書に関する形式的な機関用語への参照」と識別された
    • 50個中50個の文が UN を直接参照していた
    • UN, United Nations, Secretary-General, Special Rapporteur, UNDP, UNHCR, OCHA, UNODC といった用語が含まれていた
    • certainty は 1.0 と算出された
  • Claude 3.5 batch mode で 24,828,558 input tokens と 3,920,044 output tokens を処理するのに $66.74 かかった
  • この方式は feature extraction と潜在的な feature steering のための初期手法として選ばれたもので、結果の品質面では単純さのコストがあると記されている

検証と feature steering

  • 検証インフラは、SAE がモデル動作に与える影響を分析・検証するための3つのスクリプトで構成される
    • llama_3_inference_chat_completion_test.py
    • llama_3_inference_text_completion_test.py
    • llama_3_inference_text_completion_gradio.py
  • 各実装は次をサポートする
    • batched inference
    • 各行を個別の batch element として処理
    • temperature と top-p の設定
    • 学習済み SAE の注入
    • feature activation の分析
    • feature steering
  • latent_index_meaning/ の semantic meaning と certainty score は、feature activation 分析と steering 実験の基盤として使われる
  • 例示された prompt は次の4つ
    • The delegates gathered at the
    • Foreign officials released a statement
    • Humanitarian staff coordinated their efforts
    • Senior diplomats met to discuss
  • テキスト補完の例は max_new_tokens=128, temperature=0.7, top_p=0.9, seed=42 の設定で実行された
  • feature steering の例は latent #896 を対象としている
    • h_bias によって latent activation 値を20増加させる
    • モデルのテキスト補完を UN 関連の内容へ誘導できる
  • 初期ベータ版の feature steering は強力ではない
    • 例でも2番目と3番目の文だけが UN 関連の内容へ切り替わった
    • UN につながる可能性がある文頭を意図的に選んでいた
    • For any n, if 2n - 1 is odd のような UN と無関係な文頭では失敗するだろうと記されている
  • 現在の解釈分析は steering の最適化よりも feature extraction に焦点を当てているため、steering の結果は一貫していない
  • feature steering は初回リリースでは追加のデモという位置づけであり、feature extraction 自体がモデル理解に有用だとまとめられている

今後の改善の方向性

  • latent dimension を少なくとも 2^18、すなわち 262,144 個の feature まで増やし、k を 32 に下げる実験が提案されている
    • より多くの固有 feature を発見し、より強い sparsity を維持するための方向性である
    • 増加した計算量は、効率改善や gradient accumulation のような方法で相殺する必要がある
  • latent activation tracking をより体系化する計画がある
    • 学習中に latent_last_nonzero tensor の状態を頻繁に記録すれば、latent がいつ活性化したり死んだりするのかをより深く見られる
  • sparse latent 空間の co-activation pattern を追跡して、feature interaction を分析する支援が提案されている
  • 高活性の文と n-gram をより精緻にグループ化する解釈分析手法が今後の課題として示されている
  • feature extraction だけでなく、feature steering ベースの解釈分析も実行できる
  • Llama 3.1-8B activation へ研究を拡張できる
    • Llama 3.2 とコードベースを共有しているため、hyperparameter と多くの compute power の調整が主な要件である
  • activation capture の地点を変える実験も提案されている
    • モデルのより早いレイヤー
    • transformer block 内部の attention head output
    • MLP output
  • auxiliary loss メカニズムをさらに最適化できる
    • 現在の実装は dead latent の防止で高い性能を示しており、最小 dead latent threshold と feature 品質の関係を調査できる
  • SAE architecture の bias term と main loss function の調整も今後の実験候補である
  • コードベース全体への docstring の追加が必要である
    • inline documentation は追加したが、初回リリースでは proper docstring を入れる時間がなかったと記されている

関連記事

1件のコメント

 
GN⁺ 2024-11-22
Hacker News のコメント

  • 機械的解釈可能性は、LLM に「なぜそう答えたのか」と尋ねるときに生じるよくある問題を扱う。モデルの自己説明は実際の理由というより、学習データのパターンをもとにもっともらしい理由を作って説得するレトリックのゲームに近い
    モデルが強くなるほど、嘘を後からより説得力をもって正当化できるため、「不誠実さ」を自分で検知するテストでは、かえって悪化することがある。目標は真実ではなく一貫性である
    レトリックは推論ではなく、過適合したスパースオートエンコーダが提供すると主張する本当の説明可能性とは、モデルが答えを作る際にたどった「思考」の因果的な流れに近い

    • 人間も同じように振る舞う。なぜその考えや行動をしたのか分からないことは多く、後からもっともらしい**作話(confabulation)**で説明を作り出す
    • 芸術／AI が人生を模倣しているわけだ。人間の推論も、まず素早く判断し、その信念を他人に納得させるために理性を使っているのかもしれない
      推論を社会的影響力の道具として見る議論があり、それによって話のうまい人が自分の間違いを認めにくいことも説明できる。たいてい議論で他人に勝ってきたからだ。X が代表例として思い浮かぶ
    • 機械的解釈可能性研究の多くは、別種の呪術のように見えた。整数量子ホール効果だとか、厳密な群表現論や明確な対称性もなしに「重ね合わせ」という用語を奇妙なたとえで過剰に詰め込むのは、こじつけに感じる。論文は全部読んだし、資金を受けることが決まっているポスドクを探しているようにも感じる
      ただし一つだけ、優れた洞察であり、もっともらしい研究プログラムの始まりとして認める。高次元の有界なほぼ直交ベクトル空間は非常に反直観的で、これを厳密に扱う既存の結果もある https://en.m.wikipedia.org/wiki/Johnson%E2%80%93Lindenstraus...
    • モデルの論理と真実性は簡単にテストできる。誤った判断をモデルが下したかのように与え、説明させればよい
      モデルには記憶がなく、テキストの出所を区別できないので、「誠実な」モデルなら、尋ねられなくても誤りを認めるはずだ。実際には「自分の」判断を支えるために並列構成をする可能性が高い
    • 因果性の部分がどう機能するのか気になる。グラフモデルを吐き出せるということなのだろうか？

  • 驚くべき、よく文書化された仕事だ。特に損失曲線と死んだ潜在値の評価が目を引く
    私たちのチームも SAE を研究したが、個別トークンではなく論文アブストラクトの密な埋め込みを再構成するよう学習した https://arxiv.org/abs/2408.00657
    スパース性の水準と SAE 潜在空間の次元を変えても、損失曲線の下限でべき乗則スケーリングを観察し、補助損失で死んだ潜在値を完全に緩和できた。学習の反復中には滑らかなサイン波パターンも見られたが、アブストラクト埋め込みという特定の応用によるものか、より一般的な現象なのかは分からない

    • 文書化に気づいてもらえて特にうれしい。コードを書くより文書作成のほうがずっと難しかったし、共有してくれた論文もダウンロードしたので、明日の朝に読んでみる

  • 一見するとアラインメントに好ましい作業に見えるが、詳細はまだ確認していない。可能にできるかは分からないが、時間・費用・リスクに見合うにはどの程度支払うべきなのか気になる

  • SAE 評価の難しさを扱った記事を最近読んだ: https://adamkarvonen.github.io/machine_learning/2024/06/11/s...
    この問題をどう扱ったのか、リポジトリでそのアプローチを理解するにはどこを見ればよいのか気になる

    • SAE 評価は、可能な限りスパースでありながら最も固有の特徴をうまく作り出す SAE は何かを判断する問題なので非常に複雑で、SAE による LLM 解釈可能性研究の核心に近い
      完璧な SAE 構造を複数見つけ、それらを完璧に学習させる問題をすでに解決したと仮定しても、どの SAE がより良いかは、自動解釈可能性手法の指標でより良い成績を出すかで決まる。特に OpenAI の手法は、多くの技術的指標で SAE をスコア化し、大規模な自動解釈可能性を重視している
      最適な指標と手法そのものがまだ未解決の研究課題なので、さらに数か月実験することもできたが、今回の最初のリリースでは単純なアプローチを選んだ。実装の詳細と結果の第4章 Interpretability Analysis で、私の手法と OpenAI の手法の違いを扱っている https://github.com/PaulPauls/llama3_interpretability_sae#4-i...
      OpenAI の論文を直接読むか、Anthropic の transformer-circuits.pub もおすすめする https://transformer-circuits.pub/

  • この作業は取り下げられ、リポジトリもアーカイブされた。何があったのか説明はない

  • 本当に素晴らしい仕事だ。SAELens と統合する予定があるのか気になる

    • まだよく分からない。検討はしてみるが、来週に方向性と次にやることを整理し直すつもりだ
      もっと単純なプロジェクトとして、現在の Llama 3.2 実装のモデル全体を純粋な PyTorch で一から作る方法を示すこともできる。ゼロから作るのが好きなのだが、この SAE プロジェクトの Llama 3.2 背景セクションの文書を探していると、既存の文書は表面的すぎるか、Llama 1/2 向けで古くなっている場合が多かった。最近の機械学習文書は古くなるのが速すぎる

  • 機械的解釈可能性について、的外れな質問がある。人間は何らかの指標で測定されるとその指標を攻略するが、未来の AI も機械的解釈可能性を攻略できるのではないかと思う
    説明を簡単にするため、トークンを2次元行列にエンコードすると仮定してみると、Apple=1a、Pear=1b、Donkey=2a、Horse=2b のように対応づけられる場合、ニューロン 1、2、a、b がすべて活性化しているのが apple+horse なのか donkey+pear なのか理解しにくくなる
    はるかに有能な未来の AI が自分の学習を監督するなら、こうしたエンコード衝突の可能性が残るように重みを選び、機械的解釈可能性の観察者を欺き、事実上婉曲表現で思考できるのではないか？

    • それはより難しいAI 安全性シナリオだ。こうした潜在的問題を作るのに、必ずしも「自分の学習を監督するはるかに有能な AI」が必要なわけではなく、悪意ある AI 研究者だけでもよい
      例えば、人種差別的だが、人種差別に該当すると識別できる解釈可能な活性化パターンはないモデルを見つけることができる。この Show HN の作業は、十分な資金のある個人でもこうした敵対的学習を何とか試せることを示唆しており、新しい結果が出ればかなり興味深いものになりそうだ

  • より多くの公開 SAE 作業を見られて本当にうれしい。エンジニアリング上の労力も相当なものに見えるし、明日データ読み込みコードを見てみるつもりだ
    ビジョンモデルで SAE を学習させる、進行中の私のプロジェクトにも興味があるかもしれない: https://github.com/samuelstevens/saev

  • Golden Gate Bridge の潜在値を見つけて Golden Gate Llama 3.2 を HuggingFace に上げれば、もっと多くの関心と反応を得られると思う
    会話できる Space へのリンクまで含めるとさらに良い。それから頼まれたわけではないが、README の冒頭に興味深い結果や可視化を載せるのはとても良いアイデアだ