Observable Propagation: Uncovering Feature Vectors in Transformers

Thank you very much for your review, and for recognizing the value of observable propagation in allowing for feature vectors to be found. With regard to your comments, we address them individually below:

• “The paper should directly give a formal definition of an "observable", state how it is used/useful, and show how it is computed. Instead, the paper gives a long narrative of several motivations in words.”
  • We have added a formal definition and explanation of an observable in Appendix A. To summarize, an observable is a linear functional from $\mathbb{R}^{\mathtt{d_vocab}} \to \mathbb{R}$, where $\mathtt{d_vocab}$ is the size of the model vocabulary. We refer to the action of taking the inner product between the model’s logits outputs and an observable n as getting the output of the model with respect to observable n. Intuitively, observables correspond to tasks on which we want to measure model output. As an example, for the subject pronoun prediction task, we want to measure the extent to which the model’s logit prediction for the pronoun “she” is greater than the logit for the pronoun “he”. As such, an observable corresponding to this task is given by $e_{\mathtt{” she”}} - e_{\mathtt{” he”}}$ where $e_{\mathtt{token}}$ is a one-hot vector of size $\mathtt{d_vocab}$ with a one in the position corresponding to token. This is because the inner product between this observable and the model’s predicted logits vector precisely yields the logit for the pronoun “she” minus the logit for the pronoun “he”. We hope that this explanation has better clarified our approach.
• “The theorem about how the layer norm doesn't substantially the direction of feature vectors seems pretty specific to the narrative here; it is hard to see much general value out of the context.”
  • The reason why this theorem is important is because it implies that the general approach of finding feature vectors — something considered critical to the mechanistic interpretability community (e.g. please see Olah (2022), “finding and understanding interpretable neurons [a type of feature vector] […] isn’t just one of many interesting questions. Arguably, it’s the central task [of mechanistic interpretability]”) — is a useful means of understanding how Transformers process information, even in the presence of ubiquitous LayerNorm nonlinearities. After all, if LayerNorms substantially affected the behavior of feature vectors, then this would imply that the entire idea of using linear feature vectors is ill-founded, because nonlinear LayerNorm effects would dominate the behavior of the model, even in linear attention-only circuits. Thus, being able to show that this is not the case provides theoretical backing for this entire paradigm.

We hope that we have been able to successfully address your comments; if you have any further questions, please do not hesitate to let us know. Again, thank you for your review.

Reference: Chris Olah (2022) Mechanistic Interpretability, Variables, and the Importance of Interpretable Bases. https://transformer-circuits.pub/2022/mech-interp-essay/index.html

Thank you for your detailed review and comments. We are happy to see that you recognize the ability of observable propagation to recover meaningful feature vectors, and the broader implications that this has for the field of mechanistic interpretability.
We will address your comments individually below.

“This method seems like an oversimplification of the architecture of an LLM. The ObProp algorithm propagates through the layers by approximating almost all non-linearities by their first-order gradients.”

We find that oftentimes, useful circuits in Transformers involve either very few nonlinearities or none; additionally, these nonlinearities often display local linearity in input regions that we care about. For example, the circuit corresponding to the attn6::6 (head 6, layer 6) feature vector for the subject pronoun case contains only attention sublayers; this is also the case for the circuits corresponding to the feature vectors used in the recently-updated political party prediction experiment and the C vs. Python classification experiment.

As for circuits that contain MLP sublayers: we investigated the accuracy of the first-order linear approximation of the circuit consisting of MLP16 → MLP18 → final LayerNorm → unembeddings on the occupation prediction dataset; this is because this subcircuit is the nonlinear component of the circuit for the attn6::6 feature vector for the occupational bias observable. We found that on this dataset, the root mean squared error of the linear approximation was approximately 0.4139 logits, which is only approximately 3.5% of the mean output value of this circuit. This suggests that, at least on specific tasks with constrained input distributions, first-order approximations of MLP nonlinearities suffice. We have updated Appendix D with further information.

We recognize that the observable propagation algorithm does indeed involve a simplification, particularly in its approximation of MLP nonlinearities using first-order gradients. This approach was chosen to strike a balance between computational tractability and analytical insight. However, the method's utility extends beyond these initial simplifications. The paradigm introduced by observable propagation is a preliminary step towards more comprehensive mechanistic analyses of MLP nonlinearities. By enabling the backward propagation of feature vectors through the network, our method provides a novel framework to explore how MLP nonlinearities interact with different tasks, not just with varying inputs.

We believe that the position of observable propagation in its current state can be compared to the seminal transformers circuits work done by Elhage et al. 2021, which ignored MLP layers entirely. The assumptions made in that work were far more simplifying than our treatment of MLP linearization, but time has demonstrated that their work has born substantial fruit in pushing forward our knowledge of mechanistic interpretability of complete models. Indeed, the recent work by Bricken et al. 2023 on dictionary learning of features takes a similar approach by only analyzing one-layer transformers, a far more simplified scenario than ours, with this same aim of later providing more complete insights. We hope that our work can do the same.

“Moreover, the work abstracts out QK circuits in LLMs, i.e. the attention scores are treated as constants and not as non-linearities, therefore the problem is oversimplified and I'm not sure the theorems shown are necessarily that insightful. As the authors note, their method ObProp does not capture the action of induction heads.”

Even though we do not address QK circuits with this method, we still believe that our method’s application to OV circuits is highly valuable. If the QK circuit determines where information comes from in a Transformer, the OV circuit determines how that information is transformed throughout the model – and in many relevant tasks, these computations in the OV circuit are highly non-trivial. Using the occupational gender bias task that we address in our paper as an example: the fact that the model uses the same features to predict occupations as it does to predict gendered pronouns is certainly not obvious a priori; this is just one example of the insights that can be gained by looking at OV circuits. Furthermore, even if complex machinery in the model’s QK circuits (such as induction heads) is used to transfer information between tokens, nevertheless, whenever we want to understand what sort of information the model uses to predict one token as opposed to another, the answer to this question lies in the model’s OV circuits. And observable propagation can be used to reveal this answer.

“Theorem 1 proposes that layer norms do not affect the feature vectors, by showing that the expected cosine similarity is close to 0. However, in this theorem, the activations and observables x, n are chosen to Gaussian with mean 0 and covariance I. What's the point of this assumption? And how do the authors conclude that the theorem 1, as stated with these distributions, implies their claim about layer norm for general activations?

We appreciate the reviewer's critical assessment of Theorem 1. Upon reevaluation, we acknowledge that the Gaussian assumption for activations may not be necessary for the theorem's proof. Consequently, we have revised the theorem's statement to remove this assumption, thus broadening its applicability.

Regarding the Gaussian assumption for observables, this choice was guided by the nascent nature of the observable concept in this field. Given the lack of established intuition or precedent for the distribution of observables, we opted for a Gaussian model as a starting point, owing to its mathematical tractability and prevalence in theoretical studies. This assumption, however, is not intended to limit the theorem's scope. Empirically, we observed that Theorem 1's conclusions hold true across a range of scenarios: the cosine similarities of feature vectors obtained by directly linearly-approximating LayerNorms with feature vectors obtained by ignoring LayerNorms are all extremely high. These observations are detailed in Appendix I and provide practical evidence supporting the theorem's broader implications.

In summary, while the initial version of Theorem 1 employed specific assumptions for mathematical convenience, our empirical findings and subsequent theorem revision reinforce its relevance and potential applicability to a broader range of activation and observable distributions.

Related to the above, why did the authors look at uniform distributions on the sphere in theorem 2? Even if the vectors have bounded norm, why do the authors assume uniformity?”

Similar to our above response with regard to Theorem 1, note that we empirically observe that the coupling coefficients are accurate estimators of the constant of proportionality between the dot products of activations with two given feature vectors – please see Table 2 in Section 4.1. Thus, this theorem as well serves to theoretically motivate this quantity – that is, the coupling coefficient – which empirically bears out the behavior that is predicted by the theorem.

We hope that we have been able to demonstrate the theoretical soundness of observable propagation with regard to your questions. If you have any further comments, please do not hesitate to let us know. Again, thank you for your review.

Thank you for taking the time to review our work. We are glad to hear that you appreciate the power of observable propagation for obtaining an understanding of model internals, particularly in the gender bias case. With regard to your thoughts on evaluating our method: as per your suggestion, we ran substantial additional experiments involving different observables and performed quantitative comparisons to feature vectors obtained by linear/logistic regression. We evaluated ObProp’s performance across a broader variety of tasks, including subject pronoun prediction, identifying American politicians’ party affiliations, and distinguishing between C and Python code. In particular, feature vectors found via ObProp match or exceed performance of feature vectors found using methods that require training.

We updated the paper with these results in Section 4.3. In general, we found that for the gendered pronoun prediction setting, the feature vector obtained by linear regression was only able to match the performance of the observable propagation feature vector when the linear regression was trained on 60 embeddings. Note that, in contrast, our observable propagation feature vectors could be obtained without running any forward passes. For the Python vs. C code setting, the logistic regression feature vector only matched the performance of the observable propagation feature vector when the logistic regression was trained on 50 embeddings. And for the political party identification setting, even when trained on 3/4ths of an artificial dataset of politicians (60 embeddings in the training set), the linear regression feature vector yielded far worse results than the observable propagation feature vector (correlation coefficient $r^2\approx0.30$ versus $r^2\approx0.43$). Our findings suggest that observable propagation offers a notable advantage over traditional methods, particularly in the low-data regime.

If you have any further questions or comments, please do not hesitate to let us know. Again, thank you for your review.

(P.S. In response to your comments, we also added a brief explanation that an unembedding matrix is the matrix “which projects the model's final activations into logits space”, in the place where the term is first used.)