Not All Language Model Features Are One-Dimensionally Linear

Thank you for your thoughtful comments, we are glad you found our paper interesting!

All changes we discuss below are in the new submitted revision of our paper. Where possible, we have highlighted changes in the revised document that are specific to your comments in teal.

The clustering method was a little hard to follow for me at first, maybe pseudocode would be helpful (since it seems like it wouldn't be many lines).

Thank you for this excellent suggestion! We have added an algorithm block to Appendix F, as well as a link to this pseudocode from Section 4. Please let us know if anything is still unclear!

Is it mostly just circles? Given that they are still a one-dimensional variable, I'm wondering whether there are representations of variables with (irreducible) intrinsic dimension higher than one (like a spherical variable), which would seem to make better use of the scope of your Hypothesis 2. Otherwise, maybe the linear representation hypothesis has been violated in name but not in spirit (i.e. representation is still factorized into one-dimensional variables)?

Our method is limited to finding two-dimensional representations, so we do not identify spheres or other 3D objects. That being said, we have added the top ranked clusters to Appendix G (Figure 16) that our method discovered, and some of them look like they might have 2D intrinsic structure, although it is hard to tell. Additionally, recent work has found a globe representation in an LLM[1], so we believe these higher dimensional features certainly exist. We have added a comment about this in our discussion, thank you for the suggestion!

[1] https://www.lesswrong.com/posts/guNzr32FC6DYCaSgJ/there-is-a-globe-in-your-llm

Thank you for your helpful comments and feedback!

All changes we discuss below are in the new submitted revision of our paper. Where possible, we have highlighted changes in the revised document that are specific to your comments in blue.

Responses to questions:

As you mentioned, the title is confusing. How about 'Not All Language Model Features are One-dimensionally Linear'?

This is a great suggestion, thank you! We have made this change.

Why do you use the term 'point cloud' in Definition 1? Does it mean that the range of the map f is a discrete set?

Thank you for pointing this out, we agree this is confusing and have removed it, a feature is now simply defined as a function from the token input space to the model’s hidden space.

What is the source of randomness for $p(a,b)$ in Definition 2? Does it come from the distribution of input texts (in training set for LLMs)? Also, why do you state that statistical independence and separability are equivalent?

Regarding randomness, yes, it comes from the input text distribution $\mathcal{T}$. We have made this more clear by describing $\mathcal{T}$ explicitly at the beginning of Section 3. Regarding the equivalence of statistical independence and separability, we are just defining separability to mean that the distribution can be decomposed into (lower dimensional) independent distributions, which we believe is a common choice for the meaning of “separable distribution”. Thank you for pointing both of these out, let us know if anything is still unclear!

In Figure 2, the histogram appears to show a mixture. Why did you say that the $\epsilon$-mixture index of 0.4750 is low? What is the baseline? In Figure 10, you state that an index of 0.3621 indicates the feature is likely a mixture. This seems contradictory.

In Figure 2, we are comparing to other GPT-2 clusters that we have discovered, where an $\epsilon$-mixture index of 0.4750 is indeed low (see Figure 3). In the appendix we are working with synthetic distributions that are cleaner, so the mixture measures are lower (e.g. they do not have a circle of noisy points very near the origin). We agree this is confusing and have tweaked the example to be more of a mixture, so now the measured $\epsilon-mixture$ index in Figure 10a is $~0.64$. Thank you for pointing this out!

In line 269, why do you state that 'These dictionary elements will then have a high cosine similarity'? I believe they could have low cosine similarity even if many elements span the feature.

It is true that not all pairwise vectors in the space will have high cosine similarity, but we believe that adjacent pairs radially will likely have high cosine similarity since it is not possible to “pack in” many features in 2D all of which have pairwise low cosine similarity. Could you be more specific about the setup you are worried about?

In Figure 1, how were the points labeled? What does the label 'Monday' mean? Does it indicate that 'Monday' would be the next token in the corresponding texts?

Our apologies for not being more clear, the label corresponds to the current token, we have added this clarification to the Figure 1 caption.

In Section 5, you use specific texts and train new probes. Why didn't you use the circular representations previously identified by SAE clustering?

This is a great point, and is a result of the order in which we originally ran experiments. Empirically, the two planes (one from the probe, the other from the PCA of the SAE cluster) are highly similar. The principal angles between the planes are 45 and 50 degrees, which means that projections of basis vectors from one plane into the other maintain > cos(45) = 66% of their norm. Thus we expect the results of this experiment to be similar, although we have not run this experiment.

If our comments resolve your concerns we would appreciate if you would consider raising your score!

Thank you for your thoughtful comments, we are glad you found our paper interesting!

All changes we discuss below are in the new submitted revision of our paper. Where possible, we have highlighted changes in the revised document that are specific to your comments in orange.

Responses to weaknesses:

If I understand correctly, the proposed algorithm can extract interpretable features, however, I imagine there is a good amount of them that is not easily--if at all--interpretable. Adding some of your insight on how many potentially interesting multi-dimensional features are among the ones extracted from the algorithm could improve the article.

We agree this is an interesting potential insight! We have added some of the other features discovered in Figure 17 in Appendix G. Most of them seem interpretable in terms of the tokens that they are composed of (common words and single letters), but their structure is not entirely interpretable. We believe this is important future work!

If I am not mistaken, there is not an ablation on how much the threshold parameter T affects the extracted clusters: Having an understanding of when circular features start to emerge could be an interesting experiment to add, and may potentially lead to some useful insights.

Thank you for this excellent suggestion! We have added an experiment in Appendix F2 (Figure 16) that varies the threshold (and the k) for the Mistral clustering procedure and sees when the days of the week cluster is identifiable. We find that small T causes all features to group in one cluster, so the days of the week cluster is not found; medium T causes the days of the week cluster to become identifiable; and large T causes all features to be divided into their own clusters.

Figure 2: I would change the color for the PCAs scatter plot from black to something with less contrast (e.g. the green used in the angle subfigure of Figure 2)

Thank you for this suggestion! We’ve made the scatter points grey with alpha=0.3 with a black edgecolor. This reduces the contrast and makes the relative density of the points easier to see.

Figure 3: This is a very minor complaint, but I would change the scatter icon for the "other cluster" points to something different than the one used to indicate the clusters shown in Figure 1, for example, something like an x would probably look a bit better.

This for this suggestion too! We’ve marked them with a slightly larger x, and also made the orange color darker so they are easier to see.

Another minor issue is that the captions for Figures 5 and 6 are a bit lacking, adding some more information to help understand their respective feature would probably be helpful.

We have added more information to these captions, thank you for this suggestion!

Responses to questions:

Do you have any qualitative insights on how big the threshold parameter T should be in order to find multidimensional features?

Now we do, see our response above, thank you! Empirically, T should be set such that it is right on the boundary between everything being one large cluster and everything being in its own cluster.

Are there any extracted multi-dimensional features that show some interesting behavior, like the circular one shown in the paper? If not, do you think is it due to a structure difficult to interpret, or simply because probably most of the features are still one-dimensional?

This is a good suggestion, thank you! We did not do interventions on other multi-d features because their structure seems harder to interpret (see Figure 17 for more of these examples), but it is possible that further study of e.g. belief states (see [1]) will help us understand these structures.

[1] https://www.lesswrong.com/posts/gTZ2SxesbHckJ3CkF/transformers-represent-belief-state-geometry-in-their

Responses to questions, part 2:

p. 8, "In practice, ...": this is somewhat confusing/concerning to me. If there are alternative pathways that may affect both the circular feature and the task output, then wouldn't they "wash away" the causal implication of the feature? In other words, in what sense exactly is this a "causal" implication rather than correlational?

Could you clarify what you mean here? Certainly average ablating the other dimensions is a weaker form of subspace patching then just patching the circle and leaving the other dimensions unchanged, but it still seems that is is fundamentally a causal intervention. The idea is to show that limiting the days of the week representation to the circle maintains almost all of the performance for the task and results in an interpretable representation we can intervene on.

Broad question: are there examples of multi-dimensional features that aren't circular?

Yes, see for example this recent work from Shai et. al. finding fractal belief states [1] (although in a toy transformer trained on a Hidden Markov Model) and additional recent work finding a spherical globe in an LLM [2]. There is also recent work on “Onion-like” nested features (although in a toy RNN trained on regression). All of these representations are irreducible multi-dimensional features by our definitions. We have added a discussion of [1] to the related work, and will add a discussion of the “Onion-like” features in the non-anonymous version of the paper (they cite our work, so citing it here would break anonymity).

[1] https://www.lesswrong.com/posts/gTZ2SxesbHckJ3CkF/transformers-represent-belief-state-geometry-in-their

[2] https://www.lesswrong.com/posts/guNzr32FC6DYCaSgJ/there-is-a-globe-in-your-llm

If our comments resolve your concerns, we would appreciate if you would consider raising your score!

Responses to questions, part 1:

Figure 1: is there any intuition for why specifically we find these features in the second and third PCA components?

Yes, the first dimension is typically an “intensity” direction that manifests as the radius in the circular plots. Thus the actual multi-dimensional structure might best be thought of as a cone. We have added a discussion of this in section 4, thank you for inspiring this clarification!

Definition 1: What exactly is a df-dimensional point cloud in Rdf? Is there a mathematical definition? (Is df redundant?) If I have two points in R3, what is the "dimension" of the point cloud they form?

We intended for $d_f$-dimensional point cloud to refer to a distribution that is “full rank,” i.e. one that does not lie on a lower dimensional subspace. We agree that this is not necessary, however, and so we have omitted it in the revision.

Also, to be precise, should I understand the "probability distribution over input tokens" as that of the natural language? I also assume that this probability is restricted to n-length token sequences---is that correct? Later, in the definition of the two indices, the distribution comes up anyway as P. It may be helpful to formally define this notation/distribution up front.

Thank you for asking us to clarify this! Yes, this is correct, we have added a term $\mathcal{T}$ that represents the input text distribution and define this explicitly up front at the beginning of section 3. Note that $\mathbb{P}$ defined in equation 3 is meant to denote the “probability” operator over $\mathcal{T}$; we have made this more explicit in the current revision.

Definition 2: in the mixture definition, what is the non-bold b? Also, does "lower dimensional" simply mean $k<d_f/2$ or does it mean low-dimensional, say, $k \ll d_f$?

The non-bold b is a typo and should be bold, thank you for catching that! Lower dimensional just means $k < d_f$. For example, consider a three dimensional feature that could be decomposed as mostly lying on one or more underlying planes; then those planes should be the features.

Can you efficiently compute the separability index for features that are beyond two-dimensional?

In theory one could compute a separability index measure for higher-dimensional features. For instance, if the proposed feature was 3-dimensional, one could project onto different orthogonal lines and planes and measure the mutual information between the projections. However, in higher dimensions, both optimizing over all possible combinations of orthogonal lines and hyperplanes and computing the mutual information between higher-dimensional features may be challenging. Thus, in practice, we average the separability scores for adjacent pairs of PCA dimensions; this is not as principled but is more computationally tractable.

Hypothesis 2: while this is a stylistic point, I think the use of a possessive in "Our Superposition Hypothesis" is unnecessary and un-scientific. It could have any other objective name, say, the "Multi-Dimensional Superposition Hypothesis". See, e.g., Knuth et al., page 2 item 6.

We agree and have updated the name to your suggestion, thank you!

p. 5, lines 260--270: doesn't this claim depend heavily on (1) the two-dimensional feature being active in the LLM's representation space and (2) the SAE is large enough to capture this feature? I'm not even sure if this claim is necessary, as opposed to saying "some of SAE clusters represent irreducible features".

Thank you for pointing this out, we agree that it does rely on these things and have updated the text to make this more clear. We still think it is useful to provide intuition for why we might expect clusters of SAE features to represent multidimensional features. Does the updated text make this more clear?

Figure 3: what are some of the other features that are shown here to be more reducible than the features for Fig. 1?

We have added a selection of these in Figure 16, and an accompanying discussion in Appendix G. These clusters seem to mostly represent letters or common words. We hypothesize these may perhaps be related to the “belief states” discovered by Shai et. al. [1], although we acknowledge more work is needed to understand them.

Are there any other features that we may intuitively view as circular but not found by the SAE? Say, are color words (red-orange-yellow-green-blue-purple) represented in a circular fashion?

This is an excellent question; we have not investigated this, but we suspect there are some, as SAEs are known to have trouble representing continuous quantities. We have added some discussion of this to our conclusion.

Thank you for your thoughtful and helpful comments!

All changes we discuss below are in the new submitted revision of our paper. Where possible, we have highlighted changes in the revised document that are specific to your comments in red.

Responses to weaknesses:

First, there can be concepts that are correlated but can still be disentangled for separate interventions, so equating separability with independence can be limiting.

Thank you for pointing this out! We agree that focusing on geometric and statistical based definitions might not tell the full story, but these definitions are helpful in that they can be applied to real features across a large corpus of text, as we do in the paper (although see Appendix C where we propose an intervention-based definition). We have added this point to our limitations discussion section in the paper.

Second, I can't tell for sure if the definition of reducibility is exhaustive. (Can there be a third category?)

Our current definitions prevent multiple one-dimensional features that are either unrelated or anti-correlated from being treated as a multi-d feature (via separability index and mixture index, respectively). These cases seem like a natural breakdown of the types of statistical reducibility; if two features are neither of these things, then their co-firing pattern may have interesting structure.

While the paper's main motivation is to call the LRH into question, the proposed alternative hypothesis appears to be a stricter version of the LRH (p. 5, lines 239--240), rather than a contradiction to it. This makes the overall message a bit confusing. (I'd be happy to be corrected if this is a misunderstanding.)

This is an excellent point! As suggested by Reviewer 9dQp, we have renamed the paper “Not All Language Model Features are One-Dimensionally Linear”, which we hope makes it more clear we are not claiming the entire LRH is false. We state that the multi-d hypothesis is a stricter version of the LRH to describe how the definitions are related, but technically any possible distribution of vectors could be seen as a linear combination of one dimensional vectors (e.g. the original neuron basis), it just might not be that sparse or mono-semantic, so it is hard to “contradict” the LRH. The multi-d hypothesis might be a more “true” breakdown of the space, in that the features are monosemantic and overall sparsity is lower.

The proposed "tests" of irreducibility are not completely formal, in the sense that they are not statistical tests of significance but rather (possibly arbitrary) thresholding of a pair of metrics. This could mean that whether or not a feature is irreducible is ultimately a judgment call by the user.

We agree and this is an important point, thank you! We have added a discussion of this to the limitations section in our discussion.

I find that Section 3 is not easy to follow. The only visual aid is Figure 2, but it is not referenced or intuitive. E.g., why exactly are we looking at PCA 2/3, even for Fig. 1? What is the histogram in the middle, and what is the angle on the right? What are "typical clusters" and what are their mixture/separability indices?

Thank you for pointing this out! We have added further explanation to Figure 2 describing each part of the figure and further explanation at the end of Section 3.1 that describes the implementations of the tests. Please let us know if this is still unclear! Regarding looking at PCA 2 and 3 instead of 1 and 2, we now describe this in section 4, see our response to Q1 below for more details.

While some of the plots are great, the experiments are quite dense and unorganized. The exposition would often refer to a figure in the Appendix, while some figures (2 and 4) in the main text are not referenced. This makes it difficult to read the paper linearly.

We have now added references to figures 2 and 4, and have also added motivations to sections and subsections that were missing them. Thank you for suggesting this improvement!