A Random Matrix Theory Perspective on the Learning Dynamics of Multi-head Latent Attention

TL;DR

This paper uses random matrix theory to analyze how multi-head latent attention affects transformer capacity during training, revealing that sharing rotary embeddings across heads preserves spectral support and prevents capacity collapse.

Contribution

It introduces a spectral analysis framework for multi-head attention variants, showing that sharing rotary embeddings across heads maintains model capacity during training.

Findings

Decoupled rotary sharing prevents spectral collapse.

Capacity bottlenecks occur early and locally in training.

Sharing rotary components mitigates spectral fragmentation.

Abstract

In this work, we study how multi-head latent attention (MLA), a popular strategy for compressing key/value memory, affects a transformer's internal capacity during pretraining. Using a lightweight suite of Marchenko-Pastur (MP) diagnostics, we analyze the spectrum of the $W_{Q}W_{K}^\top$ gram matrix throughout training, comparing three variants: the standard multi-head attention (MHA) baseline, MLA-PreRoPE with rotary applied before compression, and MLA-Decoupled, which shares a single rotary sub-vector across all heads. Our random matrix analysis reveals \textbf{three key findings:} \textbf{ i)} capacity bottlenecks emerge locally: both MHA and MLA-PreRoPE exhibit sharp, early spikes in specific layers that persist and propagate, disrupting the balance between bulk and outlier directions; \textbf{ ii)} these spikes coincide with rank collapse, concentrating the model's expressivity into narrow subspaces; \textbf{ iii)} only the decoupled variant prevents this cascade, maintaining broad spectral support and suppressing outlier formation across layers. These results underscore that \emph{how} rotary embeddings are applied is just as critical as \emph{where} compression occurs. Sharing rotary components across heads mitigates spectral fragmentation and preserves representational capacity.