# Query Expansion and Key Specialization in Transformer Attention Geometry

arXiv: https://arxiv.org/abs/2609.34273 · HTML: https://mlsanity.com/paper/2609.34273/

## 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 (most cited first, 17 of 17)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [On Layer Normalization in the Transformer Architecture](https://mlsanity.com/paper/2002.04745/) (2020)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth](https://mlsanity.com/paper/2103.03404/) (2021)
- [Query-Key Normalization for Transformers](https://mlsanity.com/paper/2010.04245/) (2020)
- [Understanding the Difficulty of Training Transformers](https://mlsanity.com/paper/2004.08249/) (2020)
- [Stabilizing Transformer Training by Preventing Attention Entropy Collapse](https://mlsanity.com/paper/2303.06296/) (2023)
- [Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse](https://mlsanity.com/paper/2206.03126/) (2022)
- [Max-Margin Token Selection in Attention Mechanism](https://mlsanity.com/paper/2306.13596/) (2023)
- [Self-attention Networks Localize When QK-eigenspectrum Concentrates](https://mlsanity.com/paper/2402.02098/) (2024)
- [Dissecting Query-Key Interaction in Vision Transformers](https://mlsanity.com/paper/2405.14880/) (2024)
- [Is Random Attention Sufficient for Sequence Modeling? Disentangling Trainable Components in the Transformer](https://mlsanity.com/paper/2506.01115/) (2025)
- [SeedPrints: Fingerprints Can Even Tell Which Seed Your Large Language Model Was Trained From](https://mlsanity.com/paper/2509.26404/) (2025)
- [Transformers Are Born Biased: Structural Inductive Biases at Random Initialization and Their Practical Consequences](https://mlsanity.com/paper/2602.05927/) (2026)

## 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.
