Transformers learn factored representations

Adam Shai, Loren Amdahl-Culleton, Casper L. Christensen, Henry R Bigelow, Fernando Rosas, Alexander Boyd, Eric A. Alt, Kyle J Ray, Paul M. Riechers
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:109734-109764, 2026.

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.

Cite this Paper


BibTeX
@InProceedings{pmlr-v306-shai26a, title = {Transformers learn factored representations}, author = {Shai, Adam and Amdahl-Culleton, Loren and Christensen, Casper L. and Bigelow, Henry R and Rosas, Fernando and Boyd, Alexander and Alt, Eric A. and Ray, Kyle J and Riechers, Paul M.}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {109734--109764}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/shai26a/shai26a.pdf}, url = {https://proceedings.mlr.press/v306/shai26a.html}, 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.} }
Endnote
%0 Conference Paper %T Transformers learn factored representations %A Adam Shai %A Loren Amdahl-Culleton %A Casper L. Christensen %A Henry R Bigelow %A Fernando Rosas %A Alexander Boyd %A Eric A. Alt %A Kyle J Ray %A Paul M. Riechers %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-shai26a %I PMLR %P 109734--109764 %U https://proceedings.mlr.press/v306/shai26a.html %V 306 %X 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.
APA
Shai, A., Amdahl-Culleton, L., Christensen, C.L., Bigelow, H.R., Rosas, F., Boyd, A., Alt, E.A., Ray, K.J. & Riechers, P.M.. (2026). Transformers learn factored representations. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:109734-109764 Available from https://proceedings.mlr.press/v306/shai26a.html.

Related Material