稳定Transformer训练:防止注意力熵坍缩
Stabilizing Transformer Training by Preventing Attention Entropy Collapse
AI总结:
本文提出σReparam方法,通过谱归一化重新参数化线性层,防止注意力熵坍缩,稳定Transformer训练,并在多项任务中实现无需复杂训练技巧的竞争性能。
摘要:
训练稳定性对Transformer至关重要。在本工作中,我们通过检查注意力层的演化来研究Transformer的训练动态。具体而言,我们在训练过程中追踪每个注意力头的注意力熵,这是模型锐度的一个代理指标。我们识别出不同架构和任务中常见的模式:低注意力熵伴随高训练不稳定性,这种不稳定性可能表现为损失震荡或发散。我们将病态的低注意力熵(对应于高度集中的注意力分数)称为“熵坍缩”。作为补救措施,我们提出σReparam,一种简单高效的解决方案,其中我们使用谱归一化和一个额外的可学习标量对所有线性层进行重新参数化。我们证明σReparam成功防止了注意力层中的熵坍缩,促进了更稳定的训练。此外,我们证明了注意力熵的一个紧下界,该下界随注意力logits的谱范数以指数速度减小,为我们的方法提供了额外动机。我们在图像分类、图像自监督学习、机器翻译、语音识别和语言建模任务上使用σReparam进行了实验。我们表明,σReparam提供了对超参数选择的稳定性和鲁棒性,甚至使得以下训练成为可能:(a)无需预热、权重衰减、层归一化或自适应优化器的Vision Transformer达到竞争性能;(b)机器翻译中的深层架构;(c)无需预热和自适应优化器的语音识别达到竞争性能。代码可在https://github.com/apple/ml-sigma-reparam获取。
英文摘要:
Training stability is of great importance to Transformers. In this work, we investigate the training dynamics of Transformers by examining the evolution of the attention layers. In particular, we track the attention entropy for each attention head during the course of training, which is a proxy for model sharpness. We identify a common pattern across different architectures and tasks, where low attention entropy is accompanied by high training instability, which can take the form of oscillating loss or divergence. We denote the pathologically low attention entropy, corresponding to highly concentrated attention scores, as $\textit{entropy collapse}$. As a remedy, we propose $σ$Reparam, a simple and efficient solution where we reparametrize all linear layers with spectral normalization and an additional learned scalar. We demonstrate that $σ$Reparam successfully prevents entropy collapse in the attention layers, promoting more stable training. Additionally, we prove a tight lower bound of the attention entropy, which decreases exponentially fast with the spectral norm of the attention logits, providing additional motivation for our approach. We conduct experiments with $σ$Reparam on image classification, image self-supervised learning, machine translation, speech recognition, and language modeling tasks. We show that $σ$Reparam provides stability and robustness with respect to the choice of hyperparameters, going so far as enabling training (a) a Vision Transformer {to competitive performance} without warmup, weight decay, layer normalization or adaptive optimizers; (b) deep architectures in machine translation and (c) speech recognition to competitive performance without warmup and adaptive optimizers. Code is available at \url{https://github.com/apple/ml-sigma-reparam}.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。