Probing for Representation Manifolds in Superposition
arXiv:2605.18537v1 [cs.LG] 18 May 2026
Alexander Modell Department of Mathematics Imperial College London, U.K. [email protected]
Abstract This paper introduces the Manifold Probe, a supervised method for discovering representation manifolds in superposition. The method generalizes linear regression probes by learning the space of features of a concept that can be linearly predicted from the representations, and then learning the directions used to encode them. We demonstrate the probe on representations of time and space in Llama 2-7b, finding manifolds which linearly represent an interpretable set of features in each case. In the case of time, we show that by steering along the manifold, we can influence the model’s completions about the years in which famous songs, movies and books were released, providing evidence that the Manifold Probe can discover manifolds which are causally involved in model behaviour. § alexandermodell/maniprobe
1
Introduction
The ability to interpret the representation geometry of large language models is a fundamental goal in a larger scientific effort to understand AI systems as a whole (Bereska and Gavves, 2024). A key hypothesis in this effort is the linear representation hypothesis (Nanda et al., 2023b): that neural networks organize their internal representations so as to make semantically important features accessible via linear projections. A related hypothesis is that of superposition (Smolensky, 1990; Mikolov et al., 2013; Elhage et al., 2022): the idea that representations of distinct concepts combine additively to produce representations of their joint semantics. Initial work around these hypotheses focused on understanding representations of simple binary concepts, which are considered to be either present or absent (Elhage et al., 2022). In this setting, concepts are hypothesised to be represented by almost-orthogonal directions in representation space and the presence of multiple features is represented by summing the corresponding directions. These hypotheses motivate many contemporary interpretability methodologies, such as linear probes (Alain and Bengio, 2017; Nanda et al., 2023b), sparse autoencoders (Elhage et al., 2022; Bricken et al., 2023; Cunningham et al., 2024), and steering vectors (Li et al., 2023; Marks and Tegmark, 2024; Rimsky et al., 2024; Turner et al., 2025). More recently, there has been a push towards understanding the representation geometry of more complex, continuous concepts, which don’t fit in to the binary framework. Examples include numerics, time, space, colour, and more abstract concepts such as emotion, ideaology and phylogeny (Gurnee and Tegmark, 2024; Olah, 2024; Engels et al., 2025; Modell et al., 2025; Gurnee et al., 2025; Pearce et al., 2025; Savietto et al., 2026; Choi and Weber, 2026; Sofroniew et al., 2026; Sun et al., 2026). There is growing empirical evidence that continuous concepts are represented on manifolds which bend and twist through multiple dimensions of the representation space (Gorton, 2024; Modell et al., 2025; Yocum et al., 2025; Gurnee et al., 2025; Karkada et al., 2026). The presence of such multi-dimensional manifolds is compatible with both the linear representation and superposition hypotheses. In particular, their shape directly determines the information about the concept which can be accessed via linear projections. Preprint.
The problem of isolating multidimensional representations of concepts, and discovering the geometry of representation manifolds in superposition is relatively unexplored. For the former, we are only aware of Engels et al. (2025) who propose to do this by clustering dictionary vectors in sparse autoencoders. For the latter, Yocum et al. (2025) and Gurnee et al. (2025) propose fitting a family of linear classifier probes to a discretization of the concept space, and Modell et al. (2025) propose approximating representation manifolds with neighbour graphs. In this paper, we propose the Manifold Probe, a supervised probing methodology to discover representation manifolds which are represented in superposition with other semantic information. The probe is fitted in two stages. The first stage is to learn the space of features f (z) of the concept values z which are well-predicted by a linear function w⊤ x + b of the representations x. While a standard linear regression probe would consider a fixed feature as its target, our probe learns the features at the same time as the regression parameters. We show that under a generic statistical model of a representation manifold in superposition, these learned features approximate the geometry of the manifold with respect to some unknown basis. The second stage learns this basis using linear regression. The main methodological contribution of this paper is the formulation and optimization of the first step of this procedure. To demonstrate the Manifold Probe, we use it to explore residual-stream representations of time and space in Llama 2-7b from probing datasets curated in Gurnee and Tegmark (2024). While Gurnee and Tegmark (2024) show that the concept values themselves are linearly represented, our probe brings to light many more features which are too, some of which are represented more precisely than the concept values themselves. We show how applying factor analysis to the learned features can help us interpret them. Employing a Varimax rotation which aims to make the features approximately sparse reveals that the time manifold linearly separates decades, while the space manifold linearly separates many states in the U.S.A.. Finally, we show that the time manifold we find is not only present in the residual stream, but is used by the model. We perform an intervention experiment where, at a given layer, we steer the residual stream representations by adding steering vectors which trace the manifold. By doing this, we can influence the model to complete a prompt about the release date of a song, movie or book with a year that we target.
2
Setup and background
2.1
Concepts and representation manifolds
A concept is a topological space Z which we can attach some real world meaning to. The simplest example might be a binary concept which is considered to be either present or absent. Continuous concepts include time (a line Z = R), space (a plane Z = R2 or a sphere Z = S2 ), colour (a cylinder Z = S × R2 or cube Z = R3 ) and can include more abstract concepts such as emotion with an appropriate mathematical model (such as the valence-arousal-dominance model). We say that any injective map ϕ : Z → Rp represents the concept Z. If ϕ is also continuous, then its image M = ϕ(Z) is a representation manifold embedded in Rp which, under some mild conditions1 , is topologically equivalent to the concept Z. We’ll write U ⊆ Rp to denote the smallest subspace containing M, and d to denote its dimension. For example, if Z is an interval, then M is a curve; if Z is a circle, then M is a loop; and if Z is a rectangle, then M is a sheet, all of which might bend and twist to occupy more dimensions in the representation space than might be implied by the intrinsic dimensionality of the concept itself. 2.2
Semantics and superposition
We now turn to the question of how multiple concepts might be represented together. To this end, we will consider an abstract topological space S which we refer to as the semantic space, which we assume encodes the semantics of any input to the neural network. We will assume S can be 1 such as Z being compact.
2
factorized into a set of interpretable concepts Z1 , · · · , Zm , and a set of other semantics Ξ, so that S = Z1 × · · · × Zm × Ξ. We will be interested in hypothesizing about, and making inferences about the structural form of a map x : S → Rp which represents S. A key hypothesis in mechanistic interpretability is that of superposition: the idea that representations of concepts combine additively to produce representations of their joint semantics. Definition 1. We say that a map x : S → Rp represents the concepts Z1 , · · · , Zm in superposition if there exists maps ϕi : Zi → Ui ⊆ Rp for i = 1, · · · , m, and a map η : Ξ → V ⊆ Rp such that x(s) = ϕ1 (z1 ) + · · · + ϕm (zm ) + η(ξ) (1) for all s = (z1 , . . . , zm , ξ) ∈ S. In the special case that concept representations are one-dimensional, this superposition hypothesis has been studied extensively in the mechanistic interpretability literature. The general setting presented above, in which concept representations are allowed to occupy multidimensional subspaces, has received comparatively little attention, with the notable exception of Engels et al. (2025). It also presents an additional inference problem: not only is it of interest to estimate the subspace which the concept representation occupies, but also the geometry of the representation within that subspace. In this paper, we will be concerned with developing methodology to discover the representation ϕ := ϕ1 of a single target concept Z := Z1 . From hereon, we will absorb any additional concepts in Ξ, and assume that S = Z × Ξ. 2.3
Probing
In order to discover the representation ϕ of the target concept Z, we will employ the probing methodology (Alain and Bengio, 2016). The idea behind probing is to construct a dataset of representation-concept values pairs D = {(xi , zi )}ni=1 , and to use this in a supervised fashion to fit a statistical model which elucidates the representation geometry of interest. We assume that each representation-concept value pair (xi , zi ) ∼ P in the probing dataset D is sampled independently by sampling a semantic value si = (zi , ξi ) from a distribution PS , and then constructing the representation xi = x(si ) according to the superposition equation (1) in Definition 1. While we observe the representation-concept value pairs (xi , zi ); the nuisance semantics ξi , and the functional form of the maps ϕ and η are unobserved. We assume that the concept value zi and nuisance semantics ξi are independent, i.e. PS = PZ × PΞ . Our statistical objective is to use the probing dataset D to learn two maps which estimate ϕ(z) from either the concept value z, or a corresponding representation x: 1. a smooth non-linear map ϕ̂ : Z → M̂ ⊂ Rp from the concept values Z to a manifold M̂ in some d-dimensional subspace Û ⊂ Rp of the representation space. 2. a linear (affine) map Ψ : Rp → Û ⊂ Rp from the representation space Rp to the subspace Û ⊂ Rp . We point out at this stage that the maps ϕ and η in the superposition equation (1) are only defined up to translation, and so for the purpose of estimation, we will assume without loss of generality that E[ϕ(z)] = 0. 2.4
Manifold estimation as regression
Before discussing how we might estimate the representation map ϕ from a finite probing dataset D, it is useful to consider how we might obtain ϕ given access to the true underlying population distribution P. The following lemma, which we prove in Section C of the appendix, tells us how. Lemma 2. There exists a basis u1 , . . . , ud ∈ U and a set of mean-zero, orthonormal features f1 , . . . , fd : Z → R such that ϕ(z) = u1 f1 (z) + . . . ud fd (z) (2) 3
which also solve the following sequential population regression problems: (fk , wk , bk ) = argmin E (f (z) − w⊤ x − b)2 ,
(3)
h i 2 E ∥x − ufk (z) − c∥ .
(4)
f :Z→R, w∈Rp , b∈R E(f )=0,E(f 2 )=1 f ⊥fk−1 ,...,f1
(uk , ck ) =
argmin u,c∈Rp
where expectations are taken with respect to (x, z) ∼ P, and f ⊥ g means E(f g) = 0. Lemma 2 suggests the shape of a representation manifold is intimately connected to space of features which it linearly represents. This dual interpretation is key to our finite-sample estimation procedure, and also provides a lens through which to interpret the manifold geometry.
3
Methodology
This section is dedicated to developing a sample-based estimation procedure for estimating ϕ from the probing data D based on the population regression problems in Lemma 2. In order to fit a feature f , we parametrize it in some basis h1 , . . . , hm which we treat as known, so that it can be written as f (z) = β ⊤ h(z) ≡ β1 h1 (z) + · · · + βm hm (z)
(5)
for some unknown scalar parameters β := (β1 , . . . , βm ). We denote the space of functions of the form (5) as H. An appropriate choice of basis depends on the nature of the concept space Z. In our examples, we use cubic B-splines, or tensor products thereof, however our method can accomodate any choice of basis. In order to avoid overfitting the function f , we need some way to control its complexity. The standard approach in functional regression is to choose an overly flexible basis, and then to control the complexity of f via a penalty function J(f ) which we add to our loss function. The advantage of this approach is that it allows us to choose the level of permitted complexity using the data. In this paper, we will assume that the chosen penalty function J is quadratic which allows us to write it as a quadratic form J(f ) = β ⊤ Sβ in the basis coefficients β. In our examples, we use the integrated, squared second derivative penalty Z J(f ) = [f ′′ (z)]2 dz, Z
which is usually considered a default choice. However, our method is flexible enough to accomodate quadratic penalty, and in Section B of the appendix, we discuss how our method can be adapted to accomodate non-quadratic penalties, such as ℓ1 and mixed ℓ1 and ℓ2 -type penalties. With H and J defined, we are ready to write down our probing procedure for estimating ϕ(z). Definition 3. We write ϕ̂, Ψ = ManifoldProbe(D, d; λw , λf ) if ϕ̂(z) = û1 fˆ1 (z) + · · · + ûd fˆd (z),
Ψ(x) = û1 g1 (x) + · · · + ûd gd (x),
with gk (x) = ŵk⊤ x + b̂k , where for k = 1, . . . , d, (fˆk , ŵk , b̂k ) solves the sequential optimization problem: minimize f ∈ H, w ∈ Rp , b ∈ R subject to
n X
f (zi ) − w⊤ xi − b
2
+ λw ∥w∥22 + λf J(f )
i=1 n X
n
f (zi ) = 0,
i=1
1X [f (zi )]2 = 1, n i=1
f ⊥ fˆk−1 , . . . , fˆ1 and (ûk , ĉk ) solves the optimization problem: n X 2 minimize ∥xi − ufk (zi ) − c∥ . p u, c ∈ R i=1 4
(6)
For fixed regularization parameters λw , λf , the Manifold Probe has a closed-form solution. n×p We’ll define the centered model matrices X ∈ RP and H ∈ Rn×m with Xi,: = xi − x̄ and Prows n n Hi,: = h(zi ) − h̄ respectively, where x̄ = (1/n) i=1 xi and h̄ = (1/n) i=1 h(zi ).
The closed-form solution that we present requires that the matrix H has full-column rank, so that all of the coefficients β can be uniquely estimated. If this is not the case (which is likely given the centering), we can linearly reparametrize the basis so that it does2 . From hereon, we will assume that the basis is parametrized such that H has full-column rank. Proposition 4. Let (fˆk , ŵk , b̂k ), be the solutions to the optimization problem (6). Then, fˆk (z) = β̂k⊤ h(z) where β̂1 , . . . , β̂d are an orthonormal set of eigenvectors corresponding to the d smallest eigenvalues ν̂m , . . . , ν̂m−d of the generalized eigenvalue problem M β = νΣβ where M := H ⊤ (I − A)H + λf S,
A = X(X ⊤ X + λw I)−1 X ⊤ ,
Σ=
1 ⊤ H H. n
In addition, ŵk = (X ⊤ X + λw I)−1 X ⊤ H β̂k , Pn where x̄ = (1/n) i=1 xi .
b̂k = −ŵk⊤ x̄,
ûk =
1 ⊤ X H β̂k n
A proof of Proposition 4 is given in Section D of the appendix. 3.1
Fitting the regularization parameters
While Proposition 4 tells us how to fit the Manifold Probe for a fixed pair of regularization parameters λw , λf , it doesn’t tell us anything about how we should choose them. In practice, we will want to choose them using the data, and we will want to choose different regularization parameters for each sequentially fitted feature. One approach is to directly apply k-fold cross-validation to the objective, searching over candidate parameter values using, for example, grid search. In Section B of the appendix, we present an alternative approach which we favour in practice. We show that (6) can be optimized by solving a sequence of alternating (generalized) ridge regression problems, which we prove converge to the global minimizer under very mild conditions. We then propose to select the regularization parameters at each iteration using a closed-form criterion appropriate for ridge regression. In practice, we use either Generalized Cross-Validation (Craven and Wahba, 1978; Wood, 2004) or Restricted Maximum Likelihood (Bartlett, 1937; Wood, 2011), which have closed-forms and can be optimized very quickly using Newton’s method, without refitting the model.
4
Discovering time and space manifolds in large language models
In this section, we use the Manifold Probe to discover hidden time and space manifolds in the residual stream of Llama 2-7b (Touvron et al., 2023), an open-weights large language model. We demonstrate how we can use the probe both as an interpretability tool, to discover features which are linearly represented in the residual stream; and as a steering tool, to causally influence the model’s behaviour. We make use of two probing datasets collected by Gurnee and Tegmark (2024). The first dataset contains the names and creators of popular songs, movies and books alongside their corresponding release dates (represented as a decimal year); and the second contains the names and geographic coordinates of places in the U.S.A.. After some filtering, we have 29,503 works released in Ztime = [1950, 2020], and 17,381 places with coordinates in Zspace = [24.5, 49.5] × [−125.0, −66.5] (the bounding box of mainland U.S.A.). In both cases, we consider a 50-50 train/test split. For each entity, we construct a string such as “Queen’s Bohemian Rhapsody” or “Lake of the Ozarks” which we feed into the language model, and record the last token residual stream activations at each 2 for example, using its singular value decomposition.
5
Figure 1: A representation manifold (top left) and linear prediction (top right) from a Manifold Probe fitted the release dates of songs, books and movies from layer 16 residual stream activations of Llama 2-7b. Below: the first five fitted features (top row), corresponding linear predictions (bottom row) for representations in the test set, and test R2 coefficients.
layer. To train the manifold probe, we parametrize time features using cubic B-splines with 280 knots, and parametrize space features using a tensor product of two cubic B-splines with 40 and 80 knots respectively. We use the fitting procedure described in Section B, and use the REML criterion to select the regularization parameters. 4.1
Interpretability: exploring linearly-represented features
The bottom row of plots of Figures 1 and 2 show the first few features fˆ1 , fˆ2 , . . . fitted by the probe to the layer 16 representations from the two training sets, and the corresponding linear predictions g1 (x), g2 (x), . . . of these features from representations x in the test sets. We report the R2 coefficients of these test predictions, which measure the extent to which the features are linearly represented. Perfect predictions have an R2 coefficient of one, predicting the feature mean always has an R2 coefficient of zero, and predictions which are worse than predicting the mean have a negative R2 coefficient. The top left plots show three dimensions of the estimated manifolds M̂ and the manifold predictions Ψ(x) from the representations x in the test set, with respect to the first three fitted basis vectors û1 , û2 , û3 . Since the features f1 , . . . , fd in the decomposition (2) are only defined up to rotations of the basis u1 , . . . , ud , it can be informative to apply factor analysis to the learned features to rotate them into a basis in which they are more easily interpretable. We take the top 5 time features, and the top 32 space features, and apply Varimax rotation (Kaiser, 1958; Rohe and Zeng, 2023) which aims to make the rotated features approximately sparse. The resulting features are shown in Figure 5 in the appendix. It is of particular note that rotated space features localize on many U.S. states, showing that they are 6
Figure 2: A representation manifold (top left) and linear prediction (top right) from a Manifold Probe fitted to the geographic coordinates of places in the U.S.A. from layer 16 residual stream activations of Llama 2-7b. Below: the first three fitted features (top row), corresponding linear predictions (bottom row) for representations in the test set, and test R2 coefficients. Feature values are given by colour. approximately linearly separated in the representations. We can also interpret from the rotated time features that the decades from the 1950s to the 2010s are approximately linearly separated. To get an idea for how much information about the release dates and locations is linearly represented in the residual stream at each layer, we fitted features until the corresponding test R2 coefficients were continually below zero. In Figure 3, we plot the value of the kth ranked test R2 coefficient for each layer where this is above zero. In both datasets, the predictabilities of features increase in predictability throughout the first half of the layers before levelling off. The location representations consistently contain three features which are considerably more predictable that the rest. We include the test R2 coefficients of a ridge regression fit directly to the release dates (dotted line) from the songs, movies and books representations; and to the latitude (dotted line) and longitude (dashed line) from the U.S. places representations. These were reported in Gurnee and Tegmark (2024) as evidence that language models linearly represent space and time. In the time representations, we find that for all layers the most linearly predictable feature is very close to the identity feature, and so the test R2 coefficient for our highest ranked feature and the direct year track very closely. In the time representations, the highest ranked feature we find has a higher test R2 coefficient than the latitude and longitude features. 4.2
Steering: causally influencing the model’s understanding of time
In this section, we demonstrate that we can use the representation manifold learned by the Manifold Probe to causally influence the model’s internal belief about the release dates of songs, movies and books. 7
Figure 3: Ranked test R2 values for features fitted using the probing datasets described in Section 4 at each layer of Llama 2-7b. Left: the dotted line shows the test R2 coefficient of a ridge regression fit directly to the release dates from the songs, movies and books representations. Right: the dotted and dashed lines show the test R2 coefficients of ridge regression fits directly to the latitude and longitude from the U.S. places representations respectively. To do this, we treat the learned concept representations ϕ̂(z) as an infinite continuum of steering vectors which trace the manifold M̂. To steer the model’s internal belief about the value of a concept to a particular value z, we propose to intervene on a representation x by adding a scalar multiple of ϕ(z) to it, i.e. setting x ←− x + αϕ̂(z) for some α > 0. From the probing dataset of songs, movies and books, we fit three-dimensional manifolds ϕ̂l to the last-token residual stream activations at each layer of the model as in the previous section. We then select a stratified sample of 1,400 works from the test set, with 200 from each decade, to use for our steering experiment. For each song, movie, or book, we construct a prompt of the form “⟨creator⟩’s ⟨title⟩ was released in the year” which we feed into the model. Assuming standard temperature-one sampling, we record the probability of the model completing the prompt each year between 1945 and 2025. We do this for a clean run, and for each layer l and year z between 1950 and 2020, we do this again but intervene by steering the residual stream activations of the last token of the work’s title at layer l by the steering vector ϕ̂l (z) and α = 100. This results in a total of 32 × 70 = 2240 interventions per prompt3 . To measure efficacy of an intervention, we report the probability that the model completes the prompt with a year that is within two years of the target z. The top plots in Figure 4 show the mean efficacy of the interventions for each layer, grouped in the left-hand plot by the decade of the song, movie or book, and in the right-hand plot by the decade of the target year z. From the left-hand plot, we see that the efficacy of the interventions peak at layers 8 and 14, depending on the release decade of the work, and that this drops sharply after layer 15, and dropping to baseline performance at layer 20 and beyond. From the right-hand plot, we see that steering efficacy and the optimal layer for intervention depends quite heavily on the target year. Figure 6 in the appendix shows the probability that the model completes the prompt with a year at all. With the exception of steering to years in the 1950s, we see that these interventions have very little effect on the model’s ability to meaningfully complete the prompt. For layers 3 to 20, the bottom plot in Figure 4 shows the mean probability of each completion for each steering target. The dark diagonal streaks indicate that the interventions are having some success at influencing the model to complete the prompt with a given target year. Figure 7 in the appendix shows the standard deviations. 3 The full experiment took approximately 100 GPU hours on Nvidia RTX 4090s.
8
Figure 4: Steering experiment. Top: the mean probability a completion is within two years of the target year it was steered to at each layer, grouped by release decade (left) and target decade (right). Clean baselines are shown with dashed lines. Bottom: colour intensity (capped at 0.1) indicates the mean probability of a completion given the steering target.
5
Discussion
In this work, we introduced the Manifold Probe, a supervised method for discovering representation manifolds in superposition. We hope that our probe will serve as a useful new tool for the mechanistic interpretability community. We see potential applications as a data-driven way to discover mechanisms of continuous computation such as counting (Gurnee et al., 2025; Wu et al., 2025) and modular arithmetic (Nanda et al., 2023a; Zhong et al., 2023), and as a tool to map out abstract concepts such as emotion (Sofroniew et al., 2026; Sun et al., 2026). There are also potential implications for scientific discovery, for example, in interpreting recently-discovered phylogenetic and hematopoietic representation manifolds in biological foundation models (Pearce et al., 2025; Wu et al., 2025). One limitation of our framework is that it implicitly assumes that the number of training samples is large relative to the dimension of the representation space, so that the sample estimates concentrate around their population counterparts. For state-of-the-art foundation models, this might require in the order of tens of thousands of examples. With smaller probing datasets, for good statistical performance it may be necessary to perform a preliminary principal component analysis to reduce the dimension of the representation space before fitting the probe. Finally, while interpretability tools such as this one might be used to develop mechanistic guardrails or to steer model behaviour to serve the goals of safety, alignment and security, they might be used to learn to bypass safety guardrails or produce intentionally harmful behaviour, which as a community, we must be mindful of as we progress our scientific understanding of AI systems. 9
Acknowledgements The author would like to thank Jacob Davies, Jake Yukich, Can Rager, David Chanin, Nathalie Kirch, Patrick Rubin-Delanchy and Nick Whiteley for enlightening discussions on the topics of this paper.
References Alain, G. and Bengio, Y. (2016). Understanding intermediate layers using linear classifier probes. arXiv preprint arXiv:1610.01644. Alain, G. and Bengio, Y. (2017). Understanding intermediate layers using linear classifier probes. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Workshop Track Proceedings. OpenReview.net. Bartlett, M. S. (1937). Properties of sufficiency and statistical tests. Proceedings of the royal society of london. series a-mathematical and physical sciences, 160(901):268–282. Bereska, L. and Gavves, E. (2024). Mechanistic interpretability for ai safety - a review. Transactions on Machine Learning Research. Survey Certification, Expert Certification. Bricken, T., Templeton, A., Batson, J., Chen, B., Jermyn, A., Conerly, T., Turner, N., Anil, C., Denison, C., Askell, A., Lasenby, R., Wu, Y., Kravec, S., Schiefer, N., Maxwell, T., Joseph, N., Hatfield-Dodds, Z., Tamkin, A., Nguyen, K., McLean, B., Burke, J. E., Hume, T., Carter, S., Henighan, T., and Olah, C. (2023). Towards Monosemanticity: Decomposing Language Models With Dictionary Learning. Transformer Circuits Thread. Choi, B. J. and Weber, M. (2026). Latent structure of affective representations in large language models. arXiv preprint arXiv:2604.07382. Craven, P. and Wahba, G. (1978). Smoothing noisy data with spline functions: estimating the correct degree of smoothing by the method of generalized cross-validation. Numerische mathematik, 31(4):377–403. Cunningham, H., Ewart, A., Riggs, L., Huben, R., and Sharkey, L. (2024). Sparse autoencoders find highly interpretable features in language models. In The Twelfth International Conference on Learning Representations. Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds, Z., Lasenby, R., Drain, D., Chen, C., Grosse, R., McCandlish, S., Kaplan, J., Amodei, D., Wattenberg, M., and Olah, C. (2022). Toy Models of Superposition. Transformer Circuits Thread. Engels, J., Michaud, E. J., Liao, I., Gurnee, W., and Tegmark, M. (2025). Not All Language Model Features Are One-Dimensionally Linear. In The Thirteenth International Conference on Learning Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net. Gorton, L. (2024). Curve Detector Manifolds in InceptionV1. Gurnee, W., Ameisen, E., Kauvar, I., Tarng, J., Pearce, A., Olah, C., and Batson, J. (2025). When models manipulate manifolds: The geometry of a counting task. Transformer Circuits Thread. Gurnee, W. and Tegmark, M. (2024). Language models represent space and time. In The Twelfth International Conference on Learning Representations. Kaiser, H. F. (1958). The varimax criterion for analytic rotation in factor analysis. Psychometrika, 23(3):187–200. Karkada, D., Korchinski, D. J., Nava, A., Wyart, M., and Bahri, Y. (2026). Symmetry in language statistics shapes the geometry of model representations. arXiv preprint arXiv:2602.15029. Li, K., Patel, O., Viégas, F., Pfister, H., and Wattenberg, M. (2023). Inference-time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, 36:41451–41530. 10
Marks, S. and Tegmark, M. (2024). The geometry of truth: Emergent linear structure in large language model representations of true/false datasets. In First Conference on Language Modeling. Mikolov, T., Yih, W.-t., and Zweig, G. (2013). Linguistic regularities in continuous space word representations. In Proceedings of the 2013 conference of the north american chapter of the association for computational linguistics: Human language technologies, pages 746–751. Modell, A., Rubin-Delanchy, P., and Whiteley, N. (2025). The origins of representation manifolds in large language models. arXiv preprint arXiv:2505.18235. Nanda, N., Chan, L., Lieberum, T., Smith, J., and Steinhardt, J. (2023a). Progress measures for grokking via mechanistic interpretability. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net. Nanda, N., Lee, A., and Wattenberg, M. (2023b). Emergent linear representations in world models of self-supervised sequence models. arXiv preprint arXiv:2309.00941. Olah, C. (2024). What is a Linear Representation? What is a Multidimensional Feature? Transformer Circuits Thread. Pearce, M., Simon, E., Byun, M., and Balsam, D. (2025). Finding the tree of life in evo 2. Goodfire. Correspondence to [email protected]. Rimsky, N., Gabrieli, N., Schulz, J., Tong, M., Hubinger, E., and Turner, A. (2024). Steering llama 2 via contrastive activation addition. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 15504–15522. Rohe, K. and Zeng, M. (2023). Vintage factor analysis with varimax performs statistical inference. Journal of the Royal Statistical Society Series B: Statistical Methodology, 85(4):1037–1060. Savietto, D., Campbell, D., Panisson, A., Nurisso, M., Petri, G., Cohen, J. D., and Perotti, A. (2026). The geometry of representational failures in vision language models. arXiv preprint arXiv:2602.07025. Smolensky, P. (1990). Tensor product variable binding and the representation of symbolic structures in connectionist systems. Artificial intelligence, 46(1-2):159–216. Publisher: Elsevier. Sofroniew, N., Kauvar, I., Saunders, W., Chen, R., Henighan, T., Hydrie, S., Citro, C., Pearce, A., Tarng, J., Gurnee, W., Batson, J., Zimmerman, S., Rivoire, K., Fish, K., Olah, C., and Lindsey, J. (2026). Emotion concepts and their function in a large language model. Transformer Circuits Thread. Sun, L., Yan, L., Lu, X., Lee, A., Zhang, J., and Shao, J. (2026). Valence-arousal subspace in llms: Circular emotion geometry and multi-behavioral control. arXiv preprint arXiv:2604.03147. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. (2023). Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. Turner, A. M., Thiergart, L., Leech, G., Udell, D., Vazquez, J. J., Mini, U., and MacDiarmid, M. (2025). Steering language models with activation engineering. In The Thirteenth International Conference on Learning Representations. Wood, S. N. (2004). Stable and efficient multiple smoothing parameter estimation for generalized additive models. Journal of the American Statistical Association, 99(467):673–686. Wood, S. N. (2011). Fast stable restricted maximum likelihood and marginal likelihood estimation of semiparametric generalized linear models. Journal of the Royal Statistical Society Series B: Statistical Methodology, 73(1):3–36. Wu, W., Zhou, Z., Riley, R., Abdulqader, M., Song, X., Kautsar, S., Egan, R., Hofmeyr, S., Liu, G., Goldhaber-Gordon, S., Yu, M., Ho, H., Liu, Y., Steindorff, A. S., Liu, F., Chen, F., Morgan-Kiss, R., Shi, L., Liu, H., and Wang, Z. (2025). Uncovering the genomic manifold via scalable learning from the global microbiome. bioRxiv. 11
Yocum, J., Allen, C., Olshausen, B., and Russell, S. (2025). Neural manifold geometry encodes feature fields. In NeurIPS 2025 Workshop on Symmetry and Geometry in Neural Representations. Zhong, Z., Liu, Z., Tegmark, M., and Andreas, J. (2023). The clock and the pizza: Two stories in mechanistic explanation of neural networks. Advances in neural information processing systems, 36:27223–27250.
12
A
Additional figures
Figure 5: The top 5 time features, and top 32 space features from layer 16 of Llama 2-7b after applying a Varimax rotation to make them approximately sparse. The rotation makes the features interpretable. In particular, many of the space features localize on particular U.S. states, and the time features separate the decades from the 1950s to the 2010s.
13
Figure 6: The mean probability that the model completes the prompt with a valid year in the steering experiment, grouped by release decade (left) and target decade (right). The dashed line shows the mean probability for the clean runs. The interventions have very little effect on the model’s ability to meaningfully complete the prompt, with the exception of steering to years in the 1950s.
Figure 7: Colour intensity indicates the standard deviation of the probability of a completion given the steering target in the steering experiment.
14
B
An efficient algorithm to fit the regularization parameters
In this section, we discuss an optimization strategy which allows us to efficiently optimize the Manifold Probe objective in (6) while also selecting the regularization parameters λw and λf using a closed-form criterion such as GCV or REML which apply to linear predictors. Instead of directly employing the closed-form solution in Proposition 4, we propose the Alternating √ Least Squares procedure detailed in Algorithm 1. Here, we have used the notation ∥α∥Σ = α⊤ Σα and write α ⊥Σ β to mean α⊤ Σβ = 0. Algorithm 1: Alternating least squares optimization of (6). (0)
(0)
Input: inital parameters β1 , . . . , βd ; for k = 1, . . . , d: for t = 1, 2, . . .: w-update: w(t+1) ←−
argmin ∥y − Xw∥22 + λ̃w ∥w∥22 ,
y = Hβ (t) ;
w
β-update: 1
β (t+ /2) ←−
argmin
∥y − Hβ∥22 + λ̃f β ⊤ Sβ,
y = Xw(t+1) ;
β⊥Σ β̂k−1 ,...,β̂1 1
1
β (t+1) ←− β (t+ /2) /∥β (t+ /2) ∥Σ ;
While it may seem needlessly inefficient to optimize (6) using Algorithm 1 rather than the closedform solution in Proposition 4, when we formulate the power-iteration procedure used to solve the generalized eigenvalue problem, we see that this is exactly equivalent to the alternating least squares procedure. Power-iteration is known to converge to the global solution under mild conditions on the initial value, and therefore we can guarantee that Algorithm 1 converges to the global solution under the same conditions. (0)
Lemma 5. Suppose that ν̂k > ν̂k+1 and wk ̸⊥ ŵk for all k = 1, . . . , d, then for some λ̃w , λ̃f , Algorithm 1 converges to the global minimizer of (6). I.e. (t) lim f = fˆk , t→∞ k
and
(t)
lim wk = ŵk
t→∞
The distinct eigenvalue condition is not strictly necessary and can be relaxed to simply ν̂d > ν̂d+1 , allowing repeated eigenvalues. The stricter condition is stated for simplicity. We provide a proof based on the power-iteration argument in Section B.2. By framing the optimization in this way, we can apply closed-form criteria designed for linear predictors such as GCV or REML to select the regularization parameters λ̃w and λ̃f at each iteration of the alternating least squares procedure. Viewed this way, we are also not restricted to quadratic penalties, and can use non-quadratic penalties such as the ℓ1 or elastic-net-type penalties, provided we have an efficient off-the-shelf regression solver which accomodates it. B.1
Efficient parametrization of the ridge regression problems
To efficiently perform the required computations in Algorithm 1, we perform some reparametrizations and matrix decompositions which allow us reduce each iteration to simple matrix multiplcations of size p × p, removing the time-dependence on the number of samples and avoiding performing any matrix inversions. To do this, we first reparametrize the β-problem to enforce the linear orthogonality constraints, and then reparametrize it again so that S becomes the identity matrix. If S is rank-deficient, we simply set its zero eigenvalues to some small positive constants to make it positive-definite to allow the 15
reparametrization. We are then left with a standard ridge regression problem. We note that from here on we parametrize and solve the w-problem in exactly the same way, so we won’t discuss it separately. We next compute the singular value decomposition of H as H = U DV ⊤ , where all diagonal entries of D are positive, and reparametrize the problem again to make H = U D. This ensures that H has full-column rank and avoids some unnecessary matrix multiplications down the line. It is then straightforward to show that ∥y − Hβ∥22 + λ∥β∥22 = ∥y − Dβ∥22 + λ∥β∥22 + r
(7)
where y = U y and r y Note that y = U ⊤ Xw so as long as U X is pre-computed, this multiplication does not depend on the number of samples n. The solution to the ridge regression problem (7) is ⊤
= ∥y∥22 − ∥ ⊤
∥22 is a constant which does not depend on β.
β̂ = (D2 + λI)−1 Dy which given y can be computed in O(p) time for any λ. The GCV or REML criterion and their gradients and Hessians can also be computed efficiently using this reparametrization, allowing us to very efficiently select the regularization parameters at each iteration using Newton’s method. B.2
Proof of Lemma 5
We begin by showing the convergence of β (t) for k = 1. Once this is established, the convergence of w(t) is trivial. We’ll consider the case that k = 1 and note that the subsequent cases follow by a deflation argument. We recall that the ridge updates have closed forms w(t+1) = (X ⊤ X + λ̃w I)−1 X ⊤ Hβ (t) and β (t+ /2) = (H ⊤ H + λ̃f S)−1 H ⊤ Xw(t+1) = (H ⊤ H + λ̃f S)−1 H ⊤ X(X ⊤ X + λ̃w I)−1 X ⊤ Hβ (t) . 1
We define the matrix T = LH ⊤ AH where L = (H ⊤ H + λ̃f S)−1 and A = X(X ⊤ X + λ̃w I)−1 X ⊤ , so that 1 β (t+ /2) = T β (t) . A full β-update is then given by T β (t) β (t+1) = . ∥T β (t) ∥Σ √ where ∥a∥Σ = a⊤ Σa with Σ = H ⊤ H/n. This shows that the sequence is a power-iteration for the matrix T , and therefore as long as its largest eigenvalue is unique, and the initial value β (0) is not orthogonal to the corresponding eigenvector, it converges to the leading eigenvector of T by a standard argument4 . It remains to show that the leading eigenvector of T is the same as the eigenvector β̂ with the smallest eigenvalue ν̂ of the generalized eigenvalue problem M β̂ = ν̂Σβ̂ Plugging in M and Σ and rearranging we obtain h i H ⊤ AH β̂ = (1 − ν̂/n)H ⊤ H + λ̃f S β̂ and setting λ̃f = λf /(1 − ν̂/n), we have H ⊤ AH β̂ = (1 − ν̂/n) H ⊤ H + λ̃f S β̂ =: (1 − ν̂/n)L−1 β̂. Multiplying both sides on the left by L, we have T β̂ = (1 − ν̂/n)β̂ which shows that β̂ is the leading eigenvector of T , which completes the proof. 4 see, for example, https://en.wikipedia.org/wiki/Power_iteration.
16
C
Proof of Lemma 2
C.1
Proof of (3)
Let f1⋆ , . . . , fd⋆ be any set of mean-zero, orthonormal features which satisfy (2), and let F = span{f1⋆ , . . . , fd⋆ }. We will show that span{f1 , . . . , fd } = F. We begin with the case k = 1. By the law of iterated expectation, (3) can be written as h h h 2 i 2 ii E f (z) − w⊤ x − b = E E f (z) − w⊤ x − b | z therefore
f1 (z) = E w⊤ x + b | z = w⊤ E[x | z] + b.
Now E[x | z] = ϕ(z) + E[η] = ϕ(z) + η̄ and so f1 (z) = w⊤ [ϕ(z) + η̄]+b = w⊤ (f1⋆ (z)u1 + · · · fd⋆ (z)ud ) = (w⊤ u1 )f1⋆ (z)+· · ·+(w⊤ ud )fd⋆ (z)+w⊤ η̄+b. Now, f1 is constrained so that E[f1 ] = 0, and given that E[f1⋆ ] = · · · = E[fd⋆ ] = 0, this implies that w⊤ η̄ + b = 0. Therefore, f1 is a linear combination of f1⋆ , . . . , fd⋆ . Since E(f 2 ) = 1, at least one coefficient must be non-zero and therefore f1 ∈ F. Next, we suppose that f1 , . . . , fk−1 ∈ F and f1 ⊥ · · · ⊥ fk−1 for some k ∈ {1, . . . , d}. We will show that fk ∈ F. Minimizing (3) subject to the constraint fk ⊥ fk−1 , . . . , f1 is equivalent to minimizing (3) replacing x(z) with the deflation x(k) (z) = x(z) − (π1 f1 (z) + · · · + πk−1 fk−1 (z)),
πi = E[xfi ].
As before, this is minimized by h i fk (z) = w⊤ E x(k) | z + b = w⊤ [ϕ(z) + η̄] + b − (π1 f1 (z) + · · · + πk−1 fk−1 (z)) = (w⊤ u1 )f1⋆ (z) + · · · + (w⊤ ud )fd⋆ (z) − (π1 f1 (z) + · · · + πk−1 fk−1 (z)). Since this is a linear combination of functions in F, and fk is constrained to be non-trivial, this implies that fk ∈ F. Therefore, by induction, f1 , . . . , fd ∈ F. Now the functions f1 , . . . , fd are constrained to be orthogonal, and therefore the span F, so span{f1 , . . . , fd } = span{f1⋆ , . . . , fd⋆ }. C.2
Proof of (4)
To find the optimal (uk , ck ) that minimizes the population least-squares error i h 2 E ∥x − ufk (z) − c∥ we start by taking the gradient with respect to c and setting it equal to zero to obtain −2E [x − ufk (z) − c] = 0 which gives that c = x̄ − uE[fk ] = x̄ where x̄ = E[x] and we have used the fact that E[fk ] = 0. Substituting this back into (4), taking the gradient with respect to u and setting it equal to zero gives us −2E [fk (z)(x − x̄ − ufk (z))] = 0 which implies that
E[fk (z)(x − x̄)] = uE[fk2 ].
Since E[fk2 ] = 1, this gives us that u = E[fk (z)(x − x̄)]. 17
Now x̄ = η̄, and so u = E[fk (z)(ϕ(z) + η(ξ) − η̄)] = E[fk (z)ϕ(z)] + E[fk (z)η(ξ)] − E[fk (z)η̄] = E[fk (z)ϕ(z)] where the final two terms are zero by the assumption of independence between z and ξ, and the fact that E[fk ] = 0. Now, substituting in the decomposition of ϕ(z) gives us u = E[fk (z)(u1 f1 (z) + · · · + ud fd (z))] = u1 E[fk f1 ] + · · · + ud E[fk fd ] Now by assumption E[fk fj ] = 0 for all j ̸= k, and E[fk2 ] = 1, so we have that u = uk , as required.
D
Proof of Proposition 4
We f ∈ H, it can be written as f (z) = β ⊤ h(z) and the constraint Pn begin by observing that since ⊤ i=1 f (zi ) = 0 implies that β h̄ = 0. Therefore, we can write evaluations of f in terms on the centered model matrix H: f (zi ) = β ⊤ h(zi ) = β ⊤ (h(zi ) − h̄) = (Hβ)i . Now, to minimize (6) with respect to b, we take the gradient with respect to b and set it equal to zero to obtain n X −2 f (zi ) − w⊤ xi − b = 0. i=1
Since
P
⊤ i=1 f (zi ) = 0, we obtain that b = −w x̄. Therefore we can write
f (zi ) − w⊤ xi − b = f (zi ) − w⊤ xi − (−w⊤ x̄) = f (zi ) − w⊤ (xi − x̄) = (Hβ)i − (Xw)i . Recalling that J(f ) = β ⊤ Sβ, this means we can write (6) in matrix form as minimizing ∥Hβ − Xw∥22 + λw w⊤ w + λf β ⊤ Sβ
(8)
subject to β ⊤ Σβ = 1 and β ⊤ Σβ̂j = 0 for all j < k. For fixed β, the optimal w is given by the ridge regression solution w = (X ⊤ X + λw I)−1 X ⊤ Hβ. With A = X ⊤ (X ⊤ X + λw I)X ⊤ as in the proposition, we have 2
2
∥Hβ − Xw∥2 = Hβ − X(X ⊤ X + λw I)−1 X ⊤ Hβ 2 2
= ∥(I − A)Hβ∥2 = β ⊤ H ⊤ (I − A)2 Hβ = β ⊤ H ⊤ (I − 2A + A2 )Hβ. Now, we define K = (X ⊤ X + λw I)−1 , and noting that X ⊤ X = K −1 − λw I, we observe that A2 = XKX ⊤ XKX ⊤ = XK(K −1 − λI)KX ⊤ = XKK −1 KX ⊤ − λw XK 2 X ⊤ = XKX ⊤ − λw XK 2 X ⊤ = A − λw XK 2 X ⊤ and
∥w∥2 = ∥KX ⊤ Hβ∥22 = β ⊤ H ⊤ XK 2 X ⊤ Hβ. 18
Therefore β ⊤ H ⊤ A2 Hβ = β ⊤ H ⊤ AHβ − λw β ⊤ HXK 2 XH ⊤ β = β ⊤ H ⊤ AHβ − λw ∥w∥2 , and it follows that
∥Hβ − Xw∥22 + λw ∥w∥22 = β ⊤ H ⊤ (I − A)Hβ.
Therefore, the objective function (8) can be written as ∥Hβ − Xw∥22 + λw ∥w∥22 + λf β ⊤ Sβ = β ⊤ H ⊤ (I − A)H + λf S β =: β ⊤ M β. By the Rayleigh-Ritz theorem, minimizer of this with respect to β, subject to β ⊤ Σβ = 1 is given by the generalized eigenvector of (4) with the smallest eigenvalue. Continuing sequentially, the generalized eigenvector with the kth smallest eigenvalue gives the solution to (8) subject to β ⊤ Σβ̂j = 0 for all j < k. This completes the proof.
19