Sparse Autoencoders (SAE): вскрываем чёрный ящик LLM, чтобы понять, как они на самом деле думают

Ты смотришь на ответ GPT-4 и думаешь: "Откуда это взялось?" Промпт был простой, а на выходе — сложная логическая цепочка. Мы все привыкли, что большие языковые модели — это чёрные ящики. Пока не появились Sparse Autoencoders.

SAE — не просто очередной метод интерпретации. Это хирургический инструмент, который позволяет заглянуть внутрь скрытых слоёв трансформера и увидеть, какие "нейроны" активируются на конкретные концепции. Anthropic использовала их для анализа Claude 3. Мы разберёмся, как это работает на практике.

Почему обычные autoencoders не работают для LLM?

Стандартный autoencoder пытается сжать данные в скрытое представление и восстановить их. Для изображений — работает. Для интерпретации LLM — полный провал. Вот почему:

• Слишком плотные активации: В скрытых слоях LLM одновременно активно ~5-10% нейронов. Autoencoder без регуляризации активирует почти все скрытые нейроны — получаем шум, а не интерпретируемые признаки.
• Нет разреженности: Мы хотим, чтобы каждая скрытая фича в автоэнкодере соответствовала одной чёткой концепции ("математика", "Python-синтаксис", "ирония"). Без sparse penalty получаем суперпозицию всего и сразу.
• Переобучение на восстановление: Модель учится идеально реконструировать активации, но её скрытое пространство не имеет смысловой структуры.

Магия разреженности: как L1-регуляризация меняет всё

Sparse Autoencoder добавляет один ключевой ингредиент: штраф за активность скрытых нейронов. Математически функция потерь выглядит так:

loss = reconstruction_loss + λ * sparsity_loss
# где:
# reconstruction_loss = MSE(original_activations, reconstructed_activations)
# sparsity_loss = sum(abs(hidden_activations))  # L1 регуляризация
# λ — коэффициент разреженности, обычно 0.01-0.001

L1-регуляризация заставляет большинство скрытых нейронов быть близкими к нулю. Активируются только те, которые действительно необходимы для реконструкции. В результате каждый нейрон вынужден стать "специалистом" — кодировать одну конкретную семантическую или синтаксическую особенность.

Архитектурные хитрости: что скрыто между энкодером и декодером

Базовая архитектура SAE проста до безобразия. Но дьявол в деталях — именно они определяют, будет ли модель вообще обучаться.

Ключевые компоненты:

Самая частая ошибка новичков — забыть про bias в декодере. Без него модель вынуждена использовать веса декодера для кодирования среднего значения входных активаций. Это "крадёт" ёмкость у полезных признаков.

# НЕПРАВИЛЬНО — нет bias в декодере
class BadSAE(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.encoder = nn.Linear(input_dim, hidden_dim)
        self.decoder = nn.Linear(hidden_dim, input_dim)  # НЕТ bias!
    
    def forward(self, x):
        hidden = F.relu(self.encoder(x))
        return self.decoder(hidden), hidden

Правильная реализация с учётом нюансов:

import torch
import torch.nn as nn
import torch.nn.functional as F

class SparseAutoencoder(nn.Module):
    """
    Sparse Autoencoder для интерпретации скрытых активаций LLM.
    Ключевые особенности:
    1. Bias в декодере для учёта среднего значения
    2. Нормировка весов декодера для стабильности обучения
    3. ReLU для обеспечения разреженности
    """
    
    def __init__(self, input_dim, hidden_dim, l1_coef=0.001):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.l1_coef = l1_coef
        
        # Энкодер
        self.encoder = nn.Linear(input_dim, hidden_dim)
        # Инициализируем веса маленькими значениями
        nn.init.kaiming_normal_(self.encoder.weight, nonlinearity='relu')
        
        # Декодер С bias!
        self.decoder = nn.Linear(hidden_dim, input_dim)
        # Инициализация декодера
        nn.init.kaiming_normal_(self.decoder.weight, nonlinearity='linear')
        
        # Нормируем веса декодера (важно для стабильности)
        with torch.no_grad():
            self.decoder.weight.data = F.normalize(self.decoder.weight.data, dim=1)
    
    def forward(self, x):
        # Прямой проход
        hidden_pre = self.encoder(x)  # [batch, hidden_dim]
        hidden = F.relu(hidden_pre)   # ReLU обеспечивает разреженность
        
        # Восстановление
        reconstructed = self.decoder(hidden)
        
        # Вычисляем потери
        mse_loss = F.mse_loss(reconstructed, x, reduction='mean')
        l1_loss = hidden.abs().mean()  # L1 регуляризация
        
        total_loss = mse_loss + self.l1_coef * l1_loss
        
        return {
            'reconstructed': reconstructed,
            'hidden': hidden,
            'loss': total_loss,
            'mse_loss': mse_loss,
            'l1_loss': l1_loss
        }
    
    def normalize_decoder_weights(self):
        """Нормировка весов декодера после обновления"""
        with torch.no_grad():
            norm = self.decoder.weight.norm(dim=1, keepdim=True)
            self.decoder.weight.data = self.decoder.weight.data / norm

Собираем данные: как добыть скрытые активации из LLM

SAE не обучается на текстах. Он обучается на активациях промежуточных слоёв уже обученной LLM. Это как учить переводчика не на языках, а на мозговых волнах полиглота.

1 Выбираем слой для анализа

Не все слои одинаково полезны. Ранние слои (ближе к эмбеддингам) кодируют низкоуровневые синтаксические паттерны. Поздние слои (ближе к выходу) — семантику и логику. Средние слои — часто самый интересный компромисс.

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "gpt2"  # Начинаем с малого
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Регистрируем хук для захвата активаций
activations = []

def hook_fn(module, input, output):
    # Захватываем активации после feed-forward слоя
    # output[0] имеет размерность [batch, seq_len, hidden_size]
    activations.append(output[0].detach().cpu())

# Выбираем слой (например, 6-й из 12)
layer_idx = 6
target_layer = model.transformer.h[layer_idx].mlp  # FFN слой
handle = target_layer.register_forward_hook(hook_fn)

2 Генерируем разнообразные тексты

Чем разнообразнее данные, тем более универсальные фичи извлечёт SAE. Нужны не просто Википедия, а код, диалоги, математические рассуждения.

texts = [
    "def fibonacci(n):",
    "The capital of France is Paris.",
    "I think therefore I exist.",
    "import numpy as np",
    "The quantum state collapses when measured.",
    "SELECT * FROM users WHERE active = TRUE;",
    "The movie was surprisingly good despite the low budget.",
    "2 + 2 * 2 = 6, not 8.",
    "Once upon a time in a galaxy far, far away...",
    "The gradient descent algorithm minimizes the loss function."
]

# Токенизируем и пропускаем через модель
all_activations = []

for text in texts:
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
    with torch.no_grad():
        outputs = model(**inputs)
    
    # activations теперь содержит активации для этого текста
    # Берем активации из середины последовательности (там обычно самая интересная активность)
    seq_len = activations[-1].shape[1]
    mid_point = seq_len // 2
    hidden_states = activations[-1][:, mid_point, :]  # [batch, hidden_size]
    
    all_activations.append(hidden_states)

# Собираем в один тензор
dataset_tensor = torch.cat(all_activations, dim=0)  # [num_samples, hidden_size]

# Не забываем удалить хук
handle.remove()

Тренируем SAE: настройки, которые не описаны в статьях

Официальные статьи Anthropic умалчивают о сотне мелких деталей, которые определяют успех обучения. Вот что я узнал на собственном горьком опыте.

Оптимизатор и расписание скорости обучения:

# ИНИЦИАЛИЗАЦИЯ
input_dim = dataset_tensor.shape[1]  # Например, 768 для GPT-2 small
hidden_dim = 8192  # В 10 раз больше входной размерности — типичное соотношение

sae = SparseAutoencoder(input_dim, hidden_dim, l1_coef=0.0005)
optimizer = torch.optim.AdamW(sae.parameters(), lr=3e-4, weight_decay=1e-5)

# СОЗДАЁМ DATALOADER
from torch.utils.data import DataLoader, TensorDataset

dataset = TensorDataset(dataset_tensor)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# ЦИКЛ ОБУЧЕНИЯ
num_epochs = 1000
losses = []

for epoch in range(num_epochs):
    epoch_loss = 0
    
    for batch in dataloader:
        x = batch[0]
        
        optimizer.zero_grad()
        outputs = sae(x)
        loss = outputs['loss']
        
        loss.backward()
        optimizer.step()
        
        # КРИТИЧЕСКИ ВАЖНО: нормируем веса декодера после каждого шага
        sae.normalize_decoder_weights()
        
        epoch_loss += loss.item()
    
    avg_loss = epoch_loss / len(dataloader)
    losses.append(avg_loss)
    
    # Выводим статистику каждые 50 эпох
    if epoch % 50 == 0:
        print(f"Epoch {epoch}: Loss = {avg_loss:.6f}, "
              f"MSE = {outputs['mse_loss'].item():.6f}, "
              f"L1 = {outputs['l1_loss'].item():.6f}")
        
        # Проверяем мёртвые нейроны
        active_neurons = (outputs['hidden'].abs() > 0.01).float().mean(dim=0)
        dead_count = (active_neurons == 0).sum().item()
        print(f"Мёртвых нейронов: {dead_count}/{hidden_dim} ({dead_count/hidden_dim*100:.1f}%)")

Диагностика проблем:

Интерпретация результатов: что на самом деле нашли нейроны

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

1 Анализ весов декодера

Каждый нейрон в скрытом слое SAE связан с вектором весов в декодере. Этот вектор показывает, какие оригинальные нейроны LLM активирует данный признак.

# Берём веса декодера для конкретного нейрона
neuron_idx = 42  # Интересный нейрон
neuron_weights = sae.decoder.weight[neuron_idx]  # [input_dim]

# Находим оригинальные нейроны LLM с наибольшими весами
top_k = 10
top_indices = torch.topk(neuron_weights.abs(), k=top_k).indices
top_values = neuron_weights[top_indices]

print(f"Нейрон {neuron_idx} сильно связан с нейронами LLM:")
for idx, val in zip(top_indices, top_values):
    print(f"  Нейрон LLM #{idx.item()}: вес = {val.item():.4f}")

2 Активация на конкретных текстах

Самый интересный анализ — смотреть, на каких текстах активируется конкретный нейрон SAE.

def find_triggers_for_neuron(sae, model, tokenizer, neuron_idx, test_texts, threshold=0.5):
    """Находим тексты, которые максимально активируют нейрон"""
    triggers = []
    
    for text in test_texts:
        # Получаем активации LLM
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
        with torch.no_grad():
            outputs = model(**inputs, output_hidden_states=True)
        
        # Берём активации из целевого слоя
        hidden_states = outputs.hidden_states[layer_idx]  # [1, seq_len, hidden_size]
        
        # Пропускаем через SAE
        batch_size, seq_len, hidden_size = hidden_states.shape
        flattened = hidden_states.view(-1, hidden_size)
        sae_outputs = sae(flattened)
        
        # Активации нейрона
        neuron_activations = sae_outputs['hidden'][:, neuron_idx]  # [batch*seq_len]
        max_activation = neuron_activations.max().item()
        
        if max_activation > threshold:
            triggers.append((text, max_activation))
    
    # Сортируем по убыванию активации
    triggers.sort(key=lambda x: x[1], reverse=True)
    return triggers[:5]  # Топ-5 триггеров

# Пример использования
test_texts = [
    "import torch.nn as nn",
    "The cat sat on the mat.",
    "def calculate_loss(predictions, targets):",
    "Paris is the capital of France.",
    "The quick brown fox jumps over the lazy dog."
]

top_triggers = find_triggers_for_neuron(sae, model, tokenizer, neuron_idx=42, 
                                         test_texts=test_texts, threshold=0.3)

print(f"Тексты, активирующие нейрон 42:")
for text, activation in top_triggers:
    print(f"  Активация {activation:.3f}: {text[:50]}...")

Практическое применение: где SAE реально полезны

Интерпретация — это красиво, но что с этим делать? Вот несколько практических сценариев:

• Поиск проблемных паттернов: Если нейрон активируется на токсичные тексты — нашли потенциальный источник bias в модели. Можно использовать эту информацию для дополнительной фильтрации.
• Контроль генерации: Зная, какие нейроны отвечают за креативность vs. фактологичность, можно влиять на выход модели, подавляя или усиливая конкретные признаки.
• Сжатие представлений: Разреженные активации SAE можно использовать как более интерпретируемую альтернативу оригинальным эмбеддингам в RAG-системах.
• Диагностика fine-tuning: Сравнивая SAE до и после дообучения, можно увидеть, какие именно концепции изменились в модели.

Что дальше? Эволюция SAE и новые вызовы

Базовые Sparse Autoencoders — только начало. Сообщество уже движется в нескольких направлениях:

1. Top-k SAE: Вместо L1-регуляризации используют жёсткое ограничение — только k нейронов могут быть активны одновременно. Результаты более интерпретируемые, но обучение сложнее.
2. Динамическая разреженность: Коэффициент λ адаптируется в процессе обучения, чтобы поддерживать заданный уровень разреженности.
3. Иерархические SAE: Несколько уровней разреженных представлений для захвата концепций разного уровня абстракции.
4. SAE для attention: Анализ не только feed-forward слоёв, но и паттернов внимания — более сложная задача, но потенциально более информативная.

Самая интересная возможность — использование SAE не только для анализа, но и для контроля. Представь: ты находишь нейрон, отвечающий за "генерирование фактов vs. выдумки". Можешь искусственно его подавить, когда нужна креативность, или усилить, когда важна точность. Это следующий шаг после обычного промпт-инжиниринга — прямое вмешательство в "мыслительный процесс" модели.

Начни с GPT-2 и 8192 скрытых нейронов. Поэкспериментируй с разными типами текстов. Посмотри, какие нейроны активируются на код, а какие — на поэзию. Первый раз, когда увидишь чёткую корреляцию между нейроном и семантической концепцией, поймёшь — чёрный ящик начал открываться.