[论文解读] Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

Yuandong Tian, Yiping Wang+2 more|arXiv (Cornell University)|May 25, 2023

Neural Networks and Applications被引用 6

一句话总结

本文对一维Transformer在下一项标记预测任务中SGD训练动态进行了数学上严谨的分析，揭示了自注意力机制作为判别性扫描机制的作用：它会逐步聚焦于不同且频繁共现的标记，同时减少对常见标记的注意力。关键发现是训练动态中存在一种相变，可防止其坍缩为赢家通吃模式，从而产生稳定且稀疏的注意力模式，且该模式可通过学习率比值进行控制。

ABSTRACT

Transformer architecture has shown impressive performance in multiple research domains and has become the backbone of many neural network models. However, there is limited understanding on how it works. In particular, with a simple predictive loss, how the representation emerges from the gradient \emph{training dynamics} remains a mystery. In this paper, for 1-layer transformer with one self-attention layer plus one decoder layer, we analyze its SGD training dynamics for the task of next token prediction in a mathematically rigorous manner. We open the black box of the dynamic process of how the self-attention layer combines input tokens, and reveal the nature of underlying inductive bias. More specifically, with the assumption (a) no positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that self-attention acts as a \emph{discriminative scanning algorithm}: starting from uniform attention, it gradually attends more to distinct key tokens for a specific next token to be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, it progressively drops attention weights, following the order of low to high co-occurrence between the key and the query token in the training set. Interestingly, this procedure does not lead to winner-takes-all, but decelerates due to a \emph{phase transition} that is controllable by the learning rates of the two layers, leaving (almost) fixed token combination. We verify this extbf{\emph{scan and snap}} dynamics on synthetic and real-world data (WikiText).

研究动机与目标

• 理解自注意力在简单下一项标记预测训练过程中所隐含的归纳偏置。
• 在无位置编码的情况下，形式化注意力权重在SGD优化下的动态演化过程。
• 解释为何自注意力在稀疏性增加的情况下仍不会坍缩为赢家通吃模式。
• 表征解码器与自注意力层之间学习率比值如何控制相变的出现时机及注意力模式的密度。
• 在合成数据和真实世界数据（WikiText）上验证理论预测。

提出的方法

• 分析在SGD优化下，采用自注意力后接解码器层的一维Transformer在下一项标记预测任务中的训练动态。
• 通过非线性常微分方程（ODE）建立SGD动态的连续时间近似，以建模参数演化过程。
• 通过矩阵$Y$和$Z$对注意力权重进行重参数化，以解耦自注意力与解码器的动态。
• 在长序列与慢注意力假设下，推导注意力logits与Softmax输出的解析解。
• 利用Hoeffding不等式，对长序列中标记共现频率的经验估计进行边界约束。
• 通过在合成数据和WikiText上的实验验证结论，测量注意力熵与召回阈值曲线。

Figure 1: Overall of our setting. (a) A sequence with contextual tokens $\{x_{1},\ldots,x_{T-1}\}$ and query token $x_{T}$ is fed into 1-layer transformer (self-attention, normalization and decoding) to predict the next token $x_{T+1}$ . (b) The definition of sequence classes (Sec. 3.2 ). A sequence

实验结果

研究问题

• RQ1 在一维Transformer中，无位置编码时，自注意力在训练过程中的动态演化机制如何？
• RQ2 自注意力学习到何种归纳偏置？其与标记共现频率及独特性的关系为何？
• RQ3 为何自注意力在稀疏性增加的情况下仍能避免坍缩为赢家通吃模式？
• RQ4 解码器与自注意力层之间的学习率比值如何影响注意力模式的稀疏性与稳定性？
• RQ5 理论中提出的'扫描与捕捉'动态是否可在真实世界序列中得到实证验证？

主要发现

• 自注意力表现出判别性扫描行为：它会逐步聚焦于与查询标记不同且频繁共现的关键标记。
• 注意力机制表现出频率偏差，倾向于选择与查询在训练数据中具有高共现频率的标记。
• 尽管稀疏性增加，注意力仍不会坍缩为单热编码，这是由于训练动态中存在相变，导致梯度幅值呈对数衰减。
• 解码器层学习率越大，相变越快发生，导致注意力模式更密集；而自注意力层学习率越大，则导致更稀疏的模式。
• WikiText上的实证结果表明，更高的学习率降低了注意力图的熵并减缓了召回阈值的下降，表明注意力更稀疏。
• 通过不同学习率比值下的召回阈值曲线，验证了关于注意力稀疏性与共现依赖性的理论预测。

Figure 2: Overview of the training dynamics of self-attention map. Here $\tilde{c}_{l|m,n}:=\mathbb{P}(l|m,n)\exp(z_{ml})$ is the un-normalized attention score (Eqn. 5 ). (a) Initialization stage. $z_{ml}(0)=0$ and $\tilde{c}_{l|m,n}=\mathbb{P}(l|m,n)$ . Distinct tokens (Sec. 3.2 ) shown in blue, co