Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis
作者: Hongru Yang, Bhavya Kailkhura, Zhangyang Wang, Yingbin Liang
分类: cs.LG, cs.CL
发布日期: 2024-10-12
备注: Accepted by NeurIPS 2024
💡 一句话要点
通过梯度流分析，研究Transformer识别词共现的训练动态
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: Transformer 训练动态 梯度流分析 词共现 注意力机制 深度学习 自动平衡梯度
📋 核心要点
- 现有Transformer训练动态研究常采用简化方法，忽略了真实训练过程的复杂性，限制了对模型内在机制的理解。
- 该论文通过分析梯度流，揭示了Transformer在词共现任务中训练的两个阶段，以及注意力机制和MLP的协同演化过程。
- 论文证明了梯度流的“自动平衡梯度”特性，保证了不同样本损失的均衡下降，并实验验证了理论结果的有效性。
📝 摘要（中文）
理解Transformer的训练动态对于解释大型语言模型背后的强大能力至关重要。本文研究了浅层Transformer在识别两个指定词共现任务上的训练动态。与以往研究Transformer训练动态的文献中常见的权重重参数化、注意力线性化、特殊初始化和惰性机制等简化方法不同，我们分析了从随机初始化同时训练三个注意力矩阵和一个线性MLP层的梯度流动态，并提供了一个通过耦合动力系统分析这种动态的框架。我们建立了接近最小损失，并描述了训练后的注意力模型。我们发现梯度流作为一种内在机制，自然地将训练过程分为两个阶段。在第一阶段，线性MLP迅速与两个目标信号对齐以进行正确的分类，而softmax注意力几乎保持不变。在第二阶段，注意力矩阵和MLP共同演化，以扩大分类裕度并将损失降低到接近最小值。在技术上，我们证明了梯度流的一个新颖性质，称为梯度自动平衡，它使不同样本的损失值几乎以相同的速率下降，并进一步促进了接近最小训练损失的证明。我们还进行了实验来验证我们的理论结果。
🔬 方法详解
问题定义：论文旨在理解Transformer在识别词共现任务中的训练动态。现有研究通常采用权重重参数化、注意力线性化等简化方法，这与实际Transformer的训练方式存在差异，难以揭示其内在机制。因此，论文关注在更真实的训练场景下，Transformer如何学习词共现关系。
核心思路：论文的核心思路是通过分析梯度流的动态变化，揭示Transformer训练过程中不同模块（注意力矩阵和MLP）的协同演化过程。论文发现，梯度流存在一种“自动平衡梯度”的特性，能够使不同样本的损失值以相似的速率下降，从而促进模型收敛到接近最小损失。
技术框架：论文研究了一个浅层Transformer模型，包含三个注意力矩阵和一个线性MLP层。模型的目标是识别两个指定词的共现关系。研究通过构建一个耦合动力系统，对梯度流的动态变化进行建模和分析。整个训练过程被划分为两个阶段：第一阶段，MLP快速对齐目标信号；第二阶段，注意力矩阵和MLP协同演化，降低损失。
关键创新：论文最重要的技术创新在于发现了梯度流的“自动平衡梯度”特性。该特性保证了不同样本损失的均衡下降，避免了某些样本损失过大而影响整体训练效果的情况。此外，论文还提供了一个分析Transformer训练动态的框架，可以用于研究更复杂的Transformer模型。
关键设计：论文采用随机初始化，避免了特殊初始化带来的偏差。损失函数采用标准的交叉熵损失。通过理论分析和实验验证，论文揭示了注意力矩阵和MLP在训练过程中的演化规律。关键参数包括学习率、batch size等，这些参数的选择会影响训练的收敛速度和最终性能。
📊 实验亮点
论文通过理论分析和实验验证，揭示了Transformer训练过程中的两个阶段，并证明了梯度流的“自动平衡梯度”特性。实验结果表明，该特性能够有效促进模型收敛到接近最小损失，验证了理论分析的正确性。
🎯 应用场景
该研究成果有助于深入理解Transformer的工作机制，为设计更高效、更可解释的Transformer模型提供理论指导。潜在应用领域包括自然语言处理、计算机视觉等，例如，可以用于优化模型训练策略、提高模型泛化能力。
📄 摘要（原文）
Understanding the training dynamics of transformers is important to explain the impressive capabilities behind large language models. In this work, we study the dynamics of training a shallow transformer on a task of recognizing co-occurrence of two designated words. In the literature of studying training dynamics of transformers, several simplifications are commonly adopted such as weight reparameterization, attention linearization, special initialization, and lazy regime. In contrast, we analyze the gradient flow dynamics of simultaneously training three attention matrices and a linear MLP layer from random initialization, and provide a framework of analyzing such dynamics via a coupled dynamical system. We establish near minimum loss and characterize the attention model after training. We discover that gradient flow serves as an inherent mechanism that naturally divide the training process into two phases. In Phase 1, the linear MLP quickly aligns with the two target signals for correct classification, whereas the softmax attention remains almost unchanged. In Phase 2, the attention matrices and the MLP evolve jointly to enlarge the classification margin and reduce the loss to a near minimum value. Technically, we prove a novel property of the gradient flow, termed \textit{automatic balancing of gradients}, which enables the loss values of different samples to decrease almost at the same rate and further facilitates the proof of near minimum training loss. We also conduct experiments to verify our theoretical results.