Query Expansion and Key Specialization in Transformer Attention Geometry
Not on the map. Why?
Abstract (arXiv)
The projection of queries and keys are central to the attention mechanism in Transformer architectures. While they are mathematically symmetric, they play different roles in attention mechanisms. The question of whether there is an effect from their functional distinction on their geometric development in training remains unanswered. We investigate the problem through the training of small GPT-like Transformers on character-level WikiText-103 for three different depths (4, 6, and 8 layers), three types of initialization for queries and keys, and four random seeds, resulting in 36 runs and 54 trajectories of average layers across seeds. We track the effective dimensionality of those layers using participation ratios and discover that effective dimension of queries expand while keys shrink, and that PRQ - PRK is positive in all trajectories studied. In connection to attention, the shrinking of keys leads to a narrower spectrum of QK^top and more peaked attention weights. In order to determine if this connection is causal or coincidental, we directly control the spectrum of keys during training across five seeds: restricting it to make it shrink sharpens the attention with high directional confidence, while keeping it constant to the level of initial dispersion makes attention softer. Additional token-level checkpoint analyses show that the monotonic paired-contrast trend is not universal across pretrained families, but survives as an early-training regime that later decays over a full pretraining run, and the link between interaction-rank geometry and attention entropy remains visible in several models.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 17 of its 24 references that are in the map. First posted on arXiv 28 Sep 2026.
Builds on
17 of 24 references.
References · 17
- 2017Attention Is All You Needcited by 60,361
- 2018Neural Tangent Kernel: Convergence and Generalization in Neural Networkscited by 2,621
- 2019What Does BERT Look At? An Analysis of BERT's Attentioncited by 1,003
- 2022In-context Learning and Induction Headscited by 964
- 2020On Layer Normalization in the Transformer Architecturecited by 922
- 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 856
- 2021Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depthcited by 419
- 2020Query-Key Normalization for Transformerscited by 280
- 2020Understanding the Difficulty of Training Transformerscited by 210
- 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 161
- 2022Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapsecited by 141
- 2023Max-Margin Token Selection in Attention Mechanismcited by 58
5 more
- 2024Self-attention Networks Localize When QK-eigenspectrum Concentratescited by 19
- 2024Dissecting Query-Key Interaction in Vision Transformerscited by 12
- 2025Is Random Attention Sufficient for Sequence Modeling? Disentangling Trainable Components in the Transformercited by 12
- 2025SeedPrints: Fingerprints Can Even Tell Which Seed Your Large Language Model Was Trained Fromcited by 10
- 2026Transformers Are Born Biased: Structural Inductive Biases at Random Initialization and Their Practical Consequencescited by 6
Built on by
Not cited yet.