arXiv:2409.17335cs.LGstat.ML2024-09NeurIPS被引 15

解析单层Transformer训练动态,揭示其在数据迁移下的泛化能力。

Non-asymptotic Convergence of Training Transformers for Next-token Prediction

  • 构建基于偏序的数学框架,刻画NTP任务数据结构特性。
  • 两阶段训练中各层子线性收敛至最大间隔方向,交叉熵损失线性下降。
  • 理论解释了Transformer在数据分布变化时仍具预测能力,适合研究模型泛化者。

Transformers 在现代机器学习中因卓越的序列建模能力而取得巨大成功,尤其在下一词预测(NTP)任务中表现突出。然而,现有理论研究多聚焦于渐近性能,对训练动态的非渐近分析尚不充分。本文针对由自注意力模块与前馈层构成的单层Transformer,开展精细的非渐近分析。首先,基于偏序的数学框架刻画了适用于NTP任务的数据集本质结构特征;随后设计两阶段训练算法:预处理阶段训练前馈层、主阶段训练注意力层,两者均呈现快速收敛特性,分别以子线性速率逼近对应的最大间隔解的方向。同时,交叉熵损失实现线性收敛。此外,实验表明训练后的Transformer在数据分布发生偏移时仍具备非平凡的预测能力,为Transformer优异的泛化性能提供了理论支持。分析方法涉及注意力梯度的新性质推导及对这些性质如何促进收敛过程的深入探讨。实验进一步验证了理论发现。

原文摘要 · Abstract (English)

Transformers have achieved extraordinary success in modern machine learning due to their excellent ability to handle sequential data, especially in next-token prediction (NTP) tasks. However, the theoretical understanding of their performance in NTP is limited, with existing studies focusing mainly on asymptotic performance. This paper provides a fine-grained non-asymptotic analysis of the training dynamics of a one-layer transformer consisting of a self-attention module followed by a feed-forward layer. We first characterize the essential structural properties of training datasets for NTP using a mathematical framework based on partial orders. Then, we design a two-stage training algorithm, where the pre-processing stage for training the feed-forward layer and the main stage for training the attention layer exhibit fast convergence performance. Specifically, both layers converge sub-linearly to the direction of their corresponding max-margin solutions. We also show that the cross-entropy loss enjoys a linear convergence rate. Furthermore, we show that the trained transformer presents non-trivial prediction ability with dataset shift, which sheds light on the remarkable generalization performance of transformers. Our analysis technique involves the development of novel properties on the attention gradient and further in-depth analysis of how these properties contribute to the convergence of the training process. Our experiments further validate our theoretical findings.

Transformer理论分析泛化能力

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。