Multi-Head Attention as Ensemble Nadaraya-Watson Estimation: Variance Reduction, Decorrelation, and Optimal Head Diversity
Not on the map. Why?
Abstract (arXiv)
We develop a rigorous statistical theory of multi-head attention (MHA) as an ensemble of Nadaraya-Watson (NW) kernel regression estimators. Building on the algebraic identity between single-head softmax attention and the NW estimator, we prove that MHA is a structured ensemble of H NW estimators, each operating in a distinct learned projection subspace of the key space. We derive an explicit Bias-Variance-Covariance decomposition of the MHA mean squared error, showing that variance reduction depends not merely on the number of heads H but fundamentally on the decorrelation of head outputs. Decorrelation is governed by the principal angles between learned projection subspaces: orthogonal projections yield maximum variance reduction; aligned projections yield none. We introduce the Head Diversity Index (HDI), a computable spectral measure of inter-head decorrelation, and prove that MHA mean squared error is monotonically decreasing in HDI. This provides the first rigorous theoretical explanation for the empirically observed specialization of attention heads. Under a fixed total-dimension budget D = H * d_k, we solve the optimal head-dimension allocation problem, deriving the MSE-minimizing pair (H*, d_k*) from data distribution and regression smoothness. The solution yields a new architectural scaling law: the optimal per-head dimension grows logarithmically with training set size, while the optimal number of heads grows nearly linearly with the total budget D. Our framework unifies three strands of prior work: the NW theory of single-head attention, the general weighting theory for ensemble learning, and the decorrelation-variance-reduction isomorphism between biological and computational ensembles. Multi-head attention is the Transformer's instantiation of a universal principle: identical agents plus diversity-enforcing mechanisms yields emergent optimality.
Key facts. Cited by 1 paper in the map (data through 7 Oct 2026). Builds on 9 of its 19 references that are in the map. First posted on arXiv 18 May 2026.
Builds on
9 of 19 references.
References · 9
- 2017Attention Is All You Needcited by 59,932
- 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
- 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 844
- 2019Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernelcited by 211
- 2023Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Contextcited by 63
- 2026Efficient and Minimax Optimal In-context Nonparametric Regression with Transformerscited by 7
- 2026Decorrelation, Diversity, and Emergent Intelligence: The Isomorphism Between Social Insect Colonies and Ensemble Machine Learningcited by 3
- 2025A General Weighting Theory for Ensemble Learning: Beyond Variance Reduction via Spectral and Geometric Structurecited by 2
- 2026Isomorphic Functionalities between Ant Colony and Ensemble Learning: Part II-On the Strength of Weak Learnability and the Boosting Paradigmcited by 2
Built on by
Citing papers per year: 2026: 1.
1 paper.
- 2026The Linear Representation Hypothesis Needs a Group Action22 Sep 2026