On the Optimization and Generalization of Multi-head Attention
This work addresses a foundational problem in machine learning by providing theoretical insights into multi-head attention, which is incremental as it builds on prior analyses of single-head attention.
The paper tackles the under-explored training and generalization dynamics of multi-head attention in Transformers, deriving convergence and generalization guarantees for gradient-descent training under a realizability condition and showing these conditions hold for a tokenized-mixture model.
The training and generalization dynamics of the Transformer's core mechanism, namely the Attention mechanism, remain under-explored. Besides, existing analyses primarily focus on single-head attention. Inspired by the demonstrated benefits of overparameterization when training fully-connected networks, we investigate the potential optimization and generalization advantages of using multiple attention heads. Towards this goal, we derive convergence and generalization guarantees for gradient-descent training of a single-layer multi-head self-attention model, under a suitable realizability condition on the data. We then establish primitive conditions on the initialization that ensure realizability holds. Finally, we demonstrate that these conditions are satisfied for a simple tokenized-mixture model. We expect the analysis can be extended to various data-model and architecture variations.