PaperlayerView exact revision on arXiv

Source-linked AI summary

The Linear Representation Hypothesis and the Geometry of Large Language Models

Kiho Park, Yo Joong Choe, Victor Veitch

arXiv:2311.03658v2cs.CLcs.AIcs.LGstat.ML

TL;DR

The paper asks what linear representation means and which inner product makes geometry meaningful in language-model representation spaces. It formalizes representations with counterfactuals, connects them to probing and steering, and uses a causal inner product to unify them. LLaMA-2 experiments support linear representations across concepts and show the importance of inner-product choice.

  • Problem

    The paper asks what “linear representation” means and which inner product makes geometric operations meaningful in language-model representation spaces.

  • Method

    It uses counterfactual pairs to formalize representations in embedding and unembedding spaces, then proves connections to measurement and intervention.

  • Results

    LLaMA-2 experiments find subspace representations for varied concepts and provide evidence connecting them to measurement and intervention.

  • Takeaways & Limitations

    The causal inner product unifies embedding and unembedding representations while respecting causal separation as orthogonality.

  • Takeaways & Limitations

    The causal inner product is not unique, leaving d degrees of freedom and no stated principle for selecting a unique choice.

Abstract

from arXiv · show

Informally, the 'linear representation hypothesis' is the idea that high-level concepts are represented linearly as directions in some representation space. In this paper, we address two closely related questions: What does "linear representation" actually mean? And, how do we make sense of geometric notions (e.g., cosine similarity or projection) in the representation space? To answer these, we use the language of counterfactuals to give two formalizations of "linear representation", one in the output (word) representation space, and one in the input (sentence) space. We then prove these connect to linear probing and model steering, respectively. To make sense of geometric notions, we use the formalization to identify a particular (non-Euclidean) inner product that respects language structure in a sense we make precise. Using this causal inner product, we show how to unify all notions of linear representation. In particular, this allows the construction of probes and steering vectors using counterfactual pairs. Experiments with LLaMA-2 demonstrate the existence of linear representations of concepts, the connection to interpretation and control, and the fundamental role of the choice of inner product.

1. Introduction

The paper formalizes the linear representation hypothesis to clarify how linear representations relate to interpretation, control, and geometric operations in language-model representation spaces. It distinguishes embedding and unembedding spaces and introduces a causal inner product to unify their representations.

  • Motivation: The linear representation hypothesis proposes that high-level concepts are represented linearly in language-model representation spaces.Examples include language, tense, and gender properties.
  • Motivation: Linear representations could support simple methods for interpreting and controlling language-model behavior through linear-algebraic operations.The paper’s goal is to formalize the hypothesis and clarify its relation to interpretation and control.
  • Open questions: The paper identifies subspace, measurement, and intervention as distinct interpretations of linear representation.These correspond respectively to concept directions, linear probes, and steering vectors.
  • Open questions: It is not clear a priori how these interpretations relate or which is the right notion of linear representation.A second challenge is determining which inner product makes cosine similarity and projection meaningful.
  • Contributions: The paper introduces a causal inner product intended to respect semantic structure and unify representations across embedding and unembedding spaces.Figure 1 depicts separate input-context and output-word representations becoming coincident after transformation.
  • Contributions: Using counterfactual pairs, the paper formalizes linear representations in both embedding and unembedding spaces and connects them to measurement and intervention.Experiments with LLaMA-2 find subspace representations for a variety of concepts and provide evidence for these connections.

2. The Linear Representation Hypothesis

The paper defines concepts through counterfactual variation and formalizes linear representations separately in unembedding and embedding spaces. It then connects these representations to linear measurement and intervention, respectively.

  • Concepts: Concepts are latent factors of variation that can be changed in isolation while affecting the output.Binary concepts are ordered so representation signs encode the direction of change.
  • Concepts: A concept variable defines counterfactual outputs that differ only in the concept’s value, with output values assumed readable deterministically.Counterfactual output pairs can therefore specify concepts such as male⇒female.
  • Concepts: Causally separable concepts are those whose values can be varied freely and in isolation in well-defined counterfactual outputs.English⇒French and male⇒female are given as an example of causally separable concepts.
  • Unembedding representations: The paper formalizes unembedding representations as directions shared by output-word differences across counterfactual pairs.The representation is unique up to positive scaling because the sign of the difference is meaningful.
  • Unembedding representations and measurement: Unembedding representations connect to measurement because the target concept’s conditional output probability is logit-linear in the context embedding.The coefficient direction is shared across counterfactual pairs, although the positive scale may vary by pair.
  • Unembedding representations and measurement: The unembedding representation differs from a probe representation by excluding information about correlated but off-target concepts.Thus, it acts as an idealized probing representation for the target concept.
  • Embedding representations: Embedding representations are defined by context-embedding differences that change the target concept while preserving distributions over causally separable off-target concepts.The first condition establishes relevance to the target; the second removes relevance to off-target concepts.
  • Embedding representations and intervention: Adding an embedding representation increases the probability of its target concept without changing the probability of a causally separable off-target concept.This establishes the connection between embedding representations and intervention.

3. Inner Product for Language Model Representations

The paper argues that Euclidean geometry is not intrinsically meaningful for language-model representations because training identifies representations only up to affine transformations. It proposes causal inner products, defined by orthogonality of causally separable concepts, to unify embedding and unembedding representations and make geometric operations interpretable.

  • Unidentifiability: Training through softmax probabilities identifies representations only up to invertible affine transformations, so concept directions and geometric operations are not uniquely determined.Applying an invertible linear transformation and translation preserves the softmax distribution, while transforming concept representations changes their coordinates.
  • Causal Inner Product: A causal inner product represents causally separable concepts as orthogonal vectors, such as language and gender concepts.This supplies an additional principle for choosing an inner product beyond the geometry inherited from arbitrary coordinates.
  • Unification: The causal inner product unifies embedding and unembedding representations by mapping each unembedding concept direction to the embedding representation of the same concept.The mapping is given by the Riesz isomorphism, and in a transformed space the causal inner product becomes the Euclidean inner product.
  • Applications: The construction enables embedding representations and interventional representations to be derived from unembedding representations, avoiding the need to find many direct counterfactual prompt pairs.This provides a practical route for constructing representations used for intervention and control.
  • Explicit Construction: The paper estimates the causal inner product from statistical dependencies among unembedding vectors sampled uniformly from the vocabulary.The approach uses the assumption that causally separable concepts are independent or uncorrelated in this sampling scheme.
  • Explicit Construction: The explicit causal-inner-product family retains d degrees of freedom, and the experiments choose D = Id, yielding M = Cov(γ)^−1.Causal orthogonality imposes fewer constraints than are needed to identify a unique positive-definite matrix, although most inner products can still be ruled out.

4. Experiments

Experiments with LLaMA-2 test whether counterfactual concept directions support the subspace, measurement, and intervention notions of linear representation under a causal inner product. The results show aligned concept directions, approximate orthogonality for causally separable concepts, predictive probing, and targeted steering.

  • Concepts are represented as directions in the unembedding space: Counterfactual word differences substantially align with estimated concept directions across 26 of 27 concepts, with thing⇒part as the sole exception.The directions are estimated from normalized means of counterfactual-pair differences and evaluated with leave-one-out estimates against random word-pair differences.
  • The estimated inner product respects causal separability: The estimated causal inner product makes most concept pairs nearly orthogonal and produces a block-diagonal structure organized by semantic similarity.The heatmap evaluates absolute inner products among estimated unembedding representations; language pairs and related concepts form interpretable groups.
  • Concept directions act as linear probes: For French⇒Spanish contexts, the French⇒Spanish concept direction acts as a linear probe, whereas the off-target male⇒female direction has no predictive power.The probe experiment uses random-length French and Spanish Wikipedia contexts rather than counterfactual pairs, and analogous results are reported for all 27 concepts.
  • Concept directions map to intervention representations: Intervention representations constructed through the embedding–unembedding isomorphism are tested by adding concept directions to context representations.The experiments use causally separable concept pairs, selected next-word quadruples, and generated contexts to evaluate targeted changes in logits.
  • The estimated inner product respects causal separability: Intervening in a target direction changes its associated concept while leaving an off-target concept unchanged, across three tested target concepts.For male⇒female, the intervention raises the logit for “queen” over “king” without changing the logit for “King” over “king”; 24 additional concepts appear in the appendix.
  • The estimated inner product respects causal separability: Increasing the male⇒female intervention makes “queen” the most likely completion after “Long live the ” while “king” falls below the top five.The intervention scale ranges from 0 to 0.4, and the most likely completions increasingly reflect the target concept.

5. Discussion and Related Work

The paper formalizes linear representation and argues that its natural variants can be unified, while emphasizing the role of the inner product in interpretation and control.

  • Counterfactual pairs provide the key formalization connecting different notions of linear representation and identifying the relevant inner-product structure.
  • The paper links linear representations to interpretation and control by using concept directions to predict and change model outputs.
  • Linear subspaces in language representations: Prior work has observed linear subspaces across word, sentence, Transformer language-model, and vision-language representations.
  • Measurement, intervention, and mechanistic interpretability: The paper clarifies the relationship between probing, steering, mechanistic interpretability, and the critical role of the inner product, but excludes model parameters and intermediate-layer activations.
  • Causal representation learning: Its causal inner-product results connect to identifiability by showing that semantically respectful geometry is not fixed by usual training, but can be selected under a suitable assumption.

A. Summary of Main Results

The paper summarizes how its definitions, causal inner product, and representation types fit together into a unified account of linear representations.

  • The paper defines unembedding and embedding representations and shows that they also yield measurement and intervention representations, respectively.
  • The causal inner product unifies unembedding and embedding representations through the induced Riesz isomorphism.
  • Figure 6 provides a high-level summary of the connections between the paper’s different notions of linear representations.

B.1. Proof of Theorem 2.2

The proof derives the measurement representation from the softmax sampling distribution and the paper’s definition, yielding a logit-linear expression in the context embedding.

  • Theorem 2.2 expresses the conditional output logit as αλ⊤¯γW for a concept’s unembedding representation ¯γW.
  • The proof writes out the softmax sampling distribution and cancels the two probabilities’ normalizing constants.
  • Definition 2.1 supplies the remaining equality, while randomness in α comes from the random concept-output pair {Y (0), Y (1)}.

B.2. Proof of Lemma 2.4

The lemma proof compares how concept probabilities change across context embeddings, translates those conditions into inner-product constraints, and establishes uniqueness and equivalence under a basis assumption.

  • Conversely, under the same causal-separability and basis conditions, a representation satisfying equation (2.1) is the embedding representation for W.
  • For a causally separable concept Z, changing W leaves the conditional probability of Z unchanged, while W’s probability increases between λ0 and λ1.
  • When outputs encode combinations of W and Z, conditioning on W becomes equivalent to restricting outputs to the corresponding W-conditioned pair.
  • Theorem 2.2 converts these probability conditions into a positive alignment with ¯γW and zero alignment with ¯γZ.
  • If causally separable concept representations together with ¯γW form a basis of R^d, any representation satisfying the conditions is unique up to positive scaling.

B.3. Proof of Theorem 2.5

The proof establishes that adding the embedding representation of a concept increases the corresponding intervention probability while preserving causal separation from another concept. It then connects the causal inner product to the Riesz map from unembedding to embedding representations.

  • Adding c¯λW to λ increases the probability of the output intervention Y(W,1) over Y(W,0) as c increases.The result holds for c ∈ R.
  • The same intervention preserves the probability comparison for a causally separable concept Z, with the corresponding logit expression increasing in c.The proof treats the Z-conditioned intervention analogously.
  • The causal inner product makes the Riesz image of ¯γW nonzero on ¯γW and zero on every causally separable ¯γZ.These orthogonality relations follow from the definition of the causal inner product.
  • Under the stated basis condition, the Riesz image ϕ(¯γW) expresses the unique embedding representation ¯λW up to positive scaling.The proof identifies ϕ(¯γW) with the embedding representation using the orthogonality characterization.

B.5. Proof of Theorem 3.4

The proof derives the matrix form of the causal inner product from a basis of mutually causally separable concept representations. It also shows that the same matrix maps canonical unembedding representations to embedding representations up to positive diagonal scaling.

  • Representing the causal inner product as ¯γ^⊤M¯γ′, the proof uses mutually causally separable basis concepts to identify M as G^-⊤G^-1.Here G contains the canonical unembedding representations as columns.
  • The matrix M maps each canonical unembedding representation to an embedding representation under the theorem’s assumptions.This follows from the proof of Lemma 2.4 and Theorem 3.2.
  • The resulting matrix form is completed by substituting the diagonal positive-scaling relation into M = G^-⊤G^-1.The proof states that this substitution establishes the second half of equation (3.2).

C. Experiment Details

The experiments use LLaMA-2 with counterfactual word pairs and Wikipedia contexts to study concept measurement and intervention. The setup addresses tokenization noise and reports that most concepts appear in unembedding space, while the causal inner product improves semantic geometry across models.

  • Model: LLaMA-2-7B has seven billion parameters, two trillion pre-training tokens, 32,000 tokens, and 4,096-dimensional token embeddings.Ninety percent of the pre-training data is in English.
  • Counterfactual pairs: Counterfactual pairs combine BATS, bilingual resources, CSV data, and ChatGPT-4 generation, but token reuse across words introduces noise that cannot be fully removed.Multi-token words are excluded, while words containing tokens with other meanings remain problematic.
  • Context samples: For measurement, the study samples Wikipedia contexts following bilingual counterfactual token pairs collected from a word2word lexicon and filters pairs that are single vocabulary tokens.The contexts are organized by concept in Table 3.
  • Results: Most concepts are encoded in LLaMA-2’s unembedding representation, although concepts such as thing⇒part are not encoded there.Figure 7 compares counterfactual-pair projections with projections from randomly sampled word pairs.
  • Results: The causal inner product improves on Euclidean similarity in LLaMA-2 and still captures semantics in Gemma-2B, where Euclidean geometry does not.The comparison is based on concept-pair heatmaps, with model-specific differences in Euclidean behavior.

D.4. Additional results from the intervention experiment

Additional intervention analyses examine logit changes, probe behavior, and inner-product geometry across concepts and contexts. They show concept-consistent top-word changes and distinguish causal from non-causal separability in LLaMA-2.

  • Intervention: Adding α¯λC to λ(xj) changes logits across concepts for α ∈ [0, 0.4], with Figure 12 comparing starting and ending points of each intervention.The arrows begin at λ(xj) and end at λC,0.4(xj).
  • Intervention: After intervention in royalty contexts, top predicted words remain aligned with the queen-related concept even when “queen” is not the most likely word.Reported alternatives include “woman,” “queen,” “her,” and “female.”
  • Inner-product sanity check: Causally separable concept projections are empirically uncorrelated, whereas non-causally separable concepts such as verb⇒3pSg and verb⇒Ving are correlated.Each plotted point corresponds to an unembedding vector from the 32K-token vocabulary.
  • Geometry: For LLaMA-2, causally separable concepts are approximately orthogonal under both the estimated causal and Euclidean inner products.Figure 8 presents heatmaps of absolute pairwise inner products.
  • Geometry: For Gemma-2B, the causal inner product preserves approximate orthogonality while the Euclidean inner product does not capture semantics.Figure 9 contrasts the two inner products using concept-pair heatmaps.
Loading 2311.03658v2…