borisggg/steering-denoiser-gemma2_2b

steering-mvp denoiser — Gemma-2-2B-it, layer 12

A residual-MLP denoiser trained to repair the fluency damage activation steering does to Gemma-2-2B-it's residual stream — the external-validation counterpart to the GPT-2 checkpoint in this same lineage. This checkpoint is D4: trained on rank-1 corruptions h + rho*||h||*u, u resampled fresh every training step from Gemma Scope's pt-res SAE decoder pool (16,384 directions; no it-tuned Gemma Scope release exists, so a pt-trained SAE is applied to it-model activations — gated on measured reconstruction quality, see below).

Code, plan, and full experiment log: https://github.com/bborisggg/steering-mvp

Read this before trusting the numbers below: on this model, the result the GPT-2 checkpoint found does not replicate. D4 does not clearly beat plain additive steering here, and neither beats simply prompting the model with the trait instruction. Published because that's the honest external-validation outcome, not because this checkpoint is a recommended tool.

What it does

D(x, t) = x + scale * f(encode(x, scale), e(t)), a 2-block residual MLP with FiLM conditioning on the corruption strength t, output head zero-initialized so an untrained denoiser is exactly the identity. Applied as h_denoised = D(h + alpha*v, t(r)), where v here is a persona-vector direction (difference-in-means over eliciting vs. suppressing prompts), not an SAE feature.

Training

Base model	google/gemma-2-2b-it, layer 12 (resid_post)
Corruption	D4 — rank-1, rho in [0.05, 3.0] log-uniform, pool = 16,384 directions (Gemma Scope pt-res, layer_12/width_16k/average_l0_82)
SAE reconstruction gate	EV = 0.79 on natural corpus text through the it-model (0.6 gate) — chat-formatted persona prompts measured EV −5 to −8 and were correctly rejected as the wrong probe distribution before this pool was trusted
Objective	L = |D(x,t) - h|^2 / (|h|^2 + eps), mean over batch
Activations	12,393 cached Gemma residual-stream tokens, not centered (RMSNorm, not LayerNorm — centering here measurably perturbs the model)
Architecture	2-block residual MLP, hidden 2x, FiLM time embedding dim 128, 43.68M params
Steps / batch / lr	5000 / 256 / 0.001
Seed	0 (single seed — this is validation of a frozen GPT-2 recipe, not a second development loop)
Directions	6 persona traits (evil, sycophantic, impolite, apathetic, optimistic, humorous), extracted at this same layer, held-out separation Cohen's d 2.8–8.7

Results (judge-scored, median of 6 traits, neutral prompts — steering is the only source of the trait, never folded into the prompt for B0/B1/D1/D4)

method	r	judge coherence	judge concept
B0	0.00	100.0	0.8
B1	0.05	100.0	2.5
B1	0.10	98.3	6.7
B1	0.15	97.5	10.8
B1	0.20	93.3	18.8
B1	0.40	61.7	74.2
B1	0.60	49.2	75.8
B1	0.80	39.2	84.6
B1	1.00	20.8	67.5
D1	0.05	98.3	0.8
D1	0.10	99.2	4.2
D1	0.15	94.2	13.3
D1	0.20	93.3	21.7
D1	0.40	55.0	55.0
D1	0.60	33.3	56.7
D1	0.80	17.5	46.7
D1	1.00	7.1	35.8
D4	0.05	99.6	3.3
D4	0.10	97.9	5.0
D4	0.15	96.7	14.2
D4	0.20	92.5	20.0
D4	0.40	55.0	61.7
D4	0.60	30.8	78.3
D4	0.80	17.5	50.0
D4	1.00	9.2	39.2

D1/D4 trail B1 on coherence at every r tested and only edge ahead on concept at one point (r=0.6), at markedly worse coherence there — not a Pareto win. Prompting the model with the trait instruction directly reaches concept 77.5 at coherence 98.3 — better than any steering arm's own best concept score, at coherence none of them come close to.

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.ResidualMLPDenoiser from the repo above, or:
#   cfg = payload["config"]; model = ResidualMLPDenoiser(**cfg); model.load_state_dict(payload["state_dict"])

# h is RAW, not centered (RMSNorm) -- opposite of the GPT-2 checkpoint in this same lineage.
t = corruptions.t_for_r(payload["extra"]["corruption"], r)
denoised = model(h_steered, t=t)

What will silently break this

1. Do NOT center the activations — Gemma-2 uses RMSNorm, which never subtracts a mean, so removing one perturbs the model (this checkpoint's config["center"] is False).
2. The conditioning level t must be derived from the deployed r via the corruption family's own stored inverse map, never passed raw or fixed.

Limitations

One layer (12 of 26), one seed, 6 traits, a pt-trained SAE applied to it-model activations for the corruption pool. See the project repo's PLAN.md / REPORT.md for the full state, including where this result diverges from the GPT-2 checkpoint's.