Inductive Bias and Spectral Properties of Single-Head Attention in High Dimensions

我们研究了在由近期提出的注意力索引模型所定义的合成高维序列任务上，训练单头绑定注意力层时的经验风险最小化问题。利用随机矩阵理论、自旋玻璃物理以及近似消息传递方法，我们推导出训练误差和测试误差的精确渐近表达式，确定了插值阈值与恢复阈值，并刻画了所学得权重矩阵的极限谱分布。权重衰减会隐式地引入核范数正则化，从而倾向于得到低秩的查询矩阵和键矩阵。基于这一发现，我们将查询与键矩阵的标准分解式训练方式，与直接对它们的乘积进行逐元素参数化的训练方式进行比较，揭示了分解形式所带来的归纳偏置。值得注意的是，我们预测的谱分布特征重现了大规模Transformer模型中观察到的经验趋势，为这些现象提供了与之相符的理论视角。

论文研究在合成高维序列任务上训练的单头注意力层中的经验风险最小化问题，重点分析模型在训练和测试误差、插值与恢复阈值方面的渐近行为。该问题聚焦于理解注意力机制中权重学习的理论性质，特别是参数化方式对学习动态和泛化的影响，属于较新的理论机器学习方向。

利用随机矩阵理论、自旋玻璃物理和近似消息传递方法，推导出训练和测试误差的精确渐近结果，并揭示权重衰减隐式引入核范数正则化，倾向于低秩的查询和键矩阵。通过比较因子化训练与直接参数化其乘积的方式，揭示了因子化结构带来的归纳偏置。

理论预测的权重谱分布与大规模Transformer中观察到的经验趋势一致，为实际模型中的现象提供了理论解释。实验基于合成数据和注意力索引模型设计，虽未提及真实数据集或开源代码，但理论分析深入且具有启发性，值得进一步推广到多头注意力和真实语言任务中。

1. 'On the Expressive Power of Deep Learning: A Tensor Analysis', COLT 2016 2. 'The Implicit Bias of Gradient Descent on Separable Data', JMLR 2018 3. 'Exact Asymptotics for Phase Retrieval and Compressed Sensing with Random Measurements', IEEE IT 2013 4. 'Modeling the Distribution of Normal Data in Pre-Trained Transformers', ACL 2022 5. 'Spin Glass Theory and Beyond: An Introduction to the Replica Method', World Scientific 1987