The Linear Representation Hypothesis and the Geometry of Large Language Models
Kiho Park, Yo Joong Choe, Victor Veitch
Introduction
In the context of language models, the “Linear Representation Hypothesis” is the idea that high-level concepts are represented linearly in the representation space of a model [MYZ13, Aro+16, Elh+22, Wan+23, NLW23, e.g.]. In the context of language, a high-level concept might include: is the text in French or English? Is it in the present tense or past tense? If the text is about a person, are they male or female? The appeal of the linear representation hypothesis is that—were it true—the tasks of interpreting and controlling model behavior could exploit linear algebraic operations on the representation space. The goal of this paper is to formalize the linear representation hypothesis, and clarify how it relates to interpretation and control.
The first challenge is that it is not clear what “linear representation” actually means. There are (at least) three natural ways to interpret the idea:
Subspace: \Citep[e.g.,][]mikolov2013distributed,pennington2014glove The first idea is that each concept is represented as a subspace. For example, in the context of word embeddings, it has been argued empirically that , , and all similar pairs belong to a common subspace [Mik+13]. Then, it is natural to take this subspace to be a representation of the concept of .
Measurement: \Citep[e.g.,][]nanda2023emergent,LMSpaceTime:2023 Next is the idea that the probability of a concept value can be measured with a linear probe. For example, the probability that the output language is French is logit-linear in the representation of the input. In this case, we can take the linear map to be a representation of the concept of .
Intervention: \Citep[e.g.,][]wang2023concept,ActivationAddition:2023 The final idea is that the value a concept takes on can be changed (without changing other concepts) by adding a suitable steering vector—e.g., we change the output to French by adding a vector. In this case, we take this added vector to be the representation of the concept.
It is not clear a priori how these ideas relate to each other, nor which is the “right” notion of linear representation.
Next, suppose we have somehow found the linear representations of various concepts. The appeal of linearity is that we can now hope to use linear algebraic operations on the representation space for interpretation and control. For example, we might compute the similarity between a representation and known concept directions, or edit representations projected onto target directions. However, similarity and projection are geometric notions: they require an inner product on the representation space. The second challenge is that it is not clear what inner product is appropriate for understanding model representations.
To address these two challenges, we make the following contributions:
First, we formalize the subspace notion of linear representation in terms of counterfactual pairs, in both “embedding” (input phrase) and “unembedding” (output word) space. Using this, we prove that the unembedding notion connects to measurement, and the embedding notion to intervention.
Next, we introduce the notion of a causal inner product: an inner product with the property that concepts that can vary freely of each other are represented as orthogonal vectors. We show that such an inner product has the special property that it unifies the embedding and unembedding representations; illustrated in fig. 1. Additionally, we show how to estimate the inner product using the LLM unembedding matrix.
Finally, we study the linear representation hypothesis empirically using LLaMA-2 [Tou+23]. Using the subspace notion, we are able to find linear representations of a variety of concepts. Using these, we give evidence that the causal inner product respects semantic structure, and that subspace representations can be used to construct measurement and intervention representations.
The Linear Representation Hypothesis
We begin by formalizing the subspace notion of linear representation, one in each of the unembedding and embedding spaces of language models, and then tie the subspace notions to the measurement and intervention notions.
The first step is to formalize the notion of a concept. Intuitively, a concept is any factor of variation that can be changed in isolation. For example, we can change the output from French to English without changing its meaning, or change the output from being about a man to about a woman without changing the language it is written in.
Following [Wan+23], we formalize this idea by taking a concept variable to be a latent variable that is caused by the context , and that acts as a cause of the output . For simplicity of exposition, we will restrict attention to binary concepts. Anticipating the representation of concepts by vectors, we introduce an ordering on each binary concept—e.g., malefemale. This ordering will make the sign of a representation meaningful (so, e.g., the representation of femalemale will have the opposite sign.)
Each concept variable defines a set of counterfactual outputs that differ only in the value of . For example, for the malefemale concept, we might have
In this paper, we’ll assume that the value of concepts can be read off deterministically from the sampled output (so, e.g., the output “king” implies ). Then, can specify concepts by specifying their corresponding counterfactual outputs.
We will eventually need to reason about the relationships between multiple concepts. We say that two concepts and are causally separable if is well-defined for each . That is, causally separable concepts are those that can be varied freely and in isolation. For example, EnglishFrench and malefemale are causally separable—consider . However, EnglishFrench and EnglishRussian are not because they cannot vary freely. Also, PresentTensePastTense—verb tense—and SingularNounPluralNoun—noun plurality—are not because they do not apply to the same type of outputs.
We’ll write as when the concepts are clear from context.
2 Unembedding Representations and Measurement
We now turn to formalizing the idea of linear representation of a concept. The first observation is that there are two distinct representation spaces in play—the model representation space , and the unembedding representation space . A concept could be linearly represented in either space. We begin with the unembedding space. Defining the cone of vector as ,
We say that is an unembedding representation of concept if almost surely.
This definition captures the idea of linear representation that relies on is parallel to and so forth. We use a cone instead of subspace because the sign of the difference is significant—i.e., the difference between “king” and “queen” is in the opposite direction as the difference between “woman” and “man”. The unembedding representation (if it exists) is unique up to positive scaling, consistent with the linear subspace hypothesis that concepts are represented as directions. In other words, the unembedding representation is the unique direction that the counterfactual pairs point to in the unembedding space.
The first result is that the unembedding representation is closely tied to the measurement notion of linear representation:
Let be a concept, and let be an unembedding representation of . Then, given any context embedding ,
where a.s. is a function of .
In words: if we know the output token is either “king” or “queen” (say, the context was about a monarch), then the probability that the output is “king” is logit-linear in the language model representation with regression coefficients . The random scalar is a function of the particular counterfactual pair —e.g., it may be different for and . However, the direction used for prediction is the same for all counterfactual pairs demonstrating the concept.
Theorem 2 shows a connection between the subspace representation and the linear representation learned by fitting a linear probe to predict the concept. Namely, in both cases, we get a predictor that is linear on the logit scale. However, the unembedding representation differs from a probe-based representation in that it does not incorporate any information about correlated but off-target concepts. For example, if French text were disproportionately about men, a probe could learn this information (and include it in the representation), but the unembedding representation would not. In this sense, the unembedding representation might be viewed as an ideal probing representation.
3 Embedding Representations and Intervention
The next step is to define a linear subspace representation in the embedding space . We’ll again go with a notion anchored in demonstrative pairs. In the embedding space, each defines a distribution over concepts. We consider pairs of sentences such as and that induce different distributions on the target concept, but the same distribution on all off-target concepts. A concept is embedding-represented if the difference in all such pairs belongs to a common subspace. Formally,
We say that is an embedding representation of concept if for any context embeddings that satisfy
for each concept that is causally separable with , we have .
The first condition ensures that the direction is relevant to the target concept, and the second condition ensures that the direction is not relevant to off-target concepts.
It turns out that the embedding representation is closely tied to the intervention notion of linear representation. To get there, we’ll need the following lemma relating embedding representations to unembedding representations.
Let be the embedding representation of a concept , and let and be the unembedding representations for and any concept that is causally separable with . Then, we have
We can now give the connection to the intervention notion of linear representation.
Let be the embedding representation of a concept . Then, for any concept that is causally separable with ,
In words: adding to the language model representation of the context changes the probability of the target concept, but not the probability of off-target concepts.
Inner Product for Language Model Representations
Given linear representations, we would like to make use of them by doing things like measuring the similarity between different representations, or editing concepts by projecting onto a target direction. Similarity and projection are both notions that require an inner product. We now consider the question of which inner product is appropriate for understanding language model representations.
We define to be the space of differences between elements of . Then, is a -dimensional real vector space.Note that the unembedding space is only an affine space, since the softmax is invariant to adding a constant. We consider defining inner products on . Unembedding representations are naturally directions (unique only up to scale). Once we have an inner product, we define the canonical unembedding representation to be the element of the unembedding cone with . This lets us define inner products between unembedding representations.
Unidentifiability of the inner product
We might hope that there is some natural inner product that is picked out (identified) by the model training. It turns out that this is not the case. To understand the challenge, consider transforming the embedding and unembedding spaces according to
However, the objective function used to train the model depends on the representations only through the softmax probabilities. Thus, the representation is identified (at best) only up to some invertible affine transformation.
This also means that the concept representations are identified only up to some invertible linear transformation . The problem is that, given any fixed inner product,
in general. Accordingly, there is no obvious reason to expect that algebraic manipulations based on, e.g., the Euclidean inner product, should be preferred to manipulations using any other inner product.
1 Causal Inner Products
We require some additional principles for choosing an inner product on the representation space. The intuition we follow here is that causally separable concepts should be represented as orthogonal vectors. For example, FrenchEnglish and MaleFemale, should be orthogonal. We define an inner product with this property:
for any pair of causally separable concepts and .
This choice turns out to have the critical property that it gives a natural unification of the unembedding and embedding representations:
In the experiments, we will make use of this result to construct embedding representations from unembedding representations. In particular, this allows us to find interventional representations of concepts. This is important because it is difficult in practice to find pairs of prompts that directly satisfy Definition 3.
2 An Explicit Form for Causal Inner Product
The next problem is: if a causal inner product exists, how can we find it? In principle, this could be done by finding the unembedding representations of a large number of concepts, and then finding an inner product that maps each pair of causally separable directions to zero. In practice, this is infeasible because of the number of concepts required to find the inner product, and the difficulty of estimating the representations of each concept.
We now turn to developing a more tractable approach. Our technique is based on the following insight: knowing the value of concept expressed by a randomly chosen word tells us little about the value of that word on a causally separable concept . For example, if we learn that a randomly sampled word is French (not English), this does not give us significant information about whether it refers to a man or woman.Note that this assumption is about words sampled randomly from the vocabulary, not words sampled randomly from natural language sources. In the latter, there may well be non-causal correlations between causally separable concepts (e.g., if French text is disproportionately about men). Following Theorem 5, we formalize this idea as follows:
Suppose are causally separable concepts and that is an unembedding vector sampled uniformly from the vocabulary. Then, \bar{\lambda}_{W}^{\top}\gamma\mathchoice{\mathrel{\hbox to0.0pt{\displaystyle\perp\hss}\mkern 2.0mu{\displaystyle\perp}}}{\mathrel{\hbox to0.0pt{\textstyle\perp\hss}\mkern 2.0mu{\textstyle\perp}}}{\mathrel{\hbox to0.0pt{\scriptstyle\perp\hss}\mkern 2.0mu{\scriptstyle\perp}}}{\mathrel{\hbox to0.0pt{\scriptscriptstyle\perp\hss}\mkern 2.0mu{\scriptscriptstyle\perp}}}\bar{\lambda}_{Z}^{\top}\gamma for any embedding representations and for and , respectively.
This assumption lets us connect causal separability with something we can actually measure: the statistical dependency between words. The next result makes this precise.
for some diagonal matrix with positive entries, where is the unembedding vector of a word sampled uniformly at random from the vocabulary.
Notice that causal orthogonality only imposes constraints on the inner product, but there are degrees of freedom in defining a positive definite matrix (hence, an inner product)—thus, we expect degrees of freedom in choosing a causal inner product. Theorem 8 gives a characterization of this class of inner products, in the form of (3.6). Here, is a free parameter with degrees of freedom. Each defines the inner product. We do not have a principle for picking out a unique choice of (and thus, a unique inner product). In our experiments, we will work with the choice . Then, we have a simple closed form for the corresponding inner product:
Notice that although we don’t have a unique inner product, we can rule out most inner products. E.g., the Euclidean inner product is not a causal inner product if does not satisfy (3.6) for any .
The choice of inner product also be viewed as defining a canonical choice of representations in eq. 3.1. Namely, we define
Experiments
We now turn to empirically validating the existence of linear representations, the technique for finding the causal inner product, and the predicted relationships between the subspace, measurement, and intervention notions of linear representation. Code available at github.com/KihoPark/linear_rep_geometry.
We use the LLaMA-2 model with 7 billion parameters [Tou+23] as our testbed. This is a decoder-only Transformer LLM [Vas+17, Rad+18], trained using the forward LM objective and a 32K token vocabulary.
We start with the hypothesis that concepts are represented as directions in the unembedding representation space (Definition 1). This notion relies on counterfactual pairs of words that vary only in the value of the concept of interest. We consider 22 concepts defined in the Big Analogy Test Set (BATS 3.0) [GDM16], which provides such counterfactual pairs.We throw away any pair where one of the words is encoded as multiple tokens. We also consider 4 additional language concepts: EnglishFrench, FrenchGerman, FrenchSpanish, and GermanSpanish, where we use words and their translations as counterfactual pairs. Additionally, we consider the concept frequentinfrequent capturing how common a word is—we use pairs of common/uncommon synonyms (e.g., “bad” and “terrible”) as counterfactual pairs. In Appendix B, we list all 27 concepts we consider and example pairs.
If the subspace notion of the linear representation hypothesis holds then all counterfactual token pairs should point to a common direction in the unembedding space. In practice, this will only hold approximately for real pairs because each word can have multiple meanings (e.g., “Queen” is a female monarch, a chess piece, and a rock band). However, if the linear representation hypothesis holds, we still expect that will significantly align with a malefemale direction. So, for each concept , we look at how the direction defined by each counterfactual pair is geometrically aligned with a common direction (the unembedding representation). We estimate as the meanPrevious work on word embeddings [DGM16, FDD20] motivate taking the mean to improve the consistency of the concept direction. among all counterfactual pairs:
Figure 2 presents histograms of each projected onto with respect to the causal inner product. Because is computed using , we compute each projection using a leave-one-out (LOO) estimate of the concept direction that excludes . Across the four concepts shown (and 22 others shown in Appendix C), the differences between counterfactual pairs are substantially more aligned with than those between random pairs. The sole exception is thingpart, which does not appear to have a linear representation.
The results are consistent with the linear representation hypothesis: the directions computed by each counterfactual pair point (up to some noise) to a common direction representing a linear subspace. Further, is a reasonable estimator for that direction.
2 Concept directions act as linear probes
Next, we check the connection to the measurement notion of linear representation. We consider the concept FrenchSpanish. To construct a dataset of French/Spanish contexts, we sample contexts of random lengths from Wikipedia pages in each language. (Note: these are not counterfactual pairs.) Following Theorem 2 we expect and . Figure 3 confirms this expectation, showing that is a linear probe for the concept in . We also see that the representation of an off-target concept does not have any predictive power for this task.
3 Concept directions map to intervention representations
Theorem 5 says that we can construct an intervention representation by constructing an embedding embedding representation. Doing this directly requires finding pairs of prompt that vary only on the distribution they induce on the target concept. In preliminary experiments, we found it was difficult to construct such pairs in practice.
Here, we will instead use the isomorphism between embedding and unembedding representations (Theorem 7) to construct intervention representations from unembedding representations. We take
Theorem 5 predicts that adding to a context representation should increase the probability of , while leaving the probability of all causally separable concepts unaltered.
To test this for a given pair of causally separable concepts and , we first choose a quadruple , and then generate contexts such that the next word should be . For example, if malefemale and lowerupper, then we choose the quadruple (“king”, “queen”, “King”, “Queen”), and generate contexts using ChatGPT-4 (e.g., “Long live the”). We then intervene on using via
Figure 4 shows the results of one such experiment, confirming our expectations. We see, for example, that intervening in the malefemale direction raises the logit for choosing “queen” over “king” as the next word, but does not change the logit for “King” over “king”.
4 The estimated inner product respects causal separability
Finally, we turn to directly examining whether the estimated inner product chosen from Theorem 8,
is indeed approximately a causal inner product. In fig. 5, we plot a heatmap of the inner products between all pairs of the 27 estimated concepts. If the estimated inner product is a causal inner product, then we expect values near 0 between causally separable concepts (and large values between causally related concepts).
The first observation is that most pairs of concepts are nearly orthogonal with respect to this inner product. Interestingly, there is also a clear block diagonal structure. This arises because the concepts are grouped by semantic similarity. For example, the first 10 concepts relate to verbs, and the last 4 concepts are language pairs. The additional non-zero structure also generally makes sense. For example, lowerupper (capitalization, concept 19) has non-trivial inner product with the language pairs other than FrenchSpanish. This may be because French and Spanish obey similar capitalization rules, while English and German each have different conventions (e.g., German capitalizes all nouns, but English only capitalizes proper nouns).
Discussion and Related Work
The idea that high-level concepts are encoded linearly is appealing because—if it is true—it may open up simple methods for interpretability and controllability of LLMs. In this paper, we have formalized ‘linear representation’, and shown that all natural variants of this notion can be unified. This equivalence already suggests some approaches for interpretation and control—e.g., we show how to use collections of pairs of words to define concept directions (Section 4.1), and then use these directions to predict what the model’s output will be (Section 4.2), and to change the output in a controlled fashion (Section 4.3). A major theme is the role played by the choice of inner product.
The linear subspace hypotheses was originally observed empirically in the context of word embeddings [Mik+13, LG14, GL14, Vyl+16, GDM16, CCCP20, FDD20, e.g.,]. Similar structure has been observed in cross-lingual word embeddings [MLS13, Lam+18, RVS19, Pen+22], sentence embeddings [Bow+16, ZM20, Li+20, Ush+21], representation spaces of Transformer LLMs [Men+22, MEP23, Her+23], and vision-language models [Wan+23, Tra+23, Per+23]. These observations motivate Definition 1. The key idea in the present paper is providing formalization in terms of counterfactual pairs—this is what allows us to connect to other notions of linear representation, and to identify the inner product structure.
Measurement, intervention, and mechanistic interpretability
There is a significant body of work on linear representations for interpreting (probing) [AB17, Kim+18, nos20, RKR21, Bel22, Li+22, Gev+22, NLW23, e.g.,] and controlling (steering) [Wan+23, Tur+23, MEP23, Tra+23, e.g.,] models. This is particularly prominent in mechanistic interpretability [Elh+21, Men+22, Her+23, Tur+23, Zou+23, Tod+23, HGG23]. With respect to this body of work, the main contribution of the present paper is to clarify the linear representation hypothesis, and the critical role of the inner product. However, we do not address interpretability of either model parameters, nor the activations of intermediate layers. These are main focuses of existing work. It is an exciting direction for future work to understand how ideas here—particularly, the causal inner product—translate to these settings.
Geometry of representations
There is a line of work that studies the geometry of word and sentence representations [Aro+16, MT17, Eth19, Rei+19, Li+20, HM19, Che+21, CTB22, JAV23, e.g.,]. This work considers, e.g., visualizing and modeling how the learned embeddings are distributed, or how hierarchical structure is encoded. Our work is largely orthogonal to these, since we are attempting to define a suitable inner product (and thus, notion of distance) that respects the semantic structure of language.
Causal representation learning
Finally, the ideas here connect to causal representation learning [Hig+16, HM16, Hig+18, Khe+20, Zim+21, Sch+21, Mor+21, Wan+23, e.g.,]. Most obviously, our causal formalization of concepts is inspired by [Wan+23], who establish a characterization of latent concepts and vector algebra in diffusion models. Separately, a major theme in this literature is the identifiability of learned representations—i.e., to what extent they capture underlying real-world structure. Our causal inner product results may be viewed in this theme, showing that an inner product respecting semantic closeness is not identified by the usual training procedure, but that it can be picked out with a suitable assumption.
Acknowledgements
Thanks to Gemma Moran for comments on an earlier draft. This work is supported by ONR grant N00014-23-1-2591 and Open Philanthropy.
References
Appendix A Proofs
The proof involves writing out the softmax sampling distribution and invoking Definition 1.
In (A.3), we simply write out the softmax distribution, allowing us to cancel out the normalizing constants for the two probabilities. Equation (A.4) follows directly from Definition 1; note that the randomness of comes from the randomness of . ∎
A.2 Proof of Lemma 4
Let be a pair of embeddings such that
for any concept that is causally separable with . Then, by Definition 3,
These two conditions are also equivalent to the following pair of conditions, respectively:
The reason is that, conditional on , conditioning on is equivalent to conditioning on . And, the event is equivalent to the event . (In words: if we know the output is one of “king”, “queen”, “roi”, “reine” then conditioning on is equivalent to conditioning on the output being “king” or “roi”. Then, predicting whether the word is in English is equivalent to predicting whether the word is “king”.)
By Theorem 2, the two conditions (A.8) and (A.9) are respectively equivalent to
where ’s are positive a.s. These are in turn respectively equivalent to
In other words, also satisfies (A.11), implying that it must be the same as up to positive scaling. Therefore, for any and satisfying (A.5), . ∎
A.3 Proof of Theorem 5
Therefore, we have (2.5) since by Lemma 4.
Therefore, we have (2.6) since by Lemma 4. ∎
A.4 Proof of Theorem 7
where the second equality follows from Definition 6. By Lemma 4, expresses the unique unembedding representation (up to positive scaling); specifically, where . ∎
A.5 Proof of Theorem 8
for any causally separable concepts and . Also, is an embedding representation for each concept for by the proof of Lemma 4 and Theorem 7. Thus, by Assumption 1,
for . By applying (A.20) to the basis , we have
for some diagonal matrix with positive entries. Then, and
Appendix B Experiment Details
We utilize the llama-2-7b variant of the LLaMA-2 model [Tou+23], which is accessible online (with permission) via the huggingface library.https://huggingface.co/meta-llama/Llama-2-7b-hf Its seven billion parameters are pre-trained on two trillion sentencepiece [KR18] tokens, 90% of which is in English. This model uses 32,000 tokens and 4,096 dimensions for its token embeddings.
Counterfactual pairs
Tokenization impedes using the meaning of an exact word. First, a word can be tokenized to more than one token. For example, a word “princess” is tokenized to “prin” + “cess”, and does not exist. Thus, we cannot obtain the meaning of the exact word “princess". Second, a word can be used as one of the tokens for another word. For example, the French words “bas” and “est” (“down” and “east” in English) are in the tokens for the words “basalt”, “baseline”, “basil”, “basilica”, “basin”, “estuary”, “estrange”, “estoppel”, “estival”, “esthetics”, and “estrogen”. Therefore, a word can have another meaning other than the meaning of the exact word.
When we collect the counterfactual pairs to identify , the first issue in the pair can be handled by not using it. However, the second issue cannot be handled, and it gives a lot of noise to our results. Table 2 presents the number of the counterfactual pairs for each concept and one example of the pairs. The pairs for 13, 17, 19, 23-27th concepts are generated by ChatGPT-4 [Ope23], and those for 16th concept are based on the csv filehttps://github.com/jmerullo/lm_vector_arithmetic/blob/main/world_capitals.csv). The other concepts are based on The Bigger Analogy Test Set (BATS) [GDM16], version 3.0https://vecto.space/projects/BATS/, which is used for evaluation of the word analogy task.
Context samples
In Section 4.2, for a concept (e.g., EnglishFrench), we choose several counterfactual pairs (e.g., (house, maison)), then sample context and that the next token is and , respectively, from Wikipedia. These next token pairs are collected from the word2word bilingual lexicon [CPK20], which is a publicly available word translation dictionary. We take all word pairs between languages that are the top-1 correspondences to each other in the bilingual lexicon and filter out pairs that are single tokens in the LLaMA-2 model’s vocabulary.
Table 3 presents the number of the contexts and for each concept and one example of the pairs .
In the experiment for intervention notion, for a concept , we sample texts which (e.g., “king”) should follow, via ChatGPT-4. We discard the contexts such that is not the top 1 next word. Table 4 present the contexts we use.
Validation for Assumption 1
In Figure 6, we check that and are independent for the causally separable concepts where is estimated by (4.2). On the other hand, Figure 7 shows that and are not independent for the non-causally separable concepts.
Appendix C Additional Results
In Figure 8, we include the analog of Figure 2, where we check the causal inner product of the differences between the counterfactual pairs and an LOO estimated unembedding representation for each of the 27 concepts. While the most of the concepts are encoded in the unembedding representation, some concepts, such as thingpart, are not encoded in the unembedding space .
C.2 Additional results from the measurement experiment
We include the analog of Figure 3, where we use each of the 27 concepts as a linear probe on either FrenchSpanish (Figure 9) or EnglishFrench (Figure 10) contexts.
C.3 Additional results from the intervention experiment
In Figure 11, we include the analog of Figure 4, where we add the embedding representation (4.2) for each of the 27 concepts to and see the change in logits.