SAELens: 機械的解釈性のためのスパースオートエンコーダ

SAELensは、スパースオートエンコーダ(SAE)の訓練と分析を行うためのメインライブラリです。これは多義性を持つニューラルネットワーク活性化を、スパースで解釈可能な特徴に分解する技術です。Anthropicの単義性に関する画期的な研究に基づいています。

GitHub: jbloomAus/SAELens (1,100以上のスター)

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

ニューラルネットワークの個々のニューロンは多義的であり、複数の意味的に異なるコンテキストで活性化します。これはモデルが、ニューロン数より多くの特徴を表現するために重ね合わせを使用するために発生し、解釈性を困難にします。

SAEがこれを解決するのは、密集した活性化をスパースで単義的な特徴に分解することによってです。通常、任意の入力に対してのみ少数の特徴が活性化し、各特徴は解釈可能な概念に対応しています。

SAELensを使用するタイミング

SAELensが必要な場合:

• モデル活性化の解釈可能な特徴を発見したい
• モデルが学習した概念を理解したい
• 重ね合わせと特徴の幾何学を研究したい
• 特徴ベースのステアリングまたはアブレーションを実行したい
• 安全性関連の特徴(欺瞞、バイアス、有害なコンテンツ)を分析したい

代替手段を検討する場合:

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

インストール

pip install sae-lens

要件: Python 3.10以上、transformer-lens>=2.0.0

コア概念

SAEが学習するもの

SAEは、スパースなボトルネックを通じてモデル活性化を再構成するよう訓練されます:

入力活性化 → エンコーダ → スパース特徴 → デコーダ → 再構成活性化
  (d_model)      ↓     (d_sae >> d_model)    ↓      (d_model)
              スパース性                   再構成
              ペナルティ                    損失

損失関数: 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. モデルと事前訓練済み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. モデル活性化を取得
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. SAE特徴にエンコード
sae_features = sae.encode(activations)  # [batch, pos, d_sae]
print(f"Active features: {(sae_features > 0).sum()}")

# 4. 各位置の上位特徴を検索
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. 活性化を再構成
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. 訓練を構成
cfg = LanguageModelSAERunnerConfig(
    # モデル
    model_name="gpt2-small",
    hook_name="blocks.8.hook_resid_pre",
    hook_layer=8,
    d_in=768,  # モデルの次元

    # SAEアーキテクチャ
    architecture="standard",  # または "gated", "topk"
    d_sae=768 * 8,  # 拡張係数8
    activation_fn="relu",

    # 訓練
    lr=4e-4,
    l1_coefficient=8e-5,  # スパース性ペナルティ
    l1_warm_up_steps=1000,
    train_batch_size_tokens=4096,
    training_tokens=100_000_000,

    # データ
    dataset_path="monology/pile-uncopyrighted",
    context_size=128,

    # ロギング
    log_to_wandb=True,
    wandb_project="sae-training",

    # チェックポイント
    checkpoint_path="checkpoints",
    n_checkpoints=5,
)

# 2. 訓練
trainer = SAETrainingRunner(cfg)
sae = trainer.run()

# 3. 評価
print(f"L0 (平均活性化特徴数): {trainer.metrics['l0']}")
print(f"CE損失復旧: {trainer.metrics['ce_loss_score']}")

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

パラメータ	一般的な値	効果
d_sae	4-16× d_model	より多くの特徴、より高い容量
l1_coefficient	5e-5〜1e-4	高い = よりスパース、精度が低い
lr	1e-4〜1e-3	標準的なオプティマイザのLR
l1_warm_up_steps	500-2000	早期の特徴死滅を防ぐ

評価メトリクス

メトリクス	ターゲット	意味
L0	50-200	トークンごとの平均活性化特徴数
CE損失スコア	80-95%	元のモデルと比較して復旧した交差エントロピー
デッド特徴	<5%	活性化しない特徴
説明分散	>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"
)

# 特定の特徴を活性化させるものを検索
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):
    """SAE特徴の方向を残差ストリームに追加します。"""
    tokens = model.to_tokens(prompt)

    # デコーダから特徴方向を取得
    feature_direction = sae.W_dec[feature_idx]  # [d_model]

    def steering_hook(activation, hook):
        # すべての位置に拡大された特徴方向を追加
        activation += strength * feature_direction
        return activation

    # ステアリング付きで生成
    output = model.generate(
        tokens,
        max_new_tokens=50,
        fwd_hooks=[("blocks.8.hook_resid_pre", steering_hook)]
    )
    return model.to_string(output[0])

特徴帰属

# どの特徴が特定の出力に最も影響するか?
tokens = model.to_tokens("The capital of France is")
_, cache = model.run_with_cache(tokens)

# 最終位置の特徴を取得
features = sae.encode(cache["resid_pre", 8])[0, -1]  # [d_sae]

# 特徴ごとのロジット帰属を取得
# 特徴の寄与 = 特徴活性化 × デコーダ重み × 非埋め込み
W_dec = sae.W_dec  # [d_sae, d_model]
W_U = model.W_U    # [d_model, vocab]

# "Paris"ロジットへの寄与
paris_token = model.to_single_token(" Paris")
feature_contributions = features * (W_dec @ W_U[:, paris_token])

top_features = feature_contributions.topk(10)
print("'Paris'予測の上位特徴:")
for idx, val in zip(top_features.indices, top_features.values):
    print(f"  特徴 {idx.item()}: {val.item():.3f}")

一般的な問題と解決策

問題: デッド特徴比が高い

# 間違い: ウォームアップなし、特徴が早期に死滅
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=1e-4,
    l1_warm_up_steps=0,  # 悪い!
)

# 正しい: L1ペナルティのウォームアップ
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=8e-5,
    l1_warm_up_steps=1000,  # 段階的に増加
    use_ghost_grads=True,   # デッド特徴を復活させる
)

問題: 再構成が低品質(CE復旧が低い)

# スパース性ペナルティを削減
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=5e-5,  # 低い = より良い再構成
    d_sae=768 * 16,       # より多くの容量
)

問題: 特徴が解釈可能でない

# スパース性を増加させる(より高いL1)
cfg = LanguageModelSAERunnerConfig(
    l1_coefficient=1e-4,  # 高い = よりスパース、より解釈可能
)
# またはTopKアーキテクチャを使用
cfg = LanguageModelSAERunnerConfig(
    architecture="topk",
    activation_fn_kwargs={"k": 50},  # 正確に50個の活性化特徴
)

問題: 訓練中のメモリエラー

cfg = LanguageModelSAERunnerConfig(
    train_batch_size_tokens=2048,  # バッチサイズを削減
    store_batch_size_prompts=4,    # バッファ内のプロンプトを削減
    n_batches_in_buffer=8,         # より小さい活性化バッファ
)

Neuronpediaとの統合

neuronpedia.orgで事前訓練済みSAE特徴を参照:

# 特徴はSAE IDでインデックス化されます
# 例: gpt2-small レイヤー8 特徴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カリキュラム

論文

• Towards Monosemanticity - Anthropic (2023)
• Scaling Monosemanticity - Anthropic (2024)
• Sparse Autoencoders Find Highly Interpretable Features - Cunningham et al. (ICLR 2024)

公式ドキュメント

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

SAEアーキテクチャ

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

# TopK SAE(正確に50個の特徴が活性化)
cfg = LanguageModelSAERunnerConfig(
    architecture="topk",
    activation_fn="topk",
    activation_fn_kwargs={"k": 50},
)