Self-attention Networks Localize When QK-eigenspectrum Concentrates
This addresses a foundational issue in machine learning by clarifying contradictory failure modes in self-attention mechanisms, which is incremental as it builds on existing theories.
The paper tackles the problem of understanding attention localization in self-attention networks by linking it to the eigenspectrum of query-key matrices, showing that small eigenspectrum variance prevents both rank and entropy collapse, leading to improved model expressivity and trainability.
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.