refusal-steering-using-saes

Accepted at Nemi Conf 2026 : https://nemiconf.github.io/summer26/accepted_work.html

Recent work reports that one difference-in-means direction matches Sparse Autoencoders at refusal steering, and costs far less to run. That verdict comes from aggregate refusal rate, which does not say what the intervention breaks on the way, or where it was applied. This repo re-tests it on Gemma-3-4B-IT with a Gemma Scope 2 16k residual SAE at layer 17. It runs eight experiments and ends with a rule that turns the measurements into one recommendation.

Install and run

pip install -e .
hf auth login                       # Gemma is gated, accept the licence first
refusal-steer all --output-dir artifacts

Or open notebooks/colab_run2.ipynb and run the cells in order. It does the same thing with a table printed after each stage. notebooks/colab.ipynb is the older one-shot notebook, kept because the files in artifacts/ came from it.

Single experiments and the two CPU-only steps:

refusal-steer sweep                             # one experiment
refusal-steer plots    --output-dir artifacts   # redraw figures from saved JSON
refusal-steer decision --output-dir artifacts   # redo the verdict from saved JSON

Common flags: --model, --layer, --sae-release, --sae-width, --sae-sparsity, --num-fit, --num-eval, --coeffs 0.5 1 2 4, --ablation-coeff, --seeds 0 1 2, --batch-size, --max-new-tokens, --judge, --no-capability, --min-comprehension.

What gets compared

Five steering arms, all driven by the same coefficient so the strengths line up:

arm	what it does
diff_in_means	subtract the refusal direction from the residual stream
ace	the same direction, with the benign mean added back (affine correction)
sae_clamp_top1	clamp the top refusal feature, keep the SAE reconstruction
sae_clamp_top5	clamp the top five features
sae_residual_preserving	clamp the top feature, then add the SAE's reconstruction error back

The SAE clamp target is written in units of the feature's own mean activation on harmful prompts, so coeff means about the same strength on both sides. Clamping to zero is just coeff = 0.

The refusal direction and the feature ranking are fit on one split (num_fit, 64 prompts by default) and scored on a disjoint one (num_eval, 32). Prompts are JailbreakBench JBB-Behaviors for the harmful side and XSTest for the benign side, with the XSTest contrast_* types dropped because they are actually unsafe.

The eight experiments

Each one answers a single question.

	name	question
1	sweep	Across the whole coefficient grid, does any SAE arm beat a linear one at the same collateral cost, and does affine correction close the gap?
2	residual_leakage	Does the refusal signal survive in what the SAE fails to reconstruct, so the clamp works by throwing that error away instead of editing a feature?
3	cross_category	Can one category's feature switch refusal off for that category alone, which a single global direction cannot do?
4	positions	Every generated token, the first token only, or only the last prompt token?
5	adversarial	Which representation still spots harm under Base64, ROT13 and roleplay rewrites?
6	benchmark	What does each method really cost per generated token, with the token count held fixed?
7	capability	Does the steered model still answer MMLU and keep its perplexity, which is damage the degeneracy check cannot see?
8	seeds	Does the SAE's lead keep its sign when the split, the direction and the feature ranking are all refit?

How the run is ordered

refusal-steer all loads the model once and runs all eight in order. The sweep goes first on purpose. The three one-coefficient experiments (2, 3 and 4) then run at a coefficient picked from the sweep: the strongest one where every SAE arm still stays inside the collateral budget. decision.suggest_ablation_coeff does the picking and experiments._adopt_ablation_coeff applies it.

This matters. In the first run those three were pinned to max(coeffs), where every SAE arm is fully degenerate, so three of the decision signals were reading noise off broken arms.

If an experiment raises, the error is recorded and the rest keep going. A rented GPU is a bad place to lose seven finished experiments to the eighth one's bug. The decision rule then runs on whatever finished, and marks the rest as not measured rather than scoring them against the SAE.

The decision rule

decision.py is the only file that turns numbers into a recommendation. It reads the saved results JSON, computes ten signals, and walks an ordered list of clauses top down. The first one that fires is the verdict.

Signals: effect_lift, ace_closes_gap, leakage_honest, selectivity, robustness, latency_overhead, first_token_overhead, first_token_suffices, capability_preserved, seed_stability.

Clauses, in order:

1. pay_for_selectivity the per-category feature is genuinely local, and capability did not collapse. One direction cannot express this at any price, so it beats every cost argument.
2. linear_suffices_leakage the SAE leads on steering, but the lead does not survive adding the reconstruction error back. That was error discard, not feature editing.
3. lead_not_reproducible the SAE leads at the reported seed, but the lead does not hold its sign across seeds. That is what a lucky split looks like.
4. pay_for_effect the SAE leads, survives the residual control, and does not wreck capability.
5. pay_for_detection_only the two bases tie on steering, but SAE features keep detecting harm through re-encoded prompts. Steer with the direction, monitor with the SAE.
6. linear_suffices nothing measured pays for the SAE.

Every cutoff the rule uses lives in Config.decision (the Thresholds dataclass) and is written into decision.json and decision.md. Change one, rerun refusal-steer decision, and the verdict is rederived without a GPU.

The price the rule spends is the measured one. benchmark times every method with early stopping off, because an arm that makes the model refuse in five tokens finishes sooner and that is an effect, not a saving. It times each method twice, at every token and at the first token only. If positions shows first token steering keeps the effect, the SAE runs once per sequence instead of once per token, and the rule banks the lower price.

How the numbers are scored

Jailbreak rate, not refusal rate. A model steered into a loop ("is is is is") is not refusing, and a refusal-only metric counts that as a win. An arm counts as jailbroken only when the completion is neither a refusal nor degenerate. Every rate carries a 95% bootstrap interval.

Refusal scoring is a substring marker list by default (--judge markers), or a judge model (--judge model, with judge_model naming which one; empty means reuse the model under test). Degeneracy stays deterministic either way, since it is a statistic of the text.

Coefficients are swept, not fixed. One coefficient cannot tell "this method does not work" apart from "this method was not driven hard enough". The default grid is log spaced over the region where the SAE arms are usable, plus the two points where they collapse: 0.5, 0.71, 1.0, 1.41, 2.0, 4.0, 8.0.

Effect is scored against a collateral budget. An arm only counts for the rule if benign refusal and degeneracy stayed near where they started. Jailbreaking everything by breaking the model is not a win.

Fit and eval prompts are disjoint. The direction and the feature ranking are fit objects, so scoring them on their own fitting prompts measures memorisation.

Adversarial transforms are applied to both classes, and a comprehension control records whether the model can read the encoding at all. Splits the model scores below min_comprehension (0.5) are still reported, but they are left out of the aggregate robustness signal. A detector cannot be blamed for missing intent that never reached the model.

Layout

src/refusal_steering/
  config.py        one Config dataclass, plus Thresholds for the decision rule
  data.py          JBB-Behaviors + XSTest, disjoint split, adversarial rewrites,
                   MMLU and WikiText loaders, comprehension probes
  model.py         model and SAE loading, activations, generation, encode/decode,
                   choice logits, NLL, timing and memory helpers
  features.py      refusal direction (difference in means), SAE feature ranking, AUROC
  steering.py      the hooks: directional, ace, sae_clamp, sae_residual_preserving,
                   at_positions
  metrics.py       refusal, repetition, degeneracy, jailbreak, bootstrap CIs
  judge.py         marker judge and model judge
  experiments.py   the eight experiments, the shared Context, run_all
  decision.py      signals, clauses, verdict, all from the results JSON
  plots.py         figures, drawn from the results JSON
  cli.py           refusal-steer
notebooks/         colab_run2.ipynb (current), colab.ipynb (first run)
tests/             test_core.py, test_decision.py

Output

Results land in artifacts/<experiment>/results.json, with that experiment's figures next to them. On top of that:

• artifacts/decision.json and artifacts/decision.md, the verdict with the signals and the thresholds used
• artifacts/figures/effect_versus_cost.{png,pdf}, the cross-experiment plot
• artifacts/manifest.json, the config, which experiments finished, which failed and the verdict
• artifacts/run.json, written by the notebook, recording the GPU, the library versions and the config

Running on Colab

The notebook runs the whole suite on an A100 or L4 and then either downloads artifacts/ as a zip or pushes it back to this repo. A Colab session is temporary and a GPU run is not free, so the output should not depend on the tab staying open.

Two Colab specifics are baked in. Deps are installed directly and the package with --no-deps, because pip reads the published metadata for a git install and that still asks for a torch version Colab does not have. The repo is cloned rather than pip-installed from git for the same reason. For the push path, use a fine-grained token with Contents: read and write in the Colab secret GITHUB_TOKEN. .gitignore tracks results and figures under artifacts/ while still ignoring caches and logs, so a run commits cleanly and its scratch does not.

Model and SAE

Gemma Scope 2 SAEs for gemma-3-4b-it were trained on the instruction-tuned model itself, so unlike Gemma Scope 1 there is no base-model train/apply mismatch to argue about. The curated gemma-scope-2-4b-it-res release covers layers 9, 17, 22 and 29, and Config rejects any other layer up front instead of letting it fail as a hub 404 twenty minutes into a run. For other layers pass --sae-release gemma-scope-2-4b-it-res-all.

Width is 16k and sparsity is medium (L0 around 60) by default. small is sparse enough that refusal splits across several features, which mixes up the single-feature clamp with a feature-selection failure.

Requirements

Python 3.11 or newer, a GPU with about 16 GB (bfloat16 model, float32 SAE), and a HuggingFace token with access to google/gemma-3-4b-it.

About the files already in artifacts/

The current artifacts/ is from the first run. It used the old linear 1, 2, 4, 8 coefficient grid, pinned the ablations at 8, and has no capability or seeds directory, so those two signals are missing from its decision.md. Its verdict is pay_for_selectivity. Rerun the suite to replace it. blog/, docs/ and paper/ still describe that run in places, so do not quote their numbers next to fresh ones without checking.