[論文レビュー] Signal Propagation in Transformers: Theoretical Perspectives and the Role of Rank Collapse

Lorenzo Noci, Sotiris Anagnostidis+4 more|arXiv (Cornell University)|Jun 7, 2022

Neural Networks and Applications被引用数 4

ひとこと要約

本稿は、Transformerにおける信号伝播を調査し、初期化時のランク崩壊がクエリおよびキーにおける勾配消失の根本的原因であると特定した。本稿は、トークン表現のランクと勾配フローを安定化させるための深さに依存する残差スケーリングを提案するとともに、クエリ/キーと値のパラメータ間で勾配ノルムが非対称であることを明らかにし、Transformerにおける適応的最適化手法の有効性を説明した。

ABSTRACT

Transformers have achieved remarkable success in several domains, ranging from natural language processing to computer vision. Nevertheless, it has been recently shown that stacking self-attention layers - the distinctive architectural component of Transformers - can result in rank collapse of the tokens' representations at initialization. The question of if and how rank collapse affects training is still largely unanswered, and its investigation is necessary for a more comprehensive understanding of this architecture. In this work, we shed new light on the causes and the effects of this phenomenon. First, we show that rank collapse of the tokens' representations hinders training by causing the gradients of the queries and keys to vanish at initialization. Furthermore, we provide a thorough description of the origin of rank collapse and discuss how to prevent it via an appropriate depth-dependent scaling of the residual branches. Finally, our analysis unveils that specific architectural hyperparameters affect the gradients of queries and values differently, leading to disproportionate gradient norms. This suggests an explanation for the widespread use of adaptive methods for Transformers' optimization.

研究の動機と目的

• 自己注意層における初期化時のランク崩壊の原因と影響を理解すること。
• ランク崩壊がクエリおよびキーのパラメータにおける勾配消失を引き起こし、学習を妨げる仕組みを調査すること。
• トークン表現のランクを保持し、信号伝播を安定化させるための深さに依存する残差スケーリング戦略を提案すること。
• アーキテクチャのハイパーパrameterがクエリ、キー、値の勾配ノルムにどのように異なる影響を与えるかを分析すること。
• 深層Transformerの学習において適応的最適化手法が広く使われる理由を理論的および実験的に裏付けること。

提案手法

• 確率的行列理論を用いて、初期化時の単一自己注意層におけるクエリ、キー、値の期待勾配を理論的に導出する。
• トークン表現間のコサイン類似度を安定化させ、ランク崩壊を防ぐ深さに依存する残差スケーリング方式を構築する。
• クエリおよびキーの勾配ノルムが入力ノルムの立方に比例する一方で、値の勾配ノルムは線形に比例することを示し、勾配の大きさに顕著な不均衡が生じることを明らかにする。
• 勾配ノルムをクエリと値の間で実験的にバランスさせるために、温度スケーリング付きソフトマックスを導入する。
• 理論的分析と、トランストランスレーションタスク（IWSLT’14）およびトランショナルタスクを用いた実験的検証により、勾配ダイナミクスと安定性を検証する。
• 学習可能な残差スケーリングを用いてモデルを訓練し、相関、活性化ノルム、勾配フローをモニタリングすることで、理論的予測の妥当性を検証する。

実験結果

リサーチクエスチョン

• RQ1 初期化時のランク崩壊は、深層Transformerにおけるクエリおよびキーのパラメータの勾配にどのように影響するか？
• RQ2 ランク崩壊を防ぎ、深さにわたる信号伝播を安定化させるためのアーキテクチャ的メカニズムは何か？
• RQ3 Adamのような適応的最適化手法がTransformerで優れた性能を示す理由は何か？これは勾配ノルムの不均衡によって説明可能か？
• RQ4 入力相関、シーケンス長、埋め込み次元は、クエリ、キー、値の相対的勾配ノルムにどのように影響するか？
• RQ5 深さに依存する残差スケーリングは、トークン表現のランクを安定化させ、深層Transformerにおける学習ダイナミクスを改善できるか？

主な発見

• 初期化時のランク崩壊は、クエリおよびキーのパラメータにおける勾配消失を引き起こし、深層Transformerにおける学習を著しく損なう。
• 深さに依存する残差スケーリングは、トークン表現間のコサイン類似度を安定化させ、無限の深さに達してもランクが保持され、退化が防がれる。
• クエリおよびキーの勾配ノルムは入力ノルムの立方に比例するが、値の勾配ノルムは線形に比例するため、勾配の大きさに顕著な不均衡が生じる。
• 実験的結果から、適切なスケーリングがなければ、深さを増すほどトークンの整合性が高まり、学習性能が悪化することが示された（図1参照）。
• 温度スケーリング付きソフトマックス（$ au = 8.5$）を導入することで、クエリと値の間で勾配ノルムがバランスされ、理論的予測が妥当であることが実証された。
• 訓練済みモデルでは、残差スケーリングパラメータ $eta_1, eta_2$ が学習中に適応的に変化し、トークン相関が低く保たれるため、提案されたスケーリングの理論的安定性が裏付けられた。