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 Rep⁡(“woman”)−Rep⁡(“man”)\operatorname{Rep}(\text{``woman''})-\operatorname{Rep}(\text{``man''}), Rep⁡(“queen”)−Rep⁡(“king”)\operatorname{Rep}(\text{``queen''})-\operatorname{Rep}(\text{``king''}), 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 Male/Female\mathtt{Male/Female}.

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 English/French\mathtt{English/French}.

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 English/French\mathtt{English/French} 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 WW to be a latent variable that is caused by the context XX, and that acts as a cause of the output YY. 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., male⇒\Rightarrowfemale. This ordering will make the sign of a representation meaningful (so, e.g., the representation of female⇒\Rightarrowmale will have the opposite sign.)

Each concept variable WW defines a set of counterfactual outputs {Y(W=w)}\{Y(W=w)\} that differ only in the value of WW. For example, for the male⇒\Rightarrowfemale 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 W=0W=0). 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 WW and ZZ are causally separable if Y(W=w,Z=z)Y(W=w,Z=z) is well-defined for each w,zw,z. That is, causally separable concepts are those that can be varied freely and in isolation. For example, English⇒\RightarrowFrench and male⇒\Rightarrowfemale are causally separable—consider {“king”,“queen”,“roi”,“reine”}\{\text{``king''},\text{``queen''},\text{``roi''},\text{``reine''}\}. However, English⇒\RightarrowFrench and English⇒\RightarrowRussian are not because they cannot vary freely. Also, PresentTense⇒\RightarrowPastTense—verb tense—and SingularNoun⇒\RightarrowPluralNoun—noun plurality—are not because they do not apply to the same type of outputs.

We’ll write Y(W=w,Z=z)Y(W=w,Z=z) as Y(w,z)Y(w,z) 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 Λ\Lambda, and the unembedding representation space Γ\Gamma. A concept could be linearly represented in either space. We begin with the unembedding space. Defining the cone of vector vv as Cone⁡(v)={αv:α>0}\operatorname{Cone}(v)=\{\alpha v:\alpha>0\},

We say that γˉW\bar{\gamma}_{W} is an unembedding representation of concept WW if γ(Y(1))−γ(Y(0))∈Cone⁡(γˉW)\gamma(Y(1))-\gamma(Y(0))\in\operatorname{Cone}(\bar{\gamma}_{W}) almost surely.

This definition captures the idea of linear representation that relies on γ(“king”)−γ(“queen”)\gamma(\text{``king''})-\gamma(\text{``queen''}) is parallel to γ(“man”)−γ(“woman”)\gamma(\text{``man''})-\gamma(\text{``woman''}) 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 WW be a concept, and let γˉW\bar{\gamma}_{W} be an unembedding representation of WW. Then, given any context embedding λ∈Λ\lambda\in\Lambda,

where α>0\alpha>0 a.s. is a function of {Y(1),Y(0)}\{Y(1),Y(0)\}.

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 γˉW\bar{\gamma}_{W}. The random scalar α\alpha is a function of the particular counterfactual pair {Y(1),Y(0)}\{Y(1),Y(0)\}—e.g., it may be different for {“king”,“queen”}\{\text{``king''},\text{``queen''}\} and {“roi”,“riene”}\{\text{``roi''},\text{``riene''}\}. 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 Λ\Lambda. We’ll again go with a notion anchored in demonstrative pairs. In the embedding space, each λ(x)\lambda(x) defines a distribution over concepts. We consider pairs of sentences such as λ0=λ[“He is the monarch of England,”]\lambda_{0}=\lambda[\text{``He is the monarch of England,''}] and λ1=λ[“She is the monarch of England,”]\lambda_{1}=\lambda[\text{``She is the monarch of England,''}] 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 λˉW\bar{\lambda}_{W} is an embedding representation of concept WW if for any context embeddings λ0,λ1∈Λ\lambda_{0},\lambda_{1}\in\Lambda that satisfy

for each concept ZZ that is causally separable with WW, we have λ1−λ0∈Cone⁡(λˉW)\lambda_{1}-\lambda_{0}\in\operatorname{Cone}(\bar{\lambda}_{W}).

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 λˉW\bar{\lambda}_{W} be the embedding representation of a concept WW, and let γˉW\bar{\gamma}_{W} and γˉZ\bar{\gamma}_{Z} be the unembedding representations for WW and any concept ZZ that is causally separable with WW. Then, we have

We can now give the connection to the intervention notion of linear representation.

Let λˉW\bar{\lambda}_{W} be the embedding representation of a concept WW. Then, for any concept ZZ that is causally separable with WW,

In words: adding λˉW\bar{\lambda}_{W} 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 Γˉ\bar{\Gamma} to be the space of differences between elements of Γ\Gamma. Then, Γˉ\bar{\Gamma} is a dd-dimensional real vector space.Note that the unembedding space Γ\Gamma is only an affine space, since the softmax is invariant to adding a constant. We consider defining inner products on Γˉ\bar{\Gamma}. Unembedding representations are naturally directions (unique only up to scale). Once we have an inner product, we define the canonical unembedding representation γˉW\bar{\gamma}_{W} to be the element of the unembedding cone with ⟨γˉW,γˉW⟩=1\langle\bar{\gamma}_{W},\bar{\gamma}_{W}\rangle=1. 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 γ\gamma is identified (at best) only up to some invertible affine transformation.

This also means that the concept representations γˉW\bar{\gamma}_{W} are identified only up to some invertible linear transformation AA. 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, French⇒\RightarrowEnglish and Male⇒\RightarrowFemale, should be orthogonal. We define an inner product with this property:

for any pair of causally separable concepts WW and ZZ.

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 WW expressed by a randomly chosen word tells us little about the value of that word on a causally separable concept ZZ. 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 W,ZW,Z are causally separable concepts and that γ\gamma 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 λˉW\bar{\lambda}_{W} and λˉZ\bar{\lambda}_{Z} for WW and ZZ, 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 DD with positive entries, where γ\gamma is the unembedding vector of a word sampled uniformly at random from the vocabulary.

Notice that causal orthogonality only imposes d(d−1)/2d(d-1)/2 constraints on the inner product, but there are d(d−1)/2+dd(d-1)/2+d degrees of freedom in defining a positive definite matrix (hence, an inner product)—thus, we expect dd 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, DD is a free parameter with dd degrees of freedom. Each DD defines the inner product. We do not have a principle for picking out a unique choice of DD (and thus, a unique inner product). In our experiments, we will work with the choice D=IdD=I_{d}. 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 M=IdM=I_{d} does not satisfy (3.6) for any DD.

The choice of inner product also be viewed as defining a canonical choice of representations g,lg,l 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: English⇒\RightarrowFrench, French⇒\RightarrowGerman, French⇒\RightarrowSpanish, and German⇒\RightarrowSpanish, where we use words and their translations as counterfactual pairs. Additionally, we consider the concept frequent⇒\Rightarrowinfrequent 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 γ(“King”)−γ(“Queen”)\gamma(\text{``King''})-\gamma(\text{``Queen''}) will significantly align with a male⇒\Rightarrowfemale direction. So, for each concept WW, we look at how the direction defined by each counterfactual pair γ(yi(1))−γ(yi(0))\gamma(y_{i}(1))-\gamma(y_{i}(0)) is geometrically aligned with a common direction γˉW\bar{\gamma}_{W} (the unembedding representation). We estimate γˉW\bar{\gamma}_{W} 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 γ(yi(1))−γ(yi(0)))\gamma(y_{i}(1))-\gamma(y_{i}(0))) projected onto γˉW\bar{\gamma}_{W} with respect to the causal inner product. Because γˉW\bar{\gamma}_{W} is computed using γ(yi(1))−γ(yi(0))\gamma(y_{i}(1))-\gamma(y_{i}(0)), we compute each projection using a leave-one-out (LOO) estimate γˉW,(−i)\bar{\gamma}_{W,(-i)} of the concept direction that excludes (yi(0),yi(1))(y_{i}(0),y_{i}(1)). Across the four concepts shown (and 22 others shown in Appendix C), the differences between counterfactual pairs are substantially more aligned with γˉW\bar{\gamma}_{W} than those between random pairs. The sole exception is thing⇒\Rightarrowpart, 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, γˉW\bar{\gamma}_{W} 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 French⇒\RightarrowSpanish. 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 γˉW⊤λ(xjfr)<0\bar{\gamma}_{W}^{\top}\lambda(x_{j}^{\texttt{fr}})<0 and γˉW⊤λ(xjes)>0\bar{\gamma}_{W}^{\top}\lambda(x_{j}^{\texttt{es}})>0. Figure 3 confirms this expectation, showing that γˉW\bar{\gamma}_{W} is a linear probe for the concept WW in Λ\Lambda. We also see that the representation of an off-target concept ZZ 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 λˉW\bar{\lambda}_{W} to a context representation should increase the probability of WW, while leaving the probability of all causally separable concepts unaltered.

To test this for a given pair of causally separable concepts WW and ZZ, we first choose a quadruple {Y(w,z)}w,z∈{0,1}\{Y(w,z)\}_{w,z\in\{0,1\}}, and then generate contexts {xj}\{x_{j}\} such that the next word should be Y(0,0)Y(0,0). For example, if W=W= male⇒\Rightarrowfemale and Z=Z= lower⇒\Rightarrowupper, then we choose the quadruple (“king”, “queen”, “King”, “Queen”), and generate contexts using ChatGPT-4 (e.g., “Long live the”). We then intervene on λ(xj)\lambda(x_{j}) using λˉC\bar{\lambda}_{C} via

Figure 4 shows the results of one such experiment, confirming our expectations. We see, for example, that intervening in the male⇒\Rightarrowfemale 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, lower⇒\Rightarrowupper (capitalization, concept 19) has non-trivial inner product with the language pairs other than French⇒\RightarrowSpanish. 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 α\alpha comes from the randomness of (Y(1),Y(0))(Y(1),Y(0)). ∎

A.2 Proof of Lemma 4

Let λ0,λ1\lambda_{0},\lambda_{1} be a pair of embeddings such that

for any concept ZZ that is causally separable with WW. Then, by Definition 3,

These two conditions are also equivalent to the following pair of conditions, respectively:

The reason is that, conditional on Y∈{Y(0,0),Y(0,1),Y(1,0),Y(1,1)}Y\in\{Y(0,0),Y(0,1),Y(1,0),Y(1,1)\}, conditioning on WW is equivalent to conditioning on Y∈{Y(W,0),Y(W,1)}Y\in\{Y(W,0),Y(W,1)\}. And, the event Z=1Z=1 is equivalent to the event Y=Y(W,1)Y=Y(W,1). (In words: if we know the output is one of “king”, “queen”, “roi”, “reine” then conditioning on W=1W=1 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 α\alpha’s are positive a.s. These are in turn respectively equivalent to

In other words, λ1−λ0\lambda_{1}-\lambda_{0} also satisfies (A.11), implying that it must be the same as λˉW\bar{\lambda}_{W} up to positive scaling. Therefore, for any λ0\lambda_{0} and λ1\lambda_{1} satisfying (A.5), λ1−λ0∈Cone⁡(λˉW)\lambda_{1}-\lambda_{0}\in\operatorname{Cone}(\bar{\lambda}_{W}). ∎

A.3 Proof of Theorem 5

Therefore, we have (2.5) since λˉW⊤γˉZ=0\bar{\lambda}_{W}^{\top}\bar{\gamma}_{Z}=0 by Lemma 4.

Therefore, we have (2.6) since λˉW⊤γˉW>0\bar{\lambda}_{W}^{\top}\bar{\gamma}_{W}>0 by Lemma 4. ∎

A.4 Proof of Theorem 7

where the second equality follows from Definition 6. By Lemma 4, ϕ(γˉW)\phi(\bar{\gamma}_{W}) expresses the unique unembedding representation λˉW\bar{\lambda}_{W} (up to positive scaling); specifically, ϕ(γˉW)=λˉW⊤\phi(\bar{\gamma}_{W})=\bar{\lambda}_{W}^{\top} where λˉW⊤:γˉ↦λˉW⊤γˉ\bar{\lambda}_{W}^{\top}:\bar{\gamma}\mapsto\bar{\lambda}_{W}^{\top}\bar{\gamma}. ∎

A.5 Proof of Theorem 8

for any causally separable concepts WW and ZZ. Also, MγˉWiM\bar{\gamma}_{W_{i}} is an embedding representation for each concept WiW_{i} for i=1,⋯ ,di=1,\cdots,d by the proof of Lemma 4 and Theorem 7. Thus, by Assumption 1,

for i≠ji\neq j. By applying (A.20) to the basis G=[γˉW1,⋯ ,γˉWd]G=[\bar{\gamma}_{W_{1}},\cdots,\bar{\gamma}_{W_{d}}], we have

for some diagonal matrix DD with positive entries. Then, M=G−⊤G−1M=G^{-\top}G^{-1} 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 γ(“princess”)\gamma(\text{``princess''}) 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 γˉW\bar{\gamma}_{W}, 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 WW (e.g., English⇒\RightarrowFrench), we choose several counterfactual pairs (Y(0),Y(1))(Y(0),Y(1)) (e.g., (house, maison)), then sample context {xj0}\{x_{j}^{0}\} and {xj1}\{x_{j}^{1}\} that the next token is Y(0)Y(0) and Y(1)Y(1), 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 {xj0}\{x_{j}^{0}\} and {xj1}\{x_{j}^{1}\} for each concept and one example of the pairs (Y(0),Y(1))(Y(0),Y(1)).

In the experiment for intervention notion, for a concept W,ZW,Z, we sample texts which Y(0,0)Y(0,0) (e.g., “king”) should follow, via ChatGPT-4. We discard the contexts such that Y(0,0)Y(0,0) is not the top 1 next word. Table 4 present the contexts we use.

Validation for Assumption 1

In Figure 6, we check that λˉW⊤γ\bar{\lambda}_{W}^{\top}\gamma and λˉZ⊤γ\bar{\lambda}_{Z}^{\top}\gamma are independent for the causally separable concepts where λˉW\bar{\lambda}_{W} is estimated by (4.2). On the other hand, Figure 7 shows that λˉW⊤γ\bar{\lambda}_{W}^{\top}\gamma and λˉZ⊤γ\bar{\lambda}_{Z}^{\top}\gamma 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 thing⇒\Rightarrowpart, are not encoded in the unembedding space Γ\Gamma.

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 French⇒\RightarrowSpanish (Figure 9) or English⇒\RightarrowFrench (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 αλˉC\alpha\bar{\lambda}_{C} (4.2) for each of the 27 concepts to λ(xj)\lambda(x_{j}) and see the change in logits.