Mix-LN: Unleashing the Power of Deeper Layers by Combining Pre-LN and Post-LN
Authors: Pengxiang Li
Lu Yin
Shiwei Liu
Venue: ICLR 2025 Poster (2025)
Keywords: large language model
layer normalization
mix-Layer Normalization
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce Mix-LN that combines the benefist of Pre-LN and Post-LN to encourage a more balanced training across layers, thereby improving the overall quality of the model.

Abstract: Large Language Models (LLMs) have achieved remarkable success, yet recent findings reveal that their deeper layers often contribute minimally and can be pruned without affecting overall performance.  While some view this as an opportunity for model compression, we identify it as a training shortfall rooted in the widespread use of Pre-Layer Normalization (Pre-LN). We demonstrate that Pre-LN, commonly employed in models like GPT and LLaMA, leads to diminished gradient norms in its deeper layers, reducing their effectiveness. In contrast, Post-Layer Normalization (Post-LN) preserves larger gradient norms in deeper layers but suffers from vanishing gradients in earlier layers. To address this, we introduce Mix-LN, a novel normalization technique that combines the strengths of Pre-LN and Post-LN within the same model. Mix-LN applies Post-LN to the earlier layers and Pre-LN to the deeper layers, ensuring more uniform gradient norms across layers. This allows all parts of the network—both shallow and deep layers—to contribute effectively to training. Extensive experiments with various model sizes demonstrate that Mix-LN consistently outperforms both Pre-LN and Post-LN, promoting more balanced, healthier gradient norms throughout the network, and enhancing the overall quality of LLM pre-training. Furthermore, we demonstrate that models pre-trained with Mix-LN learn better compared to those using Pre-LN or Post-LN during supervised fine-tuning, highlighting the critical importance of high-quality deep layers. By effectively addressing the inefficiencies of deep layers in current LLMs, Mix-LN unlocks their potential, enhancing model capacity without increasing model size. Our code is available at https://github.com/pixeli99/MixLN.

Published as a conference paper at ICLR 2025
MIX-LN: UNLEASHING THE POWER OF DEEP LAYERS BY COMBINING PRE-LN AND POST-LN
Pengxiang Li
Lu Yin
Dalian University of Technology Lipengxiang/email dlut.
University of Surrey
1. y.nyinnuurree ac .k
Shiwei Liu
University of Oxford
shiwei wiwumaths....... ox.
ABSTRACT
Large Language Models (LLMs) have achieved remarkable success, yet recent findings reveal that their deeper layers often contribute minimally and can be pruned without affecting overall performance. While some view this as an op- portunity for model compression, we identify it as training shortfall rooted the widespread use of Pre-Layer Normalization (Pre-LN), We demonstrate that Pre-LN, commonly employed in models like GPT and LLaMA, leads to dimin- ished gradient norms in its deeper layers, reducing their effectiveness. In contrast, Post-Layer Normalization (Post-LNN) preserves larger gradient norms in deeper layers but suffers from vanishing gradients in earlier layers. To address this, we introduce Mix-IN, a novel normalization technique that combines the strengths of Pre-LN and Post-LN within the same model. Mix-LN applies Post-LN the earlier layers and Pre-LN to the deeper layers, ensuring more uniform gradi- across layers. This allows all parts of the network-bott shallow and deep layers-to contribute effectively to training. Extensive experiments with various model sizes from 70M 7B demonstrate that Mix-IN consistently outperforms both Pre-LN and Post-LN, promoting more balanced, healthier gradient norms throughout the network, and enhancing the overall quality of LLM pre-training. Furthermore, we demonstrate that models pre-trate pre-trained with Mix-LN learn better compared to those using Pre-LN or Post-LN during supervised fine-tuning (SFT) and reinforcement learning from human feedback (RLHF), highlighting the crit- ical importance of high-quality deep layers. By effectively addressing the inef- ficiencies of deep layers in current LLMs, Mix-LN LN unlocks their potential, en- hancing model capacity without increasing model size. Our code is available at https://gitubb.com//xx//999999
INTRODUCTION
Large Language Models (LLMs) have ushered a new era of artificial intelligence by demonstrating unprecedented capabilities in understanding and generating human like text (Brown) 2020 Achiam ct al. 2023) Touvron et al. 20233 Dubey et al.) 2024). Trained on vast datasets that span multiple languages and topics, LLMs driving advancements across industries and academia, enhancing human-computer interactions, and fostering innovation in previously unimaginable ways.
Recent studies reveal critical observation regarding the effectiveness of deeper layers in LLMs, particularly those beyond the middle layers. has been shown that these deeper layers can often be pruned significantly (Yin al. (023), even removed entirely (Gromov al., 2024 Men et al. 2024), without notably affecting the model's overall capabilities. Moreover, [Li et al. (024) demonstrated that deeper layers contribute minimally to performance during fine-tuning, further
"Equal contribution.
1 "oorrsponniig author.

---

Published as a conference paper at ICLR 2025
Layer Norm
Layer Norm
x L
Attantion
Pre-LN
FFN
Layer Norm
Layer Norm
Layer Norm
Layer Norm
Attention
+
XaL Post-LN
Attention
Layer Norm
Layer Norm
+
Attantion
(a) Post-LN
(b) Pre-LN
(c) Mix-LN
Figure I: (a) Post-LN layer; (b) Pre-LN layer; (c) Mix-LN layer.
questioning their importance. Unfortunately, this finding has been largely overlooked by the research community, where many primarily an opportunity for model compression (Siddiqui et al. 2024 Zhong et al. 2024 Sreenivas et al. 2024), rather than recognizing it as a potential shortfall in the training process.
In this paper, we seek to challenge the prevailing notion that deeper layers in LLMs are of lesser significance. The training of LLMSs is extraordinarily resource-intensive, often requiring thousands of GPUs or TPUs and several months of computation on vast datasets. For example, the training of GPT-3 reportedly incurred millions of dollars computational costs. The underutilization of deeper layers leads to inefficiencies, squandering resources that could otherwise be leveraged to enhance model performance. Ideally, all layers in model should be well-trained, with sufficient diversity in features from layer to layer, to maximize the utility of the network's 's parameters (Yang et al., 2023). This makes crucial investigate the root causes of this underutilization and to develop strategies that fully capitalize the potential of deeper layers, ensuring that the overall architecture optimized for both performance and efficiency.
We hypothesize that the inefficiency of deeper layers in LLMs primarily stems from the choice of Layer Normalization. Specifically, Pre-Layer Normalization (Pre-LN) (Dai 2019; Baevski & Auli, 2018) tends to produce smaller gradients in deeper layers, thereby diminishing their effectiveness, while Post-Layer Normalization (Post-LN) (Ba 2016) results larger gradients in deeper layers but leads to gradient vanishing in earlier ones. Most state-of-the-ar LLMs, like GPT, LLaMA, and Mistral, employ Pre-LN, which contributes to the widespread assumption that deeper layers are inherently less effective.
To validate this conjecture, we conduct experiments with the following two categories of LLMss and compare the effectiveness of layers across different depths in Pre-LN models and Post-LN models.
We select LLaMa2-7B (Touvron al. 2023) a representative Pre-LN model and BERT-large (Devlin 2018) as Post-LN model to evaluate the quality of their layers. Our findings confirm that the deeper layers LLaMa2-7B exhibit high similarity, with their removal leading to minimal impact compared to the early layers. In stark contrast, BERT shows higher similarity among its first half, which contributes less the model's output.
Open-weight large-scale LLMs:
To control for irrelevant confounding variables, we conduct sec- ond set of experiments by training small-scale LLMs ourselves, ensuring that the only difference between the models is the choice of layer normalization. Consistent trends are observed these experiments, reinforcing our earlier observations.
In-house small-scale LLMs:
Building these insights, propose novel normalization technique, dubbed Mix which synergizes Pre-LN and Post-LN to achieve more balanced and healthier gradient norms across the network. Mix-IN applies Post-LN to the earlier layers and Pre-LN to the deeper layers. The rationale behind this is that Post-LN enhances gradient flow in the deeper layers, while Pre-LN stabilizes gradients in the earlier layers. By employing Post-LN in the initial layers and Pre-LN
2

---

Published as a conference paper at ICLR 2025
in the later layers, Mix-LN promotes healthier gradient norms in the middle and deeper layers, fostering more balanced training across the entire network and ultimately improving the model's overall performance.
Our extensive experiments, spanning models from 70M to 7B parameters, demonstrate that Mix- LN consistently outperforms Pre-LN, Post-LN, and their variants. Mi M-N not only avoids the train- ing instability associated with Post-LN but also significantly improves the quality of deeper layers compared Pre-LN, leading better pre-training performance. Additionally, models pre-trained with Mix-LN -LN demonstrate superior learning during Supervised TunenT (SFT) and Reinforce- ment Learning from Human Feedback (RLHF) compared those trained with Pre-LN Post-LN, underscoring the importance of high-quality deep layers in LLMs.
2 HYPOTHESIS EVALUATION
In this section, we will evaluate our hypothesis that the inefficiency of deeper layers in LLMs stems from the choice of Pre-LN. The evaluation details are described as follows.
2.1 LAYER NORMALIZATION AND ITS GRADIENT
F (a) and (b) illustrate Post-LN and Pre-LN Transformer architectures, respectively. Formally, let us define x as the input, F(x) as either a FFN layer or multi-head attention layer, and LN(-) as the layer normalization. Post-LN applies LN after the residual addition:
(ost-I.N(ar) = LN + F(x)).
In contrast, Pre-LN applies LN(.) before the residual addition: (re-LN(x:) = + F(IN(x)).
(2)
We can calculate the derivatives of Equations and follows:
@Post-LN(())
aLN(( + F(x))
I aF(x))
(3)
dx
0l + F(x))
dx
OPr-LLN(x) dr Both the above equations involve an important component, i.g., the Jacobian matrix of layer nor- malization, J (z²) @LN(=') where x' is the input of LN(.). Following the proof of Xiong et al. (2020): Takase et al, (2023) with the assumption that x' follow a normal distribution with mean of
(F(LN(z)) OLN(x)
(4)
OLN(x) dr
2z"
0, we can have:
OLN(x')
Vd (I (I- x'a'T -=='' Ox' = ||00/12 l2
(1- || /
(5)
where oz is the standard deviations of x' and d is the hidden dimention. Hence, x'a'T
aLN(x')
Vd
22
Ox T2e vd
I
2 o.
ofd
(6)
where = (x' J/2 is the standard normal distribution obtained after layer normalization. Since d > 1 in LLMs, we can finally obtain:
aLN(x') Or'
1
In practice, we observe that ae' gradually grows larger than one during training, which means the spectral norm the Jacobian matrix of LN is smaller than 1. According to the derivative Post- LN in Equation (3). this down-scaling factor will accumulate as II over multiple layers L, leading to gradient vanishing in early layers in Post-LN Transformers.
In contrast, the derivative of the residual connection in Pre-LN is decoupled from the term associated with the derivative of LN, as shown Equation (4). This design helps prevent the vanishing gradient problem in early layers. However, since Pre-LN does not constrain the residual connection, the outputs of successive transformer blocks accumulate as the layer depth grows. Consequently, the approaches identity matrix, indicating that the entire Pre- LN operation of Equation 44) ceases to contribute effectively to learning. This explains why deeper layers in Pre-LN tend to contribute less to the model's overall learning compared earlier layers.
derivative of Pre-LN in Equation
3

---

Published as a conference paper at ICLR 2025
2.2 EMPIRICAL EVALUATION SETUP
Methods: Our evaluation methodology involves comparative analysis of two models one utiliz- ing Pre-LN and the other employing Post-LN. By empirically assessing the effectiveness layers across different depths in each model, expect to see that Pre-LN models will exhibit decrease in the effectiveness of deeper layers, whereas Post-IN models will show sustained or even improved quality in deeper layers.
LLM Models: To rigorously evaluate our hypothesis, we conduct experiments on two distinct cate- gories of LLMs: (i) Open-weight large-scale LLMs and (ii) In-house small-scale LLMs. In the open- weight category, we select LLaMa2-7B (Touvron et al. 2023) as representative Pre-LN model and BERT-large (Devlin 2018 as a Post-LN model. However, these open-weigh models differ not only in normalization but also in other factors such as training data, activation functions, and context length, complicating our ability to isolate the impact of normalization alone. To control for these confounding variables, we conduct a second set of experiments by training small-scale LLMs from scratch ourselves. The goal is to ensure that the only difference between the models is the choice of layer normalization. Specifically, we train LLaMa- 1 130M models on the C4 dataset with either Pre-LN or Post-LN, using R.MSNorm (Zhagg & Sennrich 2019) and SwiGLU activations (Shazeer, 2020 following Lialin al. (203bb) Zhao al. (2024) Please refer to Appendix A for more traiing configuration details.
Evaluation Metrics:
A critical challenge in validating our hypothesis lies in defining and selecting robust metrics that capture the effectiveness of individual layers. In this study, we employ two metrics: (i) Angular Distance and (ii) Performance Drop, which provide meaningful evaluation of the role and contribution of each layer. In addition, we report the gradient norm of each layer to demonstrate the effect of different layer normalization on the gradient flow.
(i) Angular Distance d(af)x(t)) is used in Gromov et al. 2024 to measure the angular distance between the input layer and the input to layer e + n on a neutral pre-training dataset. Formally, assuming of is the input to the layer f, and oftt is the input to the layer é + n, the angular distance between layers l and and its subsequent th layer, i.e., E n., on a single token T is given by
d(ø,, x ) = arccos
xr
IT
(8)
where ||| denotes the - and the factor of 1/TT scales d(a", (( =f) to the range [0. 1]. To eliminate the effect of randomness, the angular distance reported in this paper is averaged over 256K tokens from the C4 dataset. A smaller value of d(x', (x(*)) indicates a shorter distance, meaning that the two vectors are more similar. Layers whose representations are extremely similar to their neighboring layers mean that they can be easily removed, and therefore their weights are less effective. Ideally, representation should change substantially from layer to layer order to most effectively make use of the parameters of network (Yang al. 2023 Gromov et al. 2024
(ii) Performance Drop A refers to the difference in the performance of an LLM before and after pruning the layer f. It quantifies the performance degradation caused by the removal of that layer. Formally, it can be defined as follows:
AP
p(() prined Poriginal
where Poriginal is the performance of the model without any pruning, p(t is the performance of the model after pruning layer f. A smaller value of A(f) indicates that removing the layer causes minimal change to the model's output, suggesting the layer is less important. Specifically, for I.LaMA2-7B, we choose the commonly used MMLU (Hendrycks et al. (2200) as the evaluation task; for BERT-large, we opt for SQuAD v1.1 (Rajpurkar] (016) as the evaluation task. Given the limited capacity of our in-house trained LLMs, we choose ARC-e (Clark et al. 2018 after supervised fine-tune ten-iuning instead of MMLU, for performance drop.
2.3 EVALUATION RESULTS
2 OPEN-WEIGHT LARGE-SCALE LLMS
Figure 2 (a. c) illustrate the metric values for BERT-Large. Both metrics indicate that, as a Post-LN
model, the early layers of BERT-Large are less effective compared to the deeper layers. As shown in
4

---

Published as a conference paper at ICLR 2025
(a) BERT-large Angular Distance
(b) LLaMa2-78 Angular Distance
Subsequent Layer
Subsequent 8t Layer
20
0.4
15
0.3
10
0.2
0.2
0.1
0.1
10
10
20
30
Layer
Layer Index
(c) BERT-large Performance Drop
(d) LLaMa2-78 MB Performance Drop
0.2
6P
-02
MMLU 6PU
5
-0.4
-00
10.6
-15
0.8-
Panghaal
-20
100
15
20
20 25
30
Layer Index
Layer Index /
Figure 2: Results of open-weight large-scale LLMs. Angular Distance (a, b): Each column repre- sents the angular distance from the initial layer (x-axis) and its subsequent ntth layer (y-axis). The distance is scaled to the range [0, 1]. where yellow indicates smaller distances and purple indicates larger distances. Performance Drop (c, d): (c): SQuAD v1.1 performance drop of removing each layer from BERT-large; (d): MMLU accuracy drop of removing each layer from LLaMa2-7B.
(a) Post LN Angular Distance
Pre Angular Distance
(e) Layer Gradient Norm
0.4
Gredient Norm
0.3
Subrequest nT
Subsequeve
0.2
0.2
0.2
0.1
0.0
32
Layer Indax
Layer Index
Layer
Post-LN Performance Drop
Performance Drop
Pogna
ABC SPPA
Layer Index
Layer Index
Figure 3: Results of in-house small-scale LLaMa-130M 130M. Angular Distance (a, b): Each column represents the angular distance from the initial layer (x-axis) and its subsequent th layer (y- axis). The distance is scaled to the range [0, 11, where yellow indicates smaller distances and purple indicates larger distances. Performance Drop (c, d): ARC-e performance drop of removing each single layer from ALaMa-130M Gradient Norm (e): Gradient norm each layer in 100M.
Figure a the first half of BERT-Large tends to have a smaller angular distance (more yellow) from neighboring layers than the second half. In particular, layers 3, 4, 9, 10, and 11 show a very high similarity to their subsequent layers. In Figure 2/c, the performance drop SQuAD of removing an early layer is significantly smaller than the impact removing a deeper layer. Intriguingly, removing layers 2 and 3 can improve the performance slightly.
In contrast, Figure (b, d) display the metric values for LLaMa2-7B. As a Pre-LN model, the angular distance between neighboring layers decreases gradually (from purple yellow) from the top layers to the 30th layer as illustrated in Figure b.. Notably, the deeper layers (20th to 30th) exhibit extremely small angular distances their adjacent layers. This trend is consistent with the
5

---

Published as a conference paper at ICLR 2025
MMLU performance in Figure d where the removal of deeper layers results in almost negligible accuracy loss while removing early layers causes a substantial drop in accuracy.
In summary, observe that the least effeectve layers in LLaMa2-7B are located in the deeper layers, whereas the early layers in BERT-Large are less effective than the deeper layers. The results from the category of open-weight weight large-scale LLMs strongly support hypothesis, demonstrating a clear aligoment with our expectations.
2.3.2 IN-HOUSE SMALL-SCALE LLMS
Figure illustrates all metric values for two LLaMa- 130M models, where the only difference be- tween them is the choice of layer normalization.
(a. b) show the Angular Distance for Post-LN and Pre-LN, respectively. Without the effects of other compounding factors, this comparison provides clearer distinction between Post- LN and Pre-LN compared to open-weight large-scale LLMs. In Post-LN models, the most similar layers are concentrated in the early stages, with the first three layers showing particularly low dis- tance. As the depth increases, the layers become increasingly distinctive. In contrast, the Pre-LN LLaMa- 1AMMM-130M 130M exhibits gradual decrease in angular distance depth increases, leading to highly similar deep layers. Figures [d, e) further confirm this with the Performance Drop metric: re- moving early layers (e.g., 0-7 layers) in Post-LN results in minimal performance loss, while deeper layers (especially layers 9-1 1) are critical to preserving the original performance. However, Pre-LN LLaMa- 130M exhibits the opposite trend, where removing most layers after the first layer causes negligible performance loss, indicating that they contribute little to the model's 's output.
Figures (a,
Figure (c) shows the gradient norm of each layer for Post-LN and Pre-LN the beginning of the training. The results perfectly align with our expectations: Post-LN leads to larger gradients in deeper layers but suffers from severe gradient vanishing in early layers, whereas Pre-LN maintains healthy gradient flow in early layers but diminishes in later layers.
With the consistent findings from both open-weight LLMs and our in-house L.Mss, we can conclude that the widespread use of Pre-LN in LLMs is the root cause of the ineffectiveness of deep layers.
3 MIX-LAYER NORMALIZATION (MIX-LN)
Having validated our hypothesis that the use of Pre-LN is the root cause of the ineffectiveness of deeper layers, we propose Mix-Layer Normalization (Mix-LN), novel normalization strategy designed to enhance the effectiveness of both middle and deeper layers in LLMs.
The key idea behind Mix-IN leverage the strengths of both Pre-LN and Post-LN. Post-LN has been shown to improve the effectiveness of deeper layers, while Pre-LN is more effective for earlier layers. Therefore, we propose apply Post-LN to the initial layers and Pre-LN to the later layers, ensuring that the middle and deeper layers benefit from the advantages of both methods.
Formally, for an LLM with L layers, we apply Post-LN the first [aL] layers and Pre-LN to the remaining [(1 L. layers, where E [0, 1] is a hyperparameter controlling the transition point between the two normalization strategies. The functions [-] and [-] denote the floor and ceiling operations, respectively. Although the final layers may still experience smaller gradients due to the use of Pre-LN, the negative impact is substantially mitigated because the number of layers employ- ing Pre-LN is reduced from L. to [[1 L]. This reduction improves gradient flow in the deeper layers, enhancing their effectiveness. Additionally, we expect that Mix-IN can alleviate training instability issues caused by Post-LN (Nguyen & Salazar 2019 Wang et al. 2024). as reducing the number of layers using Post-LN leads to smaller accumulation of gradient attenuation, according to the analysis in Section 2.1
MAIN EXPERIMENTAL RESULTS
4.1 LLM PRE-TRAINING
In this section, we verify the effectiveness of Mix-LN LN by comparing it with various common nor- malization techniques, including Post-LN (Nguyen & Salazar 2019), DeepNorm (Wang et al.)

---

Published as a conference paper at ICLR 2025
2024). and Pre-LN Dai 2019) Following Lialin al. (2033a) Zhao et al. 2024). conduct experiments using the LLaMA based architecture with various sizes from 71 71M to 1B parameters, in- corporating RMSNorm (Shazeer 2020) and SwiGLU activations (Zhang & Sennrich 2019). Models are trained with Adam (Kingma 2014 using different learning rates based on model size: specifi- cally, we learning rate of Te-3 3 for models with 250M parameters and below, and learning rate of 5e-4 for the IB parameter model. All models of the same size are trained with identical config- urations except for the normalization. To determine the optimal value for the hyperparameter a in Mix-LN, we performed small hyperparameter sweep using ALaMA-250M, as shown in Table 5 We found that 0.25 provided the best performance, and therefore, we applied this value across all model sizes.
Table I: Perplexity (L) comparison of various normalization methods across various LLaMA sizes.
LLaMA-7IM
LLaMA-130M
LLaMA-250M
LLaMA-IB
Training Tokens
1.1B
2.2B
3.98
5B
Post-LN
35.18
26.95
1409.09
1411,54
DeepNorm
34.87
27.17
22.77
1410.94
Pre-LN
34.77
26.78
21.92
18.65
Mix-LN
33.12
26.07
21.39
18.18
Results are shown in Table I Post-LN generally yields the worst performance and even diverges with larger models, aligning with previous studies that indicate Post-IN suffers from training insta- bility in Transformers (Xiong et al., 2020 Takase et al. 2022). DeepNorm, as a modified version of Post-LN, achieves comparable performance to Pre-LN with smaller model sizes; however, also experiences divergence during training with IB parameter models. This observation confirms vere training instability of Post-LN, where gradients in early layers vanish, preventing proper model convergence. In contrast, Mix-LN consistently achieves the lowest perplexity across various model sizes. Mix-LN achieves a notable gain by 1.65 and 0.53 perplexity with LLaMA-7 and LLaMA- 250M, respectively, compared the widespread Pre-LN.
The above results clearly show that Mi X-IN not only overcomes the instability of Post-LN but also enhances the model quality by combining the benefits of Pre-LN and Post-LN, making it ideal choice for large-scale LLMs.
LLaMA-7B
Pre-LN
Ma-1N
Evaluation Perplextyy
7000
10300
13000
33000
Updated Steps
Figure 4: Training curve (eval perplexity) of Mix-LN and Pre-LN with LLaMa-7B
4.2 SCALING UP TO 7B MODEL
is essential evaluate if the benefits of Mix-LN also hold for larger model sizes. this end, we conducted experiments using the I.LaMa-7B architecture, which features an embedding size of 4096 and 32 total layers. All training configurations were kept identical, except for the layer normalization method. Due to computational constraints, managed to complete only 13,000 steps of training. The training curve is presented in Figure 4 We can see that it possible to scale the gain of Mix-ix LN to larger-scale training.
Nevertheless, we found that Mi Mi becomes more sensitive to a when scaling up to 7B training. In general, smaller values of a and longer warm-up period are required to stabilize the training of Mix-LN at a larger scale. For instance, we use a relatively small a 6.25% here to ensure stable training of L.aMa-BB.
7

---

Published as a conference paper at ICLR 2025
4.3
SUPERVISED FINE-TUNING
Table 2: Fine-tuning performance (1) of LLaMa with various normalizations.
Method
MMLU
BoolQ
ARC-e
PIQA
Hellaswag
OBQA Winogrande
Avg.
LLaMA-250N
Post-LN
22.95
37.83
26.94
52.72
26.17
11.60
49.56
32.54
DeepNorm
23.60
37.86
36.62
61.10
25.69
15.00
49.57
35.63
Pre-LN
24.93
38.35
40.15
63.55
26.34
16.20
49.01
36.93
Mix-LN
26.53
56.12
41.68
66.34
30.16
18.00
50.56
41.34
LLaMA-IB
Post-LN
22.95
37.82
25.08
49.51
25.04
13.80
49.57
31.96
DeepNorm
23.35
37.83
27.06
52.94
26.19
11.80
49.49
32.67
Pre-LN 26.54
62.20
45.70
67.79
30.96
17.40
50.51
Mix-LN 27.99
61.93
48.11
68.50
31.35
18.80
55.93
44.66
We believe that the superior middle and deeper layers produced by Mi X-IN are better equipped to learn during supervised fine-tuning. This advantage stems from the fact that these layers capture more diverse and rich features compared to those trained with Pre-LN. In complex downstream tasks, having access broad spectrum of features allows the model to make leading to improved generalization.
nuanced predictions,
To verify conjecture, follow Li al. (2024) and fine-tune the models obtained in Section 4.1 lon Commonsensel 70K (Hu al., 2023), evaluating them on eight downstream tasks. As shown Table 2
Mix-IN consistently outperforms other normalization techniques across all evaluated datasets. For the 1.L.MA-200M model, Mix-LN achieves significant average gain of 4.26% and a 17.31% improvement on BoolQ compared to Pre-LN. Similar trends are observed with the larger LLaMA-IB model. Even though Mix LN only slightly reduces perplexity by 0.25 compared to Pre-LN, delivers substantial performance gains in supervised fine-tuning.
4.4 REINFORCEMENT LEARNING FROM HUMAN FEEDBACK
Table RLHF comparison of final reward (1) of Pre-LN and Mix-LN with LLaMA-IB
Method
Model
Final Reward
Pre-LN
LLaMA-IB 0.75
Mix-LN LLaMA-IB
1.32
Consistently, the benefits of Mix-LN can be seamlessly transferred to RLHF. Following Adam-min (Zhang et al., 2024), we implement the RLHF workflow from InstructGPT (Ouyang et al. 2022) and train 1B models obtained Section on the ultrafeedback dataset to optimize the preference reward. Table 3 iillustrates that Mix-LN achieves notable reward gain (higher is better) over Pre- LN, i.e., 1.32 vs. 0.75.
4.5 EVALUATION WITH VISION TRANSFORMERS
Table Accuracy (t) comparison of Pre-LN and Mix-LN on ViT models.
Model
ViT-Tiny
ViT-Small
Pre-LN
67.30 67.34
75.99
Mix-LN
76.40
To evaluate Mix-LN on non-language models, we replace Pre-LN in VIT models with Mix-LN with 0.25. We train the updated model for 120 epochs ImageNet-1K KK following Liu et al. (2022a), using the ConvNeXt (Liu et al., 2022b) configurations. The results clearly demonstrate that the benefits of Mix-LN also generalize to non-language models. The results in Table 4 demonstrate that the benefits of Mix-LN extend to non-language tasks, with performance gains that are more pronounced in larger models (ViT-Small) compared to smaller ones (ViT-Tiny).
8

---

Published as a conference paper at ICLR 2025
5
ANALYSIS AND MORE EVALUATIONS
Table 5: Perplexity of LLaMA-IB with various Post-LN ratios a.
Pre-LN
Mix-LN
Post-LN
Post-LN ratios a
0
16.7%
25.0%
33.0%
41.7%
50.0%
100%
Perplexity
18.65
18.34
18.18
18.41
18.55
18.86
1434
What is the proper Post-LN ratio a for MiX-LN? Mix-LN has a hyperparameter, a, that controls the ratio of layers applying Post-LN. Specifically, = 0 means Pre-LN is applied to all layers, while a = corresponds to pure Post-LN, To determine the optimal Post-LN ratio, we conduct sweep over the values [O, 16.7%. 25.0%, 33.0%, 41.7%, 50.0%, 100%] using LLaMA-1 AB the C4 dataset. The results are shown in Table S As the normalization transitions from Pre-LN to Mix- LN. the model achieves progressively lower perplexity, reaching its best performance at 0.25. Beyond this point, performance begins to decline, although it still surpasses that pure Pre-LN until most layers apply Post-LN, where performance degrades significantly. Based these results, we choose a = 0.25 for all model sizes, although we believe there is potential to further improve the performance of Mix by searching for the optimal a for each individual model.
(a) Mix LN Angular Distance
[b) Pre-LN Angular Distance
[c] Post-LN Angular Distance
bubsequentt.o n Layer
Subsequent ni Layer
Subzaquent 6t Layer
433
9
425
0.2
0.15
Index
Index
Index
Figure 5: Angular distance from initial layer é (x-axis) with block size n (y-axis) of LLaMA-130M
Mix-LN promotes representation diversity across layers. As have claimed, hybrid ap- proach promotes a more balanced gradient flow throughout the entire network. To validate this, report the angular distance of I.LaMA-130M for Pre-LN, Post-LN, and Mi X-IN in Figure 5 Given block size n. the layers with the smallest distances are highlighted in the lightest yellow in each row. Notably, Mix-LN consistently exhibits larger distances (darker color) across layers compared to Pre-LN, except for the final two layers. This indicates that Mix LN produces more di- verse representations between layers than Pre-LN. In contrast, the smallest distances in Post-LN are concentrated in the early layers, reinforcing the notion that Post-LN tends to restrict representation diversity deeper layers.
Performance Layers
Eraten Type
ABC
MO-M
30
(a) Layer gradient norm of I.L.MA-50MM with vari- ous normalization techniques.
(b) Performance drop comparison LLaMA-130M across layers for Pre-LN, Post-LN, and Mix-LN.
Figure 6: Comparison of gradient norms and performance drops of layer pruning for different layer normalization.
Mix-LN enhances healthier gradient norms across all layers. We compare the gradient norm of different LN at initialization in Figure 6al It demonstrates that Mi x-IN maintains consistent

---

Published as a conference paper at ICLR 2025
gradient norms across all layers. This balance results in a more uniform distribution of gradient norms across layers, allowing all parts the network both shallow and deep layers to contribute effectively to model training.
Performance drop of layer pruning for Mix-IN. To further evaluate the effectiveness of Mix- LN, we compare the performance drop (AP) across layers with Pre-LN and Post-LN. [bbswwss Mix-LN achieves more significant contribution from deeper layers. Specifically, the deeper layers in Mix-LN models show larger AP compared Pre-LN, indicating that these layers contribute more effectively to the model's overall performance.
Comparing with other Layer Normalization. In addition, conducted comparisons using LLaMA-250M to evaluate Mix-LN against recently proposed normalization methods, including Ad- min (Liu et al. 2020), Sandwich-LN (Ding al. (021)) and Group-LN (Wu & He] 2018 Ma al. 2024). The results indicate that Sandwich-LN and Group-LN slightly outperform Pre-LN, while Ad- min performs worse. However, all of these methods fail to reduce perplexity below 23, falling short of Mix-LN. This result highlights the effectiveness of Mix-LN compared to other recent innovations.
Table 6: Comparison against other normalization methods on LLaMA-250M.
Model
Pre-LN
Admin
Group-IN
Sandwich-LN
Mix-LN
LLaMA-250M
23.39
24.82
23.10
23.26
22.33
CONCLUSION
In this paper, we have addressed the inefficiencies of deep layers in LLMs by identifying the widespread use of Pre-LN as the root cause. Pre-LN leads to diminished gradients deeper layers, reducing their effectiveness. While Post-LN preserves deeper gradients, it suffers from vanishing gradients earlier layers. To resolve this, we introduced Mi LNL,, a hybrid normalization tech- nique that combines the strengths of both Pre-LN and Post-LN. By applying Post-LN early layers and Pre-LN to deeper layers, Mi X-IN achieves balanced gradient norms throughout the network, enabling more effective training. Our experiments show that Mi x-IN consistently outperforms both Pre-LN and Post-LN,, enhancing pre-training and fine-tuning performance without increasing model size. By fully utilizing the potential of deep layers, Mix-IN improves the overall capacity and efficiency of LLMs.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
Jimmy Lei Ba. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. arXiv preprint arXiv:1809, 108553 2018
Ankur Bapna, Mia Xu Chen, Orhan Firat, Yuan Cao, and Yonghui Wu. Training deeper neural machine translation models with transparent attention. arXiv preprint arXiv: (tt::1808.0754 2018.
Tom B Brown. Language models are fhw-shot learners. arXiv preprint arXiv:2005. 2020.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05:: 05457,, 2018.
Zihang Dai. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv:1901.02860 2019.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding.
arXiv preprint arXiv: arXiv: 18i0.08800.0.... 2018
10

---

Published as a conference paper at ICLR 2025
Ming Ding, Zhuoyi Yang, Wenyi Hong, Wendi Zheng, Chang Zhou, Da Yin, Junyang Lin, Xu Zou, Zhou Shao, Hongxia Yang, et al. Cogview: Mastering text-to-image generation via transformers. Advances in neural information processing systems, 198229992989555 2021.
Zi-Yi Dou, Zhaopeng Tu, Xing Wang, Shuming Shi, and Tong Zhang. Exploiting deep representa- tions for neural machine translation. arXiv preprini arXiv: 1810.00181 2018.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The Ilama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A Roberts. The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403.17887 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 2020.
Zhiqiang Hu, Lei Wang, Yihuai Lan, Wanyu Xu, Ee-Peng Lim, Lidong Bing, Xing Xu, Soujanya Poria, and Roy Ka-Wei Lee. Lim-adapters: An adapter family for parameter-efficien fine-tuning of large language models. arXiv preprint 2xXi:2334401993 2023.
Tianjin Huang, Ziquan Zhu, Gaojie Jin, Lu Liu, Zhangyang Wang, and Shiwei Liu. Spam: Spike- aware adam with momentum reset for stable I1m training. arXiv preprint urXiv:2501.06845 2025.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Vedang Lad, Wes Gurnee, and Max Tegmark. The remarkable robustness of Ilms: Stages of infer- ence? arXiv preprint arXiv:2406. 19384, 2024.
Pengxiang Li, Lu Yin, Xiaowei Gao, and Shiwei Liu. Owlore: Outlier-weighed layerwise sampled low-rank projection for memory-efficient Ilm fine-tuning. arXiv preprint arXiv:2405. 18380, 2024.
Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High- rank training through low-rank updates. In The Twelfth International Conference on Learning Representations, 2023a.
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Stack more lay- ers differently: High-rank training through low-rank updates. arXiv preprint arXiv:2307.05695, 2023b.
Liyun Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the diffi- culty of training transformers. arXiv preprint arXiv:2004.082449 08249, 2020.
Shiwei Liu, Tianlong Chen, Xiaohan Chen, Xuxi Chen, Qiao Xiao, Boqian Wu, Tommi Kärkkäinen, Mykola Pechenizkiy, Decebal Mocanu, and Zhangyang Wang. More convnets in the 2020s: Scal- ing up kernels beyond 51x51 using sparsity. arXiv preprint wxXiv:207..03600 2022a.
Zhuang Liu, Hanzi Mao, Chao-Yuse Wu, Christoph Feichtenhofer. Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11976-11986, 2022b.
Xuezhe Ma, Xiaomeng Yang, Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May, Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient Ulm pretraining and inference with unlimited context length. arXiv preprint arXiv:2404.000 088001 2024.
Xin Men, Mingyu Xu. Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. Shortgpt: Layers in large language models are more redundant than you expect. arXiv preprint 24X.:22433.08833 2024.
11

---

Published as a conference paper at ICLR 2025
Toan Q Nguyen and Julian Salazar. Transformers without tears: Improving the normalization of self-attention. arXiv preprint arXiv:1910.05580 2019.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models fol- low instructions with human feedback. Advances in neural information processing systems, 35:
27730-27744, 2022.
Rajpurkar. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1006.05250, 2016.
Teven Le Scao, Thomas Wang, Daniel Hesslow, Lucile Saulnier, Stas Bekman, M Saiful Bari, Stella Biderman, Hady Elsahar, Niklas Muennighoff, Jason Phang, et al. What language model to train if you have one million gpu hours? arXiv preprint arXiv:2210.1544
2022.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 2020.
Shoaib Ahmed Siddiqui, Xin Dong, Greg Heinrich, Thomas Breuel, Jan Kautz, David Krueger, and Pavlo Molchanov. A deeper look at depth pruning of Ilms. arXiv preprint arXiv:2407.16286, 2024.
Sharath Turuvekere Sreenivas, Saurav Muralidharan, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary. Mohammad Shoeybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. Llm pruning and distillation in practice: The minitron approach. arXiv preprint arXiv:2408.11799 2024.
Wenfang Sun, Xinyuan Song, Pengxiang Li, Lu Yin, Yefeng Zheng, and Shiwei Liu. The curse of depth in large language models. arXiv preprint wrXiv:2502.05795 2025.
Sho Takase, Shun Kiyono, Sosuke Kobayashi, and Jun Suzuki. B2t connection: Serving stability and performance in deep transformers. arXiv preprint arXiv:2206.00::: 00330, 2022.
Sho Takase, Shun Kiyono, Sosuke Kobayashi, and Jun Suzuki. Spike no more: Stabilizing the pre-training of large language models. arXiv preprint arXiv:2312.16691 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Hongyu Wang. Shuming Ma, Li Dong, Shaohan Huang, Dongdong Zhang, and Furu Wei. Deep- net: Scaling transformers to 1,000 layers. IEEE Transactions Pattern Analysis and Machine Intelligence, 2024.
Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. arXiv preprint arXiv:190.001780 2019.
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3 3-19, 2018
Ruibin Xiong, Yunchang Yang. Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, pp. 10524-10533, PMLR, 2020.
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Feature learning in infinite-dephh neural networks. arXiv preprint arXiv:2310.02244, 2023.
Lu Yin, You Wu, Zhenyu Zhang, Cheng- Yu Hsieh, Yaqing Wang, Yiling Jia, Mykola Pechenizkiy, Yi Liang, Zhangyang Wang, and Shiwei Liu. Outlier weighed layerwise sparsity (owl): A missing secret sauce for pruning Ilms to high sparsity. arXiv preprint arXiv:2310.0517,. 2023.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Infor- mation Processing Systems, 32, 2019.
12

---

Published as a conference paper at ICLR 2025
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406.16799, 2024.
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient Ilm training by gradient low-rank projection. arXiv preprint arXiv:2403.03557, 2024.
Longguang Zhong, Fanqi Wan, Ruijun Chen, Xiaojun Quan, and Liangzhi Li. Blockpruner: Fine- grained pruning for large language models. arXiv preprint arXiv:2406. 10594, 2024.
13

---

Published as a conference paper at ICLR 2025
A DETAILS OF EXPERIMENTS
A.1 ARCHITECTURE AND HYPERPARAMETERS
We introduce details of the LLaMA architecture and hyperparameters used for pre-training following
Lialin et al. 2023a Zhao et al. 2024). Table S the most hyperparameters of LLaMA models sequence length of 256 for all models, with batch size of 512, and total of 131K tokens per batch. Learning rate warmup is applied to the first 10% of the training steps. We train models using Adam with a cosine annealing for the learning rate schedule, decaying to 10% of the initial learning rate. We use a learning rate of le-3 for models with 250M parameters and below, and learning rate of 5e-4 for the IB parameter model.
across model sizes. We use
Table 7: Hyperparameters of LLaMA models used in this paper.
Params
Hidden
Intermediate
Heads
Layers
Steps
Data amount
LR
Batch Size
71M
512
1368
8
10K
B
1 x 10-3
512
25%
130M
768
2048
12
12
20K
2.2 B
1 10-3
512
25%
250M
1024
2560
16
24
40K
3,9 B
1 10-3
512
25%
1B
2048
5461
32
24
100K
5.0
5 10-4
512
25%
B
4096
11008
32
32
13K
1.7
5 10
512
6.25%
B COMPATIBILITY TO ADVANCED
In this section, also evaluate if Mix LN can integrate well with the advanced techniques pro- posed to stabilize training. Specifically, we evaluate the commonly used Scaled Initialization (Nguyen & Salazar 2019 Scao al. 2022) that initializes W2 and Wo with smaller normal distribution (0, v 2/5dd/ /2N) stabilize training dynamics; and Scaled Embed Takase al, 2023) scales up embeddings to stabilize LayerNorm gradients. We observe that both Pre-IN and Mix-LN work effectively with Scaled Initialization. However, incorporating Scaled Embed on top of this setup leads to degradation in performance.
Table Perplexity of LLaMA- 130M with various normalization methods with Scaled Initialization and Scaled Embed.
Normalization
Scaled Initialization
Perplexity Scaled Embed
Pre-LN
32.18
Mix-LN
29.95
Pre-LN
30.63
Mix-LN
<
29.77
Pre-LN
V
31.28
Mix-LN
V
31.19
C RELATED WORK
C. NORMALIZATION IN LANGUAGE MODELS
Layer Normalization (LN), first proposed by Ba 2016). has become the de facto standard for nor- malizing activations modern language models. It directly estimates normalization statistics from the summed inputs to neurons within a hidden layer, ensuring that the input distribution to each layer remains stable throughout training. the original Transformer (Vaswani 2017). LN was ini- tially applied after the residual connection. configuration known Post-LN. However, subsequent studies (Bavvski & Auli 2018 Dai 2019 Nguyen & Salazar 2019) found that placing LayerNorm before the residual connection (Pre-LN) results in more stable performance, especially in large lan- guage models (Brown) 2020; Touvron et al. 2023 Jiang et al. 2023). Xiong et al. (2020) theoreti- cally demonstrated that Post-LN results in larger gradients the output layer, making the use of
14

---

Published as a conference paper at ICLR 2025
wa essential to avoid instability is necessary. Conversely, Pre-LN scales down gradients with the depth of the model, which ensures more stable gradients during initialization. Our work builds upon Xiong et al. (2020), highlighting that while Pre-LN prevents instability by reducing gradient magnitudes, smaller gradients in deeper layers can diminish the effectiveness of the corresponding weights.
To improve the effectiveness of deeper layers in language models, various LN variants have been proposed. For instance, Wang al. (2019) verified empirically that Post-LN suffers from gradient vanishing in deep Transformers, while Pre-LN facilitates stacking more layers. They consequently introduced dynamic linear combination of layers (DLCL), which connects all previous layers to improve trainability. Similar techniques have been employed in other works (Bapna et al. [2118] Dou et al. 2018). Liu al. (020) revealed that Post-LN has strong dependencies on the residual branch, often leading to instability. To address this, Adaptive Model Initialization (Admin) was in- troduced, which uses additional parameters to control residual dependencies in Post-LN, stabilizing training. DeepNorm (Wang al., 2024) further improved the trainability of deep Transformers by upscaling the residual connection before applying LN, reducing model updates, and enabling deeper architectures. Additionally, Ding al. (202I) proposed Sandwich LayerNorm, normalizing both the input and output of each transformer sub-layer. Takase et al. (2022) identified that Post-LN tends to preserve larger gradient norms in deeper layers, potentially leading to more effective training. To address the issue of gradient vanishing in early layers, they introduced B2T, a method that uses a residual connection bypass all LN except the final one in each layer. Huang al. (2025) unveiled that LN layers suffer from loss spikes and gradient spikes during LLM training. We got inspiration from Takase et al. (2022), addressing the limitations of both Pre-LN and Post-LN by combining them. We study Scaled Initialization and Scaled Embed in Appendix B
C.2 INEFFICACY OF DEEP LAYERS IN LLMS
The Inefficacy of deep layers LLMs serves as a valid indicator for LLM pruning. Yin et al. (2023) demonstrated that the deeper layers of prominent LLMs like LLaMA and Mistral can be pruned more aggressively than earlier layers, without causing a significant drop in performance. Similarly, Gromov et al. (2024) and Men et al. (2024) further explored layer pruning, identifying the deeper layers of LLMs as typically less essential. Lad et al. (2024) observed that in models like Pythia and GPT-2, deeper layers exhibit strong resilience to interventions, such as layer deletion or swapping. Our work shares similarities with Gromov et al. (024) in applying angular distance to assess the effectiveness of layers. However, while they identify the inefficacy of deeper layers, they do not offer an explanation for this phenomenon nor propose a solution to address it. Building on the insights from this paper, Sun al. (2025) introduced the concept of the "Curse of Depth" highlight the pitfalls of deeper layers in LLMs. To address this issue while avoiding the instability associated with Post-LN, they proposed LayerNorm Scaling scale down the output of LN by its depth.
15


Reviews:
[{"review_id": "GRvy4sK1mE", "reviewer": "Reviewer_iHrq", "summary": "In this paper the authors context the employment of either pre-layer normalization (Pre-LN) and post layer normalization. (Post-LN) They bring an argument showing vanishing gradient for earlier or later layers. Building on this intuition, they propose to place post-LN in early layers and pre-LN in later ones, to maximize the gradient. The evaluation is performed on LLMs, either large and small scale ones.", "questions": "- How does the performance drop when using Mix-LN?\n- Is this approach still valid for other tasks (like for example image classification, with a Transformer architecture on ImageNet-1k)?\n- Have the authors attempted to compare with other normalization techniques (like group normalization)?\n- Can the authors provide theoretical boundaries to the vanishing/exploding gradient conditions for equations (3) and (4)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The observation related to gradient norm brought by LN is interesting, clear, and straight to the point.\n- Part of the quantitative evaluation is conducted to large-scale LLMs, making this work actual.\n- All the empirical evaluation is quite realistic and likely correct.\n- Figures in general do an excellent job in providing a display on the distributions.", "weaknesses": "- The relative gains in terms of perplexity/accuracy for larger model slims down. This in a certain sense contradicts the vanishing gradient argument. Besides, the gap compared to pre-LM becomes slower and slower.\n- The success of the approach depends on $\\alpha$: there are cases like Llama-1B for BoolQ in which only pre-LM performs better, indicating that tuning $\\alpha$ properly can be determinant.\n- Metrics are not always complete: I could not find, for example, performance drop when using Mix-LN", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Wti2QjNrtV", "reviewer": "Reviewer_CgHw", "summary": "The paper proposes Mix-LN, a hybrid normalization approach that combines Pre-Layer Normalization (Pre-LN) and Post-Layer Normalization (Post-LN) within large language models (LLMs). The technique leverages Post-LN in shallow layers to mitigate gradient vanishing and Pre-LN in deeper layers to maintain gradient flow, with the aim of maximizing layer effectiveness throughout the network. While Mix-LN demonstrates some empirical gains, especially in mid-sized LLMs, the contribution remains incremental in the broader landscape of normalization techniques for LLMs. The paper’s main contributions are primarily empirical, and it lacks a strong theoretical foundation to substantiate why Mix-LN improves gradient dynamics compared to recent normalization methods.", "questions": "- Could the authors provide more insights or theoretical justifications for why Mix-LN enhances gradient flow differently across shallow and deep layers?\n\n- How does Mix-LN perform when compared to other recent normalization methods, such as Admin or Sandwich LN?\n\n- How sensitive is Mix-LN to the hyperparameter 𝛼, and is this value stable across different tasks and model scales?\n\n- Can the authors clarify Mix-LN’s computational impact during training? Does it introduce any additional training or inference overhead?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Novel Approach to Gradient Flow in LLMs: Mix-LN’s hybrid approach to layer normalization is unique, attempting to combine the advantages of Pre-LN and Post-LN across different model depths.\n\n- Good Experimental Validation: The paper includes extensive experimentation on multiple model sizes and tasks, showing consistent improvement with Mix-LN, particularly in mid-sized models.\n\n- Improved Training Stability: Mix-LN appears to mitigate some of the training instability issues commonly observed with Post-LN in large models, an advantage for practitioners.", "weaknesses": "- Lack of Theoretical Rigor: The paper does not provide a detailed theoretical framework to explain why Mix-LN achieves balanced gradient dynamics across layers, which is crucial for the paper’s validity. In the well-explored field of normalization for LLMs, incremental changes require more substantial theoretical backing to make a notable contribution.\n\n- **Limited Comparison** with State-of-the-Art Normalization Techniques: The paper lacks direct comparison with recent normalization methods, such as Admin or Sandwich LN, which also address deep-layer gradient inefficiencies. Without these comparisons, it’s challenging to assert Mix-LN’s effectiveness over other recent innovations.\n\n- **Diminished Gains on Very Large Models**: The benefits of Mix-LN become less pronounced in very large models, such as LLaMA-1B, where performance gains are smaller. This suggests potential scalability issues for Mix-LN in ultra-large models like 7B, 13B and so on, which is a limitation given the trajectory of LLM research.\n\n- Applicability Limited to LLMs: Mix-LN has only been evaluated on LLMs, and its effectiveness on non-language models or other architectures remains untested. This limits the broader applicability and impact of the approach.\n\n- Hyperparameter Sensitivity Not Thoroughly Explored: The paper does not fully address the sensitivity of the hyperparameter 𝛼 (controlling the transition point between Pre-LN and Post-LN), which could impact Mix-LN’s practical usability across different settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hsvo6iMTuv", "reviewer": "Reviewer_6By3", "summary": "The paper introduced Mix-LN, which applies Pre-LN and Post-LN in different parts of the LLMs. The author explore the diminished gradient norms of Pre-LN in deeper layers and vanishing gradient of Post-LN in earlier layers, both theoretically and experimentally. They then conduct the experiments comparing Post-LN, DeepNorm, Pre-LN and Mix-LN in pre-training and fine-tuning, with anaylsis on the results. The results show Mix-LN with certain ratios improves the overall capacity and efficiency of LLMs.", "questions": "1. Figure 2-a: Could you explain the triangle-like yellow areas (which also appears but no so significiant in Figure 2-b)? It seems to be discontinuous changes in angular distance between near or adjacent layers. However, from my understanding, the transformer blocks in BERT are the same. Therefore what makes layer 5/6 or 14/15 so special?\n2. Figure 2-b: Could you explain the dark-blue line on the right? It seems to be a significiant change of angular distance in last layer.\n3. Section 4: Minor comments. There was no introduction for perlexity?\n4. Figure 4: a&b shows exact opposite result than Figure 2/3. From my understanding, the angular distance will grow bigger (the color will thus go dark) as the distance between two layers increases. However, for example in layer 10 of Figure 4-b, the farthes has smallest angular distance. Please correct me if I am mistaken.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. This paper is very well-organized.\n2. This paper introduce the problem with theortical anaysis and strengthen with experiments, which makes the problem clear and well-motivated.\n3. Comprehensive experiments for presenting the effect of Mix-LN.\n4. The anaysis are insightfull.", "weaknesses": "1. This figures of angular distance are somewhat non-intuitive. Can't understand on the first sight. Block size more like describing the size of the model than distance between two layers under comparison. The first detailed description is actucally in line 402-403.\n2. Typo problem: Figure 3-c (in main body) or Figure 3-e (in caption of Figure 3).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TFKnG6anbv", "reviewer": "Reviewer_9E2P", "summary": "This paper addresses the inefficiency of deeper layers in Large Language Models (LLMs) by introducing Mix-LN, a novel normalization technique that combines Pre-Layer Normalization (Pre-LN) and Post-Layer Normalization (Post-LN). The authors argue that the widespread use of Pre-LN in models like GPT and LLaMA leads to diminished gradient norms in deeper layers, reducing their effectiveness. Mix-LN applies Post-LN to earlier layers and Pre-LN to deeper layers, ensuring more uniform gradient norms across all layers. Through extensive experiments across various model sizes, the authors demonstrate that Mix-LN consistently outperforms both Pre-LN and Post-LN, promoting more balanced gradient norms throughout the network and enhancing the quality of LLM pre-training and fine-tuning. This approach aims to unlock the full potential of deeper layers in LLMs, improving model capacity without increasing model size.", "questions": "1. **Scaling Properties:** How does Mix-LN perform in very large models (>7B parameters), and what theoretical guarantees can be provided for its effectiveness at scale?\n\n2. **Boundary Dynamics:** What are the theoretical considerations and potential instabilities at the transition point between Pre-LN and Post-LN layers?\n\n3. **Asymptotic Performance:** While accelerated convergence is demonstrated, how does Mix-LN's final performance compare to Post-LN when both are trained to complete convergence under same hyperparameter settings? A controlled study with matched configurations training until performance plateaus would help distinguish whether Mix-LN's benefits extend beyond computational efficiency to fundamental improvements in model capability.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This article conducts a good experimental analysis and points out that the gradient norm of Pre-LN and Post-LN shows different trends as the number of layers increases.\n2. Comprehensive experiments on both open-weight and in-house models, with multiple evaluation metrics (Angular Distance, Performance Drop, Gradient Norm).\n3. Well-structured presentation and logical flow with clear diagrams and mathematical formulations.", "weaknesses": "1. **Limited Theoretical Foundation:** The paper lacks rigorous mathematical analysis of Mix-LN's properties, particularly in contrast to well-established theoretical frameworks for methods like DeepNorm[1]. This absence of formal proofs for convergence properties and optimal ratio selection limits our understanding of the method's fundamental principles.\n\n2. **Scale Limitations:** By restricting experiments to models up to 1B parameters and primarily using LLaMA-based architectures, the work lacks demonstrating scalability to production-scale models (7B-175B parameters) or generalizability across different architecture families.\n\n3. **Incremental Innovation:** The proposed Mix-LN solution, while practical, represents a relatively straightforward combination of existing techniques rather than a fundamental advancement in normalization methodology, especially given extensive prior work referenced in related work and other forums[2,3] in this domain.\n\nReferences:\n\n[1] Wang, H., et al. (2024). \"Deepnet: Scaling transformers to 1,000 layers.\" IEEE TPAMI\n\n[2] Su, J. (2022). \"Understanding Pre-Norm vs Post-Norm in Transformers.\" Scientific Spaces\n\n[3] Raschka, S. (2023). \"Layer Normalization Variants in Transformer Architectures.\" Sebastian Raschka's ML Magazine", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vL4m1LD16X", "reviewer": "Reviewer_mk56", "summary": "This paper identifies an training shortfall of LLMs that their deep layers contribute minimally to the overall performance, which wastes the model capacity. The paper identifies the reason behind this as the use of Pre-LN leading to diminished gradients in later layers. The paper further proposes to use a mixed LN approach, where post-LNs are applied to the earlier layers and pre-LNs to the later layers. Experiments show enhancement brought by the mixed LN approach in both pretraining and supervised finetuning.", "questions": "See weakness part. I would like to see more evidence that the proposed method helps more on deeper models.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper provides novel and interesting insights on the different impact of pre/post LN in the training behavior of earlier and later layers. The observation leads to a straightforward solution of applying mixed LN, which can be easily applied in all models. The paper is well-written and easy to follow. Experiments on multiple models and tasks, as well as ablation studies show solid performance improvement of the proposed method.", "weaknesses": "One major weakness of this paper is that the provided results are not adequate to support the claim that mixed-LN helps deep model by allowing both shallow and deep layers contributes effectively. From the main results in Tab. 1 and 2, it appears that mix-LN improves more over pre-LN in that smaller 71M model (12 layers) than the deeper 1B model (32 layers), which seems to suggest the benefit of mix LN is less in deep models. It would be better to provide results of even larger and deeper models to show that the benefit of mix LN do scale up. If computation resource is limiting the further experiments on larger models, trying a slim but deep architecture may be helpful to verify the effectiveness of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)