提出新优化方法,让自注意力层快速收敛到全局最优。
Training Dynamics of Softmax Self-Attention: Fast Global Convergence via Preconditioning
- 通过矩阵分解建模自注意力训练动态。
- 实现几何级收敛,避免局部最优。
- 适合研究注意力机制优化与深度学习理论的人。
我们研究了在软最大自注意力层中进行梯度下降的训练动态,该层用于执行线性回归。结果表明,一种简单的首阶优化算法可以在无限数据极限下以几何速率收敛至全局最优的自注意力参数。分析分为两步:首先,在无限数据条件下,自注意力层所解决的回归问题等价于一个非凸矩阵分解问题;其次,利用这一联系设计了一种新型“结构感知”梯度下降算法,高效优化原始有限数据回归目标。该算法相比标准梯度下降有若干创新,包括能避开虚假驻点的预处理项与正则化项,以及基于数据的谱初始化,使参数初始值以高概率位于全局最小值流形附近。
原文摘要 · Abstract (English)
We study the training dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression and show that a simple first-order optimization algorithm can converge to the globally optimal self-attention parameters at a geometric rate. Our analysis proceeds in two steps. First, we show that in the infinite-data limit the regression problem solved by the self-attention layer is equivalent to a nonconvex matrix factorization problem. Second, we exploit this connection to design a novel "structure-aware" variant of gradient descent which efficiently optimizes the original finite-data regression objective. Our optimization algorithm features several innovations over standard gradient descent, including a preconditioner and regularizer which help avoid spurious stationary points, and a data-dependent spectral initialization of parameters which lie near the manifold of global minima with high probability.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。