Reproduce “Self-attention Networks Localize When QK-eigenspectrum Concentrates” on OpenResearch

About this paper

Han Bao, Ryuichiro Hataya, Ryo Karakida

This research introduces a theoretical characterization of self-attention localization through the eigenspectrum of the query-key parameter matrices, demonstrating that a concentrated eigenspectrum with a non-zero mean can simultaneously mitigate both rank and entropy collapse in Transformer models. An empirical regularization scheme further validates that promoting such localization leads to higher attention entropy and improved language modeling performance.

Key insights

• Attention localization is characterized by a QK-eigenspectrum concentrated around a non-zero mean, while uniformity occurs with a zero mean and finite variance.
• Minimizing the variance of the QK-eigenspectrum is a unifying mechanism that simultaneously mitigates both rank collapse (by preserving expressivity) and entropy collapse (by increasing attention entropy).
• The proposed signal propagation probability offers a quantitative measure to understand how specific tokens contribute to learning, directly linking it to the QK-eigenspectrum.

Summary by alphaXiv.