Controlling changes to attention logits

TL;DR

This paper proposes a parameter-dependent learning rate method to control attention logit changes, enhancing stability and performance in transformer models, especially when QK normalization is incompatible.

Contribution

It introduces a simple, cost-effective approach to stabilize attention logits by adjusting learning rates, outperforming existing methods like QK norm in certain settings.

Findings

Improved stability with higher learning rates

Outperforms QK norm in MLA setting

Achieves competitive performance with multi-head attention

Abstract

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.