注意力索引模型的高维学习动力学

中文摘要 AI 辅助

注意力机制是现代基础模型的核心，但其训练动力学仍鲜为人知，尤其是当注意力矩阵具有高秩时。本研究探讨注意力索引模型，这一通用框架可表征多层及多头注意力架构。首先，我们证明在合适的高维极限下，总体损失景观由有限组迹序参数刻画；与之相对，在线随机梯度下降（SGD）受无限层级的矩阵矩支配，我们表明该矩阵矩可通过有限截断系统实现指数级逼近。其次，该框架揭示注意力参数化本身可作为一种架构隐式偏置：直接优化注意力矩阵$S\in\mathbb{R}^{d\times d}$会陷入无信息状态；绑定注意力（$S=WW^\top$）会诱导自动对称破缺机制，在$\Theta(d^2\log d)$个样本上实现弱恢复；对于非绑定注意力（$S=UV^\top$），我们发现快慢机制：预激活均值先在快时间尺度上演化，而重叠在慢时间尺度上演化，当快动力学选择的状态打破初始对称性时，会在$\Theta(d^2\log d)$尺度上发生弱恢复。

英文摘要

Attention mechanisms are central to modern foundation models, yet their training dynamics remain poorly understood, especially when the attention matrices have extensive rank. In this work, we study attention-indexed models, a broad framework that can represent multi-layer and multi-head attention architectures. First, we show that, in a suitable high-dimensional limit, the population-loss landscape is characterized by a finite set of trace order parameters. In contrast, online stochastic gradient descent (SGD) is governed by an infinite hierarchy of matrix moments, which we show can be exponentially well-approximated by a finite truncated system. Second, this framework reveals that attention parameterization itself can act as an architectural implicit bias. Direct optimization of an attention matrix $S\in\mathbb{R}^{d\times d}$ can remain trapped in an uninformative state. Tied attention ($S=WW^\top$) induces an automatic symmetry-breaking mechanism and yields weak recovery in $Θ(d^2\log d)$ samples. For untied attention, $S=UV^\top$, we uncover a fast-slow mechanism: the pre-activation mean first evolves on a fast timescale, while the overlaps evolve on a slower one. Weak recovery on the $Θ(d^2\log d)$ scale occurs when the state selected by the fast dynamics breaks the initial symmetry.

发表机构

• École Polytechnique Fédérale de Lausanne (EPFL)（洛桑联邦理工学院）

机构由 AI 辅助整理，请以论文原文为准。