通过调节查询与键的参数学习率，稳定Transformer注意力机制。
Controlling changes to attention logits
- 为查询和键分配依赖参数的学习率，控制注意力logits变化。
- 在MLA场景下提升基础学习率，性能超越现有方法。
- 无需完整显式计算，适合不支持QK归一化的架构。
神经网络权重的稳定性对训练Transformer模型至关重要。查询和键权重尤其容易无限制增长，需干预。虽有QK归一化（QK norm）可解决此问题，但其要求推理时完整显式化查询和键，不适用于多潜在注意力（MLA）。本文提出：控制注意力logits的变化对稳定性至关重要。通过为查询和键权重赋予参数依赖的学习率，可有效调控logits变化。该低成本方法使网络基学习率得以提升，在MLA设置中表现优于其他方法，且在多头注意力中达到与QK归一化相当的性能。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。