Self-attention Networks Localize When QK-eigenspectrum Concentrates
Han Bao Ryuichiro Hataya Ryo Karakida
The self-attention mechanism prevails in modern machine learning. It has an
interesting functionality of adaptively selecting tokens from an input sequence
by modulating the degree of attention localization, which many researchers
speculate is the basis of the powerful model performance but complicates the
underlying mechanism of the learning dynamics. In recent years, mainly two
arguments have connected attention localization to the model performances. One
is the rank collapse, where the embedded tokens by a self-attention block
become very similar across different tokens, leading to a less expressive
network. The other is the entropy collapse, where the attention probability
approaches non-uniform and entails low entropy, making the learning dynamics
more likely to be trapped in plateaus. These two failure modes may apparently
contradict each other because the rank and entropy collapses are relevant to
uniform and non-uniform attention, respectively. To this end, we characterize
the notion of attention localization by the eigenspectrum of query-key
parameter matrices and reveal that a small eigenspectrum variance leads
attention to be localized. Interestingly, the small eigenspectrum variance
prevents both rank and entropy collapse, leading to better model expressivity
and trainability.