gpt2-steering-denoiser

gpt2-steering-denoiser

Два небольших denoiser'а (MLP + alpha-conditioning), обученные корректировать негативный побочный эффект activation steering на GPT-2 small.

Полный код, методология и разбор результатов: github.com/trtbfn/tbank_mech_interp (см. REPORT.md).

Задача

Стиринг h̃ = h + α·v (где v — вектор из SAE-декодера) усиливает нужное свойство в генерации, но при росте α разваливает связность текста. Оба чекпойнта учатся исправлять это, не убирая сам стиринг.

Чекпойнты

• denoiser.pt — коррекция на слое инъекции (слой 6 из 12): h̃ = denoiser(h + α·v, α/300). Стабильнее по всей сетке α ∈ {32, 128, 256}, разброс perplexity относительно plain steering — в пределах ±110.
• denoiser_L8.pt — коррекция ниже по потоку, на слое 8. Обучен на отдельной задаче (парные clean/corrupted активации слоя 8). Даёт лучший наблюдавшийся в проекте результат (weather, α=128: perplexity 583 против 1110 у сырого стиринга, concept-активация выше) при доминировании по обеим осям одновременно, но менее предсказуем на других доменах (разброс шире, есть выбросы).

Ни один из двух не доминирует систематически над другим — выбор зависит от того, что важнее: предсказуемость (denoiser.pt) или пиковый эффект (denoiser_L8.pt). Подробное сравнение — §6 отчёта.

Архитектура

class Denoiser(nn.Module):
    def __init__(self, d_model=768, hidden_mult=4):
        super().__init__()
        hidden = d_model * hidden_mult
        self.net = nn.Sequential(
            nn.Linear(d_model + 1, hidden),  # +1 — conditioning на alpha/ALPHA_MAX_COND
            nn.GELU(),
            nn.Linear(hidden, hidden),
            nn.GELU(),
            nn.Linear(hidden, d_model),
        )

    def forward(self, h, alpha_norm=None):
        ...
        return h + self.net(torch.cat([h, alpha_norm], dim=-1))  # residual-формулировка

Полный код в denoiser.py репозитория.

Использование

import torch
from denoiser import Denoiser, ALPHA_MAX_COND

denoiser = Denoiser()
denoiser.load_state_dict(torch.load("denoiser.pt")["state_dict"])
denoiser.eval()

# h — активация residual stream (768,) на слое 6, alpha — сила стиринга
h_corrected = denoiser(h + alpha * v, alpha / ALPHA_MAX_COND)

Для denoiser_L8.pt — та же архитектура, но применяется хуком на слое 8, при этом стиринг на слое 6 не трогается (см. eval_steering.py, метод denoised_downstream).

Обучение

• Модель: GPT-2 small (transformer_lens)
• Данные: WikiText-2 (~3120 абзацев → ~396K активаций), без пересечения с валидационными доменами
• denoiser.pt: корруптирование h + α·u (u — случайное единичное направление, α ~ U[0,300]), 300 эпох
• denoiser_L8.pt: парные данные (чистая/протёкшая активация слоя 8 при случайном возмущении на слое 6), 300 эпох

Детали и все найденные по пути ошибки (data leakage, несогласованность conditioning, невыполнимая схема зашумления) — в REPORT.md репозитория.