Geometry of Hidden State

The internal world of a large language model is a labyrinth of high-dimensional vectors. Each forward pass is a journey through this geometric space, where abstract concepts are manipulated by matrices and non-linear functions. For years, our best descriptions of this process were frustratingly opaque. But a recent thread of research suggests that within this complexity lies a surprising amount of simple, linear structure.

This isn't just a curiosity. If concepts and relations are encoded linearly, we can build more reliable tools to interpret what models know, predict how they will behave, and even perform surgical interventions to edit their knowledge or control their outputs. Three key papers-one on linguistic subspaces, one on relational maps, and one on the underlying geometry itself-form a powerful, unified narrative. This post unpacks their findings to explain what "linearity" really buys us.

What We Mean by “Linearity”

Before diving in, it’s critical to untangle a few distinct ideas often conflated under the banner of "linearity." The claim is not that transformers are secretly linear models. They are deeply non-linear. The claim is that we can often use linear tools to meaningfully describe and manipulate their internal representations.

Here are three increasingly specific senses of linearity:

Caveat: The choice of geometry-how we measure distance and angles-fundamentally changes whether these claims appear true. Most analyses default to the Euclidean inner product (and its cousin, cosine similarity). But as the work on the Linear Representation Hypothesis (LRH) shows, this can be misleading. A "causal" inner product, which respects the statistical independence of concepts in language, can reveal linear structure that is invisible to standard metrics, or show that apparent linearity was an artifact. Declaring the geometry is non-negotiable for making a rigorous claim.

Background & Notation

Let's establish a common language. We view a transformer through the lens of its residual stream. A hidden state $r^{(l)} \in \mathbb{R}^d$ at layer $l$ and a specific token position represents the accumulated information up to that point. The model's final prediction depends on the final state, typically via the unembedding matrix $W^U$ which maps $r^{(L)}$ to logits over the vocabulary.

• A concept $c$ (e.g., gender, tense) can be represented by a linear subspace $S_c \subset \mathbb{R}^d$. We can isolate the component of a hidden state that lies in this subspace using a projection operator $P_c$.
• A Linear Relational Embedding (LRE) for a relation $r$ is an affine map $\text{LRE}_r(s) = W_r s + b_r$. It takes a subject representation $s$ (a hidden state at a certain layer/position) and approximates the corresponding object representation. This map is often estimated by taking a first-order Taylor approximation of the model's computation around a reference subject $s_0$.
• The Causal Inner Product $\langle \cdot, \cdot \rangle_{\text{causal}}$ is a non-Euclidean metric designed to make concepts that are semantically independent in the world (e.g., tense and sentiment) orthogonal in the representation space. It provides a formal basis for when talking about "directions" is meaningful and connects the geometry of representations to counterfactual changes in model inputs and outputs.

With these tools, we will now examine the evidence.

Paper A (CoNLL 2021): Linear Subspaces in Practice

The work by Hernandez and Andreas (2021) provided strong, early evidence that linear structures weren't just a convenient analytical tool, but a core aspect of how models like BERT organize linguistic information.

Their method involved identifying subspaces for various linguistic features, from simple ones like part-of-speech to complex ones like the grammatical role of a word in a dependency parse. They trained simple linear probes to find these subspaces and then tested their properties.

Key Findings:

• Low-Dimensionality: Subspaces for many distinct linguistic features were surprisingly low-dimensional, often occupying just 1-10 dimensions out of the full 768 in BERT-base. This suggests information is compressed into specific, targeted geometric structures.
• Hierarchical Structure: The subspaces reflected linguistic hierarchies. For example, the subspace for verb would contain the more specific subspaces for transitive verb and intransitive verb. This geometric containment mirrors the conceptual relationship.
• Distributed Representation: These subspaces were not aligned with individual neurons. The basis vectors for a concept's subspace were distributed across many coordinates, confirming that features are stored collectively, not in single "grandmother cells."

Most importantly, the paper established a causal link between these subspaces and model behavior. They performed interventions by projecting a hidden state onto a chosen subspace or removing its component in that subspace. For example, by manipulating the subspace for a dependency arc, they could systematically change which words the model attended to and alter its masked language modeling predictions in a controlled way.

Caveat: This work demonstrates that linear signals are present and can be causally manipulated. It does not, however, prove that the model always uses that signal for every downstream task. A feature might be linearly encoded but ignored by later layers in certain contexts. The causal link is strong evidence for its functional role, but not a guarantee of its universal application.

Paper B (ICLR 2024): Linear Relation Decoding

While the CoNLL paper focused on properties of single words or tokens, Hernandez et al. (2024) investigated the geometry of relations between entities. Their central idea is the Linear Relational Embedding (LRE), which tests whether the complex computation mapping a subject to an object can be simplified to an affine transformation.

To find the relational map $(W_r, b_r)$ for a relation like country -> capital, they compute the local Jacobian of the model's transformation from the subject's hidden state to the object's hidden state. By averaging this estimate over a few example subjects, they obtain a robust linear operator that approximates the relation.

Empirical Scope and Limits:

The team tested this method on 47 different relations, spanning factual knowledge, commonsense reasoning, linguistic patterns, and even social biases.

• Successes: The linear approximation worked remarkably well for about half of the relations. For these, applying the learned LRE to a subject's hidden state could reliably produce the hidden state of the correct object, often for subjects unseen during the LRE's estimation.
• Failures: The approximation consistently failed for relations where the set of possible objects is vast and arbitrary, such as company -> CEO or person -> birth city.

Intuition: This pattern suggests a compelling hypothesis. For relations with a constrained, structured output space (like capitals or languages), the transformer can "pre-compile" the computation into a nearly-linear map within a specific processing window (i.e., across a few layers). For open-ended relations, it must rely on more complex, non-linear search and retrieval mechanisms that cannot be captured by a single LRE.

This work also introduced the Attribute Lens, a practical application of LREs. By applying a pre-computed LRE (e.g., for country -> capital) to the hidden states at every token and every layer for an input like "The government of France announced," one can visualize exactly where and when the model computes the capital "Paris," even if it is never explicitly mentioned.

Paper C (ICML 2024): Definitions & the Causal Inner Product

The Linear Representation Hypothesis (LRH) paper by Park, Choe, and Veitch (2024) provides the theoretical bedrock for these empirical findings. It asks a fundamental question: what does it even mean for a concept to be "linearly represented"?

They provide two formal definitions grounded in counterfactuals:

The central insight is that these two definitions are unified by choosing the right geometry. The standard Euclidean inner product is a poor choice because it treats all dimensions equally. Language is not like that; some concepts are related, others are independent. The authors introduce the causal inner product, a metric derived from the model's unembedding matrix that makes semantically independent concepts orthogonal.

What the Causal Inner Product Buys Us:

• Principled Directions: It provides a rigorous way to define "concept directions" that align with counterfactual meaning. A steering vector for gender should, ideally, be orthogonal to a steering vector for tense under this metric.
• Improved Separability: The paper demonstrates empirically on LLaMA-2 that concepts that are muddled under cosine similarity become cleanly separable using the causal inner product. For instance, contexts about France and Spain might seem similar in Euclidean space, but the causal metric can better isolate the countryspecific components.
• Unification: With the causal inner product, the "probe" direction (from the output-space definition) and the "steering" direction (from the input-space definition) become one and the same. This provides a single, coherent geometric object for both reading and writing a concept.

Experiments with LLaMA-2 confirm that linear representations meeting these causal definitions exist for a variety of concepts, and that the choice of inner product is not merely academic-it fundamentally determines whether these structures are found.

Synthesis - One Geometry, Many Tools

These three papers are not independent observations; they are three views of the same underlying phenomenon. Together, they paint a remarkably coherent picture of how linear geometry enables complex reasoning in non-linear models.

• The subspace picture (CoNLL '21) establishes that concepts live in specific, low-dimensional linear manifolds and that these manifolds are causally potent.
• The affine relational maps (ICLR '24) show that for many common knowledge-retrieval tasks, the multi-step computation from a subject's subspace to an object's subspace can be locally approximated by a simple affine hop.
• The causal geometry (ICML '24) provides the rigorous mathematical language to justify it all. It explains why we can talk about directions and subspaces, tells us which metric to use, and unifies the acts of observing a concept (probing) and controlling it (steering).

The unified story is this: To find the simple linear structure hidden inside a transformer, you need to look at the right slice of the representation (the right layer and token position) and measure it with the right metric (the causal inner product). In that sweet spot, complex operations often resolve into simple geometric translations and projections. Outside of that slice, or with the wrong metric, the model's behavior remains stubbornly non-linear and opaque.

Practical Guidance

For researchers and practitioners working on interpretability, this body of work suggests a more responsible way to make and evaluate claims about linearity.

My View

• Opinion: Linear tools are at their best when treated as local approximations of a model's behavior. They are powerful instruments for generating hypotheses, performing quick diagnostic tests, and executing surgical edits on model knowledge. They are not, however, a complete theory of how transformers compute. Their power lies in simplifying a small piece of a complex puzzle, but we must always remember and declare the boundaries of that simplification.
• Intuition: I suspect that transformers operate by dynamically linearizing sub-problems. When the model needs to map a country to its capital, it routes the subject representation through a sequence of blocks that, in aggregate, perform an approximately affine map. Once that attribute is computed and added to the residual stream, subsequent blocks can perform other, unrelated non-linear operations. The linearity is a temporary state in a much longer, piecewise computation.
• Caveat: The most common mistake in this area is to confuse the presence of a linearly decodable signal with its use in the model's final decision. A probe can often detect a feature at an early layer, but that doesn't mean the model uses it. Only a causal experiment-an intervention that shows a change in the representation leads to a predictable change in output-can close that loop.

Limits & Open Questions

This is a fast-moving area, and our understanding is far from complete. Key open questions include:

• Faithfulness Bounds: When, precisely, does the causal inner product lead to substantively different conclusions than standard cosine similarity? Are there classes of problems where the simpler metric is "good enough"?
• Portability: Do linear subspaces and LREs transfer across model families (e.g., LLaMA to GPT) or model sizes? Is the geometric structure of knowledge universal, or is it an idiosyncratic feature of each training run?
• Compositionality: Can we chain multiple linear operations? For example, can we compose the LRE for author -> book with the LRE for book -> main character to reason about an author's characters? Early evidence suggests this is difficult, likely because the non-linearities between steps are crucial.
• Effect of Instruction Tuning: How does fine-tuning on human preferences (RLHF) and other alignment techniques warp the pristine geometric structures found in base models? Does it consolidate concepts into clearer subspaces or shatter them into context-dependent, non-linear fragments?

The geometry of hidden states is no longer a complete black box. By choosing the right lens, we are beginning to resolve its structure, finding straight lines and simple planes in a space we once thought was impenetrably complex.

References