Transcoders Find Interpretable LLM Feature Circuits

Summary

The paper introduces the use of transcoders a tool for mechanistic interpretability as a replacement for SAE. The main difference between SAE and transcoders is that SAE uses autoencoders to take the output of an MLP and reconstruct it while transcoders (that can also be viewed as an encoder-decoder architecture) take the input of an MLP and reconstruct the output. The paper compares SAE and transcoders in terms of sparsity (L0 Norm of the features), faithfulness (matching the output of the original MLP), and human interpretability and show that they in par with SAE on multiple-size models. The paper shows how transcoders can be used to perform weights-based circuit analysis through MLP sublayers by doing the following: - First identifying attribution between feature pairs: - Each feature in a transcoder is associated with two vectors: the i-th row of $W_{enc}$ is the encoder feature vector of feature i, and the i-th column of $W_{dec}$ is the decoder feature vector of feature i. The i-th component of zTC(x) is the activation of feature i. - Assuming $l$ is in an earlier layer than $l'$, they calculate the contribution of feature i in transcoder $l$ to the activation of feature i' as the product of two terms: the earlier feature’s activation (which depends on the input to the model) $z^{(l,i)}_{TC}(x^{(l,t)})$ and the dot product of the earlier feature’s decoder vector with the later feature’s encoder vector (which is independent of the model input). - Second finding computational subgraphs: - On a given input and transcoder feature i', one can determine which earlier-layer transcoder features i are important for causing i' to activate, from the previous step. - Once the earlier-layer features i that are relevant to i' are identified, they recurse on i to understand the most important features causing i to activate by repeating this process. - Doing so iteratively yields a set of computational paths. - These computational paths can then be combined into a computational subgraph, in such a way that each node (transcoder feature), edge, and path is assigned an attribution. - They defined a de-embeddings: - A de-embedding vector for a transcoder feature is a vector that contains the direct effect of the embedding of each token in the model’s vocabulary on the transcoder feature - The de-embedding vector for feature i in the layer l transcoder is given by $W_E^Tf_{enc}^{(l,i)}$, where $W_E$ is the model’s token embedding matrix. This vector gives input-invariant information about how much each possible input token would directly contribute to the feature’s activation. - Given a de-embedding vector, looking at which tokens in the model’s vocabulary have the highest de-embedding scores tells us about the feature’s general behavior. The paper showed a blind case study of how they reverse engineer a feature in GPT2-small and analysed GPT2-small “greater-than” circuit

Strengths

### Novelty: - While transducers have been previously introduced this paper is the first to apply them to large models and analyze their performance in various settings. - The circuit analysis presented in the paper is novel in the sense that it disentangles input-invariant information from input-dependent information. ### Presentation: - The paper is very well written and easy to follow. ### Application: - The paper showed how the transducers can be used to reverse engineer features in GPT2-small and they also analyzed GPT2-small “greater-than” circuit comparing transducers to neuronal approach showing that transducers have better performance.

Weaknesses

### Method: Circuit Analysis: Assuming Transcoders perfectly emulates MLP (which it tries to but as long as faithfulness loss > 0 it doesn't) the circuit analysis graph is between Transcoders in different layers but in the actual model there is multiple-head attention in between each MLP layer which is not accounted for its effect in equation 6. ### Evaluation: The evaluation was weak at best. The paper did compare with SAE but the only actual quantitative results were the difference between SAE and transcoders in terms of sparsity and faithfulness which only shows the transcoders are on par with SAE. The more important experiment is the interpretability experiment since this is why both methods were developed in the first place for these experiments were not systematic and many of the choices were not justified or explained properly.

Questions

Section 3.2.1: - For the blind interpretably experiments why was layer 15 in particular chosen? - From the following sentence: "We recorded for each feature whether or not there seemed to be an interpretable pattern, and only after examining every feature did we look at which features came from where." - How did you identify that there is an "interpretable pattern"? - How many people performed this experiment? Section 4.1: - For attribution and graph calculation for a given feature i' at layer l' are all others feature at previous layers considered in the graph or only the same feature at the previous layer?

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

Yes, they have.

Summary

The paper compares transcoders to SAEs for interpreting parts of GPT2-small, Pythia-410M and Pythia 1.4B. Transcoders are SAEs trained to replace a particular MLP in the original model instead of implementing an identity. They find that the transcoders they train outperform the baseline SAEs they train on the pareto-curve of sparsity, measured by the L0 pseudonorm, with faithfulness, measured by the CE loss of the model recovered. They describe the attribution technique they use to attribute features to earlier features and other network components. They apply this attribution technique to interpret feature activations without access to the network inputs. They also use it to analyze the GPT2 "greater than" circuit investigated in https://arxiv.org/abs/2305.00586.

Strengths

I think transcoders are maybe the most important idea for a new interpretability technique to investigate at the moment, and this is the first experimental study of them on language models I’m aware of. SAEs have a potential issue of finding features that sparsely decompose the input data distribution rather than the computations of the neural network itself. Transcoders seem like the simplest potential solution to this problem since they essentially decompose the operations the network performs instead of its activations. I like that they showcase using the technique to reverse engineer a circuit in section 4.3.

Weaknesses

In Appendix E line 1050, they state that the baseline SAEs they compare their results against were trained on MLP inputs rather than MLP outputs. I think this is the wrong comparison point. The transcoders only need to reconstruct the new features computed by a single MLP, while an SAE trained on MLP inputs needs to reconstruct all features in the residual stream that are in the preimage of the MLP input matrix. This is potentially a lot more features, which might give the SAEs an inherent disadvantage. Training the SAEs on the MLP output activations instead would be a more appropriate comparison. I think this is the biggest weakness of the paper. In section 4.3 and figure 5, they use neuron-based analysis as a baseline to compare their transcoder feature-based analysis to. I think this is a very weak baseline. Using e.g. activation PCA components would be more appropriate. This is the main result they have comparing their technique to standard techniques outside the cluster of techniques based on SAEs/sparse decoding, so I think not using an appropriate baseline here is not great. I don't think it is a very major flaw, because this is mainly a paper that takes it as given that SAEs are interesting and attempt to improve on them, rather than a paper trying to justify the sparse decoding approach to interpretability. They describe the attribution method between transcoder features they use as a 'new technique', which I don’t think is justified. It's effectively the attribution patching technique https://arxiv.org/abs/2310.10348 , https://arxiv.org/abs/2403.00745, used e.g. in https://arxiv.org/abs/2403.19647. Transcoders just have the advantage of making it work more elegantly. Any ReLU MLP network attribution patching is applied to would yield equation 6 as well. In section 3.2.1 and table 1, manual comparison on fifty features isn’t a very large sample size for this sort of analysis. Supplementing this with e.g. a larger set of LLM generated autointerpretability scores would have been better, as was done in https://arxiv.org/abs/2309.08600. Personally, I don’t think not investing more effort into this is a major flaw, because I think reconstruction error and L0 are more important metrics than human interpretability scores. In Figure 2, they present reconstruction scores for transcoders and SAEs in terms of raw CE-loss recovered scores. I think it would have been good for the presentation to also show the recovered performance in terms of the compute required to train a model with the same CE score, as suggested in https://arxiv.org/abs/2406.04093. Raw CE scores can make the performance gap to the model look smaller than it is, since the difficulty of further reducing CE loss grows as models become more performant.

Questions

Why train the baseline SAEs on the MLP inputs instead of the outputs? This way, the SAEs need to reconstruct all features in the residual stream within the MLP preimage, while the transcoder only needs to reconstruct the new features computed by the MLP. In section 4.3, why use neuron-based analysis as a baseline? Neurons tend to be very brittle under ablations in my experience. Why not e.g. activation PCA components?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

4

Limitations

All addressed.

Summary

Sparse Autoencoders (SAEs) have been used for interpretability of transformer neural networks. SAEs take the output activations of MLPs within transformers and learn are trained to re-construct those outputs: $SAE(f(x)) \approx f(x)$. In this submission, the authors study using transcoders for a similar task. Unlike SAEs, transcoders are trained to imitate the MLPs directly: $TC(x) \approx f(x)$. The authors demonstrate via experiments and blind trials that the performance of transcoders is competitive with and sometimes exceeds that of SAEs.

Strengths

- Transcoders are a natural and unexplored extension of the SAE architecture that perform very comparably to traditional SAEs. - Deriving the decomposition into input dependent / invariant features is novel and addresses a significant challenge of interpretability. On its own this warrants future exploration. - The feature case studies are a valuable contribution to the literature. - Testing interpretability via blind feature tests is innovative.

Weaknesses

- The concrete comparisons to traditional SAEs are not sufficiently thorough. Experiments demonstrate that the KL divergence of transcoders are competitive compared to SAEs, but my understanding of the SAE literature is that it is still an open question as to how to judge the quality of an SAE. There are, however, other questions that are more concrete but are unaddressed. e.g. how easy/stable are transcoders to train compared to SAEs? Do transcoders suffer issues with "dead neurons" when the width is scaled up? Do transcoders show similar "feature splitting" effects when the width is scaled up? A more holistic comparison to SAEs would strengthen the submission significantly. - While performing a blind test of the interpretability of different features is novel, I worry that it is prematurely rigorous. - One of the implicit assumptions of the blind interpretability experiment is that the features learned by transcoders and SAEs are different. Exactly what a "feature" is is still being worked out in the literature, but if transcoders and SAEs truly do learn distinct sets of "features" that is an important fact that should not have been left out of the paper. If the authors did not investigate the extent of the overlap between the features that SAEs and transcoders learn that is unfortunate, but seems like a promising avenue for future investigations.

Questions

- What is the effect of different levels of $L_1$ penalty on the features that the transcoder learns? Do transcoders also experience "feature splitting" as the width of the transcoder goes up? - Are transcoders easy to train? Are the learning dynamics stable? Do they experience issues with "dead neurons"? - Does the input invariant / dependent construction not work for the features that SAEs learn? - When looking at - Do SAEs and transcoders trained on the same layer and same data learn the same or similar features? - Have any attempts been made to perform a similar analysis for other circuits? e.g. the [IOI](https://openreview.net/forum?id=NpsVSN6o4ul) or [Gendered Pronoun circuits](https://cmathw.itch.io/identifying-a-preliminary-circuit-for-predicting-gendered-pronouns-in-gpt-2-smal)?

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors adequately address the submission's limitations.

Summary

Motivated by prior claims on how MLP sublayers make interpretability challenging (arguably due to their extremely dense nature), this paper proposes "Transcoders", a protocol that is in spirit similar to Sparse Autoencoders (SAEs). Specifically, a Transcoder aims at "faithfully" matching the output of a layer (in this case, an MLP layer); in contrast, an SAE aims to reconstruct the input to a layer. Results show Transcoders are similarly effective when compared to SAEs (in terms of model loss). Authors follow a blind evaluation pipeline to see if Transcoders yield interpretable features, whereby hypotheses are made for what a feature means and then posthoc the hypothesis is tested by looking at input samples.

Strengths

I enjoyed reading the paper. It's written fairly clearly and the targeted problem is well described, though some more expansion on how MLPs are a challenge for interpretability would help (see weaknesses). The blind evaluation protocol was really good to see---I appreciated authors' efforts in stress testing the limits of their approach (this is how good science should go :)).

Weaknesses

While I like the paper, I think it needs a bit more work to be ready for acceptance. To that end, following are a few suggestions. - **Expand on challenges in MLP interpretability.** This is the bulk of the motivation of the paper and, from what I can gather, is grounded in off-hand references in two prior works. I don't think those works thoroughly describe what the challenges in fine-grained circuit analysis of MLP sublayers are. It is fine to use those works as motivation, but I think the paper should arguably start off with an experiment to demonstrate the challenge more thoroughly, given that prior work hasn't done that. In fact, I would argue that your own results with SAEs show similar interpretability may be achievable as found via Transcoders (since SAEs perform similar in terms of model loss). From that perspective, it isn't established if the question being addressed is sufficiently a challenge - **Analysis is correlational, lacking a quantitative confirmation and causal interventions.** While I really like the blind evaluation pipeline, there are two missing elements in the results per my opinion. These render the current analysis somewhat informal. *First*, once a hypothesis is formed, one should ideally run an algorithmic test to see on what proportion of a dataset where the hypothesis predicts a feature will activate does the feature activate. From what I gathered, the results merely involve "eyeballing" a few inputs to see whether the hypothesis holds. *Second*, the experiments are correlational, not causal. Ideally, one would systematically intervene on the input and see if the claimed hypothesis breaks in a predictable way. For example, when one identifies a feature activates for semicolons being present in the input, then if that input is kept the same but the semicolon is dropped or another relevant punctuation is added, does the feature not activate or at least activate less? - **Figure 2.** I found Figure 2 to be rather unintuitive. The axes labels are quite small, so I thought the loss increases along y-axis and hence Transcoders underperform SAEs, but then realized that the axis is decreasing along y-axis. I would have preferred a more standard plotting schematic, or at least larger labels to help avoid misinterpretation. **Post rebuttals update.** Taking into account the context provided by the authors and reading through the paper again, I am happy to raise my score. My primary concern was the paper was pitched to solve a very specific problem, i.e., interpreting MLPs. Authors' response indicates this was merely part of the problem they set out to address: identifying circuits in an input-invariant manner, since the interaction between circuit discovery and input-sensitive features makes that pipeline difficult. Overall, I understand the pitch much better now. The paper intro (and other relevant parts) should be updated to accurately reflect authors' motivation. I do not see these rewrites as a major challenge, so I'd be happy if the paper is accepted.

Questions

- From what I can gather, the "de-embedding" tool is equivalent to logit lens? Can you expand on what's different? I understand this is not a core contribution; just wondering why a different term was used. - I am still struggling to understand the benefit of Transcoders over SAEs. Can you expand on this more? It's good to define alternative tools, but unsure if this was a sufficiently different tool and led to much benefits.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

See weaknesses.

Response

Thank you to the authors for their responses and updates. The new plots comparing MLP-out SAEs fully address what I considered the primary weakness of the paper. All the other weaknesses I noted are less important. This raises my rating of the paper's soundness. > "What we consider to be a contribution of ours is our recognition that input-times-gradient, when applied to ReLU MLPs such as transcoders, does factorize into an input-dependent and an input-invariant term – i.e. the gradient is constant. Furthermore, something that we like to underscore when thinking about Equation 6 is that for transcoders (and this is not true in general for ReLU networks), the input-dependent term (the feature activation) is interpretable (because feature activations are largely interpretable)." The gradient is only constant on the subset of data where the target feature is active. Outside of that subset, it is 0. How and when the threshold is crossed and the target feature becomes active depends on the interactions between the input features. This information is important for understanding the network and cannot be straightforwardly read off from what you call the input-independent term. Meaning the input-dependent term is actually the input ‘feature’ activation times either 1 if the target feature is on, or 0 if the target feature is off. This quantity may be a lot less interpretable than a feature activation.