Dimensional Collapse in Transformer Attention Outputs: A Challenge for Sparse Dictionary Learning

Junxuan Wang, Xuyang Ge, Wentao Shu, Zhengfu He, Xipeng Qiu
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:126304-126331, 2026.

Abstract

Transformer architectures, and their attention mechanisms in particular, form the foundation of modern large language models. While transformer models are widely believed to operate in high-dimensional hidden spaces, we show that attention outputs are confined to a surprisingly low-dimensional subspace, with an effective dimensionality of only about 60% of the full space—a phenomenon that is consistently observed across diverse model families and datasets, and arises from overlap among the output subspaces of different attention heads. Critically, we find this low-rank structure as a key factor of the prevalent dead feature problem in sparse dictionary learning, where it creates a mismatch between randomly initialized features and the intrinsic geometry of the activation space. Building on this insight, we propose a subspace-constrained training method for sparse autoencoders (SAEs), initializing feature directions into the active subspace of activations. Our approach reduces dead features from 87% to below 1% in Attention Output SAEs with 1M features, and can further extend to other sparse dictionary learning methods. Our findings provide both new insights into the geometry of attention and practical tools for improving sparse dictionary learning in large language models. Code is available at https://github.com/OpenMOSS/Llamascopium.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-wang26bh, title = {Dimensional Collapse in Transformer Attention Outputs: A Challenge for Sparse Dictionary Learning}, author = {Wang, Junxuan and Ge, Xuyang and Shu, Wentao and He, Zhengfu and Qiu, Xipeng}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {126304--126331}, 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/wang26bh/wang26bh.pdf}, url = {https://proceedings.mlr.press/v306/wang26bh.html}, abstract = {Transformer architectures, and their attention mechanisms in particular, form the foundation of modern large language models. While transformer models are widely believed to operate in high-dimensional hidden spaces, we show that attention outputs are confined to a surprisingly low-dimensional subspace, with an effective dimensionality of only about 60% of the full space—a phenomenon that is consistently observed across diverse model families and datasets, and arises from overlap among the output subspaces of different attention heads. Critically, we find this low-rank structure as a key factor of the prevalent dead feature problem in sparse dictionary learning, where it creates a mismatch between randomly initialized features and the intrinsic geometry of the activation space. Building on this insight, we propose a subspace-constrained training method for sparse autoencoders (SAEs), initializing feature directions into the active subspace of activations. Our approach reduces dead features from 87% to below 1% in Attention Output SAEs with 1M features, and can further extend to other sparse dictionary learning methods. Our findings provide both new insights into the geometry of attention and practical tools for improving sparse dictionary learning in large language models. Code is available at https://github.com/OpenMOSS/Llamascopium.} }
Endnote
%0 Conference Paper %T Dimensional Collapse in Transformer Attention Outputs: A Challenge for Sparse Dictionary Learning %A Junxuan Wang %A Xuyang Ge %A Wentao Shu %A Zhengfu He %A Xipeng Qiu %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-wang26bh %I PMLR %P 126304--126331 %U https://proceedings.mlr.press/v306/wang26bh.html %V 306 %X Transformer architectures, and their attention mechanisms in particular, form the foundation of modern large language models. While transformer models are widely believed to operate in high-dimensional hidden spaces, we show that attention outputs are confined to a surprisingly low-dimensional subspace, with an effective dimensionality of only about 60% of the full space—a phenomenon that is consistently observed across diverse model families and datasets, and arises from overlap among the output subspaces of different attention heads. Critically, we find this low-rank structure as a key factor of the prevalent dead feature problem in sparse dictionary learning, where it creates a mismatch between randomly initialized features and the intrinsic geometry of the activation space. Building on this insight, we propose a subspace-constrained training method for sparse autoencoders (SAEs), initializing feature directions into the active subspace of activations. Our approach reduces dead features from 87% to below 1% in Attention Output SAEs with 1M features, and can further extend to other sparse dictionary learning methods. Our findings provide both new insights into the geometry of attention and practical tools for improving sparse dictionary learning in large language models. Code is available at https://github.com/OpenMOSS/Llamascopium.
APA
Wang, J., Ge, X., Shu, W., He, Z. & Qiu, X.. (2026). Dimensional Collapse in Transformer Attention Outputs: A Challenge for Sparse Dictionary Learning. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:126304-126331 Available from https://proceedings.mlr.press/v306/wang26bh.html.

Related Material