borisggg/steering-denoiser-gpt2

Steering Denoiser — GPT-2 small, layer 6

A small residual-MLP denoiser that repairs the damage activation steering does to the residual stream. Trained for the T-Lab 2026 Mechanistic Interpretability assignment.

Code and full report: https://github.com/bborisggg/steering-denoiser

What it does

Steering h~ = h + alpha*v pushes the residual stream off the manifold of activations the model actually produces, and fluency collapses. This module maps a steered activation back towards that manifold:

h~ = D(h + alpha * v)

D(x) = x - f(x), trained with L = ||h - D(corrupt(h))||^2 on 758,038 residual-stream activations from GPT-2 small at layer 6.

By Tweedie's identity an MSE-optimal denoiser satisfies D(x) = x + sigma^2 * grad log p(x), so this is a one-step score model over activations — a cheap stand-in for the flow-matching prior of GLP (arXiv 2602.06964).

Training

Base model	gpt2 (small), layer 6 (resid_post)
Corruption	C4 — {"corruption": "C4", "alpha_max": 3.0, "n_train_features": 256}
Objective	MSE to the clean activation
Activations	758,038 tokens, centered, attention sink excluded
Architecture	2 pre-LN residual MLP blocks, hidden 4x, 11.7M params
Steps / batch	4000 / 4096
Held-out MSE	0.024 (0.0008 of leaving the corruption alone)

That held-out MSE is measured on the training corruption distribution and is flattered by it: C4 draws from a pool of 256 directions, so the denoising task itself is easy. The number that matters is the downstream one below, measured on steering vectors this model never saw.

The evaluation steering vectors were frozen before training and never entered it (holdout fingerprint 25189650a5ecd352).

Results

Evaluated on 6 held-out SAE features, 32 prompts, GPT-2 small. Unsteered reference: perplexity 121.5.

	alpha	perplexity	concept fire rate
Plain steering h + alpha*v	1.0	193	0.161
This denoiser	0.75	151	0.176

Better on both axes: +10% concept expression at 22% lower perplexity.

Training corruption matters more than architecture. The same denoiser trained on the isotropic Gaussian noise the task originally proposed does not beat plain steering at any perplexity budget; only the structured rank-1 corruption used here does. Full ablation, the negative results, and the analysis are in the repository.

Two things that will silently break this

1. Activations must be centered along d_model. The GPT-2 SAEs this was built against were trained on TransformerLens activations with center_writing_weights=True, and the denoiser follows that convention. Feeding raw HuggingFace activations puts you in the wrong space. Centering is output-neutral — LayerNorm removes the mean anyway.

2. Pass the conditioning level matched to alpha. The model is conditioned on a corruption level t. It defaults to t=1 ("maximally corrupted"), which applies maximum denoising at every strength and wrecks lightly-steered activations. Use t_for_alpha.

Usage

import torch
from huggingface_hub import hf_hub_download

payload = torch.load(hf_hub_download("borisggg/steering-denoiser-gpt2", "denoiser.pt"), weights_only=False)
# Rebuild with steering_denoiser.denoisers.load_denoiser from the repo above,
# or read payload["config"] / payload["state_dict"] directly.

h = h - h.mean(-1, keepdim=True)          # centered space
steered = h + alpha * scale * v           # scale = median ||h|| at the layer
fixed = denoiser(steered, denoiser.t_for_alpha(alpha))

Limitations

Trained and evaluated only on GPT-2 small at layer 6, with SAE-derived steering vectors on short open-ended prompts. Transfer to other models, layers or vector families is untested here. See the repository's RESEARCH_LOG.md for negative results and open questions.