Representation Redistribution in Sparse Autoencoders

AI summary

Overview

Research area: Mechanistic interpretability of large language models, specifically sparse autoencoders (SAEs) and how prompt variation redistributes activity across learned features.

Technical level: Advanced — the paper builds formal definitions, three propositions plus supplementary corollaries, and a linear-readout decomposition, alongside large-scale empirical audits.

Scope: The paper introduces the redistribution sensitivity score (RSS) for identifying sparse feature changes that recur consistently under prompt transformations, then separates that representational recurrence from actual downstream effect on a fixed linear readout across 27 pretrained SAEs and two frozen safety probes.

What This Paper Is About

Sparse autoencoders break a language model's dense activations into many sparse features, but rewriting a prompt in a way that preserves the underlying request can shift which features fire — the paper calls this representation redistribution. The problem is twofold: a change seen on one rewritten prompt may not recur on another, and even a change that does recur may have no effect on a downstream computation that reads those activations. The authors build a scoring method that selects only recurring, directionally consistent feature changes, fix that support before evaluation, and then measure exactly how much of a linear readout's behaviour those selected changes actually explain.

Key Contributions

1. The redistribution sensitivity score (RSS). RSS builds a magnitude-based candidate pool from mean absolute feature change, requires recurrence and directional consistency within it, and keeps the largest eligible features as a "responsive support." The support and its directions are selected on one set of variants and held fixed when evaluated on disjoint variants. A companion statistic, MeanRSS, summarizes the support's response magnitude.

2. A separation of representation change from downstream effect. For a fixed linear readout, the authors give an exact decomposition showing how response magnitude, readout coefficient weight, directional alignment (including cancellation across variants), and the clean margin jointly determine whether a recurrent feature change affects a readout and whether it crosses a decision threshold.

3. Evidence of a reproducible redistribution profile. Across 27 pretrained SAEs, the tested transformations produce a consistent, reproducible redistribution profile on variants excluded from feature selection.

4. Functional validation on frozen safety readouts. Supports selected without any readout information nonetheless reconstruct held-out logit changes, reverse a large share of probe errors when restored, and improve offline error-ranking when combined with the clean margin.

Main Findings

• Context and wrapper changes redistribute more than local edits. In the representation study, 1,000 English prompts across 12 content categories and six transformation settings were used. The normalized context-and-wrapper mean exceeded the local-edit mean in all 27 SAEs, with an average difference of 0.874 and a 95% bootstrap interval of [0.867, 0.881]. After aggregation across SAEs, the pattern held for all 1,000 prompts, and every context or wrapper setting exceeded both local edits for 997 prompts.

• RSS beats comparison supports at recovering the profile on held-out variants. RSS gave stronger context-versus-local-edit separation than each comparison method in 24–26 of 27 SAEs and produced the largest contrast in 22. Comparison sets included clean-fixed features, globally fixed features, and the dictionary mean.

• Selected directions mostly persist. On held-out variants, the response rate in the selected direction was 0.952 for RSS versus 0.758 for random supports matched in size and candidate pool. The finite-sample guarantee (with n = 5 and ρ = κ = 0.8, at least four observations share the dominant direction and at most one opposes) applies only to selection observations, not to new variants.

• Magnitude alone is not enough. Against magnitude-only selection matched in support size, mean held-out response rates were 0.876 versus 0.816 for Gemma and 0.871 versus 0.794 for Llama, confirming the recurrence and consistency constraints add persistence.

• RSS features carry more readout weight than coefficient-matched controls. After matching on |w_f|, mean absolute held-out contribution was 0.240 versus 0.022 for Gemma and 0.370 versus 0.022 for Llama, with RSS larger in 96.9% and 98.0% of matched pairs. The paper's sign rule agreed with the sign of the averaged held-out contribution in 97.4% and 97.1% of cases.

• RSS retains real but not maximal contribution. At K = 20, RSS retained 43.0% and 56.2% of absolute held-out contribution, versus 32.6% and 43.5% for RSS supports selected for a different transformation. A readout-aware |e_f| ranking retained more than 90% in both models, so RSS is useful without selecting for the readout but is not the optimal retainer.

• Clean margin plus support erosion predicts recurrent failure well. Clean margin alone gave recurrent-failure AUROCs of 0.863 (Gemma) and 0.901 (Llama). Subtracting RSS support erosion from the full clean margin raised these to 0.935 and 0.958 with median supports of 20 and 19 features, compared with 0.965 and 0.964 for full erosion. Full erosion correlated with held-out erosion at 0.923 and 0.896; the RSS support versions correlated at 0.737 and 0.793.

• Errors concentrate near the decision boundary. Probe AUROC stayed above 0.90 under every transformation, yet 4.2–9.5% of decisions changed. Every case with negative full residual margin had recurrent readout failure in this cohort — an association the authors describe as transfer, not identity, since residual margin comes from selection variants and failure from held-out decisions.

• Compact supports reconstruct and localize readout change. RSS reconstructed held-out logit changes with R² of approximately 0.58 and 0.63, and restoring selected readout inputs to their clean values reversed 72.2% and 71.8% of transformed probe errors. Mismatched supports did worse; globally unstable and random features reached only about 0.01–0.02 R².

• Useful under a tight review budget. At a 20% offline review budget, the RSS residual margin captured 89.2% of Gemma errors and 92.8% of Llama errors, plus 87.3% and 94.4% of harmful false negatives, using supports with a median of 20 and 19 features against 415 and 862 nonzero coefficients in the full readouts.

Methodology in Plain English

Start with a clean prompt and a set of transformed variants that keep the request the same but change context, wrapper, punctuation, or wording. For each SAE feature, measure how much its activation changes and average the absolute change, so opposite directions do not cancel out. Keep the largest changes as candidates, then ask of each candidate: how often does it actually respond across the selection observations, and how consistently does it move in one direction? Only features that clear both thresholds survive, and the top few by magnitude become the "responsive support."

The key discipline is that selection happens first and is then frozen. Features and their chosen directions are locked in on the selection variants, and everything downstream is measured on a disjoint set of variants the selector never saw. This distinguishes genuine recurrence from fitting to the particular prompts used.

For the readout analysis, the authors fit a linear probe on SAE activations using clean training data only, calibrate a threshold on clean validation data, and freeze it. Because the readout is linear, each feature's contribution to the margin can be written exactly as the product of its readout weight, its average response magnitude, and a directional alignment factor that shrinks when responses cancel. Combining the average erosion with the clean margin gives a residual margin, which can then be compared against actual held-out failures. A counterfactual test restores selected features to their clean values inside the readout only — the language model's activations and generated text are untouched — to see whether the decision flips back.

The functional experiments use 1,800 balanced WildGuardMix prompts split 1,200/300/300 into training, validation, and test, with readouts on Gemma 3 4B layer 17 and Llama 3.1 8B layer 24. For each test prompt, one realization per operator forms the selection set and the other two form the held-out set — five selection variants and ten held-out variants. Recurrent readout failure is defined as at least two of the ten held-out margins being negative, leaving 1,524 Gemma and 1,445 Llama prompt-and-setting cases after requiring a correct clean prediction, available support, and complete outcomes.

Why This Matters

Research impact. Interpretability work often treats a feature change as a finding in itself, or summarizes representational drift as a single aggregate distance. This paper argues for a middle object — recurring, directionally consistent change on a fixed support — and then insists on a separate accounting of downstream consequence. The exact decomposition explains why large changes can be harmless and why small, well-aligned changes can flip a decision near the boundary, which is a useful corrective to equating representation change with failure.

Real-world applications:

• Safety classifier auditing. The two frozen readouts are safety probes; RSS-based residual margins give auditors a way to rank which prompt phrasings are most likely to break a classifier using only a few dozen features instead of hundreds.
• Prompt-robustness testing. Teams varying system prompts, wrappers, or formatting can use the context-versus-local-edit profile to predict which kinds of edits disturb internal representations most.
• Model regression monitoring. A fixed support and MeanRSS provide a lightweight signal for detecting when a model update changes how activations redistribute under otherwise equivalent inputs.
• Triage under limited review capacity. The offline selective-review results show how a compact risk ranking can concentrate human review on the cases most likely to be errors.

Industry relevance. The method is deliberately readout-agnostic at selection time, so a support can be computed once and reused against multiple downstream heads. The compactness result — median supports of roughly 20 features against hundreds of nonzero readout coefficients — matters for deployment settings where auditing every coefficient is impractical.

Future Directions

• Testing beyond new realizations of the same operators. The main functional holdout uses new realizations of the same transformations, so transfer to genuinely unseen transformation families remains open.
• Extending past frozen linear readouts. The exact decomposition depends on linearity; the paper does not establish how the accounting behaves for nonlinear or multi-layer downstream computations.
• Closing the gap to full erosion. Support erosion reached AUROCs of 0.935 and 0.958 versus 0.965 and 0.964 for full erosion, and the readout-aware |e_f| ranking retained over 90% of contribution while RSS retained 43.0% and 56.2% at K = 20 — leaving room for selection criteria that stay readout-free yet retain more.
• Broadening the model and transformation panel. The abstract describes eight language models while the experiments section names GPT-2, Pythia, Gemma 2, Gemma 3, Mistral, and Llama 3.1 across multiple layers, so wider coverage and more transformation categories are natural extensions.

Target Audience

Interpretability and alignment researchers who work with sparse autoencoders and want a principled way to move from "this feature changed" to "this change mattered," plus safety and evaluation engineers who need compact, auditable signals for probing classifier fragility. Readers should be comfortable with linear probes, margin-based classification, and the basic SAE setup; the formal propositions are self-contained but assume familiarity with the notation. Practitioners interested only in prompt sensitivity at the behavioural level will find the framing useful but will need to engage with the feature-level machinery to apply it.

Code is available at https://github.com/OpenSAE/Redistribution-Sensitivity-Score.

Authors’ abstract

Sparse autoencoders (SAEs) decompose language model activations into sparse features, but prompt transformations designed to preserve the underlying request can redistribute activity across them. A feature change observed on one transformed prompt need not recur, and even a recurring change need not affect a downstream computation. We introduce the redistribution sensitivity score (RSS), which identifies large feature responses that recur in a consistent direction, and evaluate the resulting supports on variants excluded from selection. Across 27 pretrained SAEs spanning eight language models, the tested prompt transformations produce a reproducible redistribution profile on these excluded variants. We then separate recurring representation change from signed readout contribution and decision failure. For a fixed linear readout, an exact decomposition shows how response magnitude, readout coefficients, directional alignment, and the clean margin determine downstream consequences. In two frozen safety readouts, combining RSS support erosion with the full clean margin yields recurrent failure AUROCs of 0.935 and 0.958 using median supports of 20 and 19 features, compared with 0.965 and 0.964 using full erosion. RSS therefore provides a way to identify recurring sparse representation changes and quantify their contribution to a downstream readout without equating representation change with failure. Code is available at https://github.com/OpenSAE/Redistribution-Sensitivity-Score