A Statistical Theory of Gated Attention through the Lens of Hierarchical Mixture of Experts
Viet Nguyen, Tuan Minh Pham, Thinh Cao, Tan Dinh, Huy Nguyen, Nhat Ho, Alessandro Rinaldo
TL;DR
This work develops a rigorous statistical theory for gated attention by mapping it to hierarchical mixtures of experts (HMoE). It shows that gated attention can be represented as a three-level HMoE with nonlinear or linear experts, enabling a clean analysis of learning as expert estimation. The authors prove that vanilla multi-head self-attention suffers exponential sample complexity for expert estimation, while gated attention with appropriately placed nonlinearities achieves polynomial rates, with concrete bounds of order $\mathcal{O}(\epsilon^{-4})$ in sample complexity. They validate the theory with numerical experiments, confirming faster convergence for gated attention compared to standard attention. Overall, the results provide a principled explanation for when and why gated attention improves performance, and offer practical guidance on where to place nonlinear gates in attention architectures.
Abstract
Self-attention has greatly contributed to the success of the widely used Transformer architecture by enabling learning from data with long-range dependencies. In an effort to improve performance, a gated attention model that leverages a gating mechanism within the multi-head self-attention has recently been proposed as a promising alternative. Gated attention has been empirically demonstrated to increase the expressiveness of low-rank mapping in standard attention and even to eliminate the attention sink phenomenon. Despite its efficacy, a clear theoretical understanding of gated attention's benefits remains lacking in the literature. To close this gap, we rigorously show that each entry in a gated attention matrix or a multi-head self-attention matrix can be written as a hierarchical mixture of experts. By recasting learning as an expert estimation problem, we demonstrate that gated attention is more sample-efficient than multi-head self-attention. In particular, while the former needs only a polynomial number of data points to estimate an expert, the latter requires exponentially many data points to achieve the same estimation error. Furthermore, our analysis also provides a theoretical justification for why gated attention yields higher performance when a gate is placed at the output of the scaled dot product attention or the value map rather than at other positions in the multi-head self-attention architecture.