Mechanistic Permutability: Match Features Across Layers

Authors: Nikita Balagansky, Ian Maksimov, Daniil Gavrilov
Primary area: interpretability and explainable AI
Keywords: Interpretability, LLM features, SAE

Abstract:
Understanding how features evolve across layers in deep neural networks is a fundamental challenge in mechanistic interpretability, particularly due to polysemanticity and feature superposition. While Sparse Autoencoders (SAEs) have been used to extract interpretable features from individual layers, aligning these features across layers has remained an open problem. In this paper, we introduce SAE Match, a novel, data-free method for aligning SAE features across different layers of a neural network. Our approach involves matching features by minimizing the mean squared error between the folded parameters of SAEs, a technique that incorporates activation thresholds into the encoder and decoder weights to account for differences in feature scales. Through extensive experiments on the Gemma 2 language model, we demonstrate that our method effectively captures feature evolution across layers, improving feature matching quality. We also show that features persist over several layers and that our approach can approximate hidden states across layers. Our work advances the understanding of feature dynamics in neural networks and provides a new tool for mechanistic interpretability studies.

Decision: accept (poster)
Reviewer scores: 8, 5, 5, 8

Reviews and discussion:
<thread>
creation date:
1730585386216

thread_type:
official_review

<thread_object>
creator: reviewer_ukcq

<content>
summary:
The paper presents a technique to match SAE latents from two SAEs across different layers called SAE Match. The technique finds permutations of latents which minimize MSE between the encoder and/or decoder representation of the feature in both SAEs. SAE Match addresses different layer norms by folding in the jumprelu threshold parameter into encoder and decoder matrices. The paper evaluates the technique on the Gemma Scope SAEs for Gemma 2.

strengths:
Using MSE loss between SAE decoders is simple and fast. The idea of folding the jumprelu threshold parameter into weights to account for differences in layer norm is a really nice idea as well. Aligning SAE latents across layers is important for understanding how features progress throughout the foward pass of the model, so this work is useful to the field.

weaknesses:
The paper only focused on a single model (Gemma) and a single type of SAE (JumpReLU). The paper evaluates on Gemma Scope, but uses only a single SAE per layer despite Gemma Scope containing multiple SAEs per layer with different sparsities. It would be a good check to see how the method performs comparing different SAEs trained on the same layer as well. The method seems like it does not address feature splitting, where a single latent in one SAE becomes multiple latents in a different SAE - the method seems like it will only pair latents 1-to-1.

questions:
- In Figure 3, there are different results for folder and unfolded b_enc, but b_enc is not affected by folding according to equation 4. How is b_enc different due to folding?
- In Figure 3, vanilla matching performs not much worse than the actual matching techniques. Is vanilla matching essentially equivalent to randomly pairing decoder rows? If so, why is the MSE loss so low?
- In Section 4, it says SAEs with regularization coefficient near 70 are used. Is this referring to the L0 of the SAE in Gemma Scope? Neuronpedia only supports Gemma Scope SAEs with L0 closest to 100 - if Neuronpedia is used for all SAEs, does this mean that the L0 near 100 Gemma Scope SAEs are used as well?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730657206296

thread_type:
official_review

<thread_object>
creator: reviewer_hmcz

<content>
summary:
This paper proposes a new method for comparing and matching features of SAEs of nearby layers in transformers. They propose scaling features according to the activation threshold $\theta$ of a JumpReLU activation function in order to match features in a more natural norm for the underlying activations.

strengths:
- Proposes a novel and interesting strategy for pairing features between layers.
    - Studies some of the shortcomings (e.g. long tail of pairing 'failures') of this strategy.
- The presentation is very clear and understandable.

weaknesses:
- It would be good to spend more time justifying the hypothesies of Section 3. I do not think that the results in Figure 3 constitute much evidence for Hypothesis 2, since the reasoning here seems slightly circular - you propose parameter folding based off the observation that $\theta$ tracks the activation norms, but then evaluate feature similarity using the same objective that you are explicitly trying to minimize. Therefore, it is trivially true that 'folding+matching' outperforms 'matching'.
    - The evidence for Hypothesis 2 could be strengthened by including an analysis of how the scale of $\theta_i$ tracks with the scale of feature $i$ *on the same layer* (and similarly measuring MSE by scaling the features according to data-dependent statistics instead of $\theta$, e.g. by comparing MSE in the norm induced by some whitening transformation?); the observation in Figure 1 shows that *mean* $\theta$ tracks with *mean* activation norm, but your hypothesis rests on the assumption that $\theta_i$ is predictive of the scale of feature $i$ for features of the same layer.
    - It would be more convincing if you backed up your claims in Figure 2 that "reconstructions of matched features are closer to each other than in the unfolded variant of the algorithm" by comparing the reconstruction loss (perhaps again under a norm induced by a whitening transformation?) of 'folded+matched' and 'matched' permutations when 'skipping' a layer, as in Section 5.5.
- There was little analysis of the results of the experiments in Section 5.5. The loss differences resulting from using feature-matched activation reconstructions were not compared to any nontrivial baseline (e.g. perhaps something like a linear approximation of the layer?) and it was hard to see what the desired conclusion/hypothesis or suggestion at a direction for future study was here, and the paper might benifit from expanding on this a bit.
    - It might also be useful to baseline the figures in Figure 10 using the reconstructions of an SAE at layer $t+1$, i.e. substituting $x^{(t+1)}$ with $W_\mathrm{dec}^{(t+1)} \sigma ( W_\mathrm{enc}^{(t+1)} x^{(t+1)} + b_\mathrm{enc}^{(t+1)} ) + b_\mathrm{dec}^{(t+1)}$; it is known that SAEs achieve nonzero reconstruction loss and that this incurs a performance penalty when using the reconstructions in place of the original activations. Intuitively the penalty incured from using the matched activations is 'mixed' with this base reconstruction penalty.
    - As stated before, it would also be good to include the 'matched' baseline in this section as well.

questions:
Could you provide more direct evidence for Hypothesis 2, or provide a clearer argument for why your current result provide evidence for it? Could you provide more discussion of what the takeaways for Section 5.5 were intended to be?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730704819265

thread_type:
official_review

<thread_object>
creator: reviewer_mxq3

<content>
summary:
The authors define a similarity metric between two learned encodings, based on permutations. They find that, specifically for the JumpReLU architecture, θ tracks the growing residual stream norm. Accounting for this observation by normalizing weight vectors with θ improves permutation matching. They find that W_dec MSE decreases substantially for high layers, while not making differences for early layers. They further employ an LLM to judge the semantic similarity of matched features. Interestingly, they find that feature matching works well in early layers (<10) as per MSE but show that feature descriptions do not match. Additionally, they compare matching features across multiple layers by exact match and layer-wise composition, and investigate SAE match as layer pruning technique.

strengths:
Main finding: Cosine similarity alone is not a great proxy for late layers, as residual stream norms increase. The authors propose parameter folding, which effectively addresses this problem for JumpReLU SAEs. Current work relies on cosine similarity, and I am convinced the field should adopt this proposed technique.

weaknesses:
### Critiques that can be addressed in this paper
- I am unsure whether the original hypothesis of permutation is answered. The term "matching" implies a binary measure of whether a feature mapping is true or false. This might require the introduction of a cutoff threshold, or applying a clustering technique. Otherwise, the framing of similarity measures might be clearer than permutation. I'm curious about the authors' opinion on whether there is a binary criterion for whether features do/don't match.
- To quantify the fraction of "true matches," the MSE over all features is too coarse; only the matching score distribution provides a clear picture of which fraction of features are "well enough" matches.
- Line 315: "...unfolded matching showed higher MSE in the scale of hidden state representations, supporting behaviour described with Hypothesis 2." Comment: This is not true for benc MSE.
- I'm unsure about the statistical significance of the LLM evaluation. I understand that 100 out of 16k (<0.1%) SAE latents (aka features) were chosen. Did the authors choose these features at random? Increasing the number of SAE features would increase the significance of their findings. Using an LLM to judge the coherence of max activating examples of features would be an insightful additional metric that would track the suspected polysemanticity in early layers.
- I do not agree that Figure 7 right reflects the findings in Figure 7 left, where the matching score suddenly increases at layer 10.
- The quantification of "too far apart" in Hypothesis 3 would be a useful improvement.

### Further suggestions (that can be addressed in future work)
How does the mechanistic permutability transfer to other SAE architectures? The only difference will be adapting scaling factors defined in Equation 4. It would be great to provide a generalized definition of parameter folding for all SAE architectures.

questions:
### Questions
- I'm curious about the takeaways of the experimental results from Section 5. How do I interpret a scale of ∆L about 1? Do the authors believe their results indicate that layers between 10 and 20 can be collectively pruned with the SAE matching method?
- I'd be interested in a discussion of why encoder matching performs worse than decoder matching.

### Further notes
- Line 154: Repetition of meaning in two following sentences.
- Why is Hypothesis 2 formatted as a hypothesis? It seems like a definition of a method to me.
- Line 234: Calling "average ℓ0" a regularization coefficient is misleading; calling it (average) sparsity is clearer.
- Line 235: Do the authors refer to Equation 3 when mentioning MSE? A reference of that equation or a different naming would be useful, since MSE is often used for the reconstruction loss in the context of SAEs.
I- 'm very curious how these results compare to findings with crosscoders (https://transformer-circuits.pub/2024/crosscoders/index.html). Matched features might share decoder vectors of a single crosscoder feature.
- Section 5.3: Gurnee operates on MLP neurons, which are not incentivized to be sparse. I expect SAE features to be (more) monosemantic, so Gurnee's results ideally shouldn't apply to SAE features.
- I'd be curious about a deeper investigation of the sharp increase in semantic feature similarity in Figure 4 left.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730755080779

thread_type:
official_review

<thread_object>
creator: reviewer_3w8e

<content>
summary:
The paper introduces SAE Match, a technique that produces a bijection from [sparse autoencoder (SAE) features in layer n] and [SAE features in layer n+1], with the objective of reducing the average squared distance (i.e., mean squared error (MSE)) between a feature and its target. They also introduce parameter folding, which seeks to augment SAE Match to work on features with different scales. The paper shows via experimentation on Gemma Scope SAEs  that features have more similar LLM-derived interpretations if they have lower MSE, and that SAE Match with parameter folding results in more similar interpretations.

strengths:
- The paper's SAE Match technique appears to work well at finding corresponding SAE features between layers, which can contribute to the goal of mapping a "feature circuit" (as in Marks et al. (2024) https://arxiv.org/pdf/2403.19647). This technique is data-free, meaning only the SAE weights are required, and not any model or SAE activations.

- The paper also provides useful empirical evidence that SAEs find features between layers that are simultaneously 1) close in the space of parameters, and 2) interpreted to have similar meaning.

weaknesses:
- The authors get far worse results on layers 0-9 of the model than on layers 10-25, indicating that the technique may not fully generalize. The authors claim that this is to be expected, saying "This phenomenon aligns with findings from previous research. Gurnee et al. (2023) also reported increased polysemanticity in the early layers of neural networks." This explanation is unsatisfactory because Gurnee et al. (2023) were working with LLM neurons, not SAE features. Additionally, Cunningham et al. (2023) found that earlier layers were more interpretable (see Figure 2 in https://arxiv.org/pdf/2309.08600).

- The proposed technique of parameter folding is only defined for SAEs using the JumpReLU activation function, and it is not clear how could be adapted to over activation functions like ReLU or TopK.

questions:
Questions for the authors:

1. Parameter folding serves as a form of normalization for encoder/decoder weights (as is mentioned in lines 156-157). What happens if one instead matches feature with encoder/decoder weights normalized to be unit vectors? This would be equivalent to maximizing cosine similarity instead of minimizing MSE, as ||x-y||^2=||x||^2+||y||^2 -2 ||x||*||y||*cossim(x,y).

2. One might expect the set of features to change across layers as the model processes information. In that case, the "correct" form for a matching might not be a bijection. Could one instead define a non-bijective "matching" via P(f)=argmin_i ||f-g_i||_2 where f is a feature in layer n, and {g_i}_i=1^k is the set of features in layer n+1? How does this compare to Feature Matching?

3. In Section 5.4, Figure 8, the y-axis is labelled "GPT Score". What is that metric and how was it calculated? Previous metrics have been a "GPT Score" on a scale of Different/Maybe/Same, or Matching Score in a range of 0-1, but this does not appear to be either of those.

4. In Section 5.5, it seems Layer Pruning may introduce error in two ways: 1) feature activations in layer N do not result in perfect reconstruction of the residual stream even for layer N, and 2) features in layer N do not perfectly match features in layer N+1. Did any experiment disentangle those effects? For instance, what is the \Delta L from replacing x with \hat x? This would provide useful context for the quantities shown in Figure 10.

5. Many SAEs have "dead" features, which presumably contribute little to the matching process. Could SAE Match be modified to exclude "dead" features, and if so how?



If it is permitted for authors to make revisions before the final submission, there are several small changes that could improve the quality of the paper:

- Line 93: The loss function as written is incorrect. The L2 term is squared, and what is written as L0 should be L1. See e.g. Equation 4 in (Cunningham et al, 2023) (https://arxiv.org/pdf/2309.08600). 

- Line 132-133 (equation 3): The L2 norm in the argmin needs to be squared to get the mean *squared* distance.

- Lines 147-149 (equation 4): It appears that b_dec should be b_enc.

References:

Wes Gurnee, Neel Nanda, Matthew Pauly, Katherine Harvey, Dmitrii Troitskii, and Dimitris Bertsimas. Finding neurons in a haystack: Case studies with sparse probing. Trans. Mach. Learn. Res., 2023, 2023. URL https://openreview.net/forum?id=JYs1R9IMJr.

Hoagy Cunningham, Logan Riggs Smith, Aidan Ewart, Robert Huben, and Lee Sharkey. Sparse autoencoders find highly interpretable features in language models. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum? id=F76bwRSLeK.

Samuel Marks, Can Rager, Eric J. Michaud, Yonatan Belinkov, David Bau, Aaron Mueller. Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models. 2024. https://arxiv.org/abs/2403.19647
</content>
</thread_object>

</thread>
<thread>
creation date:
1732188392457

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We would like to highlight the changes made in the rebuttal revision of our paper:

- **Updated the list of used SAEs** due to changes in Gemma Scope. This update affected only the two initial layers, and all relevant results have been reevaluated accordingly.
- **Reformulated the explanation for performance on initial layers in Section 5.3** to attribute it to differences in L0 norms between matched SAEs, and we have provided supporting results.
- **Added experiments evaluating different thresholds of matching based on MSE values**, which further support Hypothesis 1.
- **Expanded Section 5.5 with new experiments** to enhance our findings.
- **Included evaluation with LLaMa**, demonstrating that the proposed method can be applied not only to the Gemma model.
- **Made other minor improvements in wording** as suggested by the reviewers.

All changes are highlighted in red for clarity.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper introduces a novel method for mapping Sparse Autoencoder features across adjacent layers, which demonstrates significant improvements in feature correspondence and interpretability compared to traditional cosine similarity metrics. Empirical results showcase its effectiveness with reduced MSE and semantic similarity evaluations, while visualizations and thorough discussions enhance the method’s clarity and accessibility. Reviewers highlight the potential applicability of SAE Match across diverse tasks, particularly in advancing sparse representation and interpretability research. On the other hand, focus on a single model (Gemma) and activation type (JumpReLU) limits generalization. Some claims, such as those related to Hypothesis 2, would benefit from stronger empirical evidence and comparison with alternative matching methods. Nevertheless, SAE Match is recognized as a valuable and innovative contribution, with significant potential to inform future advancements in feature alignment and complex model interpretability.