[論文レビュー] 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における学習不安定化の根本的要因として特定した。本研究では、スペクトル正則化に基づく再パラメータ化手法である$ abla$ Reparamを提案し、可 learnable スカラー$ abla$を用いることで、視覚、音声、言語処理のあらゆるタスクにおいて学習を安定化させ、ウォームアップや重み減衰、適応的最適化手法を一切不要とすることを可能にした。

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における学習不安定化の根本的要因を解明すること。
• 低エントロピーな注目スコアが損失の振動や発散と相関する共通のパターンを同定すること。
• 多様なアーキテクチャーやタスクに一般化可能な、効率的なエントロピー収縮防止の解決策を開発すること。
• $\sigma$ Reparamが最小限のハイパーパramータチューニングで安定した学習を可能にすることを実証すること。

提案手法

• アテンションヘッド、レイヤー、学習ステップを横断して、モデルの鋭さの代理指標として注目エントロピーを追跡する。
• $ abla$ Reparamの導入：すべての線形層をスペクトル正則化と可学習スカラー$ abla$を用いて再パラメータ化する。
• 自己注意機構のクエリ、キー、バリューの投影に$ abla$ Reparamを適用する。
• 注目エントロピーのタイトな下界を証明し、それが注目ログイットのスペクトルノルムに指数関数的に減少することを示す。
• 温度スケーリングなどの制御された干渉を用いて、エントロピー収縮と学習不安定化の因果関係を確立する。
• 画像分類、自己教師付き学習、機械翻訳、音声認識、言語モデリングの各タスクで$ abla$ 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 $\sigma$ Reparamは多様なタスクにおいてエントロピー収縮を効果的に防止し、学習を安定化させることができるか？
• RQ4 $\sigma$ Reparamは、重み減衰、ウォームアップ、適応的最適化手法といった一般的な正則化手法の必要性を排除できるか？
• RQ5 注目ログイットのスペクトルノルムと注目エントロピーの理論的関係は何か？

主な発見

• 注目エントロピー収縮（注目スコアにおけるほぼゼロのエントロピー）は、ViT、NLP、ASRモデルの両方で学習の発散または振動と一致する。
• $\sigma$ Reparamはエントロピー収縮を効果的に防止し、ウォームアップ、重み減衰、層正則化、適応的最適化手法を一切不要としないVision Transformerの安定した学習を可能にした。
• 機械翻訳のタスクでは、$ abla$ Reparamにより、元々収束しなかった深層Transformerアーキテクチャの学習が可能になった。
• 音声認識タスクでは、$ abla$ Reparamによりウォームアップや適応的最適化手法を不要としながら、ベースラインと同等の性能を達成した。
• 理論的解析により、注目エントロピーにはタイトな下界が存在し、それが注目ログイットのスペクトルノルムに指数関数的に減少することが示された。
• 実験により、$ abla$ Reparamはハイパーパramータの選択に対して頑健であり、視覚、音声、言語処理のあらゆるタスクで簡素化された学習レシピの実現に成功した。

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