Understanding the Difficulty of Training Transformers [1] は、Post-LN Transformer の学習の不安定さが勾配消失によるものでないことを解析と対照実験によって示し、サブ層の出力の分散のうち残差接続の出力が占める割合が小さいと学習が不安定になり、逆に大きいと層を深くした効果が得られないことを解析し、これを調整するために残差接続に対してトークン埋め込みベクトルの成分ごとに係数をかける (初期値はデータから決める) 手法 Admin: Adaptive model initialization を提案する論文です。University of Illinois at Urbana-Champaign の Liyuan Liu らによって著されました (他に Microsoft Research、Microsoft Dynamics 365 AI 所属の著者も)。2020 年の EMNLP に採択されました。
参考文献
- Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, Jiawei Han. Understanding the Difficulty of Training Transformers. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP 2020), 2020.
https://aclanthology.org/2020.emnlp-main.463/
https://arxiv.org/abs/2004.08249
関連記事
【文献メモ】RAdam: 2 次モーメントによる正規化の分散が一定な Adam (ICLR 2020) — 同じ第一著者の過去研究についての文献メモです。
以下の内容は [1] の arXiv 版に基づきます。
文献概要
研究背景
Transformer の学習にはオプティマイザや学習率スケジューラの注意深い設計が必要で、Post-LN を Pre-LN に替えると安定するものの到達性能は劣ると指摘されていました。具体的に、IWSLT'14 De-En で、RAdam (関連記事) の 2 次モーメントの減衰率と学習率の組み合わせを 15 通り試すと、Pre-LN は 15 通りすべてで収束する一方 Post-LN は 7 通りで発散し、収束した場合は 8 通り中 7 通りで Pre-LN を上回っていました。著者らは、Post-LN の学習を不安定にし、かつ Pre-LN の性能を制限している要因を明らかにすることを目的としています。
初期化直後の勾配消失の解析
著者らは、初期化直後に勾配消失が生じるのは Post-LN のデコーダだけで、Post-LN のエンコーダには生じないと解析によって主張しました (定理 1。LayerNorm のパラメータの初期値や、勾配と各モジュールの微分の独立性などを仮定しています)。そのうえで、Post-LN エンコーダと Pre-LN デコーダを組み合わせて勾配消失を取り除いても 18 層モデルの学習は発散したと報告し、勾配消失は学習の不安定さの直接の原因ではないと結論しました。また、勾配の大きさが層ごとに不均衡でも、オプティマイザがパラメータごとの学習率を調整して更新量を揃えていると述べています。
サブ層の出力の分散のうち残差接続の出力が占める割合の解析と実測
初期化直後、サブ層の出力の分散のうち残差接続の出力が占める割合は、Post-LN では層によらずほぼ一定なのに対し、Pre-LN では上の層ほど大きいです。Pre-LN でこの割合が上の層ほど大きくなるのは、層の出力が正規化されないまま足し合わされ、上の層ほど層の出力の分散が大きくなるためです。著者らは、このために Pre-LN の各サブ層は自分の変換部分に強く依存できず、層 i と層 i+1 の出力がほとんど区別できなくなるため、層を 1 つ足しても、同じ入力に並列な変換を足すこと、つまり最終層の幅を 2 倍にすることと変わらないと論じ、これが Pre-LN の到達性能を制限していると説明しました。実測でも、この割合は学習が進むにつれて小さくなるものの、学習後の Pre-LN では Post-LN より大きいままでした。また、残差接続の出力への依存をさらに強める ReZero を 6 層モデルに適用すると、IWSLT'14 De-En の BLEU が 1〜2 下がりました。
パラメータ変化に対する出力の変化の解析と Admin の提案・実験
他方、割合がほぼ一定に保たれる Post-LN の側には別の難点があります。パラメータをわずかに変化させたときの最終出力の変化は、Post-LN では層数に比例して大きくなり、Pre-LN では層数の対数程度にとどまります (定理 2)。(実験でも、この関係は決定係数 0.99 で成り立っていました。) Post-LN の最終出力の変化を抑えたいです。そこで Admin では、Post-LN の残差接続の出力に係数をかけて、最終出力の変化を Pre-LN と同じ層数の対数程度に抑えます。係数の決め方は 2 段階に分かれます。まず標準的な初期化のもとで最初の 1 バッチを順伝播し、各サブ層で残差接続される部分の出力の分散を測ります ("Profiling")。次に各サブ層の係数を、それ以前のサブ層について測った分散の総和の平方根に設定します ("Initialization")。係数は学習後に他のパラメータへ吸収でき、通常の Post-LN に戻せます。実験の結果、Admin は WMT'14 En-De の 6・12・18 層すべてで Post-LN・Pre-LN を上回り (12 層・18 層の Post-LN は発散)、60 層エンコーダ + 12 層デコーダの構成で WMT'14 En-Fr の BLEU 43.8 を達成しました。