Peri-LN: Revisiting Normalization Layer in the Transformer Architecture

1) “$o$ seems intermediate; softmax is applied in an unnatural way.”

We believe there may be a misunderstanding regarding the reviewer’s concern that $o$ is an intermediate representation and that softmax is applied in an unnatural manner. As noted in Section 3.4 and Section D, our theoretical analysis focuses on the final MLP layer. This choice is motivated by two reasons: (1) the gradient norm at the final layer is empirically known to be the most unstable (see Figure1); and (2) it allows for a rigorous mathematical analysis without requiring approximations or assumptions. Analyzing the final layer for theoretical insight is an established practice in the literature (see Theorem 1 in [1]). Importantly, our empirical observations show that other layers exhibit similar trends in gradient behavior and hidden-state variance (see Section 4.4), supporting the generality of our theoretical insights.

[1] Xiong et al. "On layer normalization in the transformer architecture." ICML 2020.

2) “No theoretical proof for MHSA”

As discussed in our Response #1, we focus on the final MLP sub-layer rather than the intermediate MHSA layers. For these reasons, we introduce a proposition centered on the MLP layer, aiming to explain the phenomena observed in our experiments.

3) “Summation is necessary in multivariates”

We respectfully note that the comment regarding explicit summation in multivariate expressions appears to stem from a misunderstanding: in our matrix notation, the summation is inherently handled by matrix multiplication. In componentwise notation, the multivariate chain rule involves a summation over the relevant indices. However, in matrix notation, the summation is taken care of by matrix multiplication. Please refer to Theorem 5.3 in [2].

[2] Colley, Susan Jane. Vector calculus. PEARSON EDUCATION LIMITED, 2012.

4) “$\gamma$ is assumed positive, yet it can become negative in practice.”

We appreciate this point. In theory, $\gamma$ is a scaling parameter primarily intended to adjust magnitude, so the mathematical derivation naturally assumes $\gamma > 0$. In this paper, we assume $\gamma$ remains positive for theoretical simplicity, and we will include a clarification of this assumption in the revised manuscript. Empirically, we further verified $\gamma$ during by monitoring the magnitude across all checkpoints for a 30B-token run; it stayed strictly positive in every layer. We would like to emphasize that, as discussed in Section 4.4.1, even when we freeze $\gamma$ to 1, Peri-LN still retains its main benefits.

5) “RMSNorm <-> LN or ReLU <-> GeLU might not match exactly”

We understand the concern that RMSNorm and ReLU may feel “approximate” compared to LN and GeLU. Our motivation was analytical tractability:

• RMSNorm omits mean-centering, which simplifies the Jacobian while still capturing the main effect of normalization (rescaling by the vector’s norm). LayerNorm introduces an additional term subtracting the mean, but in large dimensions, mean removal has a smaller relative effect—so the bounding principle remains largely the same [3].
• ReLU is piecewise-linear, making it straightforward to compute partial derivatives explicitly. Since we are only considering the gradient with respect to $W^{(2)}$ in the final layer, it makes no difference whether the preceding hidden activation function is ReLU or GeLU. In either case, the hidden activation $h$ is treated as a constant when differentiating with respect to $W^{(2)}$, so the conclusions regarding stability with respect to $\||h\||$ remain the same.

Hence, these substitutions allow us to cleanly show how placing LN at different points can dampen or amplify the backpropagated gradients. We agree it is not a full 1:1 equivalence to LN or GeLU, but the essential mechanics are preserved. In the revised text, we will add disclaimers that our derivation is an instructive idealization—RMSNorm vs. LayerNorm and ReLU vs. GeLU do not qualitatively change the conclusion about how LN placement impacts model dynamics.

[3] Zhang et al. "Root mean square layer normalization." Neurips 2019.

Concluding Remarks

It seems there has been a misunderstanding, and we kindly ask you to reevaluate our work in light of its theoretical and experimental contributions. As Reviewer 1xqQ and AtH3 highlighted, we do believe that our theoretical contribution is solid. Your comments have prompted us to refine our exposition and add more detailed supporting materials. While we respectfully maintain that our core derivations—even with simplified assumptions (RMSNorm instead of LN, ReLU instead of GeLU)—provide a valid perspective on how LN placement influences Transformers, we will include additional clarifications and empirical validations in a revised manuscript. We hope these efforts will address your concerns and more clearly convey the findings of our work. Thank you again for your time.

1) Extending Analysis to Other Layers

Thank you for highlighting this. As noted in Section 3.4 and Section D, our analysis focuses on the final layer. Following Theorem 1 in [1], we analyze the last layer because its gradients are often the largest in magnitude. We chose $W^{(2)}$ (the final linear projection in the MLP) as a representative example since it most directly feeds into the residual connection $(x + a)$, making it more transparent to illustrate how gradient norms can explode or vanish. This direct link to the residual path can significantly impact gradient stability in subsequent layers. Nonetheless, for a more comprehensive understanding of LLMs training dynamics, extending this theoretical foundation to other components (such as attention) is indeed important. We appreciate the reviewer’s insight on this matter and plan to pursue this direction as part of our future research.

[1] Xiong, Ruibin, et al. "On layer normalization in the transformer architecture." ICML 2020.

2) Extending Exploration to Vision or Multimodal Tasks

Due to time and resource constraints, we could not run additional experiments in vision or multimodal settings. However, existing literature lends some support to the broader applicability of our findings. For instance, Sun et al. [2] reports that in ViT (Vision Transformer) architectures, massive activations can also emerge under a Pre-LN setup, paralleling what we observe in language models. This similarity suggests that the insights from our Peri-LN analysis could extend beyond pure language tasks. Given the trend of integrating LLMs into large-scale vision-language models, the reviewer’s question is indeed highly relevant. As outlined in our paper, we focused on large language models as the primary use case for Peri-LN. However, we see great potential in exploring vision or multimodal tasks in future work, building on the theoretical and empirical observations presented here.

[2] Sun, Mingjie, et al. "Massive activations in large language models." COLM 2024.

3) Weight Initialization: Additional Experiments & Clarification

• Additional Experiments: In response to the reviewer’s question, we conducted additional experiments to explore different weight initialization methods. In this study, for both Pre-LN and Peri-LN architectures, we apply Xavier initialization [3]. As shown in the table below, Xavier initialization yields better performance compared to our previous weight initialization configurations. We also confirm that our main observation still holds: large variance occurs in Pre-LN Transformers but not in Peri-LN Transformers. We will provide detailed results that gradient and loss spikes still occur in Pre-LN training curves. Thank you for your insightful guidance on improving the experimental quality of the paper.
• Experimental Settings: We pre-train the 400M-parameter Transformers on 30B tokens each under the controlled same training seed. We measure the training loss and averaged benchmark score for these experiments under the same evaluation settings used in Table 2 of the paper. Other configurations follow those outlined in Section 4.1.
• Clarification on the Weight Initialization: We acknowledge that we did not provide sufficient detail about initialization methods in the original manuscript. In the experiments discussed in the paper, we initialized the weights using a zero-mean Gaussian distribution with a standard deviation of 0.02. We will clarify these details in the revised manuscript.

[3] Glorot, Xavier, and Yoshua Bengio. "Understanding the difficulty of training deep feedforward neural networks." international conference on artificial intelligence and statistics. JMLR Workshop and Conference Proceedings, 2010.

Concluding Remarks

Once again, we sincerely appreciate your insightful feedback. Your questions on extending Peri-LN’s theoretical analysis to other components and applying it to tasks beyond language have highlighted valuable directions for our future work. We are committed to further investigating these avenues—particularly how Peri-LN might generalize to vision or multimodal settings—and to incorporating additional details on initialization strategies to ensure that our results remain transparent and consistent. We hope these efforts will address your questions and more clearly convey the findings of our work. We will make sure to incorporate your valuable suggestions into the revised manuscript. If you have any further questions or topics you would like to discuss, please feel free to let us know.

1) Considering $W^{(2)}$ in the MLP

Following Theorem 1 in [1], we analyze the last layer, as the gradient norm at the final layer is empirically known to be the most unstable (see Figure 1 in the paper). We choose $W^{(2)}$ (the final linear projection in the MLP) as a representative example because it most directly feeds into the residual connection $x+a$, making it clearer to illustrate how gradient norms can explode or vanish.

[1] Xiong et al. "On layer normalization in the transformer architecture." ICML2020.

2) Post-LN and Vanishing Gradients

In Proposition 3.1, the bound $\frac{\||h\||}{\||x+a\||}$ involves not only $\||h\||$ but also $\||x+a\||$. Since $h$, $x$, and $a$ are interrelated, what ultimately matters is their relative magnitudes. The theory indicates that the Post-LN structure introduces $\||x+a\||$ into the bound, whereas Peri-LN introduces only $\||a\||$. Experimentally, we confirm that Post-LN exhibits vanishing gradients (see Figure 6(b, d) and [2]). As shown in Appendix D, the presence of a normalization layer significantly influences the gradient scale not only in Peri-LN but also in Post-LN—an observation consistent with [1] and [2]. For more discussion, we kindly refer to Response #2 of Reviewer AtH3.

[2] Kedia et al. "Transformers get stable: an end-to-end signal propagation theory for language models." ICML2024.
[3] Fishman et al. "Scaling FP8 training to trillion-token LLMs." ICLR2025.

3) Pre-LN and Exploding Gradients

As discussed in our paper (Figures 1 and 6) and by Sun et al. [4], we observe that Pre-LN architectures can produce large activation values. Fishman et al. [3] and Wortsman et al. [5] note that quadratic computations in both the Attention [5] and MLP [3] modules can yield large activations. In Pre-LN architectures, these outputs are not regulated by normalization at the sub-layer level. Consequently, high-variance spikes in intermediate activations can lead to training instabilities because, in Proposition 3.1, $\||h\||$ appears in the numerator of the gradient bound, posing a risk of gradient explosion. Although Adam adaptively rescales gradients, raw gradient spikes may still destabilize updates or force Adam to make abrupt learning rate adjustments (see Section C.1 in the supplementary, where Post-LN uses a lower LR). Indeed, as our paper shows, Pre-LN experiences gradient spikes and training instability despite Adam’s adaptive rescaling. Our main point is that Peri-LN mitigates these large activations at the sub-layer output, reducing the risk of high-magnitude raw gradients before the optimizer’s moment-based rescaling takes effect.

[4] Sun et al. "Massive activations in large language models." COLM 2024.
[5] Wortsman et al. "Small-scale proxies for large-scale transformer training instabilities." ICLR 2024.

4) Weight Initialization & Decay

• Weight Initialization: Due to limited space, we kindly refer you to Response #3 of Reviewer 1xqQ.
• Weight Decay: We conduct additional studies for various weight decay condition for both Pre-LN and Peri-LN architectures. As shown in the table, Peri-LN continues to offer better performance than Pre-LN under the same settings. We will include this detailed ablation studies in the revised manuscript to solidify that our findings (especially large hidden state variance) still hold under varied weight decay and initialization methods. Experimental settings are in Response #3 of Reviewer 1xqQ. |400M||Decay=0|0.0033|0.033|0.33| |-|-|-|-|-|-| |Loss|Pre-LN|3.03|3.03|3.03|3| ||Peri-LN|2.94|2.94|2.93|2.90| |Avg.|Pre-LN|49.26|49.18|49.01|49.51| ||Peri-LN|51.41|51.14|50.68|52.13|

5) Theory references SGD, but experiments use Adam

The key difference between Adam and SGD lies not in the gradients themselves but in how the learning rates are adjusted afterwards—Adam employs adaptive learning rates, while SGD uses a fixed one. Our theoretical analysis focuses on the structural characteristics of the raw gradients rather than assuming any specific optimizer behavior. Therefore, theoretical analyses on the gradients themselves remain valid regardless of whether SGD or Adam is used in the experiments.

6) Over-claimed novelty & Provide the background of LN

Since prior work has analyzed gradient behaviors of Post-&Pre-LN, our central contribution is to consolidate and extend these insights to a third placement—Peri-LN—that recent models (e.g., Gemma2 & 3, Olmo2) employ with limited understanding. We will moderate the overall expressions and provide a concise background on normalization layers.

Concluding Remarks

Your feedback has been instrumental in this process, and we sincerely extend our gratitude for your invaluable insights. Should you have any inquiries or require clarifications about our rebuttal, please don't hesitate to reach out. We are eager to address any concerns and elucidate potential ambiguities in greater depth.

1) Confirming That Pre + Post LayerNorm is Worse Than Peri-LN

“It seems that pre + post layernorm should be worse than periLN, perhaps that can be confirmed experimentally?”

In response to the reviewer's comment, we additionally conduct further experiments on LN placements to compare different combinations (referred to as A, B, and C positions in Figure 2 of the paper). We add configurations where LN is placed at both A + C (akin to combining Pre- and Post-LN), as well as only at B, to compare them with Peri-LN at final training loss under the controlled same training seed. We pre-train the 400M-parameter Transformers on 30B tokens each, using the same training configurations described in the paper. As aligned with Xiong et al. [1], our new results confirm that placing LN exclusively at C leads to training instability or suboptimal performance. In particular, the A + C configuration inherits characteristics of Post-LN (large gradient norm shifts), forcing the use of smaller learning rates and still resulting in lower overall performance than Peri-LN architecture. We will include additional learning rate sweep results and detailed training loss curves for the additional A+C and B experiments in the revised manuscript to more comprehensively illustrate these differences.

[1] Xiong, Ruibin, et al. "On layer normalization in the transformer architecture." ICML 2020.

2) Reconciling Gradients Blowup vs. Vanishing Gradients in Post-LN

“My takeaway from Xiong et al. [1] was that Postlayernorm will lead to gradient norm blowup, but from this paper it seems like the issue is actually gradient norm vanishing. How should I reconcile this?”

Thank you for raising this point. Our layer-wise observations in Post-LN (Figure 6 in the paper) indeed show signs of gradient vanishing through layers, yet we also observe strong gradient spikes (i.e., blowups in the total gradient summation) at various stages of training. This aligns with Xiong et al. [1], where the large shift in Post-LN gradients causes instabilities that lead to sudden spikes. In essence, both Xiong et al. and our Proposition 3.1 suggest that the gradient scale in Post-LN can swing dramatically. When examining training iteratively (step by step), we observe occasional gradient spikes (blowups). However, across the broader span of training, Proposition 3.1 and [2] show that Post-LN gradients ultimately exhibit a vanishing tendency overall—consistent with Figure 6. We appreciate the chance to clarify further. In the revised manuscript’s supplementary material, we will include additional plots demonstrating both the micro-scale spikes in both gradient and loss over the course of training.

[2] Kedia, Akhil, et al. "Transformers get stable: an end-to-end signal propagation theory for language models." ICML 2024.

3) QK-Norm Discussion

We agree that QK-Norm plays an increasingly important role in modern Transformers. As you suggest, we will move or more prominently feature the QK-Norm discussion from the supplementary section into the main paper.

4) Clarifying Olmo2 vs. Peri-LN

Yes, as noted in Appendix G, the Olmo2 architecture slightly differs from Peri-LN. Peri-LN uses both a Pre-LN and an Output-LN, whereas Olmo2 relies on QK-Norm plus the output LN (No Pre-LN). From our experiments, applying only an output LN (or only a QK-Norm) proved insufficient to stabilize training under challenging hyperparameter settings, which is consistent with remarks in the Olmo2 paper [3].

[3] OLMo, Team, et al. "2 OLMo 2 Furious." arxiv 2024.

5) Is QK-Norm Unnecessary for Peri-LN?

While Peri-LN alone provides robust training dynamics, QK-Norm can still enhance performance. In response to reviewers comment, we conducted additional experiments that confirm combining Peri-LN with QK-Norm yields slight improvements in training loss. We pre-train the 1.5B-parameter Transformers on 30B tokens each, using the same training configurations described in the paper. This observation is consistent with prior work [4] indicating that QK-Norm can synergize with various LN placements. We will include these new results in the revised manuscript with detail, along with references to recent works like gemma3 [5], which successfully integrate Peri-LN and QK-Norm.

[4] Wortsman, Mitchell, et al. "Small-scale proxies for large-scale transformer training instabilities." ICLR 2024.
[5] Team, Gemma, et al. "Gemma 3 Technical Report." arXiv 2025.

Concluding Remarks

We believe our additional experiments and clarifications—regarding LN placements, QK-Norm, gradient behavior—will strengthen the paper significantly. Your insightful guidance has been instrumental in refining our analysis. We will incorporate your valuable comments into the revised manuscript.