[论文解读] In-Context Convergence of Transformers

Huang Yu, Yuan Cheng+1 more|arXiv (Cornell University)|Oct 8, 2023

Domain Adaptation and Few-Shot LearningComputer Science被引用 3

一句话总结

本文首次对一维Transformer中基于Softmax注意力机制的梯度下降训练动态进行了理论分析，适用于上下文学习线性函数的情形。研究在特征分布均衡与非均衡两种情况下均建立了有限时间收敛性，揭示了分阶段的学习过程：主导特征首先被学习，随后是代表性不足的特征，其核心在于对竞争性注意力权重动态的创新分析。

ABSTRACT

Transformers have recently revolutionized many domains in modern machine learning and one salient discovery is their remarkable in-context learning capability, where models can solve an unseen task by utilizing task-specific prompts without further parameters fine-tuning. This also inspired recent theoretical studies aiming to understand the in-context learning mechanism of transformers, which however focused only on linear transformers. In this work, we take the first step toward studying the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent in order to in-context learn linear function classes. We consider a structured data model, where each token is randomly sampled from a set of feature vectors in either balanced or imbalanced fashion. For data with balanced features, we establish the finite-time convergence guarantee with near-zero prediction error by navigating our analysis over two phases of the training dynamics of the attention map. More notably, for data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process, where the transformer first converges to a near-zero prediction error for the query tokens of dominant features, and then converges later to a near-zero prediction error for the query tokens of under-represented features, respectively via one and four training phases. Our proof features new techniques for analyzing the competing strengths of two types of attention weights, the change of which determines different training phases.

研究动机与目标

• 理解通过梯度下降训练的基于Softmax的Transformer在上下文学习中的学习动态。
• 分析在部分特征代表性不足的非均衡特征分布下，收敛行为的表现。
• 将训练过程表征为由竞争性注意力权重驱动的相变过程。
• 在均衡与非均衡设置下，均建立有限时间收敛性保证。

提出的方法

• 分析一个在一维Transformer中使用Softmax注意力机制，并在结构化数据模型上通过梯度下降进行训练，其中标记从K个特征向量中以概率{pk}k=1K采样。
• 将上下文学习建模为一个线性回归任务，使用包含输入-标签对的提示序列和一个查询标记。
• 针对均衡特征提出两阶段分析：快速注意力对齐，随后是损失最小化。
• 提出一种新颖的证明技术，通过追踪“目标”与“非目标”注意力权重之间的相互作用，识别出不同的训练阶段。
• 利用浓度不等式和递归界来控制梯度下降下注意力图的演化。
• 采用概率论论证，证明在样本量为poly(K)且误差容限依赖于log(1/ε)的条件下，收敛具有高概率性。

实验结果

研究问题

• RQ1 基于Softmax的Transformer在梯度下降训练动态下，于上下文学习中如何收敛？
• RQ2 目标特征与非目标特征之间的注意力权重竞争，在塑造训练阶段中起到何种作用？
• RQ3 为何在非均衡数据分布中，即使某些特征代表性不足，上下文学习仍能成功？
• RQ4 是否能为均衡与非均衡特征分布均保证有限时间收敛？
• RQ5 当特征分布非均衡时，学习动态中的不同训练阶段是什么？

主要发现

• 对于均衡特征，模型在两个阶段内实现接近零的预测误差：快速注意力对齐，随后是损失最小化。
• 对于非均衡特征，模型表现出分阶段收敛：主导特征被快速学习，随后在四个不同的训练阶段后，代表性不足的特征被学习。
• 主导特征的预测误差在O(log(ε^(-1/2)) / (ηε))个梯度下降步数内收敛，误差被限制在ε以内。
• 经过足够迭代后，正确特征的注意力得分以高概率1 - e^(-Ω(poly(K)))集中到接近1。
• 分析表明，‘目标’与‘非目标’注意力权重的主导性随时间变化，从而定义了训练阶段。
• 基于竞争性注意力权重动态的证明技术具有通用性，可能适用于其他Transformer相关问题。