{"notes":[{"id":"transformers/pre-ln","category":"transformers","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"Pre-LN Transformer","subtitle":"ウォームアップはもういらない","author":"Z. Yuan","paperDate":"2020-02-12","publishedOn":"2026-03-18","world":"external/transformer-layernorm","paperTitle":"On Layer Normalization in the Transformer Architecture","paperUrl":"https://arxiv.org/abs/2002.04745","bibliography":{"title":"On Layer Normalization in the Transformer Architecture","authors":["Ruibin Xiong","Yunchang Yang","Di He","Kai Zheng","Shuxin Zheng","Chen Xing","Huishuai Zhang","Yanyan Lan","Liwei Wang","Tie-Yan Liu"],"url":"https://arxiv.org/abs/2002.04745","pdfUrl":"https://arxiv.org/pdf/2002.04745v2","codeUrl":null,"preprint":{"id":"2002.04745","submittedOn":"2020-02-12","url":"https://arxiv.org/abs/2002.04745"},"publication":null,"citation":"@misc{transformerspreln2020,\n  title = {{On Layer Normalization in the Transformer Architecture}},\n  author = {Ruibin Xiong and Yunchang Yang and Di He and Kai Zheng and Shuxin Zheng and Chen Xing and Huishuai Zhang and Yanyan Lan and Liwei Wang and Tie-Yan Liu},\n  year = {2020},\n  eprint = {2002.04745},\n  archivePrefix = {arXiv},\n  url = {https://arxiv.org/abs/2002.04745}\n}\n","citationPath":"/content/papers/citations/transformers/pre-ln.419d25cfd84b.bib"},"thumbnail":{"src":"/content/papers/assets/f44d59bcadead95b9b93a42fc681bc473b038c13ed0c5c9e1effc7ee90409fc3.png","width":436,"height":531},"locale":"ja","path":"/content/papers/ja/transformers/pre-ln.149dd51245a0.md","headings":[{"id":"問題の定義","text":"問題の定義","depth":2},{"id":"二つの置き方","text":"二つの置き方","depth":3},{"id":"ウォームアップはどれほど重要か","text":"ウォームアップはどれほど重要か？","depth":3},{"id":"問題の解決","text":"問題の解決","depth":2},{"id":"初期化の瞬間を見る","text":"初期化の瞬間を見る","depth":3},{"id":"ステップ1主経路の信号はどれくらい大きいか","text":"ステップ1：主経路の信号はどれくらい大きいか？","depth":3},{"id":"ステップ2layernormは勾配を入力の大きさで割る","text":"ステップ2：LayerNormは勾配を入力の大きさで割る","depth":3},{"id":"ステップ3つなぎ合わせる","text":"ステップ3：つなぎ合わせる","depth":3},{"id":"実際に測る","text":"実際に測る","depth":3},{"id":"考察","text":"考察","depth":2},{"id":"翻訳ウォームアップなしでしかも速い","text":"翻訳：ウォームアップなしで、しかも速い","depth":3},{"id":"bert4割の節約","text":"BERT：4割の節約","depth":3},{"id":"結論","text":"結論","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"transformers/transformer","relation":"referenced","topic":""},{"id":"transformers/sparse-transformer","relation":"referenced","topic":""},{"id":"classic-cnns/resnet-v2","relation":"referenced","topic":""}]},{"id":"transformers/transformer","category":"transformers","legacyPath":"transformers/1706-transformer","legacySlug":"1706-transformer","legacyPaths":[],"title":"Transformer","subtitle":"新世界の起点","author":"Z. Yuan","paperDate":"2017-06-12","publishedOn":"2024-05-05","paperTitle":"Attention Is All You Need","paperUrl":"https://arxiv.org/abs/1706.03762","bibliography":{"title":"Attention Is All You Need","authors":["Ashish Vaswani","Noam Shazeer","Niki Parmar","Jakob Uszkoreit","Llion Jones","Aidan N. Gomez","Lukasz Kaiser","Illia Polosukhin"],"url":"https://arxiv.org/abs/1706.03762","pdfUrl":"https://arxiv.org/pdf/1706.03762","codeUrl":null,"preprint":{"id":"1706.03762","submittedOn":"2017-06-12","url":"https://arxiv.org/abs/1706.03762"},"publication":{"venue":"NeurIPS (NIPS)","year":2017,"kind":"conference","url":"https://proceedings.neurips.cc/paper_files/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html"},"citation":"@misc{transformerstransformer2017,\n  title = {{Attention Is All You Need}},\n  author = {Ashish Vaswani and Noam Shazeer and Niki Parmar and Jakob Uszkoreit and Llion Jones and Aidan N. Gomez and Lukasz Kaiser and Illia Polosukhin},\n  year = {2017},\n  eprint = {1706.03762},\n  archivePrefix = {arXiv},\n  url = {https://arxiv.org/abs/1706.03762}\n}\n","citationPath":"/content/papers/citations/transformers/transformer.86c0068d96fd.bib"},"thumbnail":{"src":"/content/papers/assets/29a55cb05497bdd9bbb3f3a96b7ed6656fda3552fa6a0fa5eb29e8e899c1acc1.jpg","width":756,"height":1080},"locale":"ja","path":"/content/papers/ja/transformers/transformer.ed6af0e8846d.md","headings":[{"id":"問題の定義","text":"問題の定義","depth":2},{"id":"問題の解決","text":"問題の解決","depth":2},{"id":"モデル設計","text":"モデル設計","depth":3},{"id":"入力層","text":"入力層","depth":3},{"id":"位置エンコーディング","text":"位置エンコーディング","depth":3},{"id":"自己注意力機構","text":"自己注意力機構","depth":3},{"id":"多頭注意力機構","text":"多頭注意力機構","depth":3},{"id":"クロスアテンション機構","text":"クロスアテンション機構","depth":3},{"id":"フィードフォワードネットワーク","text":"フィードフォワードネットワーク","depth":3},{"id":"レイヤ正規化layer-normalization","text":"レイヤ正規化（Layer Normalization）","depth":3},{"id":"討論","text":"討論","depth":2},{"id":"なぜ自己注意力機構を使うのか","text":"なぜ自己注意力機構を使うのか？","depth":3},{"id":"実験結果機械翻訳","text":"実験結果：機械翻訳","depth":3},{"id":"結論","text":"結論","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"transformers/flashattention","relation":"next","topic":"Attention のメモリと計算量"},{"id":"transformers/bert","relation":"next","topic":"双方向言語表現の事前学習"},{"id":"transformers/gpt_1","relation":"next","topic":"言語事前学習から文脈内学習へ"}]},{"id":"transformers/sparse-transformer","category":"transformers","legacyPath":"transformers/1904-sparse-transformer","legacySlug":"1904-sparse-transformer","legacyPaths":[],"title":"Sparse Transformer","subtitle":"疎な注意機構","author":"Z. Yuan","paperDate":"2019-04-23","publishedOn":"2024-06-09","paperTitle":"Generating Long Sequences with Sparse Transformers","paperUrl":"https://arxiv.org/abs/1904.10509","bibliography":{"title":"Generating Long Sequences with Sparse Transformers","authors":["Rewon Child","Scott Gray","Alec Radford","Ilya Sutskever"],"url":"https://arxiv.org/abs/1904.10509","pdfUrl":"https://arxiv.org/pdf/1904.10509","codeUrl":null,"preprint":{"id":"1904.10509","submittedOn":"2019-04-23","url":"https://arxiv.org/abs/1904.10509"},"publication":null,"citation":"@misc{transformerssparsetransformer2019,\n  title = {{Generating Long Sequences with Sparse Transformers}},\n  author = {Rewon Child and Scott Gray and Alec Radford and Ilya Sutskever},\n  year = {2019},\n  eprint = {1904.10509},\n  archivePrefix = {arXiv},\n  url = {https://arxiv.org/abs/1904.10509}\n}\n","citationPath":"/content/papers/citations/transformers/sparse-transformer.961330a17a72.bib"},"thumbnail":{"src":"/content/papers/assets/847fd8e03beeb4395ae629700c5daaf1fa75dd302949008b9c248f7d77a8f9bc.jpg","width":1392,"height":694},"locale":"ja","path":"/content/papers/ja/transformers/sparse-transformer.bbb5722de6c8.md","headings":[{"id":"問題の定義","text":"問題の定義","depth":2},{"id":"最前面の数層","text":"最前面の数層","depth":3},{"id":"1920-層","text":"19〜20 層","depth":3},{"id":"特定されない層","text":"特定されない層","depth":3},{"id":"64128-層","text":"64〜128 層","depth":3},{"id":"解決問題","text":"解決問題","depth":2},{"id":"解法アプローチ-1-固定ステップサイズ","text":"解法アプローチ 1: 固定ステップサイズ","depth":3},{"id":"解法アプローチ-2-特定のインデックスの-token-を固定","text":"解法アプローチ 2: 特定のインデックスの Token を固定","depth":3},{"id":"修正された注意機構","text":"修正された注意機構","depth":3},{"id":"その他の最適化","text":"その他の最適化","depth":3},{"id":"訓練設定","text":"訓練設定","depth":3},{"id":"討論","text":"討論","depth":2},{"id":"cifar-10","text":"CIFAR-10","depth":3},{"id":"enwik8","text":"EnWik8","depth":3},{"id":"imagenet-64x64","text":"ImageNet 64x64","depth":3},{"id":"結論","text":"結論","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"transformers/longformer","relation":"comparison","topic":"長い系列に対する Attention の計算コスト"},{"id":"transformers/bigbird","relation":"comparison","topic":"長い系列に対する Attention の計算コスト"}]},{"id":"classic-cnns/resnet-v2","category":"classic-cnns","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"ResNet v2","subtitle":"ショートカットはきれいに","author":"Z. Yuan","paperDate":"2016-03-16","publishedOn":"2026-02-14","world":"external/resnet-v2","paperTitle":"Identity Mappings in Deep Residual Networks","paperUrl":"https://arxiv.org/abs/1603.05027","bibliography":{"title":"Identity Mappings in Deep Residual Networks","authors":["Kaiming He","Xiangyu Zhang","Shaoqing Ren","Jian Sun"],"url":"https://arxiv.org/abs/1603.05027","pdfUrl":"https://arxiv.org/pdf/1603.05027v3","codeUrl":null,"preprint":{"id":"1603.05027","submittedOn":"2016-03-16","url":"https://arxiv.org/abs/1603.05027"},"publication":null,"citation":"@misc{classiccnnsresnetv22016,\n  title = {{Identity Mappings in Deep Residual Networks}},\n  author = {Kaiming He and Xiangyu Zhang and Shaoqing Ren and Jian Sun},\n  year = {2016},\n  eprint = {1603.05027},\n  archivePrefix = {arXiv},\n  url = {https://arxiv.org/abs/1603.05027}\n}\n","citationPath":"/content/papers/citations/classic-cnns/resnet-v2.3d311a21e308.bib"},"thumbnail":{"src":"/content/papers/assets/7eda4acdde3019e92c5b44f59917b1e493a5b11b5b8a63280d03b15787934ef5.png","width":967,"height":469},"locale":"ja","path":"/content/papers/ja/classic-cnns/resnet-v2.2eccc1d938b8.md","headings":[{"id":"問題の定義","text":"問題の定義","depth":2},{"id":"問題の解決","text":"問題の解決","depth":2},{"id":"二つの恒等写像","text":"二つの恒等写像","depth":3},{"id":"ショートカットに何かを足すとどうなるか","text":"ショートカットに何かを足すとどうなるか？","depth":3},{"id":"加算の後のreluはどこに置くべきか","text":"加算の後のReLUはどこに置くべきか？","depth":3},{"id":"見方を変えるpre-activation","text":"見方を変える：pre-activation","depth":3},{"id":"考察","text":"考察","depth":2},{"id":"深いほど差は大きくなる","text":"深いほど差は大きくなる","depth":3},{"id":"ほかの手法と比べる","text":"ほかの手法と比べる","depth":3},{"id":"結論","text":"結論","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"classic-cnns/resnet","relation":"referenced","topic":""},{"id":"classic-cnns/highway-networks","relation":"referenced","topic":""},{"id":"classic-cnns/inception-v4","relation":"referenced","topic":""}]}],"text":"Transformerの学習には、奇妙な習慣があります。学習率を最初から本来の値にできないのです。\n\n元のTransformerの論文が使ったスケジュールはこうです。最初の4000ステップで学習率を0近くから最大値までゆっくり上げ、その後少しずつ減衰させます。この上昇の部分を**学習率のウォームアップ**（warm-up）と呼びます。畳み込みネットワークは普通、最初から大きめの学習率を使い、だんだん小さくしていきます。ウォームアップが必要になるのは、超大規模バッチでの学習のような特殊な場合だけです。ところがTransformerは、ほぼ毎回ウォームアップを必要とします。\n\nウォームアップは厄介です。学習が遅くなるうえ、調整すべきハイパーパラメータが二つ増えます。最大学習率とウォームアップのステップ数で、しかも最終的な成績はこの二つの値にとても敏感です。BERTのような大きなモデルを学習させるときは、設定を一つ試すたびに大きなコストがかかります。\n\nTransformerは、なぜウォームアップなしでは済まないのでしょうか。2020年2月、中国科学院、北京大学、Microsoft Research Asia、南開大学のRuibin Xiong、Yunchang Yang、Di Heらがこの論文を公開し、のちにICML 2020で発表しました。彼らの答えは、問題はLayerNormを置く位置にある、というものです。\n\n\n> [!TIP]\n> \n> 元のTransformerの各サブ層の構造は、先に私たちのノートで確認できます。そのノートのLayerNormの節でも、LayerNormの位置はさらにNorm-FirstとNorm-Lastに分けられることに触れています。\n> \n> - **[[17.06] Transformer：新世界の起点](/blog/papers/transformers/transformer)**\n\n\n## 問題の定義\n\n### 二つの置き方\n\nTransformerの各層には、自己注意と順伝播ネットワーク（FFN）の二つのサブ層があります。どのサブ層にも残差接続とLayerNormがあり、この三つの順序には異なる並べ方があります。\n\n[![左右二つのTransformer層の模式図。(a)入力x_lがMulti-Head Attentionを通り、ショートカットと加算された後にLayer Norm。さらにFFNを通り、ショートカットと加算された後にもう一度Layer Normを行い、x_l+1になる。(b)入力x_lはまずLayer Normを通ってからMulti-Head Attentionに入り、ショートカットと加算される。続いてLayer Normを通ってからFFNに入り、ショートカットと加算されてx_l+1になる。どちらもショートカットは灰色の太い矢印で描かれている](/content/papers/assets/f44d59bcadead95b9b93a42fc681bc473b038c13ed0c5c9e1effc7ee90409fc3.png)](/content/papers/assets/f44d59bcadead95b9b93a42fc681bc473b038c13ed0c5c9e1effc7ee90409fc3.png)\n\n*原論文Figure 1。[**v2 PDFの2ページ**](https://arxiv.org/pdf/2002.04745v2#page=2)から抜粋。(a)Post-LN Transformer層、(b)Pre-LN Transformer層。*\n\n左は元のTransformerとBERTのやり方で、サブ層、残差の加算、LayerNormの順に進みます。LayerNormは残差ブロックの**あいだ**に置かれ、著者らはこれを**Post-LN**と呼んでいます。FFNサブ層を例にとると次のとおりです。\n\n$$\nx_{l+1} = \\mathrm{LN}\\big(x + \\mathrm{FFN}(x)\\big)\n$$\n\n右はLayerNormを残差分岐の中に移し、サブ層の**前**に置いたもので、**Pre-LN**と呼びます。\n\n$$\nx_{l+1} = x + \\mathrm{FFN}\\big(\\mathrm{LN}(x)\\big)\n$$\n\nPre-LNの主経路には加算しかありません。主経路では最後まで正規化しないので、出力の前にLayerNormをもう1層加える必要があります。\n\nPre-LNはこの論文以前から使われていました。BaevskiとAuliの言語モデルや、OpenAIの[Sparse Transformer](/blog/papers/transformers/sparse-transformer)が採用し、tensor2tensorやOpenNMTといった翻訳ツールにも実装されています。Wangらは、層が多くなるとPre-LNがPost-LNより良いことも見出しました。ただ、これらの研究の学習手順には、どれもウォームアップが残っていました。\n\n### ウォームアップはどれほど重要か？\n\n著者らはまず、IWSLT14の独英翻訳タスクで、Post-LN Transformerを使ってウォームアップの影響を確かめました。ウォームアップのスケジュールは次のとおりです。\n\n$$\n\\mathrm{lr}(t) = \\frac{t}{T_{\\text{warmup}}}\\,\\mathrm{lr}_{\\max}, \\quad t \\le T_{\\text{warmup}}\n$$\n\nその後は逆平方根減衰を続けます。Adamと通常のSGDのそれぞれで学習させ、ウォームアップの有無と長さを比べました。\n\n[![2組の折れ線グラフで、横軸はどちらもエポック。(a)Adamで学習：左が検証損失、右がBLEU。ウォームアップなしの2本の破線は、損失が7前後で止まり、BLEUは10未満。ウォームアップ500ステップ、最大学習率1e-3の橙色の線も悪い。ウォームアップ4000ステップの2本の線は、損失が約4まで下がり、BLEUは30以上に上がる。(b)SGDで学習：ウォームアップありでは損失が10から6〜7のあいだまで下がり、BLEUは約7。ウォームアップなしでは損失が8.5以上で止まり、BLEUは0に近い](/content/papers/assets/82e2f8ace58e613addf613e637786b7479429ee4db56d81756a47b8554ff3968.png)](/content/papers/assets/82e2f8ace58e613addf613e637786b7479429ee4db56d81756a47b8554ff3968.png)\n\n*原論文Figure 2。[**v2 PDFの4ページ**](https://arxiv.org/pdf/2002.04745v2#page=4)から抜粋。IWSLT14独英の検証損失とBLEU：(a)Adam、(b)SGD。*\n\n結果ははっきりしています。\n\n- **ウォームアップなしでは学習できない**。Adamでウォームアップなしだと、BLEUは8.45にとどまり、ウォームアップしたモデルは約34です。SGDの成績はもともとAdamよりずっと悪いものの、傾向は同じで、ウォームアップなしでは15エポックのあいだBLEUが0をわずかに上回るだけです。\n- **ウォームアップの長さにも敏感**。ウォームアップが500ステップだけだと、最大学習率5e-4と1e-3のBLEUはそれぞれ31.16と2.77にとどまります。\n\nウォームアップの働きについては、以前にも説明が試みられていました。RAdamの著者らは、ウォームアップは学習の初期にAdamの適応的な学習率の分散を抑えるものだと考えました。しかしここではSGDでもウォームアップが必要なので、著者らは、ウォームアップの利点は特定の最適化手法とは関係なく、原因はモデル自体に求めるべきだと考えています。\n\n## 問題の解決\n\n### 初期化の瞬間を見る\n\nウォームアップは最初の数千ステップでしか起きないので、著者らは初期化に焦点を当てました。ネットワークを作ったばかりで、パラメータがすべてランダムな値のとき、勾配はどんな姿をしているのでしょうか？\n\n彼らは平均場理論で分析し、いくつかの簡略化を行いました。注意はシングルヘッドだけにする。重み行列はすべて$d\\times d$で、Xavier初期化を使う。注意の$W^Q$、$W^K$を0で初期化し、最初は注意が一様分布になるようにする。入力ベクトルもガウス分布とみなす。\n\n分析は3段階です。まず隠れ層の大きさを見て、次にLayerNormが勾配をどう扱うかを見て、最後に両者をつなぎます。\n\n### ステップ1：主経路の信号はどれくらい大きいか？\n\nPost-LNでは各層の最後にLayerNormを行うので、各層の出力のノルムは$\\sqrt{d}$前後に固定されます。そこで著者らは、最後の正規化の直前の$x^{\\text{post},5}$、つまりFFNの出力に残差を足した後のベクトルを見ます。Pre-LNの主経路は正規化されず、層ごとに何かが足されていきます。著者らは次を証明しました。\n\n$$\n\\mathbb{E}\\big\\|x^{\\text{post},5}_{l}\\big\\|_2^2 = \\tfrac{3}{2}d, \\qquad\n\\Big(1+\\tfrac{l}{2}\\Big)d \\le \\mathbb{E}\\big\\|x^{\\text{pre}}_{l}\\big\\|_2^2 \\le \\Big(1+\\tfrac{3l}{2}\\Big)d\n$$\n\nPost-LNはどの層でも同じ大きさで、Pre-LNは層数に比例して大きくなります。論文の簡略化した設定どおりに一度動かして、本当にそうなるかを見てみましょう。\n\n```python\nimport torch\nimport torch.nn.functional as F\n\ntorch.manual_seed(0)\nd, n, L = 512, 8, 12                       # 隠れ次元、系列長、層数\nln = lambda v: F.layer_norm(v, (d,))\nxavier = lambda: torch.randn(d, d) / d ** 0.5   # Xavier：d×dの行列、分散は2/(d+d) = 1/d\n\ndef attention(h, Wv):                      # 論文の簡略化した設定：W_Q = W_K = 0で、注意は一様分布\n    return (h @ Wv).mean(dim=0, keepdim=True).expand_as(h)\n\ndef ffn(h, W1, W2):\n    return F.relu(h @ W1) @ W2\n\nx_post = x_pre = torch.randn(n, d)\nfor l in range(1, L + 1):\n    Wv, W1, W2 = xavier(), xavier(), xavier()\n    # Post-LN：加算の後に正規化\n    x3 = ln(x_post + attention(x_post, Wv))\n    x5 = x3 + ffn(x3, W1, W2)\n    x_post = ln(x5)\n    # Pre-LN：正規化してからサブ層に入り、主経路では直接加算\n    x_mid = x_pre + attention(ln(x_pre), Wv)\n    x_pre = x_mid + ffn(ln(x_mid), W1, W2)\n    if l in (1, 2, 4, 8, 12):\n        post = (x5 ** 2).sum(1).mean() / d\n        pre = (x_pre ** 2).sum(1).mean() / d\n        print(f\"第{l:2d}層の後：Post-LNの‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = {post:.2f}   \"\n              f\"Pre-LNの‖x‖²/d = {pre:5.2f}（理論上の範囲 {1 + l / 2:.1f}～{1 + 3 * l / 2:.1f}）\")\n```\n\n```text\n第 1層の後：Post-LNの‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.51   Pre-LNの‖x‖²/d =  1.62（理論上の範囲 1.5～2.5）\n第 2層の後：Post-LNの‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.41   Pre-LNの‖x‖²/d =  2.27（理論上の範囲 2.0～4.0）\n第 4層の後：Post-LNの‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.55   Pre-LNの‖x‖²/d =  4.41（理論上の範囲 3.0～7.0）\n第 8層の後：Post-LNの‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.54   Pre-LNの‖x‖²/d =  9.73（理論上の範囲 5.0～13.0）\n第12層の後：Post-LNの‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.32   Pre-LNの‖x‖²/d = 15.44（理論上の範囲 7.0～19.0）\n```\n\nPost-LNはずっと1.5付近にあります（1回のサンプリングにすぎないので、多少のばらつきはあります）。Pre-LNは1層ごとに大きくなり、どれも理論上の範囲に収まっています。\n\n$\\tfrac{3}{2}$はどこから来るのでしょうか。$x^{\\text{post},3}$はLayerNormを通った直後なので、二乗ノルムは$d$です。これがFFNを通ります。Xavier初期化した$W^1$を掛けると、各要素の分散は1になります。ReLUで半分が切り落とされ、二乗ノルムは$d/2$が残ります。さらに$W^2$を掛けても$d/2$のままです。両者を足すと$d + d/2 = \\tfrac{3}{2}d$です。Pre-LNでも各層で少なくともFFNの$d/2$が加わり、これが下限の$l/2$になっています。\n\n### ステップ2：LayerNormは勾配を入力の大きさで割る\n\nLayerNormは入力から平均を引き、標準偏差で割ります。入力を$k$倍にしても出力は変わらないので、逆に出力の入力に対する微分は$k$分の1になります。著者らはこれを次のように書いています。\n\n$$\n\\big\\|J_{\\mathrm{LN}}(x)\\big\\|_2 = O\\!\\left(\\frac{\\sqrt{d}}{\\|x\\|_2}\\right)\n$$\n\n$J_{\\mathrm{LN}}$はLayerNormのヤコビ行列です。直接計算してみましょう。\n\n```python\nimport torch\nimport torch.nn.functional as F\nfrom torch.autograd.functional import jacobian\n\ntorch.manual_seed(0)\nd = 256\nz = torch.randn(d)\nfor k in (1, 2, 4, 8):\n    x = k * z                                          # 同じ方向のまま、長さをk倍にする\n    J = jacobian(lambda v: F.layer_norm(v, (d,)), x)   # 入力に対するLayerNormのヤコビ行列、形は[d, d]\n    print(f\"‖x‖ = {x.norm():6.1f}   ヤコビ行列のスペクトルノルム = {torch.linalg.matrix_norm(J, ord=2):.4f}\"\n          f\"   √d / ‖x‖ = {d ** 0.5 / x.norm():.4f}\")\n```\n\n```text\n‖x‖ =   15.0   ヤコビ行列のスペクトルノルム = 1.0651   √d / ‖x‖ = 1.0647\n‖x‖ =   30.1   ヤコビ行列のスペクトルノルム = 0.5325   √d / ‖x‖ = 0.5323\n‖x‖ =   60.1   ヤコビ行列のスペクトルノルム = 0.2663   √d / ‖x‖ = 0.2662\n‖x‖ =  120.2   ヤコビ行列のスペクトルノルム = 0.1331   √d / ‖x‖ = 0.1331\n```\n\n入力の長さが2倍になるたびにヤコビ行列は半分になり、$\\sqrt{d}/\\|x\\|$とほぼ完全に一致します。勾配が逆向きにLayerNormを通るとき、掛けられるのがこのヤコビ行列です。入力が大きいほど、戻っていく勾配は小さくなります。\n\n### ステップ3：つなぎ合わせる\n\n前の二つのステップをつなぎ合わせたものが、この論文の主定理です。著者らが分析しているのは最後の層のFFNの2番目の重み行列$W^{2,L}$で、結論は（高い確率で成り立つ条件のもとで）次のとおりです。\n\n$$\n\\text{Post-LN：}\\left\\|\\frac{\\partial \\tilde{\\mathcal{L}}}{\\partial W^{2,L}}\\right\\|_F \\le O\\big(d\\sqrt{\\ln d}\\big), \\qquad\n\\text{Pre-LN：}\\left\\|\\frac{\\partial \\tilde{\\mathcal{L}}}{\\partial W^{2,L}}\\right\\|_F \\le O\\left(d\\sqrt{\\frac{\\ln d}{L}}\\right)\n$$\n\n理屈はこうです。Post-LNの最後のLayerNormの入力は$x^{\\text{post},5}$で、大きさは$\\sqrt{3d/2}$前後に固定され、層数とは無関係なので、最後の層の勾配も層数と無関係になります。Pre-LNの最後のLayerNormの入力は、ずっと足し合わされてきた主経路で、その二乗ノルムは層数$L$に比例して増え、ノルムはおよそ$\\sqrt{L}$のオーダーです。勾配はここを通るときに$\\sqrt{L}$で割られます。\n\n著者らはさらに予備的な一般化の分析も行い、Post-LNでは出力層に近いパラメータほど勾配が大きく、入力側へ行くほど小さくなり、Pre-LNでは各層の勾配がほぼ同じになると結論づけています。\n\n### 実際に測る\n\n理論はかなり簡略化しているので、著者らは実際のモデルで確かめました。まず隠れ層のノルムの集中性が成り立つことを確認し、それから勾配を測ります。\n\n[![4枚の図。(a)(b)はエンコーダー6層、デコーダー6層のTransformerで、各層のFFNのW1とW2の勾配の大きさを示す棒グラフ。青のPre-LN（初期化時）は層ごとの高さがあまり変わらない。橙のPost-LN（初期化時）は層とともに増え、第6層が最も高い。緑のPost-LN（ウォームアップ後）はほぼ0。(c)Pre-LNの最後の層のW2の勾配は、全体の層数が6から14に増えるにつれて0.3から約0.13まで下がる。(d)Post-LNの最後の層のW2の勾配は1.6前後のままで、層数によって変わらない](/content/papers/assets/1365260baf3a0bcbf5df3b21d3b56df47b2abd8ebd257a7c0472af1a24dd360a.png)](/content/papers/assets/1365260baf3a0bcbf5df3b21d3b56df47b2abd8ebd257a7c0472af1a24dd360a.png)\n\n*原論文Figure 3。[**v2 PDFの7ページ**](https://arxiv.org/pdf/2002.04745v2#page=7)から抜粋。(a)(b)6-6層Transformerの各層のFFNの重みの勾配の大きさ、(c)(d)異なる大きさのモデルにおける最後の層の$W^2$の勾配の大きさ。各実験は10個の乱数シードで繰り返しています。*\n\n理論とよく合っています。\n\n- 右の2枚：Post-LNの最後の層の勾配は、モデルの大きさにかかわらず1.6前後です。Pre-LNは層数が増えるにつれて下がります。\n- 左の2枚：Post-LNの勾配は第1層から第6層までずっと大きくなっていき、出力に最も近い層が最大です。Pre-LNは層ごとの差がずっと小さくなっています。\n\nつまりPost-LNでは、初期化時に出力に近いパラメータの勾配がとても大きくなっています。ここでいきなり大きな学習率を使うと、更新の歩幅が大きすぎて学習が不安定になります。ウォームアップは、ちょうどこの区間を避けてくれます。勾配が小さくなってから学習率を上げるのです。図の緑の棒（ウォームアップ後のPost-LN）はほぼ0で、勾配は確かにずっと小さくなっています。\n\n著者らは別に、ウォームアップなしで、最初から最後まで1e-4という固定の小さな学習率でPost-LNを学習させることも試しました。これでもある程度まで学習できますが、収束はずっと遅くなります。\n\n逆に言えば、Pre-LNは初期化の時点で勾配がおとなしいので、ウォームアップを外せるはずです。\n\n## 考察\n\n### 翻訳：ウォームアップなしで、しかも速い\n\n著者らは二つの翻訳タスクで検証しました。\n\n- **IWSLT14独英**：エンコーダー6層、デコーダー6層、埋め込み次元512、FFN 1024、注意ヘッド4個。\n- **WMT14英独**：Transformer baseの設定で、対訳文ペアは約450万組。\n\nPost-LNは元の論文どおり4000ステップのウォームアップの後、逆平方根減衰を続けます。Pre-LNはウォームアップを一切しません。IWSLT14では5e-4から始めて8エポック目に0.1を掛け、WMT14では7e-4と1.5e-3からそれぞれ始めて6エポック目に減衰させ、その後は逆平方根減衰を続けます。両者の最大学習率は同じにしています。ウォームアップを外せるとうたうRAdamとも、IWSLT14で一緒に比較しました。\n\n[![4枚の折れ線グラフで、横軸はどれもエポック。(a)(b)はIWSLT14の検証損失とBLEU：Pre-LN（AdamでもRAdamでも、ウォームアップなし）の赤線と橙線は最初の数エポックで低いところまで下がり、BLEUもすぐに30以上に上がる。Adamとウォームアップを使うPost-LNの緑線と、RAdamでウォームアップなしの青線は、どちらも上がり方が遅く、最後にようやく追いつく。(c)(d)はWMT14の検証損失とBLEU：二つの学習率のどちらでも、Pre-LNは対応するPost-LNより早く収束し、最終的なBLEUはどれも26前後](/content/papers/assets/26ae58e24cef2c40c73b93b2c17d74d07d2a6fafe71d071e70ed0a490d9d585f.png)](/content/papers/assets/26ae58e24cef2c40c73b93b2c17d74d07d2a6fafe71d071e70ed0a490d9d585f.png)\n\n*原論文Figure 4。[**v2 PDFの8ページ**](https://arxiv.org/pdf/2002.04745v2#page=8)から抜粋。IWSLT14独英（a、b）とWMT14英独（c、d）の検証損失とBLEUです。*\n\n三つのことがわかります。\n\n1. **ウォームアップはもう必要ない**。IWSLT14では、ウォームアップなしのPre-LNの最終的なBLEUは約34、検証損失は約4で、ウォームアップしたPost-LNと同等です。\n2. **収束が速い**。IWSLT14では、Pre-LNの9エポック目の成績が、Post-LNの15エポック目とほぼ同じです。WMT14でも同様です。\n3. **LayerNormの位置のほうが最適化手法より重要**。RAdamは確かに、Post-LNをウォームアップなしで学習させられます。しかしPre-LNにすると、RAdamとAdamのどちらを使ってもほとんど差がありません。\n\n著者らは、ハイパーパラメータをうまく調整すれば（初期学習率7.5e-4、dropout 0.3など）、Pre-LNはIWSLT14でテストBLEU 35.5に達するとも述べています。\n\n### BERT：4割の節約\n\n事前学習にはBERT base（12層、隠れ次元768、ヘッド12個）を使い、データは英語版Wikipediaと、著者らが自分で収集したBookCorpusで、合計約34億語です。\n\nPost-LN BERTは元の論文の設定どおり、ウォームアップ10,000ステップ、最大学習率1e-4です。著者らがPost-LNの学習率を3e-4に上げてみたところ、学習はそのまま発散しました。Pre-LN BERTはウォームアップなしで、3e-4から線形に減衰させます。\n\n[![3枚の折れ線グラフで、横軸はどれも事前学習のステップ数（千ステップ）。(a)BERTの検証損失：橙のPre-LNは最初から青い破線のPost-LNより低く、2本の線はほぼ平行に下がる。Pre-LNは約500kステップで1.69まで下がり、Post-LNは約700kステップかかる。(b)MRPCの正解率と(c)RTEの正解率：どの事前学習段階でもPre-LNがPost-LNを上回り、差は事前学習の初期に最もはっきりしている](/content/papers/assets/9707b62474bc0eba4a5ff16f9578eb437856e931b5127793ec370b6a1b0d0a06.png)](/content/papers/assets/9707b62474bc0eba4a5ff16f9578eb437856e931b5127793ec370b6a1b0d0a06.png)\n\n*原論文Figure 5。[**v2 PDFの8ページ**](https://arxiv.org/pdf/2002.04745v2#page=8)から抜粋。(a)BERT事前学習の検証損失、(b)(c)さまざまな事前学習段階のモデルをMRPCとRTEでファインチューニングしたときの正解率で、エラーバーは6回のファインチューニングの95%信頼区間です。*\n\nPost-LNは700kステップまで学習させてようやく検証損失が1.69まで下がりますが、Pre-LNは500kステップで到達し、約40%速くなっています。節約できた200kステップは、ウォームアップそのものの10kステップよりはるかに多いので、速くなった理由はウォームアップがなくなったことだけでなく、Pre-LNがずっと大きな学習率に耐えられることにあります。\n\nさまざまな段階のモデルをMRPCとRTEでファインチューニングすると、Pre-LNは下流タスクでも速く収束しました。\n\n## 結論\n\nこの論文は、エンジニアリング上の一つの習慣を、アーキテクチャ上の原因までたどりました。Post-LNでは初期化時に出力に近いパラメータの勾配がとても大きく、ウォームアップで避ける必要があります。LayerNormを残差分岐の中に移してサブ層の前に置くと、最後の層の勾配は主経路の大きさによって$1/\\sqrt{L}$の割合で縮み、各層の勾配もより均一になるので、ウォームアップを外せ、学習率もさらに大きくできます。\n\n理論分析の範囲は覚えておく必要があります。見ているのは初期化の瞬間で、一様な注意やシングルヘッドといった簡略化もしています。「Post-LNは出力に近いほど勾配が大きい」という部分は著者ら自身も予備的な結果としていて、主にFigure 3の実測が裏づけています。\n\n「正規化を分岐の中に移し、ショートカットをきれいに保つ」という考え方は、畳み込みネットワークでも登場しています。GPT-2も、LayerNormを各サブ層の入力の位置に移しました。\n\n\n> [!TIP]\n> \n> - **[[16.03] ResNet v2：ショートカットはきれいに](/blog/papers/classic-cnns/resnet-v2)**\n> - **[[19.02] GPT-2：48 層デコーダ](/blog/papers/transformers/gpt_2)**\n\n\nTransformerをゼロから学習させるなら、この論文の勧めは明快です。LayerNormはサブ層の前に置き、出力の前にLayerNormをもう1層加えるのを忘れないこと。論文の実験規模では、これでウォームアップなしでも学習でき、学習率も少し大きくできます。\n\nPre-LNはのちに多くの大規模言語モデルで標準の配置になりましたが、代償もあります。それはステップ1の結果からもわかります。主経路の二乗ノルムは層数に比例して増えるのに、各層が足し込む量はほぼ同じなので、後ろの層ほど主経路に与える相対的な影響が小さくなります。2021年の[NormFormer](https://arxiv.org/abs/2110.09456)は、Pre-LNでは事前学習時に前の方の層の勾配が後ろの層よりずっと大きいことを指摘し、各層に正規化をいくつか追加しました。2022年の[DeepNet](https://arxiv.org/abs/2203.00555)は逆にPost-LNを改造し、Post-LNの性能とPre-LNの安定性を両立させようとして、Transformerを1000層まで深くしました。\n\nLayerNormを結局どこに置くべきかという話題はまだ終わっていません。また機会があれば取り上げましょう。\n","imageSizes":{"/content/papers/assets/f44d59bcadead95b9b93a42fc681bc473b038c13ed0c5c9e1effc7ee90409fc3.png":{"width":436,"height":531},"/content/papers/assets/82e2f8ace58e613addf613e637786b7479429ee4db56d81756a47b8554ff3968.png":{"width":1348,"height":391},"/content/papers/assets/1365260baf3a0bcbf5df3b21d3b56df47b2abd8ebd257a7c0472af1a24dd360a.png":{"width":1332,"height":263},"/content/papers/assets/26ae58e24cef2c40c73b93b2c17d74d07d2a6fafe71d071e70ed0a490d9d585f.png":{"width":1342,"height":276},"/content/papers/assets/9707b62474bc0eba4a5ff16f9578eb437856e931b5127793ec370b6a1b0d0a06.png":{"width":1011,"height":369}}}
