[论文解读] Stabilizing Transformer Training by Preventing Attention Entropy Collapse

Shuangfei Zhai, Tatiana Likhomanenko+6 more|arXiv (Cornell University)|Mar 11, 2023

Advanced Neural Network Applications被引用 6

一句话总结

本文将注意力熵崩溃——即自注意力分数中病理性的低熵——识别为Transformer训练不稳定的根源。提出了一种基于谱归一化的重参数化方法$σ$ Reparam，通过引入可学习标量对线性层进行重参数化，可在视觉、语音和语言任务中稳定训练，实现无需预热、权重衰减或自适应优化器的鲁棒训练。

ABSTRACT

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 $ extit{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}.

研究动机与目标

• 通过分析注意力熵动态，探究Transformer训练不稳定的根源原因。
• 识别出低注意力熵与损失振荡或发散之间存在一致关联的模式。
• 开发一种通用且高效的解决方案，以防止在多种架构和任务中出现注意力熵崩溃。
• 证明$σ$ Reparam可实现极简超参数调优下的稳定训练。

提出的方法

• 将注意力熵作为模型尖锐度的代理指标，追踪其在注意力头、层和训练步骤中的变化。
• 提出$σ$ Reparam：使用谱归一化和可学习标量$σ$对所有线性层进行重参数化。
• 将$σ$ Reparam应用于自注意力机制中的查询、键和值投影。
• 证明注意力熵存在一个紧的下界，其随注意力logits谱范数的增加呈指数下降。
• 通过受控干预（如温度缩放）建立熵崩溃与训练不稳定性之间的因果关系。
• 在图像分类、自监督学习、机器翻译、语音识别和语言建模任务中评估$σ$ Reparam的性能。

Figure 1: Transformers are sensitive to hyperparameters. Increasing the learning rate easily causes attention entropy collapse and training divergence. Left: baseline Vision Transformer with default hyperparameters from Touvron et al. ( 2021 ) ; right: $2\times$ learning rate ( $5\times 10^{-4}\maps

实验结果

研究问题

• RQ1 在Transformer中，是否存在低注意力熵与训练不稳定性之间的一致相关性？
• RQ2 注意力熵崩溃是否可被因果性地关联到损失发散或振荡？
• RQ3 $σ$ Reparam是否能有效防止熵崩溃，并在多种任务中实现训练稳定？
• RQ4 $σ$ Reparam是否能消除对权重衰减、预热和自适应优化器等常见正则化技术的依赖？
• RQ5 注意力logits的谱范数与注意力熵之间存在何种理论关系？

主要发现

• 注意力熵崩溃——即注意力分数中接近零的熵——与ViT、NLP和ASR模型的训练发散或振荡同时发生。
• $σ$ Reparam成功防止了熵崩溃，使视觉Transformer的训练无需预热、权重衰减、层归一化或自适应优化器即可稳定进行。
• 在机器翻译任务中，$σ$ Reparam使原本无法收敛的深层Transformer架构得以成功训练。
• 在语音识别任务中，$σ$ Reparam在无需预热或自适应优化器的情况下实现了具有竞争力的性能，与基线结果相当。
• 理论分析表明，注意力熵存在一个紧的下界，其随注意力logits谱范数的增加呈指数下降。
• 实验结果证实，$σ$ Reparam对超参数选择具有鲁棒性，并可在视觉、语音和语言任务中实现简化的训练方案。

Figure 2: Training can become unstable due to rapid change in attention logit magnitude. We train a Vision Transformer, sharply reducing its temperature in the attention logits by $10\times$ at different intervention epochs. (Blue) Intervention during warmup – at epoch 10 – induces a sharp drop in t