自注意力机制在下一词元预测中的机理

基于 Transformer 的语言模型通过大规模数据集训练，以根据输入序列预测下一个词元。尽管训练目标简单，它们却引领了自然语言处理的革命性进展。这一成功的背后是自注意力机制。在这项工作中，我们提出如下问题：$\textit{单个自注意力层从下一词元预测中学到了什么？}$ 我们证明，通过梯度下降训练自注意力会学习到一个自动机，该自动机以两个不同步骤生成下一个词元：$\textbf{(1)}$ $\textbf{硬检索：}$ 给定输入序列，自注意力精确地选择与最后一个输入词元相关联的 $\textit{高优先级输入词元}$。$\textbf{(2)}$ $\textbf{软组合：}$ 然后它创建这些高优先级词元的凸组合，从中可以采样下一个词元。在适当条件下，我们通过从训练数据中提取的词元上的有向图，严格刻画了这些机理。我们证明梯度下降隐式地发现了该图的强连通分量（SCC），并且自注意力学会了检索属于上下文窗口内可用的最高优先级 SCC 的词元。我们的理论依赖于将模型权重分解为方向性分量和有限分量，它们分别对应于硬检索和软组合步骤。这也形式化了 [Tarzanagh et al. 2023] 中猜想的一个相关隐式偏差公式。我们希望这些发现能阐明自注意力如何处理序列数据，并为揭示更复杂架构的奥秘铺平道路。