mlsanityData to 14 Oct 2026

Controlling changes to attention logits

arXiv 2511.21377 ↗First posted 26 Nov 2025cs.LGcited by 3

Not on the map. Why?

Abstract (arXiv)

Stability of neural network weights is critical when training transformer models. The query and key weights are particularly problematic, as they tend to grow large without any intervention. Applying normalization to queries and keys, known as `QK norm', fixes stability issues in practice, but is not always applicable. For example, QK norm is not compatible with Multi Latent Attention (MLA) because QK norm requires full materialization of queries and keys during inference, which is not done in MLA. In this paper we suggest that controlling the changes to logits is important for stability. We show that these changes are controllable by assigning parameter-dependent learning rates to the query and key weights. We find that our cheap intervention allows us to increase the base learning rate of the network, outperform other methods in the MLA setting, and achieve performance competitive with QK norm when using Multi-head Attention.

Key facts. Cited by 3 papers in the map (data through 14 Oct 2026). Builds on 27 of its 30 references that are in the map. First posted on arXiv 26 Nov 2025.

Builds on

27 of 30 references.

References · 27

  1. 2017Attention Is All You Needcited by 60,361
  2. 2025Qwen3 Technical Reportcited by 8,820
  3. 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
  4. 2024Gemma 2: Improving Open Language Models at a Practical Sizecited by 1,998
  5. 2019On the Variance of the Adaptive Learning Rate and Beyondcited by 772
  6. 2023Scaling Vision Transformers to 22 Billion Parameterscited by 694
  7. 2019Large Batch Optimization for Deep Learning: Training BERT in 76 minutescited by 674
  8. 2016Neural Combinatorial Optimization with Reinforcement Learningcited by 608
  9. 2025Kimi K2: Open Agentic Intelligencecited by 567
  10. 2017Large Batch Training of Convolutional Networkscited by 535
  11. 2022Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfercited by 288
  12. 2020Query-Key Normalization for Transformerscited by 280
15 more
  1. 2020Understanding the Difficulty of Training Transformerscited by 210
  2. 2023Small-scale proxies for large-scale Transformer training instabilitiescited by 183
  3. 2023Stabilizing Transformer Training by Preventing Attention Entropy Collapsecited by 161
  4. 2023A Spectral Condition for Feature Learningcited by 109
  5. 2025Don't be lazy: CompleteP enables compute-efficient deep transformerscited by 78
  6. 2023Spike No More: Stabilizing the Pre-training of Large Language Modelscited by 52
  7. 2020On the distance between two neural networks and the stability of learningcited by 46
  8. 2023LipsFormer: Introducing Lipschitz Continuity to Vision Transformerscited by 34
  9. 2025Peri-LN: Revisiting Normalization Layer in the Transformer Architecturecited by 34
  10. 2025The Sharpness Disparity Principle in Transformers for Accelerating Language Model Pre-Trainingcited by 30
  11. 2019Learning an Adaptive Learning Rate Schedulecited by 29
  12. 2025Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMscited by 27
  13. 2024Methods of improving LLM training stabilitycited by 13
  14. 2025Taming Transformer Without Using Learning Rate Warmupcited by 9
  15. 2025Function-Space Learning Ratescited by 4

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 3.

3 papers.

  1. 2026Learning Less Is More: Premature Upper-Layer Attention Specialization Hurts Language Model Pretrainingcited by 2
  2. 2026PoLoRA: A Preconditioned Orthogonalized LoRA Optimizercited by 2
  3. 2026One QK Channel, Many Sources: Tracing Low-Precision Attention Collapsecited by 0