R/01
As training progresses, queries use more directions while keys concentrate into fewer. The attention strip illustrates how attention becomes more focused. This is a schematic, not experimental data.
Query Expansion and Key Specialization in Transformer Attention Geometry
How query and key representations evolve during Transformer training.
Decoder-only GPT-style Transformers trained on character-level WikiText-103 show that queries and keys develop differently: query representations expand, while keys contract and become more specialized. The study compares Xavier, QK-identical and disjoint-support Xavier initialization across model depths.
Across 54 initialization-seed-depth trajectories, query participation increased in 49 cases and key participation decreased in 50. Key specialization was associated with compression of the QK interaction spectrum and reduced attention entropy. Initialization left a persistent signature even when models achieved similar validation losses.
- Setup
- The researchers trained GPT-style models with 4, 6 or 8 layers and representation sizes of 64, 128 or 256 dimensions, using 10% subsets of the WikiText-103 training and validation data. They compared three initialization methods to track how query and key geometry evolved.
- Finding
- Queries spread across more representational directions, while keys concentrated into fewer. Key specialization was associated with more focused attention, and initialization continued to influence the learned geometry.
authors Gupta V. · Nadkarni S. · Chaudhari M. · Sawant V. · Satam P.