borisggg/steering-denoiser-gemma2_2b

Steering Denoiser — gemma-2-2b-it, layer 12

A residual-MLP denoiser over Gemma-2-2B residual-stream activations, trained for the T-Lab 2026 Mechanistic Interpretability assignment.

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

Read this first: on this model the denoiser does not help, and at higher steering strengths it hurts. It is published because that negative result is the point, and because reproducing it should not require re-running the training.

What it is

D(x) = x - f(x), trained with L = ||h - D(corrupt(h))||^2 on 958,722 residual-stream activations at layer 12. 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.

Base model	google/gemma-2-2b-it, layer 12 (resid_post)
Corruption	C3 — {"corruption": "C3", "alpha_max": 3.0}
Activations	958,722 tokens, not centered (RMSNorm — see below), top 5% by norm excluded
Architecture	2 pre-LN residual MLP blocks, hidden 4x, 95.1M params
Steps / batch	6000 / 4096
Held-out MSE	3.523 (0.0988 of leaving the corruption alone)

Results: a predicted null, and worse than null

Evaluated with persona vectors (difference-in-means trait directions) on gemma-2-2b-it, trait "evil", 10 held-out questions x 4 generations, both axes scored 0-100 by an LLM judge:

method	alpha=0.5 trait / coherence	alpha=1.0 trait / coherence
Plain steering h + alpha*v	14.3 / 56.9	13.0 / 21.8
Norm-matched	8.5 / 53.1	12.6 / 22.4
This denoiser, D(h + alpha*v)	10.5 / 51.6	1.5 / 12.3
Concept-preserving	11.9 / 35.9	0.0 / 1.3

Plain steering dominates both denoiser applications on both axes at every alpha, and the denoised variants collapse faster past alpha=1.

This was predicted before training. A 30-minute diagnostic that trains nothing measures how the model degrades under steering: this one becomes improbable (perplexity 25x baseline at alpha=1) while repetition stays at exactly 0.000 and dist-2 stays above the unsteered baseline. The same denoiser idea works on GPT-2, which degrades by repeating. Measure the failure mode before training anything.

Under-training is excluded as an explanation: held-out relative MSE 0.0988 against GPT-2's own C3 denoiser at 0.1016, i.e. the same relative quality on the same corruption family.

Why it fails, measured at the geometry of the correction

Decomposing D(h~) - h~ into norm change and direction change, probed on a direction the denoiser never trained on, gives three regimes -- and the regime is set by the training corruption, not the architecture. C3 on GPT-2 and C3 on Gemma trace numerically the same curves across different architectures, 3x different width, and different normalisation:

| denoiser | delta@a=1 | ||D||/||h||@a=3 | cos@a=3 | regime | outcome | |-------------|-----------|-----------------|---------|-------------------|----------| | GPT-2 C4 | 0.232 | 0.934 | 0.979 | bounded rotation | works | | Gemma C4 | 0.025 | 0.994 | 0.974 | identity collapse | null | | this one | 0.608 | 0.249 | 0.441 | mean regression | harmful |

C1/C3 denoisers strip 75% of the activation's norm at alpha=3: an MSE-optimal response to a far off-manifold input approaches the global mean, so the module deletes the steered activation instead of repairing it. That is why the denoised methods collapse faster than plain steering past alpha=1.

The natural suspicion -- that RMSNorm renormalises the correction away -- was tested and refuted. The correction's survival through the next block's norm is 0.95-1.14 here and 1.03-1.12 on GPT-2's LayerNorm: it passes through intact and is amplified at high alpha.

Practical consequence: a low training loss is not evidence of a usable denoiser. Gemma's C4 sibling reached relative MSE 0.00049 and is the identity on unseen directions, having memorised its 256-direction pool -- which spans 1/9 of this model's 2304-dim space against 1/3 of GPT-2's 768-dim one. Check what the module does to a direction it has not seen.

Two things that will silently break this

1. Do NOT center the activations. Gemma-2 uses RMSNorm, which divides by the RMS but never subtracts the mean, so the mean component is observable by the model and removing it changes the output (measured: 3.5e-02 relative logit change, ~14000x GPT-2's 2.5e-06). The GPT-2 denoiser in the sibling repo requires the opposite. Centering is a property of the normalisation layer, not a preference.

2. Pass the conditioning level matched to alpha. The model is conditioned on a corruption level t and 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-gemma2_2b", "denoiser.pt"), weights_only=False)
# Rebuild with steering_denoiser.denoisers.load_denoiser from the repo above.

steered = h + alpha * scale * v           # raw space, NOT centered; scale = median ||h||
fixed = denoiser(steered, denoiser.t_for_alpha(alpha))

Limitations

One trait, one layer, one corruption family. The layer was fixed at mid-depth by the assignment, and a later sweep found layer 8 dominates layer 12 on both judge axes for this model — so this checkpoint was trained and evaluated at a layer that is not the best one available. The denoiser is also trained on corpus-text activations but applied to dialogue activations. See RESEARCH_LOG.md for the full set of negative results.