arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2403.08081cs.LGcs.AIcs.CLmath.OC
收藏

自注意力下下一词预测的机制

Mechanics of Next Token Prediction with Self-Attention

University of Michigan(密歇根大学);Google Research NYC(谷歌研究院纽约分部)

Yingcong Li, Yixiao Huang, M. Emrullah Ildiz, Ankit Singh Rawat, Samet Oymak

AI总结:

本文揭示自注意力层通过梯度下降学习自动机,以硬检索和软组合两步实现下一词预测,并严格刻画其机制,发现隐式发现强连通分量并检索最高优先级分量。

摘要:

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

英文摘要:

Transformer-based language models are trained on large datasets to predict the next token given an input sequence. Despite this simple training objective, they have led to revolutionary advances in natural language processing. Underlying this success is the self-attention mechanism. In this work, we ask: $\textit{What}$ $\textit{does}$ $\textit{a}$ $\textit{single}$ $\textit{self-attention}$ $\textit{layer}$ $\textit{learn}$ $\textit{from}$ $\textit{next-token}$ $\textit{prediction?}$ We show that training self-attention with gradient descent learns an automaton which generates the next token in two distinct steps: $\textbf{(1)}$ $\textbf{Hard}$ $\textbf{retrieval:}$ Given input sequence, self-attention precisely selects the $\textit{high-priority}$ $\textit{input}$ $\textit{tokens}$ associated with the last input token. $\textbf{(2)}$ $\textbf{Soft}$ $\textbf{composition:}$ It then creates a convex combination of the high-priority tokens from which the next token can be sampled. Under suitable conditions, we rigorously characterize these mechanics through a directed graph over tokens extracted from the training data. We prove that gradient descent implicitly discovers the strongly-connected components (SCC) of this graph and self-attention learns to retrieve the tokens that belong to the highest-priority SCC available in the context window. Our theory relies on decomposing the model weights into a directional component and a finite component that correspond to hard retrieval and soft composition steps respectively. This also formalizes a related implicit bias formula conjectured in [Tarzanagh et al. 2023]. We hope that these findings shed light on how self-attention processes sequential data and pave the path toward demystifying more complex architectures.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑