Representation redistribution in sparse autoencoders

A sparse autoencoder, or SAE, turns one layer of a language model into sparse features that can be read one by one. When a prompt is reworded without changing the request, which features change, and when does it matter?

Ziye Mei, Haotian Ma University of Wisconsin–Madison

A safety monitor can look accurate while individual decisions flip

Prompt

→

LLM layer

→

SAE features

→

Linear probe

→

Harmful?

The same request after one harmless rewrite

Click a kind of rewrite. The sentences are our own examples of each kind.

What changed inside the model, and which of those changes reached the decision?

Frozen probe on Gemma 3 4B

—

AUROC after the rewrite. It still looks healthy.

—

of answers change after the rewrite.

AUROC is the chance that a harmful prompt gets a higher probe score than a benign one.

AUROC = \Pr (z (x_{harmful}) > z (x_{benign}))

Each dot is 0.5% of decisions. Red dots changed their answer.

The monitor being audited

The monitor is a logistic probe on standardized SAE activations. It is trained once on clean prompts and then frozen. We use Gemma 3 4B layer 17 with 415 nonzero weights and Llama 3.1 8B layer 24 with 862. The data are 1,800 balanced WildGuardMix prompts, split into 1,200 for training, 300 for validation and 300 for testing.

Each test prompt gets 15 rewrites per transformation, five operators with three realizations each. The flip rate is the share of rewritten predictions that differ from the clean prediction. The slide shows Gemma, and the table also lists Llama.

Transformation	Gemma AUROC	Flip	Llama AUROC	Flip

The rewritten sentences on the slide are our own examples. The study used generated variants that keep the target phrase intact.

Paper Table S19.

~50 s

Talk track

This is the setup. A prompt goes into the model, we read one layer through the SAE, and a linear probe on those features decides whether the prompt is harmful. The probe is trained once and then frozen.
Now we rewrite the request without changing what it asks. The six buttons are the six kinds of rewrite we test, and the sentence shows our own example of each.
This number is the AUROC of the probe after the rewrite. AUROC is the chance that a harmful prompt scores above a benign one. It stays above 0.90 for every kind of rewrite, so nothing looks wrong.
This number and the dots count answers that changed. It is 4 to 8 percent on this model and up to 9.5 percent on the second model. Each dot is half a percent of decisions, and red means the same request now gets a different answer.
So the rest of the talk asks what changed inside the model, and which of those changes reached the decision.

Rewrite the prompt five ways. Keep features that respond again and again, in the same direction.

feature rose feature fell no change kept

Toy numbers. The selection rule is the one in the paper.

The change of feature f under rewrite v and its average size are

δ_{v, f} = a_{f} (x_{v}) - a_{f} (x_{0})

h_{f} = \frac{1}{| V |} \sum_{v \in V} | δ_{v, f} |

Of the n rewrites, the feature goes up in

n_{f}^{+}

and down in

n_{f}^{-}

. The share of rewrites in which it responds and the share of those that go the dominant way are

r_{f} = \frac{n_{f}^{+} + n_{f}^{-}}{n}, c_{f} = \frac{\max (n_{f}^{+}, n_{f}^{-})}{n_{f}^{+} + n_{f}^{-}}

Kept features

S

?

MeanRSS

?

MeanRSS = \frac{1}{| S |} \sum_{f \in S} h_{f}

The RSS selector

For a clean prompt and its rewrites, the change of feature f is measured at the aligned target token.

δ_{v, f} = a_{f} (x_{v}) - a_{f} (x_{0})

h_{f} = \frac{1}{| V |} \sum_{v \in V} | δ_{v, f} |

Over the n selector observations, $n_{f}^{+}$ counts changes above τ and $n_{f}^{-}$ changes below −τ.

r_{f} = \frac{n_{f}^{+} + n_{f}^{-}}{n}, c_{f} = \frac{\max (n_{f}^{+}, n_{f}^{-})}{n_{f}^{+} + n_{f}^{-}}

Among the M features with the largest $h_{f}$ , the eligible features are

ℰ_{M} = {f \in C_{M} : r_{f} \geq ρ, c_{f} \geq κ}

RSS keeps the K eligible features with the largest $h_{f}$ as the support $S$ , and MeanRSS is their mean $h_{f}$ . The paper uses ρ = κ = 0.8, τ = 10⁻⁵, M = 400 and K = 20. With n = 5 observations a kept feature moves in its dominant direction at least four times, so one rewrite cannot reverse it. Biggest movers is selection by magnitude alone with the same pool and size.

On the slide there are six toy features, K = 3 and τ = 0.05.

Paper Section 3.1, Propositions 1 and 2, and Table S2.

~70 s

Talk track

To find where the change lives, we rewrite the same prompt five ways. Each column is one SAE feature and each row one rewrite. The colour says whether the feature went up or down. These are toy numbers so that the pattern is easy to see.
Which three would you keep? Click three column headers, then press Reveal.
After the reveal three rows appear. They show in how many rewrites each feature responds, how often it moves the same way, and its average size. These are r f, c f and h f, defined on the right.
Now the rule switch appears. Switch it and watch the gold columns. Biggest movers keeps f1, f2 and f4. But f2 flips between up and down, and f4 moves only once. The RSS rule asks for a response in at least four of five rewrites, mostly in one direction, and then takes the largest. That gives f1, f3 and f5.
The kept features form the responsive support S. MeanRSS is the mean of their h f, so it says how strongly the representation moved. In the paper we keep twenty features per prompt.

Context and wrappers reshuffle features far more than local edits in every SAE we tested

Each dot is one SAE, a model and a layer. Diamonds are means over the 27 SAEs.

Normalized MeanRSS

R_{g} = \underset{i}{mean} {MeanRSS}_{i, g}

{\tilde{R}}_{g} = \frac{R_{g}}{\frac{1}{6} \sum_{g^{'}} R_{g^{'}}}

For each SAE, MeanRSS is averaged over prompts i for each setting g and divided by the average of the six settings. Every SAE then has mean one.

Context and wrapper mean minus local edit mean

0.874

95% interval 0.867 to 0.881. Positive in all 27 SAEs and for all 1,000 prompts.

Highlight one model

Reading the atlas

There are 27 pretrained SAEs from eight models, GPT-2 Small, Pythia 70M, Gemma 3 270M, 1B and 4B, Gemma 2 2B, Mistral 7B and Llama 3.1 8B, with several layers each. Each SAE sees 1,000 prompts under six transformations with 15 variants each.

Normalization divides each SAE's six MeanRSS values by their own average, so the plot compares profiles across architectures rather than raw sizes. Every context or wrapper setting exceeds both local edits for 997 of 1,000 prompts. Redistribution also grows with transformation strength. The median prompt level Spearman correlation between intensity and MeanRSS is 0.95 for punctuation, 0.74 for irrelevant context and 0.50 for the format wrapper.

The profile does not come from the selection itself. On rewrites never used for selection, RSS features respond in their selected direction at a rate of 0.952, compared with 0.758 for random supports of the same size. RSS separates context from local edits more strongly than each baseline in 24 to 26 of 27 SAEs. What is stable is the profile across transformations, not the exact identity of the selected features.

Transformation	Mean	95% CI

Paper Figure 2, Section 4.1 and Table S7.

Restoring about twenty features undoes most flipped decisions

How a linear probe decides

z (x) = 𝐰^{⊤} \tilde{𝐚} (x) + b, m (x) = y [z (x) - t]

The margin m is positive when the decision is correct. A rewrite moves it one feature at a time.

e_{f} = - y w_{f} \frac{1}{| V |} \sum_{v} {\tilde{δ}}_{v, f}, E_{S} = \sum_{f \in S} e_{f}

e_{f}

is how much feature f erodes the margin over the rewrites.

E_{S}

is the part carried by the RSS features.

72%

of the probe's wrong answers on the whole test set are reversed when only the RSS features are restored, on Gemma 3 4B.

One flipped decision, feature by feature. Illustrative numbers.

Restoring recomputes the readout with those features at their clean values. The model itself is not changed.

From feature changes to the margin

The probe reads standardized activations. A positive margin is a correct decision, and $m_{0}$ is the margin on the clean prompt.

z (x) = 𝐰^{⊤} \tilde{𝐚} (x) + b, m (x) = y [z (x) - t]

Over the selection rewrites, feature f erodes the margin by $e_{f}$ . For a selected feature this factors exactly into the weight, the response size and the agreement with the RSS direction.

e_{f} = - y w_{f} s_{f} A_{f} η_{f}

Restoring the support gives the counterfactual margin below. It changes only the readout calculation.

m_{cf} (x_{v}) = m_{0} + y \sum_{f \notin S} w_{f} {\tilde{δ}}_{v, f}

Restoring the RSS support in the readout reverses 72.2% of transformed probe errors for Gemma and 71.8% for Llama. RSS supports reconstruct held out logit changes with R² of 0.58 and 0.63, and unstable or random features with 0.01 to 0.02. After matching probe weights, RSS features contribute 0.240 against 0.022 for Gemma and 0.370 against 0.022 for Llama. The RSS direction predicts the sign of the contribution in 97.4% and 97.1% of cases.

At K = 20, RSS keeps 43.0% and 56.2% of the absolute held out contribution, compared with 32.6% and 43.5% for the support of another transformation. Ranking features by $| e_{f} |$ keeps more than 90%, but it needs the probe. RSS never looks at the probe or the labels. The size of the move alone is a poor guide, and MeanRSS alone ranks recurrent failures at an AUROC of 0.53 for Gemma and 0.51 for Llama.

Paper Sections 3.2 to 3.4 and 4.2 to 4.4, Proposition 3 and Table S29.

~60 s

Talk track

Why do some rewrites flip the monitor? A linear probe adds up evidence one feature at a time. The margin is the signed distance of the score from the threshold, positive when the decision is correct. A rewrite moves the margin by each feature's weight times its change. So a feature matters only if the probe weights it and it moves the wrong way.
Here is one flipped decision broken down that way, with illustrative numbers. The gold bar is the margin on the original prompt. Each red bar is what one feature takes away. The RSS features take most of it, the other features take little, and the answer crosses the line.
Press restore. We put the 20 RSS features back to their original values in the probe's calculation, and the answer comes back. Restoring 20 random features does almost nothing.
Across the whole test set, restoring only the RSS features reverses 72 percent of the probe's wrong answers. This is done in the readout and not inside the model, so it is accounting and not a causal claim about the model.

With the clean margin, about twenty features rank the prompts that keep failing

Gemma 3 4B. Each point is one test prompt under one kind of rewrite. Red points fail on at least two of ten new rewrites.

RSS residual margin

m_{res}^{S} = m_{0} - E_{S}

The clean margin minus the erosion predicted from the RSS features on five rewrites. Failures are measured on ten other rewrites.

Review the cases with the lowest

Review budget

Probe errors caught at this budget

Probe confidence

RSS residual, about 20 features

All 415 probe features

Residual margin and review

The RSS residual margin is the clean margin minus the erosion predicted from the RSS support on the five selection rewrites. Outcomes are measured on ten different rewrites.

m_{res}^{S} = m_{0} - E_{S}

A case is a recurrent failure when at least two of its ten held out margins are negative. This happens for 6.6% of Gemma cases, 101 of 1,524, and for 8.4% of Llama cases, 121 of 1,445. Every case with a negative full residual margin is a recurrent failure. The median RSS support has 20 and 19 features, against 415 and 862 probe weights.

MeanRSS alone, the size of the move, catches 18.8% and 23.3% of errors at a 20% review, no better than random review.

At 20% review	Errors	Harmful FN	AUROC

Paper Figure 5, Table S27 and Sections 4.3 and 4.4.

~70 s

Talk track

Which prompts keep failing, and can a few features rank them? Every point is one real test prompt under one kind of rewrite. Left to right is how sure the probe was on the original prompt. Bottom to top is that margin minus the damage predicted from the twenty RSS features. Red points really fail on new rewrites that we never used. Solid red is caught by the review and hollow red is missed.
This is the score on the vertical axis, the RSS residual margin. It starts from the clean margin and subtracts the erosion carried by the RSS features.
A person can check some of the cases. Which ones? Rank them by the probe's own confidence, or by the RSS residual.
Drag the budget. The gold line sweeps across, and everything in the shaded area is reviewed. Ranked by confidence, the review is a strip on the left and misses failures that looked safe. Ranked by the RSS residual, it is a strip along the bottom, where the red points are.
The bars compare the three rankings at the current budget. At 20 percent, confidence alone catches 69 percent of the errors, the RSS residual 89 percent, and all 415 probe features 94 percent.

Takeaways

Where features move under rewording is systematic, measurable and traceable to decisions

Rewording reshuffles SAE features predictably

Context and wrappers move features far more than local edits, in all 27 SAEs.

RSS finds the features that move reliably

They respond again and again in the same way, also on rewrites never used to pick them.

About 20 of them explain and rank failures

Restoring them reverses 72% of errors. With the clean margin, reviewing 20% of cases catches 89% of errors.

What these results do not show

A causal mechanism inside the model. Restoring is done in the readout.
What the moved features mean.
That the size of the move, MeanRSS, detects failures by itself. It does not.