Transformers learn factored representations

TL;DR

Transformers pretrained with next token prediction tend to develop factored, orthogonal subspace representations of their input factors, balancing between efficiency and accuracy depending on factor independence.

Contribution

This paper formalizes hypotheses about the geometric structure of transformer representations, demonstrating their tendency to learn factored subspaces and analyzing the tradeoffs involved.

Findings

Transformers learn factored representations when factors are conditionally independent.

Models favor factored representations early in training, even with dependencies or noise.

Factored representations persist in complex data, explaining part decomposition in transformers.

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.