[论文解读] On the Optimization and Generalization of Multi-head Attention

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

Neural Networks and Applications被引用 4

一句话总结

该论文首次为使用逻辑损失进行梯度下降训练的单层多头自注意力模型提供了收敛性和泛化性保证。在温和的可实现性和NTK可分性条件下，当注意力头数为多项式对数级（H = Ω(log⁶n)）时，模型实现了 ˜O(1/n) 的训练损失和泛化差距，通过算法稳定性与自有界损失分析，建立起与过参数化MLP理论的联系。

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.

研究动机与目标

• 为填补有限宽度模型下多头注意力训练动态理论理解的空白，特别是针对有限宽度模型。
• 将此前仅限于单头注意力的分析扩展至具有实际头数的多头机制。
• 利用过参数化神经网络理论工具，建立多头注意力的有限时间收敛与泛化边界。
• 识别出可检查的初始化条件，确保可实现性并实现紧致的泛化边界。
• 在分词混合数据模型上验证框架的适用性，表明在一次梯度更新后即可实现快速数据分离。

提出的方法

• 推导经验损失 ̂L(θ) 的自有界性与弱凸性性质，通过依赖于模型权重的参数 κ 轻微量化 Hessian 曲率。
• 利用过参数化MLP中的算法稳定性工具（Taheri & Thrampoulidis, 2023），以训练损失和与初始化的距离来界定泛化差距。
• 提出一种新颖的分析框架，将多头注意力视为类似于过参数化隐藏层的并行结构，从而实现现有泛化技术的迁移。
• 形式化初始化条件，以确保在模型的NTK特征下数据可实现，特别要求初始化时的模型输出为训练集大小 n 的对数级。
• 分析一个分词混合数据模型，发现在从零初始化出发的一轮随机梯度更新后，NTK 特征可实现常数边际可分性，边际为 γ⋆。
• 利用逻辑损失的自有界性与曲率边界，推导出在标准梯度下降（步长 η = ˜O(1)）下的有限时间优化与泛化速率。

实验结果

研究问题

• RQ1 多头注意力能否使用为过参数化全连接网络开发的理论框架进行分析？
• RQ2 何种初始化条件可确保在模型的NTK特征下数据可实现，从而实现收敛与泛化？
• RQ3 在给定误差率下，实现有限时间泛化边界的最小注意力头数 H 是多少？
• RQ4 在梯度下降下，多头注意力的泛化差距如何随训练集大小 n 变化？
• RQ5 所提出的框架能否在具体数据模型（如分词混合模型）上实例化，且是否能预测在极少优化步骤后数据即实现快速分离？

主要发现

• 当 H = Ω(log⁶n) 时，在常数边际 NTK 可分性与 η = ˜O(1) 步长条件下，模型实现了 ˜O(1/n) 的训练损失与泛化差距。
• 经验损失 ̂L(θ) 满足自有界弱凸性条件：λmin(∇²̂L(θ)) ≳ −κ/√H ⋅ ̂L(θ)，其中 κ 对参数向量的依赖较弱。
• 一个简单的初始化条件可确保可实现性：若初始化时模型输出为 O(log n)，且数据在 NTK 下具有边际 γ，则边界成立。
• 在从零初始化出发的一轮随机梯度更新后，MHA 模型的 NTK 特征可实现对分词混合数据的边际 γ⋆ 分离。
• 泛化差距受一个随 ˜O(1/n) 缩放的项界定，表明即使相对于 n 的头数较少，多头注意力仍能实现良好泛化。
• 该分析框架可扩展至多种数据-模型与架构变体，表明其在当前设定之外也具有广泛适用性。