基于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.
School of Mathematical Sciences, Shanghai Jiao Tong University(上海交通大学数学科学学院);School of Artificial Intelligence, Shanghai Jiao Tong University(上海交通大学人工智能学院);Institute for Advanced Algorithms Research, Shanghai(上海高级算法研究院)
Identifiable Token Correspondence for World Models
Interdisciplinary Program in Artificial Intelligence, Seoul National University(人工智能交叉学科项目,首尔国立大学);Department of Computer Science(计算机科学系);Engineering, Seoul National University(工程系,首尔国立大学)
Next-token prediction capacity: general upper bounds and a lower bound for transformers
Department of Computer Science, University of British Columbia(计算机科学系,不列颠哥伦比亚大学);Department of Electrical and Computer Engineering, University of British Columbia(电气与计算机工程系,不列颠哥伦比亚大学)
Next Token Prediction Towards Multimodal Intelligence: A Comprehensive Survey
Peking University(北京大学);Beihang Univerisy(北京航空航天大学);University of Hongkong(香港大学);Shenzhen Institute of Advanced Technology, China Academy of Sciences(中国科学院深圳先进技术研究院);Tsinghua University(清华大学);M-A-P;The Chinese University of Hong Kong(香港中文大学);Alibaba Group(阿里巴巴集团);University of Cambridge(剑桥大学);Microsoft Research(微软研究院);UIUC(伊利诺伊大学厄巴纳-香槟分校);Humanify Inc.(Humanify公司);Zhejiang University(浙江大学)
Is Next Token Prediction Sufficient for GPT? Exploration on Code Logic Comprehension
Microsoft Cloud and AI(微软云与人工智能);School of Artificial Intelligence, Jilin University(吉林大学人工智能学院);Mohamed bin Zayed University of Artificial Intelligence(穆罕默德·本·扎耶德人工智能大学)
I Predict Therefore I Am: Is Next Token Prediction Enough to Learn Human-Interpretable Concepts from Data?
Responsible AI Research Centre(负责任人工智能研究中心);Australian Institute for Machine Learning(澳大利亚机器学习研究所);Adelaide University(阿德莱德大学);School of Computer Science and Engineering(计算机科学与工程学院);The University of New South Wales(新南威尔士大学);Halıcıoğlu Data Science Institute(哈利奇奥格卢数据科学研究所);University of California San Diego(加州大学圣地亚哥分校);School of Mathematics and Statistics(数学与统计学学院);The University of Melbourne(墨尔本大学)