PaperlayerView exact revision on arXiv

Source-linked AI summary

Not All Language Model Features Are One-Dimensionally Linear

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

arXiv:2405.14860v3cs.LG

TL;DR

The paper asks whether some language-model representations are inherently multidimensional rather than one-dimensional features. It defines irreducible multidimensional features and uses sparse autoencoders to find them in GPT-2 and Mistral 7B. It discovers interpretable circular representations used for modular arithmetic, while acknowledging uncertainty about feature discovery and the adequacy of its theory.

  • Problem

    The paper addresses whether language models use inherently multidimensional representations, a question left insufficiently unified by prior one-dimensional feature work.

  • Method

    The authors define irreducible multidimensional features and develop a sparse-autoencoder-based method to automatically identify them in language models.

  • Results

    The method finds interpretable circular representations for days and months, and Mistral 7B and Llama 3 8B use them for modular addition tasks.

  • Takeaways & Limitations

    Understanding multidimensional features is necessary to mechanistically decompose some model behaviors.

  • Takeaways & Limitations

    The paper leaves unclear why few interpretable multidimensional features were found and whether its relaxed statistical definitions best describe model representations.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Mechanistic interpretability seeks to reverse-engineer language models into features and circuits, but existing work leaves the nature of their features insufficiently unified. This paper formalizes and searches for multidimensional features, finding interpretable circular representations and showing their use in modular arithmetic.

  • Mechanistic interpretability reverse-engineers trained models into variables called features and programs called circuits.
  • Prior work found multidimensional circles in toy models and one-dimensional high-level representations in large language models.
  • The paper bridges these directions by formalizing, investigating, and systematically searching for multidimensional language-model features.
  • 1.1 CONTRIBUTIONS: The authors generalize the one-dimensional feature definition and update the superposition hypothesis for multidimensional features.
  • 1.1 CONTRIBUTIONS: A sparse-autoencoder test automatically identifies multidimensional representations in GPT-2 and Mistral 7B, including circular day and month representations.
  • 1.1 CONTRIBUTIONS: Mistral 7B and Llama 3 8B use these circular representations for modular addition involving days of the week and months of the year.

2 RELATED WORK

Prior interpretability research established linear features and studied multidimensional structures mainly in toy or specialized settings. This paper builds on that literature while distinguishing its circular concept representations from earlier positional or noncausal findings.

  • Linear Representations: Word embeddings and sequence models provide evidence for one-dimensional directions representing semantic concepts, board positions, truth values, and quantities.
  • Linear Representations: Sparse autoencoders have been used to decompose model feature spaces into overcomplete bases of linear features under superposition assumptions.
  • Multi-Dimensional Representations: Multidimensional representation research is comparatively limited, with examples including fractal next-token structures and helical GPT-2 position vectors.
  • Multi-Dimensional Representations: The paper distinguishes its causal circular concept features from GPT-2 positional circles and a noncausal numerical U-shape found only within a specific prompt distribution.
  • Circuits: Circuits research explains specific behaviors using components such as layers, neurons, or sparse-autoencoder features.
  • Interpretability for Arithmetic Problems: Earlier modular-arithmetic studies found circular representations and clock- or pizza-style algorithms, but focused on small models trained only on modular arithmetic.

3 DEFINITIONS

The paper defines irreducible multidimensional features by ruling out decomposition into independent or non-co-occurring lower-dimensional features. It then relaxes these definitions into empirical indices and incorporates irreducible low-dimensional features into superposition.

  • 3.1 MULTI-DIMENSIONAL FEATURES: A feature is irreducible when it cannot be decomposed into statistically independent or non-co-occurring lower-dimensional features.
  • 3.1 MULTI-DIMENSIONAL FEATURES: The separability condition requires the transformed joint distribution to factorize as p(a, b) = p(a)p(b).
  • 3.1 MULTI-DIMENSIONAL FEATURES: The mixture condition represents the distribution as disjoint components, so the lower-dimensional features cannot be active simultaneously.
  • 3.1 MULTI-DIMENSIONAL FEATURES: The separability index minimizes mutual information over allowed decompositions, with smaller values indicating greater separability.
  • 3.1 MULTI-DIMENSIONAL FEATURES: The ϵ-mixture index measures how often an active feature can be projected near zero, with larger values indicating greater mixture structure.
  • 3.1 MULTI-DIMENSIONAL FEATURES: Empirical procedures sweep rotations for mutual information and use gradient descent to find the ϵ-band containing the largest feature fraction.
  • 3.2 MULTI-DIMENSIONAL SUPERPOSITION HYPOTHESIS: The multidimensional superposition hypothesis represents hidden states as sums of sparse, low-dimensional irreducible features with pairwise δ-orthogonal subspaces.

4 SPARSE AUTOENCODERS FIND MULTI-DIMENSIONAL FEATURES

The paper uses sparse autoencoders to cluster hidden-state dictionary elements, reconstruct cluster-restricted activations, and test them for irreducible multi-dimensional structure. Applied to GPT-2 and Mistral 7B, this method recovers interpretable circular representations of calendrical concepts.

  • Method: Sparse autoencoders deconstruct hidden states into sparse sums over an over-complete dictionary, providing the basis for searching multi-dimensional features.The dictionary contains m vectors in R^d, and the training objective includes a sparsity penalty.
  • Method: Clustering dictionary elements by pairwise cosine similarity yields approximately T-orthogonal subspaces that may correspond to irreducible multi-dimensional features.The method builds a graph from dictionary elements, prunes edges below threshold T, and uses connected components or spectral clustering.
  • Method: For each cluster, the method reconstructs activations using only that cluster’s dictionary elements, then inspects PCA projections or applies irreducibility tests.Points with no active dictionary elements from a cluster are ignored.
  • Evaluation: The method is applied to GPT-2 SAEs at every layer and to Mistral 7B SAEs trained on layers 8, 16, and 24.The procedure also succeeds on synthetic toy datasets with irreducible multi-dimensional features.
  • Discovered features: GPT-2 contains interpretable circular representations in which days, months, and years are arranged in order, while Mistral 7B contains circular day and month representations.The clearest circular structure usually appears in the second and third PCA dimensions; the first often represents intensity as radius.

5 CIRCULAR REPRESENTATIONS IN LARGE LANGUAGE MODELS

The paper tests whether circular representations of weekdays and months are used causally for modular arithmetic in Mistral 7B and Llama 3 8B. Interventions show that these circular subspaces substantially account for the models’ computations, and SAE-discovered circles remain effective across layer shifts.

  • Circular representations: Mistral 7B and Llama 3 8B represent weekdays and months circularly in the α token across multiple layers.PCA projections show the circle as the highest-varying two components of α’s representation.
  • Intervention method: The circular intervention replaces α’s projected representation with a clean circular point while ablating other activation dimensions.The method learns a probe from the top five PCA components and intervenes on the resulting two-dimensional subspace.
  • Causal tests: Across models and tasks, early-layer circular-subspace interventions nearly match whole-layer patching and usually outperform patching the top PCA dimensions.This effect is especially strong for the Weekdays task.
  • Experimental scope: Intervention effects drop on layers 15–17 because α is copied to the final token, while SAE-cluster circles also support intervention.The experiments cover 49 Weekday prompts and 144 Month prompts with alternative β values used as clean runs.
  • Causal tests: Off-distribution interventions imply that Mistral encodes α in the angle of the circle, because sweeping radius and angle changes the predicted γ.The sweep evaluates positions across the circular subspace rather than only observed circumference points.
  • SAE-plane intervention: The SAE-discovered circle slightly reduces intervention performance but is substantially more robust to layer shifts than a manually trained circular probe.On layer 6, the layer 8 probe gives average logit difference 0.029, whereas the layer 8 SAE probe gives -2.32; on layer 8, performance changes from -2.58 to -2.01.

6 DISCUSSION

The discussion presents multi-dimensional representations as a refinement of the one-dimensional feature hypothesis and argues that they matter for mechanistic decomposition. It also emphasizes that the evidence is preliminary and that the best theory of model representations remains unresolved.

  • Discussion: The paper argues that irreducible multi-dimensional representations refine the one-dimensional feature hypothesis and may be necessary to uncover model algorithms.Sparse autoencoders can discover such representations, but treating individual SAE features as computational mediators remains an open question.
  • Limitations: The authors report uncertainty about why their search found relatively few interpretable multi-dimensional features.Possible explanations include missed interpretations, higher-dimensional structure, insufficient clustering, or genuinely few such features.
  • Limitations: The multi-dimensional superposition hypothesis remains preliminary because it is unclear whether it best describes the representations models use.Future work may examine new decomposition methods, higher-dimensional representations, and whether model representations are necessarily non-linear.

A MULTI-DIMENSIONAL FEATURE CAPACITY

The paper shows that models can represent exponentially many higher-dimensional features, while defining and testing irreducibility through separability and mixture criteria.

  • Capacity bounds: Theorem 1 gives exponential lower and upper bounds on the number of pairwise δ-orthogonal higher-dimensional feature subspaces a model can represent.The lower bound scales as 1/d_max times e^(C1(d/d′^2)δ^2), while the upper bound has a corresponding exponential form.
  • Capacity bounds: These bounds imply that models can still represent an exponential number of higher-dimensional features.The lower bound may be reasonably tight, which would favor encoding features in the smallest dimensional spaces possible.
  • Caveat: The capacity analysis assumes the worst case that all features have dimension d_max, so its effect on real-model capacity is unlikely to be as extreme.In practice, many features may be one-dimensional or low-dimensional.
  • Irreducibility criteria: The irreducibility definition targets representations that fill a multi-dimensional space and cannot be decomposed into independent or non-co-occurring lower-dimensional features.The mixture condition addresses whether the distribution fills the space, while separability rules out independent co-occurring components such as latitude and longitude.
  • Irreducibility criteria: The authors operationalize reducibility with a separability index S(f) and an ϵ-mixture index M_ϵ(f).The separability index searches over rotations that split a representation and evaluates mutual information between the resulting components.

D TOY CASE OF TRAINING SAES ON CIRCLES

Toy experiments show how sparse autoencoders reconstruct irreducible circular features and how clustering can recover multiple circles from mixed activations.

  • Synthetic setup: The toy datasets contain multiple 2D irreducible circles, motivated by circular representations learned for modular arithmetic.The synthetic activations combine circular features in separate orthogonal planes.
  • Single-circle reconstruction: With m = 2 dictionary elements, the SAE uses both features for each point, whereas with m = 10, features spread around the circle and only a subset activates per input.The m = 2 model also uses the decoder bias to center the reconstructed circle at the origin.
  • Multiple-circle reconstruction: For mixed activations in R10, points are sampled from one of two orthogonal unit-circle planes with equal probability.The SAE uses m = 64 dictionary elements with the same hyperparameters as the single-circle experiment.
  • Automatic rediscovery: The clustering procedure recovers the two original circles almost exactly from the SAE dictionary elements and gated reconstructions.The clusters align with the two planes, and reconstruction suppresses SAE features outside each cluster.

E TRAINING MISTRAL SAES

The paper trains large sparse autoencoders and applies clustering and reconstruction analyses to search for multi-dimensional representations in GPT-2 and Mistral 7B.

  • SAE training: Mistral SAEs are trained on over one billion tokens from subsets of the Pile and Alpaca datasets.The SAEs cover layers 8, 16, and 24 of Mistral 7B.
  • SAE training: Each Mistral SAE uses a 16× expansion factor with 65536 dictionary elements.Training uses an L_p sparsity penalty with p = 1/2 and coefficient λ = 0.012.
  • Clustering method: The high-level method clusters SAE dictionary elements, reconstructs activations using only each cluster, and returns the resulting candidate features.The algorithm computes cosine similarities, forms clusters, gates encodings by cluster membership, and reconstructs selected activations.
  • GPT-2 analysis: For GPT-2-small, spectral clustering is applied to roughly 25k layer-7 SAE features, with projections of reconstructed activations inspected across principal components.The analysis uses 1000 clusters and manually examines roughly 500 of them.
  • GPT-2 analysis: Principal-component projections are used to visualize days-of-week, months-of-year, and twentieth-century-years representations.These projections reveal dimensions beyond those shown in the initial figure.

F.2 MISTRAL 7B METHODS AND RESULTS

For Mistral 7B, a graph-based clustering method rediscovers circular representations of weekdays and months, while results vary with the similarity threshold and the clustering approach remains improvable.

  • Clustering method: The Mistral clustering algorithm connects each dictionary element to its k = 2 nearest neighbors, prunes edges below τ = 0.5, and returns connected components.Directed nearest-neighbor edges are first made undirected before pruning.
  • Results: Roughly 2700 clusters were found in the Mistral 7B layer-8 SAE, and manual inspection rediscovered circular representations of weekdays and months.The inspected clusters contained between 2 and 1000 dictionary elements.
  • Results: The analysis did not find other obviously interesting and clearly irreducible features among the inspected Mistral clusters.This bounds the reported discovery result to the circular weekday and month representations identified by the authors.
  • Sensitivity analysis: Varying k has minimal effect, whereas τ produces regimes where features either merge into one cluster, reveal the weekday cluster, or separate into individual clusters.The weekday cluster is identifiable only in the intermediate threshold regime.
  • Limitations: The authors identify better graph-based clustering and combined cosine–Jaccard edge weights as future improvements.The current graph-based technique is therefore presented as improvable rather than definitive.

H.3 ERROR BAR CALCULATION

This section describes how intervention error bars are computed and summarizes the associated Weekdays and Months analyses. It also reports circular structure in model representations and the prompt coverage used for evaluation.

  • H.3 ERROR BAR CALCULATION: 96% error bars are computed as confidence intervals for each intervention-method and layer pair across prompts.The calculation assumes normally distributed errors.
  • H.3 ERROR BAR CALCULATION: The standard error formula uses the sample mean, a z score slightly above 2, and the standard error of the samples.The standard error is the standard deviation divided by the square root of the sample count.
  • H.3 ERROR BAR CALCULATION: Months analyses have smaller error bars because they contain 1584 intervention-effect values, compared with 294 for Weekdays.These counts come from 12 ∗12 ∗11 Months values and 7 ∗7 ∗6 Weekdays values.
  • I.0.1 BASIC PLOTS: PCA projections show circular α structure across many, but not all, layers, while linear structure is attributed largely to β.Mistral Weekdays activations projected into the SAE-discovered plane are also clearly circular.
  • J PATCHING: Attention and MLP patching results are averaged over runs varying β with fixed α and varying α with fixed β.The experiments use 20 runs of each type, and patching evaluates 40 problem pairs in total.

J PATCHING

This section uses patching and EVR to localize and characterize computation of γ. The analysis places γ computation in late MLPs and reveals a circular representation after removing α and β interference.

  • J PATCHING: γ is computed on the final token around layers 20 to 25 in Llama 3 8B Months.The corresponding deconstruction is compared with Mistral 7B Weekdays results.
  • J PATCHING: Attention heads add almost no explanatory power when γ is included before layer 25, indicating they do not participate in computing γ.This conclusion follows after γ has already been computed by that point.
  • J PATCHING: Layer-wise patching identifies a γ circuit consisting of MLPs on α and β, a copy to the token before γ, and further MLPs there.Fine-grained patching finds only a few responsible attention heads for writes to the preceding token.
  • J PATCHING: Patching localizes where to seek γ but cannot determine how γ is represented.The representation is therefore analyzed with Explanation via Regression.
  • K EXPLANATION VIA REGRESSION (EVR): EVR regresses interpretable input functions onto hidden states, using r2 to measure explained variance and residuals to expose unexplained representation components.This provides a representation analysis beyond direct PCA inspection.
  • K.1 USING EVR TO UNCOVER A CIRCULAR REPRESENTATION FOR γ: After removing components explained by one-hot α and β functions, layer 25 Mistral 7B residuals form a clear circle in γ.The result suggests a trigonometry-based algorithm such as the clock or pizza algorithm.

K.2 MORE EXPERIMENTS WITH EVR

This section iteratively applies EVR to decompose Weekdays and Months hidden states into interpretable functions. Residual visualizations track how α, β, and γ-dependent structure is explained across layers.

  • K.2 MORE EXPERIMENTS WITH EVR: EVR greedily adds interpretable functions by regressing current features, inspecting residual errors, and adding functions that represent those errors.The process stops when most variance is explained.
  • K.2 MORE EXPERIMENTS WITH EVR: Treating each EVR function as a feature makes regression coefficients identify the directions those features occupy in activation space.The same analysis indicates what can and cannot be extracted by a linear probe.
  • K.2 MORE EXPERIMENTS WITH EVR: Residual RGB plots encode the top three PCA components of errors with red, green, and blue across α- and β-indexed heatmaps.The plots visualize how residual errors depend on α, β, or γ.
  • K.2 MORE EXPERIMENTS WITH EVR: A γ circle develops and grows in explanatory power across Mistral 7B Weekdays layers 17 to 29.Layer 25 residuals after explaining one-hot α and β functions show the especially clear γ circle.
  • K.2 MORE EXPERIMENTS WITH EVR: The section includes EVR results for attention-head interventions and iterative hidden-state deconstruction figures for Llama 3 8B Months.The supplied figure references identify the corresponding tables and visualizations.
Loading 2405.14860v3…