揭示多头注意力的统计机制,解释为何分散头能提升模型性能。
Multi-Head Attention as Ensemble Nadaraya-Watson Estimation: Variance Reduction, Decorrelation, and Optimal Head Diversity
- 将多头注意力视为多个核回归估计器的组合,各头在不同投影空间中独立工作。
- 证明头间去相关性越高,误差方差越低,最优配置依赖于训练数据规模。
- 提出头多样性指数,为模型设计提供理论指导,适合研究架构优化者。
我们建立了多头注意力(MHA)作为纳达拉亚-沃森(NW)核回归估计器集合的严格统计理论。基于单头softmax注意力与NW估计器的代数等价性,证明了MHA是H个在键空间不同学习投影子空间中运行的NW估计器的结构化集成。我们推导出MHA均方误差的显式偏差-方差-协方差分解,表明方差降低不仅取决于头数H,更根本上依赖于头输出的去相关性。去相关性由学习投影子空间间的主角决定:正交投影实现最大方差降低;对齐投影则无降低效果。我们引入可计算的头多样性指数(HDI),证明MHA均方误差随HDI单调递减。这首次为注意力头的实证专属性提供了理论解释。在固定总维度预算D = H * d_k下,求解最优头维度分配问题,从数据分布和回归平滑性中推导出使均方误差最小的最优对(H*, d_k*)。结果揭示新的架构缩放规律:最优每头维度随训练集大小对数增长,最优头数近似线性随总预算D增长。该框架统一了三类前期工作:单头注意力的NW理论、集成学习的通用加权理论,以及生物与计算集成中的去相关-方差降低同构性。多头注意力是Transformer对普适原则的实现:相同代理加多样性机制,催生涌现最优性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。