How Query–Key Initialization Shapes Gradients and Training in Transformers

We study how query and key (QK) initialization shapes gradient propagation and training in Transformers. We derive an explicit gradient-norm recurrence for softmax attention and Pre-Norm residual blocks, revealing an asymmetry between large and small initialization scales. Large initialization can cause gradient explosion: the expected input-gradient gain diverges in the long-sequence limit. For normalized, weakly correlated inputs, this yields $(2\ln T{)}^{1/4}$ as a practical upper reference for ${\sigma }_{qk}$, where QK weight entries have variance ${\sigma }_{qk}^2/d$, $d$ is the input dimension, and $T$ is the sequence length. Small, nonzero initialization preserves a nonvanishing gradient-to-weight ratio, while gradient normalization in Adam and Muon permits substantial updates despite nearly uniform attention. Experiments with $100$-layer attention stacks and 124M-parameter language models support the analysis, with AdamW and Muon exhibiting a broad low-loss plateau down to ${\sigma }_{qk}=10^{-5}$. These findings motivate a practical recipe: favor small, nonzero QK initialization; and for broader searches, stay below the upper reference $(2\ln T{)}^{1/4}$.