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

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 \url{https://github.com/OpenMOSS/Llamascopium}.

Modern language models, such as those used in chatbots and writing assistants, are built from Transformer architectures. A central part of these models is the attention mechanism, which helps the model decide which words or pieces of information to focus on. Although these models are usually thought to use very large and complex internal spaces, we find that the outputs of attention actually occupy a much smaller space than expected. This pattern appears consistently across many different models and datasets. This observation helps explain a common problem in methods that try to understand what language models are doing internally. These methods often learn many small components, or features, that are meant to capture meaningful patterns in the model. However, many of these features never become useful, a problem known as ``dead features.'' We show that this happens partly because standard training methods start from random directions that do not match the true shape of the attention outputs. To address this, we propose a simple change: initialize these learned features inside the part of the space where attention outputs actually live. This greatly reduces the number of dead features, from 87\% to below 1\% in large-scale experiments. Our work provides a clearer picture of how attention behaves inside language models and offers a practical way to improve tools for interpreting them.