{"notes":[{"id":"transformers/pre-ln","category":"transformers","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"Pre-LN Transformer","subtitle":"No More Warm-Up","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":"en","path":"/content/papers/en/transformers/pre-ln.3fc6fded75ba.md","headings":[{"id":"defining-the-problem","text":"Defining the Problem","depth":2},{"id":"two-arrangements","text":"Two Arrangements","depth":3},{"id":"how-much-does-warm-up-really-matter","text":"How Much Does Warm-Up Really Matter?","depth":3},{"id":"solving-the-problem","text":"Solving the Problem","depth":2},{"id":"looking-at-the-moment-of-initialization","text":"Looking at the Moment of Initialization","depth":3},{"id":"step-1-how-big-is-the-signal-on-the-main-path","text":"Step 1: How Big Is the Signal on the Main Path?","depth":3},{"id":"step-2-layernorm-divides-the-gradient-by-the-size-of-its-input","text":"Step 2: LayerNorm Divides the Gradient by the Size of Its Input","depth":3},{"id":"step-3-putting-it-together","text":"Step 3: Putting It Together","depth":3},{"id":"measuring-it-for-real","text":"Measuring It for Real","depth":3},{"id":"discussion","text":"Discussion","depth":2},{"id":"translation-no-warm-up-and-faster-too","text":"Translation: No Warm-Up, and Faster Too","depth":3},{"id":"bert-40-saved","text":"BERT: 40% Saved","depth":3},{"id":"conclusion","text":"Conclusion","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":"The Dawn of a New Era","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":"en","path":"/content/papers/en/transformers/transformer.420ab06f9af4.md","headings":[{"id":"problem-definition","text":"Problem Definition","depth":2},{"id":"solving-the-problem","text":"Solving the Problem","depth":2},{"id":"model-design","text":"Model Design","depth":3},{"id":"input-layer","text":"Input Layer","depth":3},{"id":"positional-encoding","text":"Positional Encoding","depth":3},{"id":"self-attention-mechanism","text":"Self-Attention Mechanism","depth":3},{"id":"multi-head-attention-mechanism","text":"Multi-Head Attention Mechanism","depth":3},{"id":"cross-attention-mechanism","text":"Cross-Attention Mechanism","depth":3},{"id":"feed-forward-network","text":"Feed-Forward Network","depth":3},{"id":"layer-normalization","text":"Layer Normalization","depth":3},{"id":"discussion","text":"Discussion","depth":2},{"id":"why-use-self-attention","text":"Why Use Self-Attention?","depth":3},{"id":"experimental-results-machine-translation","text":"Experimental Results: Machine Translation","depth":3},{"id":"conclusion","text":"Conclusion","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"transformers/flashattention","relation":"next","topic":"Attention memory and computation"},{"id":"transformers/bert","relation":"next","topic":"Pretraining bidirectional language representations"},{"id":"transformers/gpt_1","relation":"next","topic":"From language pretraining to in-context learning"}]},{"id":"transformers/sparse-transformer","category":"transformers","legacyPath":"transformers/1904-sparse-transformer","legacySlug":"1904-sparse-transformer","legacyPaths":[],"title":"Sparse Transformer","subtitle":"Sparse Attention Mechanism","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":"en","path":"/content/papers/en/transformers/sparse-transformer.0a8593e9aca3.md","headings":[{"id":"defining-the-problem","text":"Defining the Problem","depth":2},{"id":"the-initial-layers","text":"The Initial Layers","depth":3},{"id":"layers-19-20","text":"Layers 19-20","depth":3},{"id":"various-layers","text":"Various Layers","depth":3},{"id":"layers-64-128","text":"Layers 64-128","depth":3},{"id":"solving-the-problem","text":"Solving the Problem","depth":2},{"id":"solution-1-fixed-stride","text":"Solution 1: Fixed Stride","depth":3},{"id":"solution-2-fixed-specific-index-tokens","text":"Solution 2: Fixed Specific Index Tokens","depth":3},{"id":"modified-attention-mechanism","text":"Modified Attention Mechanism","depth":3},{"id":"other-optimizations","text":"Other Optimizations","depth":3},{"id":"training-setup","text":"Training Setup","depth":3},{"id":"discussion","text":"Discussion","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":"conclusion","text":"Conclusion","depth":2}],"anchorAliases":{},"guide":null,"relatedReading":[{"id":"transformers/longformer","relation":"comparison","topic":"Reducing attention cost on long sequences"},{"id":"transformers/bigbird","relation":"comparison","topic":"Reducing attention cost on long sequences"}]},{"id":"classic-cnns/resnet-v2","category":"classic-cnns","legacyPath":"","legacySlug":"","legacyPaths":[],"title":"ResNet v2","subtitle":"Keep the Shortcut Clean","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":"en","path":"/content/papers/en/classic-cnns/resnet-v2.1b56c20d8b4c.md","headings":[{"id":"defining-the-problem","text":"Defining the Problem","depth":2},{"id":"solving-the-problem","text":"Solving the Problem","depth":2},{"id":"two-identity-mappings","text":"Two Identity Mappings","depth":3},{"id":"what-if-the-shortcut-carries-something-extra","text":"What If the Shortcut Carries Something Extra?","depth":3},{"id":"where-should-the-after-addition-relu-go","text":"Where Should the After-Addition ReLU Go?","depth":3},{"id":"a-different-angle-pre-activation","text":"A Different Angle: Pre-Activation","depth":3},{"id":"discussion","text":"Discussion","depth":2},{"id":"the-deeper-the-wider-the-gap","text":"The Deeper, the Wider the Gap","depth":3},{"id":"comparison-with-other-methods","text":"Comparison with Other Methods","depth":3},{"id":"conclusion","text":"Conclusion","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":"Training a Transformer comes with an odd habit: you can't set the learning rate where you want it right from the start.\n\nThe original Transformer paper uses this schedule: for the first 4000 steps, the learning rate climbs slowly from near 0 to its maximum, then gradually decays. This climb is called **learning-rate warm-up**. Convolutional networks usually start with a fairly large learning rate and lower it over time; only special cases such as very-large-batch training need warm-up. Transformers, though, need it almost every time.\n\nWarm-up is a nuisance. It slows training down and adds two more hyperparameters to tune, the maximum learning rate and the number of warm-up steps, and the final results are sensitive to both. When training a large model like BERT, every setting you try is expensive.\n\nWhy do Transformers insist on warm-up? In February 2020, Ruibin Xiong, Yunchang Yang, Di He, and colleagues from the Chinese Academy of Sciences, Peking University, Microsoft Research Asia, and Nankai University released this paper, later published at ICML 2020. Their answer: the problem lies in where LayerNorm is placed.\n\n\n> [!TIP]\n> \n> For the structure of each sub-layer in the original Transformer, you can start with our note; its LayerNorm section also notes that the original Transformer uses Post-LN, with LayerNorm after each residual addition:\n> \n> - **[[17.06] Transformer: The Dawn of a New Era](/blog/papers/transformers/transformer)**\n\n\n## Defining the Problem\n\n### Two Arrangements\n\nEach Transformer layer has two sub-layers: self-attention and a feed-forward network (FFN). Each sub-layer has a residual connection, and each has a LayerNorm. These three pieces can be arranged in different orders:\n\n[![Diagrams of two Transformer layers side by side. (a) The input x_l goes through Multi-Head Attention, is added to the shortcut, and then goes through Layer Norm; it then goes through the FFN, is added to the shortcut, and goes through Layer Norm again to give x_l+1. (b) The input x_l first goes through Layer Norm, then into Multi-Head Attention, and is added to the shortcut; next comes Layer Norm, then the FFN, then the addition with the shortcut, giving x_l+1. In both, the shortcut is drawn as a thick gray arrow](/content/papers/assets/f44d59bcadead95b9b93a42fc681bc473b038c13ed0c5c9e1effc7ee90409fc3.png)](/content/papers/assets/f44d59bcadead95b9b93a42fc681bc473b038c13ed0c5c9e1effc7ee90409fc3.png)\n\n*Figure 1 from the paper, excerpted from [**page 2 of the v2 PDF**](https://arxiv.org/pdf/2002.04745v2#page=2). (a) A Post-LN Transformer layer; (b) a Pre-LN Transformer layer.*\n\nOn the left is what the original Transformer and BERT do: sub-layer, residual addition, then LayerNorm, in that order. LayerNorm sits **between** the residual blocks, and the authors call this **Post-LN**. Taking the FFN sub-layer as an example:\n\n$$\nx_{l+1} = \\mathrm{LN}\\big(x + \\mathrm{FFN}(x)\\big)\n$$\n\nOn the right, LayerNorm moves inside the residual branch, **before** the sub-layer; this is called **Pre-LN**:\n\n$$\nx_{l+1} = x + \\mathrm{FFN}\\big(\\mathrm{LN}(x)\\big)\n$$\n\nPre-LN's main path has nothing but additions; since the main path is never normalized along the way, one more LayerNorm has to be added before the final output.\n\nPre-LN was already in use before this paper: the language models of Baevski and Auli and OpenAI's [Sparse Transformer](/blog/papers/transformers/sparse-transformer) both adopted it, and translation toolkits such as tensor2tensor and OpenNMT implemented it; Wang et al. even found that Pre-LN beats Post-LN as the number of layers grows. Still, all of these works kept warm-up in their training pipelines.\n\n### How Much Does Warm-Up Really Matter?\n\nThe authors first confirm the effect of warm-up with a Post-LN Transformer on the IWSLT14 German-to-English translation task. The warm-up schedule is:\n\n$$\n\\mathrm{lr}(t) = \\frac{t}{T_{\\text{warmup}}}\\,\\mathrm{lr}_{\\max}, \\quad t \\le T_{\\text{warmup}}\n$$\n\nfollowed by inverse square root decay. They train with Adam and with plain SGD, comparing runs with and without warm-up and with different warm-up lengths:\n\n[![Two groups of line charts, all with epochs on the horizontal axis. (a) Training with Adam: validation loss on the left, BLEU on the right. The two dashed lines without warm-up stall at a loss of around 7 with BLEU below 10; the orange line with 500 warm-up steps and a maximum learning rate of 1e-3 is also poor; the two lines with 4000 warm-up steps bring the loss down to about 4 and BLEU up above 30. (b) Training with SGD: with warm-up, the loss falls from 10 to between 6 and 7, with BLEU around 7; without warm-up, the loss stays above 8.5 and BLEU is near 0](/content/papers/assets/82e2f8ace58e613addf613e637786b7479429ee4db56d81756a47b8554ff3968.png)](/content/papers/assets/82e2f8ace58e613addf613e637786b7479429ee4db56d81756a47b8554ff3968.png)\n\n*Figure 2 from the paper, excerpted from [**page 4 of the v2 PDF**](https://arxiv.org/pdf/2002.04745v2#page=4). Validation loss and BLEU on IWSLT14 German-to-English: (a) with Adam, (b) with SGD.*\n\nThe results are clear:\n\n- **Without warm-up, it won't train.** With Adam and no warm-up, BLEU reaches only 8.45; models with warm-up get about 34. SGD does much worse than Adam to begin with, but the trend is the same: without warm-up, BLEU stays barely above 0 over 15 epochs.\n- **The warm-up length is a sensitive setting too.** With only 500 warm-up steps, maximum learning rates of 5e-4 and 1e-3 give BLEU scores of just 31.16 and 2.77, respectively.\n\nOthers had explained what warm-up does before. The authors of RAdam argued that warm-up reduces the variance of Adam's adaptive learning rate early in training. But here SGD needs warm-up too, so the authors conclude that the benefit of warm-up isn't tied to any particular optimizer, and the cause has to be sought in the model itself.\n\n## Solving the Problem\n\n### Looking at the Moment of Initialization\n\nWarm-up only happens during the first few thousand steps, so the authors focus on initialization: when the network has just been built and all its parameters are random, what do the gradients look like?\n\nThey analyze it with mean field theory and make a few simplifications: single-head attention only; every weight matrix is $d\\times d$ with Xavier initialization; the attention matrices $W^Q$ and $W^K$ are initialized to 0, so attention starts out uniform; and the input vectors are treated as Gaussian as well.\n\nThe analysis has three steps: first the size of the hidden states, then how LayerNorm treats gradients, and finally the two put together.\n\n### Step 1: How Big Is the Signal on the Main Path?\n\nPost-LN ends every layer with LayerNorm, so each layer's output norm is fixed at around $\\sqrt{d}$. The authors look instead at $x^{\\text{post},5}$, just before the last normalization: the vector after the FFN output is added to the residual. Pre-LN's main path isn't normalized, and every layer adds something on top. The authors prove:\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\nFor Post-LN it's the same size at every layer; for Pre-LN it grows linearly with the number of layers. Let's run the paper's simplified setup once and see whether that holds:\n\n```python\nimport torch\nimport torch.nn.functional as F\n\ntorch.manual_seed(0)\nd, n, L = 512, 8, 12                       # hidden dimension, sequence length, number of layers\nln = lambda v: F.layer_norm(v, (d,))\nxavier = lambda: torch.randn(d, d) / d ** 0.5   # Xavier: a d×d matrix with variance 2/(d+d) = 1/d\n\ndef attention(h, Wv):                      # the paper's simplification: W_Q = W_K = 0, so attention is uniform\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: normalize after the addition\n    x3 = ln(x_post + attention(x_post, Wv))\n    x5 = x3 + ffn(x3, W1, W2)\n    x_post = ln(x5)\n    # Pre-LN: normalize before each sub-layer; the main path just adds\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\"After layer {l:2d}: Post-LN ‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = {post:.2f}   \"\n              f\"Pre-LN ‖x‖²/d = {pre:5.2f} (theoretical range {1 + l / 2:.1f} to {1 + 3 * l / 2:.1f})\")\n```\n\n```text\nAfter layer  1: Post-LN ‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.51   Pre-LN ‖x‖²/d =  1.62 (theoretical range 1.5 to 2.5)\nAfter layer  2: Post-LN ‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.41   Pre-LN ‖x‖²/d =  2.27 (theoretical range 2.0 to 4.0)\nAfter layer  4: Post-LN ‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.55   Pre-LN ‖x‖²/d =  4.41 (theoretical range 3.0 to 7.0)\nAfter layer  8: Post-LN ‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.54   Pre-LN ‖x‖²/d =  9.73 (theoretical range 5.0 to 13.0)\nAfter layer 12: Post-LN ‖x⁽ᵖᵒˢᵗ,⁵⁾‖²/d = 1.32   Pre-LN ‖x‖²/d = 15.44 (theoretical range 7.0 to 19.0)\n```\n\nPost-LN stays around 1.5 (this is a single random draw, so there's some fluctuation); Pre-LN grows layer by layer, always within the theoretical range.\n\nWhere does the $\\tfrac{3}{2}$ come from? $x^{\\text{post},3}$ has just gone through LayerNorm, so its squared norm is $d$. Then it passes through the FFN: after multiplying by the Xavier-initialized $W^1$, each element has variance 1; ReLU cuts out half, leaving a squared norm of $d/2$; multiplying by $W^2$ keeps it at $d/2$. Add the two together and you get $d + d/2 = \\tfrac{3}{2}d$. Each Pre-LN layer likewise adds at least the FFN's $d/2$, which is where the $l/2$ in the lower bound comes from.\n\n### Step 2: LayerNorm Divides the Gradient by the Size of Its Input\n\nLayerNorm subtracts the mean from its input and divides by the standard deviation. Scale the input up by $k$, and the output stays the same; so, conversely, the derivative of the output with respect to the input shrinks by a factor of $k$. The authors write this as:\n\n$$\n\\big\\|J_{\\mathrm{LN}}(x)\\big\\|_2 = O\\!\\left(\\frac{\\sqrt{d}}{\\|x\\|_2}\\right)\n$$\n\n$J_{\\mathrm{LN}}$ is LayerNorm's Jacobian matrix. Let's compute it directly:\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                                          # same direction, length scaled by k\n    J = jacobian(lambda v: F.layer_norm(v, (d,)), x)   # Jacobian of LayerNorm with respect to its input, shape [d, d]\n    print(f\"‖x‖ = {x.norm():6.1f}   spectral norm of the Jacobian = {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   spectral norm of the Jacobian = 1.0651   √d / ‖x‖ = 1.0647\n‖x‖ =   30.1   spectral norm of the Jacobian = 0.5325   √d / ‖x‖ = 0.5323\n‖x‖ =   60.1   spectral norm of the Jacobian = 0.2663   √d / ‖x‖ = 0.2662\n‖x‖ =  120.2   spectral norm of the Jacobian = 0.1331   √d / ‖x‖ = 0.1331\n```\n\nEach time the input's length doubles, the Jacobian halves, matching $\\sqrt{d}/\\|x\\|$ almost exactly. When the gradient propagates back through LayerNorm, it's multiplied by this Jacobian: the larger the input, the smaller the gradient passed back.\n\n### Step 3: Putting It Together\n\nPutting the first two steps together gives the paper's main theorem. The authors analyze the second weight matrix of the last FFN layer, $W^{2,L}$, and conclude (with high probability):\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\nHere's the reasoning. In Post-LN, the input to the last LayerNorm is $x^{\\text{post},5}$, whose size is fixed at around $\\sqrt{3d/2}$ regardless of the number of layers, so the last layer's gradient doesn't depend on the number of layers either. In Pre-LN, the input to the final LayerNorm is the main path, accumulated all the way up; its squared norm grows linearly with the number of layers $L$, so its norm is on the order of $\\sqrt{L}$, and the gradient passing through it is divided by $\\sqrt{L}$.\n\nThe authors also give a preliminary extension of the analysis, concluding that in Post-LN, the closer parameters are to the output layer, the larger their gradients, shrinking toward the input; in Pre-LN, gradients are roughly the same across layers.\n\n### Measuring It for Real\n\nThe theory makes quite a few simplifications, so the authors check it on real models. They first confirm that the concentration property of the hidden-state norms holds, then measure the gradients:\n\n[![Four plots. (a) and (b) are bar charts of the gradient magnitudes of W1 and W2 in each FFN layer of a Transformer with a 6-layer encoder and a 6-layer decoder: the blue bars, Pre-LN (at initialization), vary little across layers; the orange bars, Post-LN (at initialization), increase with the layer and are highest at layer 6; the green bars, Post-LN (after warm-up), are almost 0. (c) The gradient of the last-layer W2 in Pre-LN falls as the total number of layers increases from 6 to 14, from 0.3 to about 0.13. (d) The gradient of the last-layer W2 in Post-LN stays around 1.6, unchanged by the number of layers](/content/papers/assets/1365260baf3a0bcbf5df3b21d3b56df47b2abd8ebd257a7c0472af1a24dd360a.png)](/content/papers/assets/1365260baf3a0bcbf5df3b21d3b56df47b2abd8ebd257a7c0472af1a24dd360a.png)\n\n*Figure 3 from the paper, excerpted from [**page 7 of the v2 PDF**](https://arxiv.org/pdf/2002.04745v2#page=7). (a)(b) Gradient magnitudes of the FFN weights in each layer of a 6-6-layer Transformer; (c)(d) gradient magnitude of the last-layer $W^2$ in models of different sizes. Each experiment is repeated with ten random seeds.*\n\nThis matches the theory:\n\n- The two plots on the right: Post-LN's last-layer gradient is around 1.6 no matter how big the model, while Pre-LN's falls as the number of layers grows.\n- The two plots on the left: Post-LN's gradients grow steadily from layer 1 to layer 6, largest in the layer closest to the output; Pre-LN's vary much less across layers.\n\nSo at initialization, Post-LN has large gradients on the parameters near the output. Using a large learning rate right away makes the update steps too big, and training becomes unstable. Warm-up sidesteps exactly this stretch: it waits until the gradients have shrunk before turning up the learning rate. The green bars in the figure (Post-LN after warm-up) are almost all 0; the gradients really have become much smaller.\n\nThe authors also tried training Post-LN without warm-up, using a fixed small learning rate of 1e-4 throughout. That does train the model to some extent, but convergence is much slower.\n\nLooking at it the other way around, Pre-LN's gradients are well-behaved right from initialization, so it should be able to do without warm-up.\n\n## Discussion\n\n### Translation: No Warm-Up, and Faster Too\n\nThe authors test this on two translation tasks:\n\n- **IWSLT14 German-to-English**: a 6-layer encoder and a 6-layer decoder, embedding dimension 512, FFN 1024, 4 attention heads.\n- **WMT14 English-to-German**: the Transformer base setting, with about 4.5 million parallel sentence pairs.\n\nPost-LN follows the original paper with 4000 warm-up steps, then inverse square root decay. Pre-LN has no warm-up at all: on IWSLT14 it starts at 5e-4 and is multiplied by 0.1 at epoch 8; on WMT14, two runs start at 7e-4 and 1.5e-3, decay at epoch 6, and then follow inverse square root decay. The two models' maximum learning rates are set to be the same. Since RAdam claims to remove the need for warm-up, the authors include it in the IWSLT14 comparison as well.\n\n[![Four line charts, all with epochs on the horizontal axis. (a) and (b) are validation loss and BLEU on IWSLT14: the red and orange lines for Pre-LN (with Adam or RAdam, both without warm-up) drop very low within the first few epochs, and BLEU quickly rises above 30; the green line for Post-LN with Adam and warm-up, and the blue line for Post-LN with RAdam and no warm-up, both climb more slowly and only catch up at the end. (c) and (d) are validation loss and BLEU on WMT14: at both learning rates, Pre-LN converges earlier than the corresponding Post-LN, and BLEU ends up around 26 for all](/content/papers/assets/26ae58e24cef2c40c73b93b2c17d74d07d2a6fafe71d071e70ed0a490d9d585f.png)](/content/papers/assets/26ae58e24cef2c40c73b93b2c17d74d07d2a6fafe71d071e70ed0a490d9d585f.png)\n\n*Figure 4 from the paper, excerpted from [**page 8 of the v2 PDF**](https://arxiv.org/pdf/2002.04745v2#page=8). Validation loss and BLEU on IWSLT14 German-to-English (a, b) and WMT14 English-to-German (c, d).*\n\nThree observations:\n\n1. **Warm-up is no longer necessary.** On IWSLT14, Pre-LN without warm-up ends at a BLEU of about 34 and a validation loss of about 4, on par with Post-LN with warm-up.\n2. **Faster convergence.** On IWSLT14, Pre-LN at epoch 9 performs about as well as Post-LN at epoch 15; WMT14 shows something similar.\n3. **Where LayerNorm goes matters more than the optimizer.** RAdam does let Post-LN train without warm-up; but once you switch to Pre-LN, whether you use RAdam or Adam makes almost no difference.\n\nThe authors also mention that with well-tuned hyperparameters (an initial learning rate of 7.5e-4, dropout of 0.3, and so on), Pre-LN can reach a test BLEU of 35.5 on IWSLT14.\n\n### BERT: 40% Saved\n\nPre-training uses BERT base (12 layers, hidden dimension 768, 12 heads), on English Wikipedia plus a BookCorpus the authors collected themselves, about 3.4 billion words in total.\n\nPost-LN BERT follows the original paper's settings: 10,000 warm-up steps and a maximum learning rate of 1e-4. The authors tried raising Post-LN's learning rate to 3e-4, and training simply diverged. Pre-LN BERT has no warm-up and decays linearly from 3e-4.\n\n[![Three line charts, all with pre-training steps (in thousands) on the horizontal axis. (a) BERT validation loss: the orange Pre-LN line is below the blue dashed Post-LN line from the start, and the two fall roughly in parallel; Pre-LN reaches 1.69 at about 500 thousand steps, while Post-LN takes until about 700 thousand. (b) MRPC accuracy and (c) RTE accuracy: Pre-LN is above Post-LN at every stage of pre-training, with the gap most pronounced early on](/content/papers/assets/9707b62474bc0eba4a5ff16f9578eb437856e931b5127793ec370b6a1b0d0a06.png)](/content/papers/assets/9707b62474bc0eba4a5ff16f9578eb437856e931b5127793ec370b6a1b0d0a06.png)\n\n*Figure 5 from the paper, excerpted from [**page 8 of the v2 PDF**](https://arxiv.org/pdf/2002.04745v2#page=8). (a) Validation loss during BERT pre-training; (b)(c) accuracy on MRPC and RTE after fine-tuning models from different pre-training stages, with error bars showing 95% confidence intervals over six fine-tuning runs.*\n\nPost-LN has to train until step 700k for its validation loss to reach 1.69; Pre-LN gets there by step 500k, about 40% faster. The 200k steps saved far exceed the 10k steps of warm-up itself, so the speed-up comes from more than skipping warm-up: above all, Pre-LN can handle a much larger learning rate.\n\nFine-tuning models from different stages on MRPC and RTE shows that Pre-LN also converges faster on downstream tasks.\n\n## Conclusion\n\nThis paper traces an engineering habit back to an architectural cause: at initialization, Post-LN has large gradients on the parameters near the output, which warm-up is needed to avoid. Move LayerNorm inside the residual branch, before the sub-layer, and the size of the main path shrinks the last layer's gradient by $1/\\sqrt{L}$, the gradients become more even across layers, warm-up can be dropped, and the learning rate can be turned up higher.\n\nKeep the scope of the theoretical analysis in mind: it looks only at the moment of initialization and makes simplifications such as uniform attention and a single head; the part about Post-LN gradients growing larger closer to the output is labeled by the authors themselves as a preliminary result, supported mainly by the measurements in Figure 3.\n\nThe idea of “move normalization into the branch and keep the shortcut clean” had also appeared in convolutional networks; GPT-2 likewise moved LayerNorm to the input of each sub-layer:\n\n\n> [!TIP]\n> \n> - **[[16.03] ResNet v2: Keep the Shortcut Clean](/blog/papers/classic-cnns/resnet-v2)**\n> - **[[19.02] GPT-2: Forty-Eight Layers of Decoders](/blog/papers/transformers/gpt_2)**\n\n\nIf you're training a Transformer from scratch, this paper's advice is straightforward: put LayerNorm before each sub-layer, and remember to add one more LayerNorm before the output. At the scale of the paper's experiments, this trains without warm-up, and you can use a somewhat larger learning rate.\n\nPre-LN later became the default arrangement for many large language models, but it comes at a cost. You can see it from the result of step 1: the main path's squared norm grows linearly with the number of layers, while each layer adds roughly the same amount, so the later the layer, the smaller its relative influence on the main path. In 2021, [NormFormer](https://arxiv.org/abs/2110.09456) pointed out that during Pre-LN pre-training, the early layers' gradients are much larger than the later layers', and added a few extra normalizations to each layer; in 2022, [DeepNet](https://arxiv.org/abs/2203.00555) went back and reworked Post-LN, aiming to keep both Post-LN's performance and Pre-LN's stability, and pushed Transformers to a thousand layers.\n\nWhere LayerNorm should go is a question that isn't settled yet; we'll come back to it when we get the chance.\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}}}