Max-Margin Token Selection in Attention Mechanism
Max-Margin Token Selection in Attention Mechanism
复制标题
DOI：
10.48550/arxiv.2306.13596
复制
发表时间：
2023-06
期刊：
影响因子：
--
通讯作者：
Davoud Ataee Tarzanagh;Yingcong Li;Xuechen Zhang;Samet Oymak
中科院分区：
文献类型：
--
作者：
Davoud Ataee Tarzanagh;Yingcong Li;Xuechen Zhang;Samet Oymak
注意机制是变压器结构的核心组成部分，它导致了大语言模型的惊人成功。但是，注意力机制的理论原理知之甚少，尤其是其非凸优化动力学。在这项工作中，我们探讨了开创性的软性注意模型$ f（\ boldsymbol {x}）= \ langle \ boldsymbol {xv}，\ texttt {softmax}（\ boldsymbol {xwp}） $（\ boldsymbol {v}，\ boldsymbol {w}，\ boldsymbol {p}）$是可训练的参数。我们证明，在$ \ boldsymbol {p} $上运行梯度下降，或等效地$ \ boldsymbol {w} $，将方向收敛到将$ \ textit {local-optimal-optimal} $ subles toection complienct over toession。这显然将注意力正式为最佳的代币选择机制。值得注意的是，我们的结果适用于通用数据，并精确地表征了代币的$ \ textit {optimality} $，以值嵌入$ \ boldsymbol {xv} $和问题几何形状来表征。我们还提供了更广泛的正规化路径分析，即使对于非线性预测头，也可以确定关注的最大程度。当优化$ \ boldsymbol {v} $和$ \ boldsymbol {p} $同时与逻辑损失时，我们确定了正则化路径方向将其定向收敛到其各自的硬质量SVM解决方案的条件，其中$ \ boldsymbol {v} $分隔了基于其标签的输入功能。有趣的是，$ \ boldsymbol {p} $的SVM公式受$ \ boldsymbol {v} $的支持向量几何形状的影响。最后，我们通过数值实验来验证我们的理论发现并提供见解。
Attention mechanism is a central component of the transformer architecture which led to the phenomenal success of large language models. However, the theoretical principles underlying the attention mechanism are poorly understood, especially its nonconvex optimization dynamics. In this work, we explore the seminal softmax-attention model $f(\boldsymbol{X})=\langle \boldsymbol{Xv}, \texttt{softmax}(\boldsymbol{XWp})\rangle$, where $\boldsymbol{X}$ is the token sequence and $(\boldsymbol{v},\boldsymbol{W},\boldsymbol{p})$ are trainable parameters. We prove that running gradient descent on $\boldsymbol{p}$, or equivalently $\boldsymbol{W}$, converges in direction to a max-margin solution that separates $\textit{locally-optimal}$ tokens from non-optimal ones. This clearly formalizes attention as an optimal token selection mechanism. Remarkably, our results are applicable to general data and precisely characterize $\textit{optimality}$ of tokens in terms of the value embeddings $\boldsymbol{Xv}$ and problem geometry. We also provide a broader regularization path analysis that establishes the margin maximizing nature of attention even for nonlinear prediction heads. When optimizing $\boldsymbol{v}$ and $\boldsymbol{p}$ simultaneously with logistic loss, we identify conditions under which the regularization paths directionally converge to their respective hard-margin SVM solutions where $\boldsymbol{v}$ separates the input features based on their labels. Interestingly, the SVM formulation of $\boldsymbol{p}$ is influenced by the support vector geometry of $\boldsymbol{v}$. Finally, we verify our theoretical findings via numerical experiments and provide insights.