Saelens

Обучение разреженных автоэнкодеров для интерпретации признаков модели.

Метаданные навыка

Справочник: полный SKILL.md

SAELens: разреженные автоэнкодеры для механистической интерпретируемости

SAELens — основная библиотека для обучения и анализа разреженных автоэнкодеров (SAE). Техника разлагает полисемантические активации нейронных сетей на разреженные, интерпретируемые признаки. Основана на фундаментальных исследованиях Anthropic о моносемантичности.

GitHub: jbloomAus/SAELens (1100+ звёзд)

Проблема: полисемантичность и суперпозиция

Отдельные нейроны в нейронных сетях полисемантичны — они активируются в нескольких семантически различных контекстах. Причина в том, что модели используют суперпозицию: кодируют больше признаков, чем у них нейронов, что серьёзно затрудняет интерпретацию.

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” эксперты-оценщики признали 70% признаков SAE подлинно интерпретируемыми. Среди обнаруженных признаков:

• последовательности ДНК, юридический язык, HTTP-запросы
• текст на иврите, информация о питательной ценности, синтаксис кода
• тональность, именованные сущности, грамматические структуры

Рабочий процесс 1: загрузка и анализ предобученных SAE

Пошаговая инструкция

from transformer_lens import HookedTransformer
from sae_lens import SAE

# 1. Загрузить модель и предобученный SAE
model = HookedTransformer.from_pretrained("gpt2-small", device="cuda")
# In sae-lens v6, SAE.from_pretrained() returns JUST the SAE (not a tuple).
sae = SAE.from_pretrained(
    release="gpt2-small-res-jb",
    sae_id="blocks.8.hook_resid_pre",
    device="cuda"
)
# If you also need the cfg dict and feature sparsity, use:
# sae, cfg_dict, sparsity = SAE.from_pretrained_with_cfg_and_sparsity(...)

# 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

Чеклист

• Загрузить модель через TransformerLens
• Загрузить подходящий SAE для целевого слоя
• Закодировать активации в разреженные признаки
• Определить топ-активирующиеся признаки для каждого токена
• Проверить качество восстановления

Рабочий процесс 2: обучение пользовательского SAE

Пошаговая инструкция

from sae_lens import (
    LanguageModelSAETrainingRunner,
    LanguageModelSAERunnerConfig,
    StandardTrainingSAEConfig,
    LoggingConfig,
)

# 1. Настроить обучение (в v6 конфиг ВЛОЖЕННЫЙ: настройки самого SAE живут
#    в под-конфиге `sae=`, настройки логирования — в под-конфиге `logger=`).
#    Обратите внимание: `architecture`, `d_sae`, `l1_coefficient` и прочее теперь
#    в под-конфиге SAE, а плоские устаревшие опции `hook_layer`, `activation_fn`,
#    `log_to_wandb` убраны.
cfg = LanguageModelSAERunnerConfig(
    # Архитектура SAE и разреженность (вложенно)
    sae=StandardTrainingSAEConfig(
        d_in=768,          # Размерность модели
        d_sae=768 * 8,     # Коэффициент расширения 8
        l1_coefficient=8e-5,  # Штраф за разреженность
        apply_b_dec_to_input=True,
        normalize_activations="expected_average_only_in",
    ),

    # Источник данных (модель и точка перехвата)
    model_name="gpt2-small",
    hook_name="blocks.8.hook_resid_pre",  # слой выводится из hook_name (hook_layer больше нет)

    # Обучение
    lr=4e-4,
    l1_warm_up_steps=1000,
    train_batch_size_tokens=4096,
    training_tokens=100_000_000,

    # Данные
    dataset_path="monology/pile-uncopyrighted",
    context_size=128,

    # Логирование (вложенно)
    logger=LoggingConfig(
        log_to_wandb=True,
        wandb_project="sae-training",
    ),

    # Контрольные точки
    checkpoint_path="checkpoints",
    n_checkpoints=5,
)

# 2. Обучить
trainer = LanguageModelSAETrainingRunner(cfg)  # SAETrainingRunner still works as an alias
sae = trainer.run()

# 3. Оценить результат
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 для желаемой разреженности
• Включить L1 warm-up для предотвращения гибели признаков
• Следить за метриками во время обучения (W&B)
• Проверить L0 и восстановление CE loss
• Проконтролировать долю мёртвых признаков

Заметка о переходе на v6: для других типов SAE меняется под-конфиг sae= — GatedTrainingSAEConfig, TopKTrainingSAEConfig (значение k задаётся напрямую) или JumpReLUTrainingSAEConfig (работает с l0_coefficient). Плоские устаревшие опции (architecture, expansion_factor, hook_layer, activation_fn/activation_fn_kwargs, use_ghost_grads, ghost grads, параметры инициализации b_dec и декодера) в v6 удалены.

Рабочий процесс 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(  # v6 returns just the SAE
    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 к residual stream."""
    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]

# Рассчитать атрибуцию логитов по признакам
# Вклад признака = activation × decoder_weight × unembedding
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("Top features for 'Paris' prediction:")
for idx, val in zip(top_features.indices, top_features.values):
    print(f"  Feature {idx.item()}: {val.item():.3f}")

Типичные проблемы и решения

Все примеры ниже используют вложенный конфиг v6: настройки самого SAE идут в под-конфиг sae= (StandardTrainingSAEConfig / TopKTrainingSAEConfig и прочие), а ручки обучения остаются на верхнем уровне LanguageModelSAERunnerConfig.

Проблема: высокая доля мёртвых признаков

from sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig

# НЕПРАВИЛЬНО: без warm-up признаки гибнут рано
cfg = LanguageModelSAERunnerConfig(
    sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4),
    l1_warm_up_steps=0,  # Плохо!
)

# ПРАВИЛЬНО: постепенно увеличивать штраф L1 (ghost grads в v6 убрали — рычаг теперь warm-up)
cfg = LanguageModelSAERunnerConfig(
    sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5),
    l1_warm_up_steps=1000,  # Постепенное увеличение
)

Проблема: плохое восстановление (низкий CE recovery)

# Снизить штраф за разреженность и/или добавить ёмкости (и то и другое — в под-конфиге SAE)
cfg = LanguageModelSAERunnerConfig(
    sae=StandardTrainingSAEConfig(
        d_in=768,
        d_sae=768 * 16,       # Больше ёмкости
        l1_coefficient=5e-5,  # Ниже = лучше восстановление
    ),
)

Проблема: признаки плохо интерпретируются

from sae_lens import LanguageModelSAERunnerConfig, StandardTrainingSAEConfig, TopKTrainingSAEConfig

# Повысить разреженность (больший L1)
cfg = LanguageModelSAERunnerConfig(
    sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=1e-4),
)
# Или взять TopK SAE (в v6 k задаётся напрямую, а не через activation_fn_kwargs)
cfg = LanguageModelSAERunnerConfig(
    sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50),  # Ровно 50 активных признаков
)

Проблема: ошибки памяти при обучении

cfg = LanguageModelSAERunnerConfig(
    sae=StandardTrainingSAEConfig(d_in=768, d_sae=768*8, l1_coefficient=8e-5),
    train_batch_size_tokens=2048,  # Уменьшить размер батча
    store_batch_size_prompts=4,    # Меньше промптов в буфере
    n_batches_in_buffer=8,         # Меньший буфер активаций
)

Интеграция с Neuronpedia

Просматривайте признаки предобученных SAE на neuronpedia.org:

# Признаки индексируются по ID SAE
# Пример: gpt2-small слой 8, признак 1234
# → neuronpedia.org/gpt2-small/8-res-jb/1234

Справочник по классам

Справочная документация

Подробная документация по API, туториалы и продвинутые сценарии использования — в папке references/:

Внешние ресурсы

Туториалы

• Basic Loading & Analysis
• Training a Sparse Autoencoder
• ARENA SAE Curriculum

Статьи

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

Официальная документация

• SAELens Docs
• Neuronpedia — браузер признаков

Архитектуры SAE

from sae_lens import LanguageModelSAERunnerConfig, TopKTrainingSAEConfig

# TopK SAE (ровно 50 активных признаков) — в v6 `k` задаётся в под-конфиге SAE
cfg = LanguageModelSAERunnerConfig(
    sae=TopKTrainingSAEConfig(d_in=768, d_sae=768*8, k=50),
)