sparse-autoencoder-training

SAELens：メカニスティック解釈可能性のためのスパースオートエンコーダー

SAELens は、スパースオートエンコーダー（SAE）の学習と分析に用いる主要なライブラリです。SAE は、多義的なニューラルネットワークの活性化を、スパースで解釈可能な特徴に分解する手法です。単義性に関する Anthropic の画期的な研究に基づいています。

GitHub: jbloomAus/SAELens

問題：多義性と重ね合わせ

ニューラルネットワークの個々のニューロンは多義的であり、意味的に異なる複数の文脈で活性化します。これは、モデルがニューロン数を超える特徴を表現するために重ね合わせを用いており、解釈可能性を難しくしているためです。

SAE はこの問題を解決します。密な活性化をスパースで単義的な特徴に分解します。通常、特定の入力に対して活性化する特徴は少数であり、各特徴は解釈可能な概念に対応します。

SAELens を使う場面

次のような場合に SAELens を使います：

• モデルの活性化から解釈可能な特徴を見つける
• モデルがどのような概念を学習したか理解する
• 重ね合わせと特徴の幾何学を研究する
• 特徴に基づく誘導やアブレーションを行う
• 安全性に関連する特徴（欺瞞、バイアス、有害なコンテンツ）を分析する

次のような場合は、別の方法を検討してください：

• 基本的な活性化分析が必要 → TransformerLens を直接使います
• 因果的介入実験を行いたい → pyvene または TransformerLens を使います
• 本番環境での誘導が必要 → 活性化の直接的なエンジニアリングを検討します

インストール

pip install sae-lens

要件：Python 3.10 以降、transformer-lens>=2.0.0

基本概念

SAE が学習するもの

SAE は、スパースなボトルネックを通じてモデルの活性化を再構成するように学習します。

Input Activation → Encoder → Sparse Features → Decoder → Reconstructed Activation
    (d_model)       ↓        (d_sae >> d_model)    ↓         (d_model)
                 sparsity                      reconstruction
                 penalty                          loss

損失関数：MSE(original, reconstructed) + L1_coefficient × L1(features)

主な検証結果（Anthropic の研究）

「Towards Monosemanticity」では、人間の評価者がSAE の特徴の 70% は真に解釈可能であると判断しました。見つかった特徴には、次のようなものがあります。

• DNA 配列、法律文書の言葉遣い、HTTP リクエスト
• ヘブライ語のテキスト、栄養に関する記述、コード構文
• 感情、固有表現、文法構造

ワークフロー 1：事前学習済み SAE の読み込みと分析

手順

from transformer_lens import HookedTransformer
from sae_lens import SAE

# 1. Load model and pre-trained SAE
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, cfg_dict, sparsity = SAE.from_pretrained(
    release="gpt2-small-res-jb",
    sae_id="blocks.8.hook_resid_pre",
    device="cuda"
)

# 2. Get model activations
tokens = model.to_tokens("The capital of France is Paris")
_, cache = model.run_with_cache(tokens)
activations = cache["resid_pre", 8]  # [batch, pos, d_model]

# 3. Encode to SAE features
sae_features = sae.encode(activations)  # [batch, pos, d_sae]
print(f"Active features: {(sae_features > 0).sum()}")

# 4. Find top features for each position
for pos in range(tokens.shape[1]):
    top_features = sae_features[0, pos].topk(5)
    token = model.to_str_tokens(tokens[0, pos:pos+1])[0]
    print(f"Token '{token}': features {top_features.indices.tolist()}")

# 5. Reconstruct activations
reconstructed = sae.decode(sae_features)
reconstruction_error = (activations - reconstructed).norm()

利用可能な事前学習済み SAE

リリース	モデル	レイヤー
gpt2-small-res-jb	GPT-2 Small	複数の残差ストリーム
gemma-2b-res	Gemma 2B	残差ストリーム
HuggingFace 上の各種モデル	タグ saelens を検索	各種

チェックリスト

• TransformerLens でモデルを読み込む
• 対象レイヤーに対応する SAE を読み込む
• 活性化をスパースな特徴にエンコードする
• トークンごとに最も強く活性化する特徴を特定する
• 再構成の品質を検証する

ワークフロー 2：カスタム SAE の学習

手順

from sae_lens import SAE, LanguageModelSAERunnerConfig, SAETrainingRunner

# 1. Configure training
cfg = LanguageModelSAERunnerConfig(
    # Model
    model_name="gpt2-small",
    hook_name="blocks.8.hook_resid_pre",
    hook_layer=8,
    d_in=768,  # Model dimension

    # SAE architecture
    architecture="standard",  # or "gated", "topk"
    d_sae=768 * 8,  # Expansion factor of 8
    activation_fn="relu",

    # Training
    lr=4e-4,
    l1_coefficient=8e-5,  # Sparsity penalty
    l1_warm_up_steps=1000,
    train_batch_size_tokens=4096,
    training_tokens=100_000_000,

    # Data
    dataset_path="monology/pile-uncopyrighted",
    context_size=128,

    # Logging
    log_to_wandb=True,
    wandb_project="sae-training",

    # Checkpointing
    checkpoint_path="checkpoints",
    n_checkpoints=5,
)

# 2. Train
trainer = SAETrainingRunner(cfg)
sae = trainer.run()

# 3. Evaluate
print(f"L0 (avg active features): {trainer.metrics['l0']}")
print(f"CE Loss Recovered: {trainer.metrics['ce_loss_score']}")

主なハイパーパラメーター

パラメーター	一般的な値	効果
d_sae	4-16× d_model	特徴数と容量が増加します
l1_coefficient	5e-5 to 1e-4	値が大きいほどスパースになりますが、精度は低下します
lr	1e-4 to 1e-3	標準的なオプティマイザーの学習率
l1_warm_up_steps	500-2000	学習初期に特徴が死ぬのを防ぎます

評価指標

指標	目標	意味
L0	50-200	トークンごとの平均アクティブ特徴数
CE Loss Score	80-95%	元のモデルに対する交差エントロピーの回復率
Dead Features	<5%	一度も活性化しない特徴
Explained Variance	>90%	再構成の品質

チェックリスト

• 対象レイヤーとフックポイントを選ぶ
• 拡張係数を設定する（d_sae = 4-16× d_model）
• 目標とするスパース性に合わせて L1 係数を調整する
• 特徴が死ぬのを防ぐために L1 ウォームアップを有効にする
• 学習中に指標を監視する（W&B）
• L0 と CE 損失の回復率を検証する
• 死んだ特徴の割合を確認する

ワークフロー 3：特徴の分析と誘導

個々の特徴の分析

from transformer_lens import HookedTransformer
from sae_lens import SAE
import torch

model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
sae, _, _ = SAE.from_pretrained(
    release="gpt2-small-res-jb",
    sae_id="blocks.8.hook_resid_pre",
    device="cuda"
)

# Find what activates a specific feature
feature_idx = 1234
test_texts = [
    "The scientist conducted an experiment",
    "I love chocolate cake",
    "The code compiles successfully",
    "Paris is beautiful in spring",
]

for text in test_texts:
    tokens = model.to_tokens(text)
    _, cache = model.run_with_cache(tokens)
    features = sae.encode(cache["resid_pre", 8])
    activation = features[0, :, feature_idx].max().item()
    print(f"{activation:.3f}: {text}")

特徴による誘導

def steer_with_feature(model, sae, prompt, feature_idx, strength=5.0):
    """Add SAE feature direction to residual stream."""
    tokens = model.to_tokens(prompt)

    # Get feature direction from decoder
    feature_direction = sae.W_dec[feature_idx]  # [d_model]

    def steering_hook(activation, hook):
        # Add scaled feature direction at all positions
        activation += strength * feature_direction
        return activation

    # Generate with steering
    output = model.generate(
        tokens,
        max_new_tokens=50,
        fwd_hooks=[("blocks.8.hook_resid_pre", steering_hook)]
    )
    return model.to_string(output[0])

特徴の帰属分析

# Which features most affect a specific output?
tokens = model.to_tokens("The capital of France is")
_, cache = model.run_with_cache(tokens)

# Get features at final position
features = sae.encode(cache["resid_pre", 8])[0, -1]  # [d_sae]

# Get logit attribution per feature
# Feature contribution = feature_activation × decoder_weight × unembedding
W_dec = sae.W_dec  # [d_sae, d_model]
W_U = model.W_U    # [d_model, vocab]

# Contribution to "Paris" logit
paris_token = model.to_single_token(" Paris")
feature_contributions = features * (W_dec @ W_U[:, paris_token])

top_features = feature_contributions.topk(10)
print("Top features for 'Paris' prediction:")
for idx, val in zip(top_features.indices, top_features.values):
    print(f"  Feature {idx.item()}: {val.item():.3f}")

よくある問題と解決策

問題：死んだ特徴の割合が高い

# WRONG: No warm-up, features die early
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=1e-4,
    l1_warm_up_steps=0,  # Bad!
)

# RIGHT: Warm-up L1 penalty
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=8e-5,
    l1_warm_up_steps=1000,  # Gradually increase
    use_ghost_grads=True,   # Revive dead features
)

Issue：再構成の精度が低い（CE の回復率が低い）

# Reduce sparsity penalty
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=5e-5,  # Lower = better reconstruction
    d_sae=768 * 16,       # More capacity
)

Issue：特徴を解釈できない

# Increase sparsity (higher L1)
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=1e-4,  # Higher = sparser, more interpretable
)
# Or use TopK architecture
cfg = LanguageModelSAERunnerConfig(
    architecture="topk",
    activation_fn_kwargs={"k": 50},  # Exactly 50 active features
)

Issue：トレーニング中にメモリエラーが発生する

cfg = LanguageModelSAERunnerConfig(
    train_batch_size_tokens=2048,  # Reduce batch size
    store_batch_size_prompts=4,    # Fewer prompts in buffer
    n_batches_in_buffer=8,         # Smaller activation buffer
)

Neuronpedia との統合

neuronpedia.org で、事前学習済みの SAE 特徴を閲覧できます：

# Features are indexed by SAE ID
# Example: gpt2-small layer 8 feature 1234
# → neuronpedia.org/gpt2-small/8-res-jb/1234

主要クラスのリファレンス

クラス	目的
SAE	スパースオートエンコーダーモデル
LanguageModelSAERunnerConfig	トレーニング設定
SAETrainingRunner	トレーニングループマネージャー
ActivationsStore	アクティベーションの収集とバッチ処理
HookedSAETransformer	TransformerLens と SAE の統合

リファレンスドキュメント

詳細な API ドキュメント、チュートリアル、高度な使い方については、references/ フォルダーを参照してください：

ファイル	内容
references/README.md	概要とクイックスタートガイド
references/api.md	SAE、TrainingSAE、設定の完全な API リファレンス
references/tutorials.md	トレーニング、分析、ステアリングの手順を追ったチュートリアル

外部リソース

チュートリアル

• 基本的な読み込みと分析
• スパースオートエンコーダーのトレーニング
• ARENA SAE カリキュラム

論文

• 単義性に向けて - Anthropic（2023）
• 単義性のスケーリング - Anthropic（2024）
• スパースオートエンコーダーが高い解釈性を持つ特徴を発見 - Cunningham ら（ICLR 2024）

公式ドキュメント

• SAELens ドキュメント
• Neuronpedia - 特徴ブラウザー

SAE アーキテクチャ

アーキテクチャ	説明	用途
Standard	ReLU + L1 ペナルティ	汎用
Gated	学習可能なゲーティング機構	より優れたスパース性の制御
TopK	有効な特徴が正確に K 個	一貫したスパース性

# TopK SAE (exactly 50 features active)
cfg = LanguageModelSAERunnerConfig(
    architecture="topk",
    activation_fn="topk",
    activation_fn_kwargs={"k": 50},
)

エージェントの作業手順

1. 環境を確認します。 モデル、チェックポイント、ベンチマークのバージョン、評価ハーネスのバージョンを確認します。
2. 入力を特定します。 データまたはユーザーから、形式、識別子、パラメーターを確認します。結果を変える値は推測せず、質問してください。
3. まず小規模に実行します。 小さなサブセットで評価し、プロンプト、解析、スコアリングを確認します。
4. 上記の指示と参考資料に従って、タスク全体を実行します。
5. 結果を検証します。 ベンチマークの公式スプリットと指標を使用し、データ汚染がないか確認して、シードやプロンプト間のばらつきを報告します。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、まだ不確かな点を明記します。

このような場合	対応
スコアが公開値と異なる	結論を出す前に、プロンプト形式、少数例設定、ハーネスのバージョンを確認します。
この手順に記載された関数、フラグ、エンドポイントがインストール済みバージョンにない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認して適宜対応し、ユーザーに伝えます。APIを捏造してはいけません。
必須の入力、識別子、パラメーターが曖昧	ユーザーに質問するか、実行前に前提を明示します。

誠実性に関するルール

• 結果、パラメーター、識別子、引用、統計を捏造してはいけません。実行または検証できない場合は、その旨を明確に伝えます。
• 評価設定をすべて報告します。プロンプト、シード、サブセットを都合よく選んではいけません。
• ここに記載されたバージョン固有の情報は古くなっている可能性があるものとして扱い、インストール済みバージョンの公式ドキュメントで確認します。
• 費用が発生する、共有GPUやクラウドのクォータを消費する、個人データや患者データを扱う、または取り消せない操作を行う前に、必ず許可を求めます。

関連スキル

• nnsight-remote-interpretability：オプションのNDIFリモート実行を含め、nnsightを使用してニューラルネットワークの内部を解釈・操作するためのガイダンスを提供します。
• unsloth：LoRAとQLoRAを使用した、より高速でメモリ使用量の少ないトレーニング用ライブラリUnslothで、大規模言語モデルをファインチューニングするためのガイダンスを提供します。base…
• transformer-lens-interpretability：HookP…を介してTransformerの内部を調査・操作するTransformerLensを使用した、メカニズム解釈研究のためのガイダンスを提供します。