上下文学习中多头 Softmax 注意力的训练动态：涌现、收敛与最优性

我们研究了用于多任务线性回归上下文学习（in-context learning）的多头 softmax 注意力模型的梯度流动态。我们在合适的初始化选择下建立了梯度流的全局收敛性。此外，我们证明了在梯度流动态过程中会出现一种有趣的“任务分配”现象，即每个注意力头专注于解决多任务模型中的单个任务。具体而言，我们证明梯度流动态可分为三个阶段：预热阶段，此时损失下降较为缓慢，注意力头逐渐建立起对各个任务的倾向；涌现阶段，此时每个头选择一个单一任务且损失迅速下降；以及收敛阶段，此时注意力参数收敛至极限。进一步地，我们证明了梯度流的最优性，即梯度流所学到的极限模型在常数因子范围内与最佳可能的多头 softmax 注意力模型相当。我们的分析还 delineates 了单头与多头注意力模型在上下文学习（ICL）预测精度上的严格差异。我们收敛分析的关键技术是将参数空间中的梯度流动态映射到谱域中的一组常微分方程，其中注意力权重半奇异值的相对大小决定了任务分配。据我们所知，我们的工作提供了首个针对多头 softmax 注意力模型的收敛性结果。