On the Origins of Linear Representations in Large Language Models (2403.03867v1)

Abstract: Recent works have argued that high-level semantic concepts are encoded "linearly" in the representation space of LLMs. In this work, we study the origins of such linear representations. To that end, we introduce a simple latent variable model to abstract and formalize the concept dynamics of the next token prediction. We use this formalism to show that the next token prediction objective (softmax with cross-entropy) and the implicit bias of gradient descent together promote the linear representation of concepts. Experiments show that linear representations emerge when learning from data matching the latent variable model, confirming that this simple structure already suffices to yield linear representations. We additionally confirm some predictions of the theory using the LLaMA-2 LLM, giving evidence that the simplified model yields generalizable insights.

• The paper establishes that linear representations in LLMs arise from log-odds matching and the implicit bias of gradient descent during next-token prediction.
• It introduces a latent variable model that encodes human-interpretable concepts using a probabilistic Markov random field for context mapping.
• Empirical validations show that gradient descent aligns similar concept vectors while maintaining orthogonal separations, enhancing model interpretability.

On the Origins of Linear Representations in LLMs

Introduction

The paper "On the Origins of Linear Representations in LLMs" (2403.03867) explores the phenomenon where high-level semantic concepts are encoded linearly within the vector space of LLMs. This linearity, observed empirically in models like LLaMA-2, raises fundamental questions about how such representations originate during the training process. The authors introduce a novel mathematical framework to explore the dynamics of next-token prediction, proposing that both the objective functions used in training and the implicit biases in gradient descent contribute to this linear structuring.

Latent Variable Model for Concept Dynamics

The study introduces a latent variable model where context sentences and next tokens reflect latent binary concept variables. This model encapsulates human-interpretable concepts within the latent space, allowing for an analytical exploration of how concepts are represented in embeddings and unembeddings.

Latent Space and Mapping: The model defines a set of binary variables representing different concepts, structured as a Markov random field. Each context maps to these concept variables, introducing a probabilistic component to next-token prediction. The injective mapping from latent space to token space ensures each concept combination is uniquely represented in the output space.

Linearity through Log-Odds and Gradient Descent

The research establishes that linear representations stem from two primary sources:

Log-Odds Matching: The paper initially demonstrates that when the log-odds of learned probabilities across concepts match, a linear structure inherently arises in subspaces. This finding is backed by prior evidence from studies on word embeddings.

Implicit Bias of Gradient Descent: Extending beyond log-odds, the authors link the linear encoding of representations to the biases induced by gradient descent. Through theoretical underpinnings, they illustrate how gradient descent inherently promotes the alignment and linearity of concept vectors over successive training iterations.

Figure 1: Unembedding steering vectors of the same concept in LLaMA-2 have nontrivial alignment, but steering vectors of different concepts are represented almost orthogonally.

Orthogonality and Semantic Structures

The study further examines how unrelated concepts, those not linked in the Markov random field, exhibit orthogonal representations. This orthogonality, unaccounted for by standard training objectives, emerges as a byproduct of the learned Euclidean geometry in LLMs, which inadvertently respects semantic distances.

Euclidean Geometry: The implicit bias of gradient descent not only causes concept alignment within representations but also ensures that disparate concepts maintain orthogonal separations, refining how semantics are geometrically encoded.

Empirical Validation and Theoretical Insights

Extensive experiments using simulated data from the proposed latent model confirm the emergence of linear and orthogonal representations. These results are consistent even when varying the set of contexts or latent dimensions, highlighting the robustness of the theoretical claims.

In practical scenarios with LLMs like LLaMA-2, empirical observations support the presence of these geometric structures. Notably, the context and concept vectors align as predicted, substantiating the theoretical models.

Conclusion

This research provides a comprehensive explanation for the linearity observed in LLM representations, attributing it to systemic biases in training processes. The introduction of a latent variable model offers a powerful framework for interpreting such phenomena, with implications extending to improving model interpretability and potentially guiding future AI developments. Through its synthesis of theory and empirical validation, the study enhances our understanding of the geometric principles governing large-scale AI models.