DeepNet: Scaling Transformers to 1,000 Layers [1] は、深い Post-LN Transformer の学習が不安定になるのは学習初期のモデル更新が大きすぎるためであることを示し、残差接続に定数 $\alpha > 1$ をかける DeepNorm と、一部の重みの初期値を $\beta < 1$ 倍する初期化を組み合わせた DeepNet を提案して、1000 層の Transformer の学習を可能にした論文です。Microsoft Research の Hongyu Wang らによって著されました。IEEE TPAMI に掲載されました (arXiv 版は 2022 年)。

参考文献

Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, Furu Wei. DeepNet: Scaling Transformers to 1,000 Layers. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(10), 2024.
https://dl.acm.org/doi/abs/10.1109/TPAMI.2024.3386927
https://arxiv.org/abs/2203.00555

以下の内容は arXiv 版に基づきます。

文献概要

深い Post-LN が不安定になる理由

著者らは、深い Post-LN では学習のごく初期にモデル更新が爆発し、その後ほとんど更新されなくなることを観察しました。大きな更新で LN の入力が大きくなると、LN を通る勾配の大きさは入力の大きさに反比例するため勾配消失が起き、モデルが悪い局所解にとらわれると説明しています。

DeepNet: DeepNorm + 初期化の縮小

DeepNet は、Post-LN の代わりに DeepNorm を用いた Transformer です。DeepNorm では、$l$ 番目のサブ層 (エンコーダ層のセルフアテンション・FFN、デコーダ層のセルフアテンション・クロスアテンション・FFN) の Post-LN の残差接続を

x_{l+1} = \mathrm{LN}\bigl(\alpha x_l + G_l(x_l, \theta_l)\bigr)

とします ($G_l$ はアテンションまたは FFN)。さらに初期化時に、FFN の重みと、アテンションの V 変換・出力変換の重みを $\beta$ 倍します (Q・K の変換は通常の Xavier 初期化)。エンコーダのみ ($N$ 層) なら $\alpha = (2N)^{1/4}$, $\beta = (8N)^{-1/4}$ とします。通常の Post-LN ではモデル更新の大きさが全サブ層のパラメータ変化の和のオーダーで積み上がるのに対し、DeepNet では層数によらない $O(1)$ に抑えられると著者らは解析しています。

$\alpha$ と $\beta$ はどちらもアーキテクチャ (エンコーダ・デコーダの層数) だけで決まる定数で、学習するパラメータではありません。

実験

1000 層 (エンコーダ 500 層・デコーダ 500 層で、サブ層は計 2500 ¹) の DeepNet を問題なく学習できました。

また、対訳コーパスで 200 層・3.2B パラメータの DeepNet を学習したところ、48 層・12B パラメータの M2M-100 (Facebook の多言語翻訳モデル) を WMT・OPUS・TED・Flores-101 のすべての多言語翻訳評価データセットにおいて BLEU スコアで上回りました。

エンコーダではエンコーダ層 1 層あたりセルフアテンション・FFN の 2 サブ層、デコーダではデコーダ層 1 層あたりセルフアテンション・クロスアテンション・FFN の 3 個の計 2500。 ↩