Transformers learn factored representations

TL;DR

Transformers trained with next-token prediction naturally factor their world into independent parts, representing each factor in orthogonal subspaces and achieving substantial dimensionality reduction. The authors formalize this via generalized hidden Markov models, defining predictive vectors and two representation geometries (joint vs factored), and prove that a lossless factorization exists under conditional independence through a linear mapping to a direct-sum embedding. Through synthetic experiments with five factors (Mess3 and Bloch Walk) and GPT‑2–style transformers, they show that the factored form emerges early, with dimensionality ≈ ∑_n (d_n−1) and approximate orthogonality between factor subspaces, while still allowing a joint representation when needed. They also demonstrate an inductive bias toward factoring, observed even when cross-factor correlations are introduced or noise corrupts observations, and extend findings to RNNs, underscoring potential interpretability benefits from persistent low-dimensional structure in residual streams.

Abstract

Transformers pretrained via next token prediction learn to factor their world into parts, representing these factors in orthogonal subspaces of the residual stream. We formalize two representational hypotheses: (1) a representation in the product space of all factors, whose dimension grows exponentially with the number of parts, or (2) a factored representation in orthogonal subspaces, whose dimension grows linearly. The factored representation is lossless when factors are conditionally independent, but sacrifices predictive fidelity otherwise, creating a tradeoff between dimensional efficiency and accuracy. We derive precise predictions about the geometric structure of activations for each, including the number of subspaces, their dimensionality, and the arrangement of context embeddings within them. We test between these hypotheses on transformers trained on synthetic processes with known latent structure. Models learn factored representations when factors are conditionally independent, and continue to favor them early in training even when noise or hidden dependencies undermine conditional independence, reflecting an inductive bias toward factoring at the cost of fidelity. This provides a principled explanation for why transformers decompose the world into parts, and suggests that interpretable low dimensional structure may persist even in models trained on complex data.

Figures (15)

• Figure 1: Transformers learn to factor the world into parts, representing each factor in orthogonal subspaces. (a) Consider a data-generating process with 4 hidden states (top-left). Although the joint process appears complex, it may admit a representation as a product of simpler independent processes running in parallel (top-right). Below, we illustrate how this factorization affects the geometry of representations: the full simplex over joint latent states (a tetrahedron for four joint states) decomposes into the product of simplices for each factor (two line segments). Even in cases where parts influence each other's dynamics, as long as conditioning on observed tokens renders the factors independent, the same geometric decomposition applies: uncertainty over each factor can be tracked separately in orthogonal subspaces. (b) When factors remain correlated even after conditioning on observations, the joint state cannot be written as a product state---it lies off the product-state manifold. Projecting onto the factored representation (red arrow) is lossy: some predictive information is sacrificed for the dimensional savings. (c) This decomposition yields dramatic representational savings. For $N$ three-state factors, the joint representation requires $3^N - 1$ dimensions (red), while the factored representation requires only $2N$ dimensions (green)---an exponential reduction. (d) To determine if transformers learn factored representations, we construct training data that is latently composed of sub-tokens generated from 5 factors. Each of these factors has a geometric prediction from our theory. Our analysis methods consist of quantifying the dimensionality of the activations using PCA and explained variance, as well as determining the geometric arrangement of factor subspaces. (e) Our main results are that transformers learn factored representations, placing these representations in orthogonal subspaces of the residual stream. Crucially, transformers do this even when it hurts prediction.
• Figure 2: Transformers discover latent factored structure. We train a transformer on sequences generated by five independent factors (F0–F4); the model sees only tokens corresponding to the Cartesian product of factor outputs. (a) Left: ground-truth belief geometry for each factor. Right: linear regression from activations recovers each geometry with high fidelity. (b) Cumulative explained variance (CEV) over training. Dashed lines show predictions factored versus joint representations; the model converges to the factored prediction. (c) Effective dimensionality (dimensions for 95% variance) collapses to about $10$ during training, far below the joint requirement. Dashed lines show dimensions required to explain 95% of the variance for the factored (10 at 95% CEV) and joint (135 at 95% CEV) predictions.
• Figure 3: Transformers represent factors in orthogonal multi-dimensional subspaces. (a) The number of dimensions needed to reach a CEV of 95% in the residual stream at the output layer of the last transformer block when driven by data from the vary-one dataset. At initialization, the activations associated with each factor use approximately the same number of dimensions as the union suggesting the subspaces are effectively fully overlapping. Over training the sum of the factor's dimensions approaches the dimension of the union, which suggests orthogonality. (b) Indeed, the first two principal components (PCs) associated with each factor evolve over training to become orthogonal to those associated with other factors. Thin lines show the normalized subspace overlap between the PCs belonging to different factors, while the heavy lines are the average over all pairs at each training step. The highlight shows a 90% confidence interval for initial factor overlap, generated from an ensemble of randomly initialized models.
• Figure 4: Transformers have an inductive bias towards factored representations. Even when the generative process doesn't admit a lossless factorization (as is the case for nonzero values of $\varepsilon$), the transformer quickly learns one anyways. For small $\varepsilon$, this factored representation persists throughout training. As $\varepsilon$ grows, transformers find the factored form early before expanding their dimensionality for greater predictive fidelity.
• Figure 5: A minimal example (two SNS factors) to compare two candidate representations (joint vs. factored). The joint representation (bottom left) is a curved 2d submanifold of a 3-simplex, and varies in three dimensions. The factored representation (bottom right) contains belief updates over each factor in an orthogonal subspace, and so only varies in two dimensions.

Theorems & Definitions (3)

• Definition 2.1
• Proposition 2.2
• Theorem 2.3