Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Not All Language Model Features Are One-Dimensionally Linear

Joshua Engels, Eric J. Michaud, Isaac Liao, Wes Gurnee, Max Tegmark

TL;DR

Not All Language Model Features Are One-Dimensionally Linear addresses whether language models rely exclusively on one-dimensional feature directions or also use irreducible multi-dimensional representations. The authors formalize multi-dimensional features, define reducibility measures (S and M_epsilon), and propose a multi-dimensional superposition hypothesis. They introduce a scalable approach using sparse autoencoders to automatically discover irreducible multi-dimensional features in GPT-2 and Mistral 7B, uncovering circular representations corresponding to days of the week and months of the year. They provide causal evidence via activation patching and EVR that these circular features are used in modular arithmetic tasks and exhibit continuity across time. The work suggests that understanding multi-dimensional representations is crucial for mechanistic decomposition and challenges the dominance of the one-dimensional linear representation hypothesis.

Abstract

Recent work has proposed that language models perform computation by manipulating one-dimensional representations of concepts ("features") in activation space. In contrast, we explore whether some language model representations may be inherently multi-dimensional. We begin by developing a rigorous definition of irreducible multi-dimensional features based on whether they can be decomposed into either independent or non-co-occurring lower-dimensional features. Motivated by these definitions, we design a scalable method that uses sparse autoencoders to automatically find multi-dimensional features in GPT-2 and Mistral 7B. These auto-discovered features include strikingly interpretable examples, e.g. circular features representing days of the week and months of the year. We identify tasks where these exact circles are used to solve computational problems involving modular arithmetic in days of the week and months of the year. Next, we provide evidence that these circular features are indeed the fundamental unit of computation in these tasks with intervention experiments on Mistral 7B and Llama 3 8B, and we examine the continuity of the days of the week feature in Mistral 7B. Overall, our work argues that understanding multi-dimensional features is necessary to mechanistically decompose some model behaviors.

Not All Language Model Features Are One-Dimensionally Linear

TL;DR

Not All Language Model Features Are One-Dimensionally Linear addresses whether language models rely exclusively on one-dimensional feature directions or also use irreducible multi-dimensional representations. The authors formalize multi-dimensional features, define reducibility measures (S and M_epsilon), and propose a multi-dimensional superposition hypothesis. They introduce a scalable approach using sparse autoencoders to automatically discover irreducible multi-dimensional features in GPT-2 and Mistral 7B, uncovering circular representations corresponding to days of the week and months of the year. They provide causal evidence via activation patching and EVR that these circular features are used in modular arithmetic tasks and exhibit continuity across time. The work suggests that understanding multi-dimensional representations is crucial for mechanistic decomposition and challenges the dominance of the one-dimensional linear representation hypothesis.

Abstract

Recent work has proposed that language models perform computation by manipulating one-dimensional representations of concepts ("features") in activation space. In contrast, we explore whether some language model representations may be inherently multi-dimensional. We begin by developing a rigorous definition of irreducible multi-dimensional features based on whether they can be decomposed into either independent or non-co-occurring lower-dimensional features. Motivated by these definitions, we design a scalable method that uses sparse autoencoders to automatically find multi-dimensional features in GPT-2 and Mistral 7B. These auto-discovered features include strikingly interpretable examples, e.g. circular features representing days of the week and months of the year. We identify tasks where these exact circles are used to solve computational problems involving modular arithmetic in days of the week and months of the year. Next, we provide evidence that these circular features are indeed the fundamental unit of computation in these tasks with intervention experiments on Mistral 7B and Llama 3 8B, and we examine the continuity of the days of the week feature in Mistral 7B. Overall, our work argues that understanding multi-dimensional features is necessary to mechanistically decompose some model behaviors.
Paper Structure (40 sections, 3 theorems, 20 equations, 26 figures, 4 tables, 1 algorithm)

This paper contains 40 sections, 3 theorems, 20 equations, 26 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Contributions
  3. Related Work
  4. Definitions
  5. Multi-Dimensional Features
  6. Superposition
  7. Sparse Autoencoders Find Multi-Dimensional Features
  8. Circular Representations in Large Language Models
  9. Intervening on Circular Day and Month Representations
  10. Intervening with the SAE Plane
  11. Continuity of Circular Representations
  12. Discussion
  13. Limitations:
  14. Multi-Dimensional Feature Capacity
  15. More on Reducibility
  16. ...and 25 more sections

Key Result

Theorem 1

For any $d'$ and $\delta$, it is possible to choose $\frac{1}{d_{\max}}e^{C_1 (d / d'^2) \delta^2}$ pairwise $\delta$-orthogonal matrices ${\mathbf A}_i \in \mathbb{R}^{n_i \times d'}$ for some constant $C_1$. Furthermore, it is not possible to choose more than $e^{C_2(d - d_{\max}\delta \log(\frac{

Figures (26)

  • Figure 1: Circular representations of days of the week, months of the year, and years of the 20th century in layer 7 of GPT-2-small colored by the token they fire on. These representations were discovered via clustering SAE dictionary elements, described in \ref{['sec:sae_search']}. Points are colored according to the token which created the representation. See \ref{['fig:gpt2-3projectionsperrep']} for other axes and \ref{['fig:mistral-cluster-reconstructions']} for similar plots for Mistral 7B.
  • Figure 1: Aggregate model accuracy on days of the week and months of the year modular arithmetic tasks. Performance broken down by problem instance in \ref{['appendix:more_circular_plots']}.
  • Figure 2: Empirical $\epsilon$-mixture index and separability index for the "days of the week" cluster along PCA components 2 and 3. Left: The $\epsilon$ band parameterized by $\mathbf{v}$ and $c$ that the optimization procedure found contained the highest fraction of points. Mid: Dot products of points in the feature distribution with the $\epsilon$ band; $M_{\epsilon}({\bm{f}})$ is the percent of dot products within $\epsilon = 0.1$ of $0$. Right: Estimated mutual information for different rotations of the space; $S({\bm{f}})$ is the minimum over all rotations. This point cloud has a lower $\epsilon$-mixture index and higher separability index than PCA projections within typical clusters (see \ref{['fig:mixture_vs_separability_scatter']}), indicating that it is more likely to be an irreducible multi-dimensional feature.
  • Figure 3: Mixture index and separability index of GPT-2 features. Features from \ref{['fig:gpt2-nonlinears']}, which we had manually identified, score highly as candidate multidimensional features with these metrics.
  • Figure 4: Top two PCA components on the $\alpha$ token. Colors show $\alpha$. Left: Layer $30$ of Mistral on Weekdays. Right: Layer $5$ of Llama on Months.
  • ...and 21 more figures

Theorems & Definitions (12)

  • Definition 1: Feature
  • Definition 2
  • Definition 3: Separability Index and $\epsilon$-Mixture Index
  • Definition 4: $\delta$-orthogonal matrices
  • Theorem 1
  • Lemma 1
  • proof
  • Theorem 1
  • proof
  • Definition 5: Representation Space
  • ...and 2 more