Source-linked AI summary
Where Decoder Cosine Similarity Fails for SAE Feature Flow Discovery
Hendrik Droste, Christian Medeiros Adriano, Kathrin Korte, Holger Giese
TL;DR
The paper asks how residual-state and MLP-update features interact to produce downstream residual features, a question relevant to understanding model adaptation. It builds and ablation-validates a transition atlas, finding that most strong Pythia transitions are below both pairwise cosine thresholds; Gemma provides qualified, weaker cross-model support.
Problem
Sparse autoencoders provide interpretable feature dictionaries, but how state and update features interact to produce downstream residual features remains unclear.
Source: §1 · p. 1Method
Results
88.0% of strong Pythia transitions have both state-target and update-target decoder cosine below 0.7; the corresponding share is 53.6% in Gemma's top-30,000 validation subset.
Source: §3 · pp. 3–4Takeaways & Limitations
Feature-flow atlases offer a diagnostic of representation-update mechanisms, while pairwise decoder cosine alone misses many strong transitions in the full Pythia atlas.
Source: §3 · p. 3, §4 · p. 4Gemma causally validates only the top 30,000 ranked triples, so it does not establish full-atlas recall; dependence on model, layer, SAE, pruning, and validation budget remains open.
Source: §4 · p. 4
Abstract
from arXiv · showhide
Foundation models are increasingly adapted through fine-tuning, model editing, and alignment procedures while retaining previously acquired capabilities. Understanding the internal computations that support these adaptations is therefore becoming increasingly important for continual model evolution. Sparse autoencoders (SAEs) provide interpretable feature dictionaries for residual-stream activations and sublayer outputs, but it remains unclear how state features and update features interact to produce downstream residual features. In this work, we focus on MLP updates as a first test case. We construct a transition atlas of triples $s_k + u_j \rightarrow t_\ell$, where a residual-state feature and an MLP-update feature jointly predict a target residual feature, and validate candidate triples by ablating the decoded update feature. In a 20M-token Pythia-160M $L_7 \rightarrow L_8$ run, we find 38,125 strong ablation-effect transitions, but 88.0% have both state-target and update-target decoder cosine similarity below 0.7. As a preliminary cross-model check, a run of 20M-token Gemma-3-4B $L_{21} \rightarrow L_{22}$ causally validates only the top 30,000 ranked candidate triples by ablating the decoded update feature, and 53.6% of strong-effect triples have both state-target and update-target decoder cosine similarity below 0.7. The Gemma result is directionally consistent with Pythia, but weaker, since update-target cosine recovers many of the strongest Gemma effects and the run is not a full-atlas causal validation. Ultimately, our results suggest that feature flow atlases can serve as diagnostics of representation-update mechanisms and thereby inform tools for steering model updates. Future work will validate more complex patterns across layers, models, and SAE families.
1 INTRODUCTION
Foundation models are increasingly adapted while retaining prior capabilities, but existing work offers limited insight into causal feature-level transformations. This paper studies how residual-state and MLP-update features interact across layers.
1 INTRODUCTION
- Existing work localizes factual associations and tracks concepts that shift or emerge, but offers less insight into causal feature-level transformations during computation.These transformations matter in settings where models incorporate new information while preserving prior capabilities.Source: §1 · p. 1
- Sparse autoencoders provide interpretable feature decompositions, motivating feature-flow discovery as a way to study interactions between residual-state and MLP-update features.The paper positions this as a mechanistic step toward understanding how features interact during computation.Source: §1 · p. 1
- The paper asks whether strong cross-layer transitions can be recovered from pairwise decoder geometry alone or depend on the joint presence of prior state and update features.It frames the question through triples of a residual-state feature, MLP-update feature, and downstream residual feature.Source: §1 · p. 1
2 FEATURE FLOW DISCOVERY
The paper builds an atlas of state–update–target feature triples and tests them through decoded-update ablation. It compares this joint, activation-based approach with pairwise decoder cosine similarity.
2 FEATURE FLOW DISCOVERY
- The atlas represents each transition as a source residual-state feature and target-layer MLP-update feature jointly predicting a downstream residual feature.Pairwise state-target and update-target decoder cosines serve as the geometric baseline.Source: §2 · pp. 1–2
2.1 MATCHING PIPELINE
- Three SAE feature spaces define state, update, and target features; activation-frequency scores rank triples, which are then checked by decoded-update ablation.The check asks whether removing the decoded update changes the target feature by at least two target-feature standard deviations.Source: Figure 1 · p. 2
- The pipeline uses TopK SAEs trained on Pythia-160M residual streams after layers 7 and 8 and the layer-8 MLP output, storing the top 64 active features per token.Stage one collects and encodes source residual, MLP-sublayer, and target residual activations.Source: §2.1 · p. 2
- Candidate triples are ranked against state-target and update-target pairwise frequency baselines, then filtered by minimum activation and co-occurrence thresholds.The baselines prevent ranking triples solely because their individual features are common.Source: §2.1 · p. 2
- For validation, the decoded contribution of the active update feature is ablated, and the target feature's normalized effect is measured in its own standard deviations.Ablation effects are normalized because target-feature activations differ in scale and variability.Source: §2.1 · p. 3
- A complementary experiment ablates selected state features in high-ranked state-update pairs and measures the effect on the corresponding update feature.This tests the state-to-update relation rather than the update-to-target effect used in the main validation.Source: §2.1 · p. 3
2.2 COSINE MATCHING
- Decoder cosine compares feature write directions without testing shared-token activity, causal change, or dependence on prior residual state.Consequently, two individually low-cosine features may jointly move the residual stream into a target feature's activation region.Source: §2.2 · p. 3
3 PRELIMINARY RESULTS
In the Pythia full-atlas run, most strong causal transitions fall below both pairwise cosine thresholds, while the Gemma subset gives a weaker cross-model check. The results indicate that pairwise decoder geometry alone misses many strong transitions, with important scope limits for Gemma.
3 PRELIMINARY RESULTS
- 38,125 of 262,735 filtered Pythia triples reach |Az| ≥2, representing 14.5% of the filtered triples.These validated strong-effect triples are used to test recovery from pairwise decoder geometry.Source: §3 · p. 3
- High update-target cosine captures direct MLP write relations, while state-feature ablation also produces strong effects on many corresponding update features.Across the 4,096 highest-ranked unique state-update pairs, 3,765 with at least ten active tokens reach Az ≥2 (91.9%), with median update-feature effect Az = 40.6.Source: §3 · p. 3
- 88.0% of Pythia strong-effect triples have both state-target and update-target cosine below 0.7, so pairwise cosine misses most strong causal transitions.At τ = 0.4, 62.1% still have both scores below threshold; Figure 2 shows the full-atlas Pythia distribution.Source: §3 · p. 3, Figure 2 · p. 4
- For Gemma, 53.6% of strong-effect candidates have both cosines below 0.7, but update-target cosine overlaps with 68.2% of strong effects.This is a top-30,000-candidate validation subset, not a full-atlas recall measurement, and the result is weaker than Pythia's.Source: §3 · p. 4
4 LIMITATIONS AND ONGOING WORK
The evidence is strongest for the completed Pythia L7 → L8 full-atlas run, while the Gemma check validates only the top 30,000 ranked triples. The authors outline further validation and suggest feature-flow analysis could offer a mechanistic perspective on continual adaptation.
4 LIMITATIONS AND ONGOING WORK
- Gemma-3-4B validates only the top 30,000 ranked triples, and how cosine recall depends on model scale, layer, SAE architecture, pruning, and validation budget remains open.The completed Pythia L7 → L8 full-atlas run is the paper’s strongest evidence.Source: §4 · p. 4
- Future work will compare triple-ranking scores, validate more Gemma Scope layer transitions, and explore pretrained SAEs to improve scalability and feature insight.The authors currently have only estimates of how pruning parameters affect the atlas.Source: §4 · p. 4
- Comparing feature-flow atlases across adaptation stages may complement parameter-update and representation-change analyses with a mechanistic view of transfer, interference, and knowledge integration.The study focuses on inference dynamics within a fixed model.Source: §4 · p. 4
A.1.1 SYMBOLS
The symbols name residual-stream and MLP outputs, SAE encoders, state, update, and target features, and decoder-cosine and ablation quantities.
A.1.1 SYMBOLS
- post denotes the residual stream after the target layer; r(i) denotes the residual stream before it, while a(i) and m(i) are layer-i attention and MLP outputs.em(i) is the MLP output after removing the decoded update feature.Source: §A.1.1 · p. 6
- E_res and E_mlp denote the residual-stream and MLP-output SAE encoders for layer i; s_k, u_j, and t_ℓ denote state, update, and target features.The state feature is from layer i−1; the update and target features are from layer i.Source: §A.1.1 · p. 6
- The decoder directions are d_s, d_u, and d_t; A_z is normalized ablation effect, c_s,t and c_u,t are decoder-cosine baselines, and τ is the cosine threshold.ϵ is a small numerical-stability constant.Source: §A.1.1 · p. 6
A.1.2 LAYER DECOMPOSITION
The layer decomposition approximates each token’s residual stream after layer i as the sum of the preceding residual stream, attention output, and MLP output.
A.1.2 LAYER DECOMPOSITION
- For each token, the post-layer-i residual stream is approximated by adding the residual stream before the target layer, attention output, and MLP output.Source: §A.1.2 · p. 6
A.1.3 MLP-INDUCED TARGET CHANGE
The MLP-induced target-feature change is defined using encodings in a shared layer-i residual SAE basis.
A.1.3 MLP-INDUCED TARGET CHANGE
- The method encodes both the original residual stream and the no-MLP residual stream in the same layer-i residual SAE basis to define the MLP-induced target-feature change.Source: §A.1.3 · p. 7
A.1.4 DECODER-COSINE BASELINES
The section establishes pairwise decoder-cosine baselines and intervention-based validation, then describes state-to-update checks, candidate pruning, and SAE training configurations.
A.1.4 DECODER-COSINE BASELINES
- Decoder-cosine baselines compare state- and update-feature decoder directions with the target-feature direction.Source: §A.1.4 · p. 7
A.1.5 CAUSAL VALIDATION
- A triple is validated when ablating its decoded update contribution produces a normalized target-feature change with magnitude |Az| ≥2.The intervention is evaluated at positions where both the state and update features are active.Source: §A.1.5 · p. 7
A.1.6 STATE-TO-UPDATE CHECK
- A complementary check ablates the decoded state-feature contribution and measures the resulting change in the corresponding MLP update feature.This tests whether the update feature depends on the prior residual-state feature.Source: §A.1.6 · p. 7
A.2 PRUNING
- The search retains the top 64 active features per token for Pythia and prunes candidate triples, so the atlas is precision-oriented rather than exhaustive.Token-level pruning is distinct from the later budget for interventionally testing ranked triples.Source: §A.2 · p. 7
A.3 PYTHIA SAE TRAINING HYPERPARAMETERS
- Table 2 lists training hyperparameters for custom Pythia-160M SAEs; Gemma uses pretrained width-16k, L0-big Gemma Scope 2 residual SAEs at layers 21 and 22 and a layer-22 MLP SAE.Source: §A.3 · p. 8, Table 2 · p. 8