A Statistical Theory of Gated Attention through the Lens of Hierarchical Mixture of Experts
This provides a theoretical foundation for gated attention, potentially improving efficiency in training large language models, though it is incremental as it builds on existing gated attention models.
The paper tackles the lack of theoretical understanding of gated attention in Transformers by showing it can be expressed as a hierarchical mixture of experts, demonstrating that gated attention is more sample-efficient, requiring only polynomial data points compared to exponential for standard multi-head self-attention to achieve the same estimation error.
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.