[Paper Review] On the Optimization and Generalization of Multi-head Attention

Puneesh Deora, Rouzbeh Ghaderi+2 more|arXiv (Cornell University)|Oct 19, 2023

Neural Networks and Applications4 citations

TL;DR

This paper provides the first convergence and generalization guarantees for gradient descent training of a single-layer multi-head self-attention model with logistic loss. It establishes that with a polylogarithmic number of heads (H = Ω(log⁶n)), and under mild realizability and NTK separability conditions, the model achieves ˜O(1/n) training loss and generalization gap, leveraging connections to overparameterized MLP theory via algorithmic stability and self-bounded loss analysis.

ABSTRACT

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.

Motivation & Objective

• To close the gap in theoretical understanding of multi-head attention training dynamics, especially for finite-width models.
• To extend existing analyses—previously limited to single-head attention—toward multi-head mechanisms with practical head counts.
• To establish finite-time convergence and generalization bounds for multi-head attention using tools from overparameterized neural network theory.
• To identify primitive, checkable conditions on initialization that ensure realizability and enable tight generalization bounds.
• To demonstrate the applicability of the framework on a tokenized-mixture data model, showing rapid data separation after one gradient step.

Proposed method

• Derive self-bounded and weak-convexity properties of the empirical loss ̂L(θ), quantifying the Hessian curvature via a parameter κ that depends mildly on the model weights.
• Leverage algorithmic stability tools from overparameterized MLPs (Taheri & Thrampoulidis, 2023) to bound the generalization gap in terms of training loss and distance to initialization.
• Introduce a novel analysis framework that treats multi-head attention as a parallelized structure analogous to overparameterized hidden layers, enabling transfer of existing generalization techniques.
• Formalize conditions on initialization ensuring realizability of the data under the model’s NTK features, particularly requiring model output at initialization to be logarithmic in training set size n.
• Analyze a tokenized-mixture data model where, after one randomized gradient step from zero, the NTK features achieve constant-margin separability with margin γ⋆.
• Use the self-boundedness of the logistic loss and curvature bounds to derive finite-time optimization and generalization rates under standard gradient descent with step size η = ˜O(1).

Experimental results

Research questions

• RQ1 Can multi-head attention be analyzed using frameworks developed for overparameterized fully-connected networks?
• RQ2 What conditions on initialization ensure that the data are realizable under the model’s NTK features, enabling convergence and generalization?
• RQ3 What is the minimum number of attention heads H required to achieve finite-time generalization bounds with a given error rate?
• RQ4 How does the generalization gap scale with training set size n under gradient descent for multi-head attention?
• RQ5 Can the proposed framework be instantiated on concrete data models such as the tokenized-mixture model, and does it predict fast separation of data after minimal optimization steps?

Key findings

• With H = Ω(log⁶n) heads, the model achieves ˜O(1/n) training loss and generalization gap under constant-margin NTK separability and η = ˜O(1) step size.
• The empirical loss ̂L(θ) satisfies a self-bounded weak-convexity condition: λmin(∇²̂L(θ)) ≳ −κ/√H ⋅ ̂L(θ), where κ depends mildly on the parameter vector.
• A simple initialization condition ensures realizability: if the model output at initialization is O(log n), and data are NTK-separable with margin γ, then the bounds hold.
• After one randomized gradient step from zero initialization, the NTK features of the MHA model separate the tokenized-mixture data with margin γ⋆.
• The generalization gap is bounded by a term that scales as ˜O(1/n), demonstrating that multi-head attention can generalize well even with a small number of heads relative to n.
• The analysis framework is extendable to various data-model and architecture variations, suggesting broad applicability beyond the current setup.