Transformers learn factored representations
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. Both track context-induced uncertainty over the latent parts, but the factored representation sacrifices fidelity when factors are not conditionally independent. 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. When factors are conditionally independent, models learn factored representations; when noise or dependencies break this structure, models gradually expand their effective dimensionality over training to recover 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.
Lay Summary
It’s easy to take for granted that our world is not monolithic, but is made of parts. But for ourselves and AI models like transformers, we merely perceive the world via a stream of undifferentiated inputs. Humans take this stream of raw data, and decompose it into discrete objects—chairs, coffee mugs, other agents—in order to understand the world. Do transformers also learn to factor their world into parts? How can they from a mere stream of tokens, and why would they? In this paper, we show that neural networks like transformers indeed learn and have an inductive bias for factored representations as a natural consequence of next-token pretraining. They internally represent coexisting multi-dimensional fractal geometries that correspond to belief updates over these parts of the world.