注意力机制中的最大间隔 token 选择

注意力机制是 transformer 架构的核心组件,它带来了大语言模型的惊人成功。然而,注意力机制背后的理论原理,尤其是其非凸优化动力学,仍知之甚少。本工作中,我们探究开创性的 softmax 注意力模型 $f(\boldsymbol{X})=\langle \boldsymbol{Xv}, \texttt{softmax}(\boldsymbol{XWp})\rangle$,其中 $\boldsymbol{X}$ 为 token 序列, $(\boldsymbol{v},\boldsymbol{W},\boldsymbol{p})$ 为可训练参数。我们证明,对 $\boldsymbol{p}$ 或等价地对 $\boldsymbol{W}$ 运行梯度下降,会在方向上收敛到一个最大间隔解,该解将局部最优 token 与非最优 token 分离。这明确地将注意力形式化为一种最优 token 选择机制。值得注意的是,我们的结果适用于一般数据,并根据值嵌入 $\boldsymbol{Xv}$ 与问题几何精确刻画了 token 的最优性。我们还提供了更宽泛的正则化路径分析,确立了即使对于非线性预测头,注意力也具有间隔最大化性质。当使用 logistic 损失同时优化 $\boldsymbol{v}$ 和 $\boldsymbol{p}$ 时,我们识别出若干条件,在这些条件下正则化路径在方向上收敛到各自的硬间隔 SVM 解,其中 $\boldsymbol{v}$ 根据其标签分离输入特征。有趣的是, $\boldsymbol{p}$ 的 SVM 表述受 $\boldsymbol{v}$ 的支撑向量几何影响。最后,我们通过数值实验验证了理论发现并提供了见解。