GPAS: Accelerating Convergence of LLM Pretraining via Gradient-Preserving Activation Scaling
Authors: Tianhao Chen
Xin Xu
Zijing Liu
Pengxiang Li
Xinyuan Song
AJAY KUMAR JAISWAL
Fan Zhang
Jishan Hu
Yang Wang
Hao Chen
Shizhe Diao
Shiwei Liu
Yu Li
Lu Yin
Can Yang
Venue: NeurIPS 2025 poster ()
Keywords: Large Language Models
Pretraining
Normalization
Primary area: deep_learning
TL;DR: We propose GPAS, a simple method that scales activations without scaling gradients to accelerate pretraining convergence of LLMs.

Abstract: Modern Large Language Models, such as the LLaMA, Qwen and DeepSeek series, predominantly adopt the Pre-LayerNorm (Pre-LN) Transformer architecture. While being stable during pretraining and scalable to large model sizes, Pre-LN suffers from an exponential growth in activation variance across layers, causing the shortcut to dominate over sub-layer outputs in the residual connection and limiting the learning capacity of deeper layers. To mitigate this issue, we propose Gradient-Preserving Activation Scaling (GPAS), a simple technique that can be used in combination with existing approaches. GPAS works by scaling down the intermediate activations while keeping their gradients unchanged. This leaves information in the activations intact, and avoids the gradient vanishing problem associated with gradient downscaling. Extensive experiments across various model sizes from 71M to 1B show that GPAS achieves consistent performance gains. Beyond enhancing Pre-LN Transformers, GPAS also shows promise in improving alternative architectures such as Sandwich-LN and DeepNorm, demonstrating its versatility and potential for improving training dynamics in a wide range of settings. Our code is available at https://github.com/dandingsky/GPAS.

GPAS: Accelerating Convergence of LLM Pretraining via Gradient-Preserving Activation Scaling
Tianhao Chen Xin Xul Zijing Liu², Pengxiang Li Xinyuan Song'
Ajay Kumar Jaiswal'. Fan Zhang'. Jishan Hu Yang Wang'. Hao Chen'.
Shizhe Diao Shiwei Liu', Li, Lu Yin Can Yang"!
The Hong Kong University of Science and Technology FInternational Digital Economy Academy
Dalian University of Technology *Emory University Universsty of Texas at Austin
'NVIDIA Unniversity of Oxford UUniversity of Surrey
(tchenbb, xxuca)@connect, ust.hk, 1. 1.yingsurrey. ac. ac.uk, macyangeust. hk
Abstract
Modern Large Language Models, such as the LLaMA, Qwen and DeepSeek series, predominantly adopt the Pre-LayerNorm (Pre-LN) Transformer architecture. While being stable during pretraining and scalable to large model sizes, Pre-LN suffers from an exponential growth activation variance across layers, causing the shortcut to dominate over sub-layer outputs in the residual connection and limiting the learning capacity of deeper layers. To mitigate this issue, we propose Gradient-PP Preeeettt-reeeiinggggg Activation Scaling (GPAS), simple technique that can used in combination with existing approaches. GPAS works by scaling down the intermediate activations while keeping their gradients unchanged. This leaves information in the activations intact, and avoids the gradient vanishing problem associated with gradient downscaling, Extensive experiments across various model sizes from 71M to IB show that GPAS achieves consistent performance gains. Beyond enhancing Pre-LN Transformers, GPAS also shows promise in improving alternative architectures such as Sandwich-LN and DeepNorm, demonstrating its versatility and potential for improving training dynamics in wide range of settings. Our code is available at pettps:/ggthhe com/ dad/ddnggyy//GPAS
Introduction
The pursuit of more cost-effective and performant architectures is a central theme the development of Large Language Models (LLMs). The introduction of the Transformer architecture [I] laid solid foundation for modern language modeling and has since become the backbone of most state-of- the-art LLMs. To facilitate stable and scalable training, recent models such as the LLaMA, Qwen and DeepSeek series B 4 adopt the Pre-LayerNorm (Pre-LN) Transformer architecture 5 6 enabling training models with up to hundreds of billions of parameters.
Although Pre-LN Transformers offer strong scalability and stability, they remain suboptimal in terms of convergence speed and parameter efficiency. Prior studies [8] 9 10] have shown that deeper layers of Pre-LN Transformers can often be pruned or removed entirely with minimal impact on performance. While this creates opportunities for reducing inference costs, it also highlights significant inefficiency in the training process- where most of the computational budget is spent.
One of the core reasons for the underutilization of deep layers lies in the accumulation of activa- tion variance in Pre-LN Transformers. While the original Transformer architecture places Layer
Equal contribution. Corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Normalization [1]] after the residual connection (commonly referred to as Post-LNN, Pre-LN places before the attention and FFN modules, leaving the shorteut branch unnormalized (Figure Tap. As result, activation variance tends to accumulate with layer depth-ofteen at an exponential rate. Recent works [12 13 have shown that this variance growth causes the signal from attention and FFN modules to be increasingly overpowered by the shortcut branch, limiting the effectiveness of deeper layers in transforming hidden states.
GPAS
Attention/FFN
Scale
8 Scale
Allention/FFN
Normalization
Stop Gradient
Stop Gradient
Normalization
(a) Pre-LN
(b) Pre GPAS
(c) GPAS block
(d) GPAS backprop
Figure I: Architectures of Pre-LN, Pre + GPAS, and the GPAS block. Red lines indicate gradient
flow. Red dotted line means the scale parameter can be either learnable or fixed.
In this work, we propose Cradient-Preservinn Activation Scaling (GPAS), simple yet effective
solution to mitigate variance growth in Pre-LN Transformers. The idea of GPAS is straightforward: reduce variance growth by directly scaling down intermediate activations. However, naively scaling down activations in the forward pass leads to downscaled gradients during backpropagatton, which could lead to vanishing gradients and slow down the learning process. As shown in Figure GPAS works around this issue by scaling down forward activations while preserving the magnitude of backward gradients. Experiments on models ranging from 71M to 1B parameters show that Pre-LN with GPAS achieves faster convergence and improved downstream performance. Detailed analysis reveals that GPAS-augment models exhibit a more compact distribution of activation variance across layers and more uniform layerwise importance.
Beyond enhancing Pre-LN Transformers, GPAS also serves as general-purpose plugin for other architectural variants. We apply GPAS to architectures including DeepNorm [14, Sandwich-LN [I5], Mix-LN [12] and LayerNorm Scaling [33, and observe consistent improvements in convergence speed and performance.
Our main contributions summarized as follows:
Activation Scaling to alleviate activation variance growth in Pre-LN Transformers while preserving gradient magnitudes.
We propose
We validate the effectiveness of GPAS through pretraining experiments across various model sizes, and show that its benefits carry over downstream supervised finetuning tasks.
We further apply GPAS to other architectures such as DeepNorm, Sandwich-LN, Mix-LN and LayerNorm Scaling, and observe consistent performance gains.
We conduct thorough analysis of the training dynamies and model properties of models with and without GPAS.
2 Related Work
Normalization layers. Normalization is crucial for training deep neural networks, especially for modern LLMs with increasing depth and parameter count. Layer Normalization [IT] plays a vital role in the success of the original Transformer []]. Despite its effectiveness, improving normalization remains an active area of research. GroupNorm [66 divides channels into subgroups and normalizes within each group. RMSNorm 17 which omits mean centering and uses root mean square instead of variance, offers a simpler and more efficient alternative to LayerNorm, and has gained popularity [18 proposed new transformation function to replace the gain
in recent LLMs 2 3 AdaNorm

---

and bias in LayerNorm. Dynamic Tanh [99 eliminates normalization altogether by replacing it with a non-linear activation.
LLM normalization schemes. Modern LLM architectures can be broadly characterized by three key components: normalization schemes, attention mechanisms, and FFN designs. Among these, normalization is most relevant to our work, while attention and FFN variants are largely orthogonal.
Transformer normalization has evolved around two main paradigms-PPstt-ee and Pre-LN, which have inspired numerous refinements to address their respective limitations. Post-LN was proposed in the vanilla Transformer M. while Pre-LN was proposed in several works S
to address the optimization issue of Post-LN, theoretical perspective from [20] showed that Post-LN leads to larger gradients near the output layer, necessitating warm-up avoid divergence. Conversely, Pre-LN yields smaller gradients in deeper layers, mitigating instability at initialization but also potentially curtailing the effectiveness of stte- parameters. B2T
6
7
21 bypasses normalization in 14 modifies Post-LNN by scaling up
early layers to combat gradient decay of Post-LN. DeepNorm the residual connection before normalization and downscaling layer parameters during initialization. Sandwich-LA. LN [5] 22 applies normalization both before and after the attention and FFN modules to further stabilize Pre-LN. Mix-LN [12] combines Post-LN and Pre-LN layers in the same model to alleviate both the variance growth of Pre-LN and the gradient vanishing issue of Post-LN. LayerNorm Scaling 133 scales output of normalization layers by inverse square root of layer depth, which facilitates linear variance growth in Pre-LN.
As for attention and FFN architectures, existing variants mainly focus expressiveness, scalability and hardware efficiency [23 24 25 26). While these directions are orthogonal to our work, we briefly mention them here to provide more complete picture.
3 Gradient-Preserring Activation Scaling
3.1 Pre-LN with GPAS
We first demonstrate how to incorporate GPAS into the prevailing Pre-LN architecture, before extending other baselines. Equation 11) shows a typical forward pass of Pre-IN sub-layer (of layer l). GPAS introduces simple modification after each sub-layer as Equation (2). Each layer now
has an additional learnable scalar aq E IR. S-LU(-) is the Sigmoid Linear Unit (27 sg(-) the stop gradient operator, which acts as an identity mapping during forward pass, but gives zero gradient during backpropagation. A visualization of this modification is shown in Figure
The learnable gate aj controls the amount of scaling that is applied to the intermediate activation 4'41- Moreover, by scaling 2'+1- ag controls the mixing ratio between shortcut ø1+1 and the output of the next sub-layer f (LN((1++1) effectively balancing information from these 2 variables. SiLU activation is used to encourage positive a values to reduce activation variance, while also allowing the network to learn negative values when necessary. Notice that the input the next sub-layer's attention or FFN is not scaled regardless of the value of az. This because Pre-LN architecture applies LayerNorm right before the attention or FFN module, which cancels out any scaling applied to 2/41 Overall, Equation
(2) scales activation by (1 SiLU(a)), while the Jacobian 0I141/0z(41 =I.
Pre-LN:
I1 ((LN(x))), E (Attn, FFN), 4/411
+ GPAS: 11+1 4/41 Si. sg(x(+1)
Preserving gradient with stop gradient. We analyze the gradient preservation of GPAS using the first layer's gradient as example. For brevity only one sub-layer per layer is considered. Let
3 1 - SiLU(aa)) be the scaling factor for each layer. If we replace sg(-) with identity, Equation (2)
becomes naive scaling: 2 2'41 3.. Suppose we have L transformer layers in total, with
= C(zz. 3) being the loss. Then the gradient 0c/Ox, is:
Oc
Oc
I-
Oxh
L-1
naive scaling:
II Br,
gz
deL
Os
1=1
ac
ac
GPAS:
Om
1=1
3

---

Since 3, < 1, the product I 3i decays exponentially with depth, causing gradient vanishing problem in early layers. Applying GPAS eliminates the scaling terms 3r during backpropagation and preserves gradient magnitude. Detailed derivations are provided in Appendix
3.2 Apply GPAS to Other Normalization Schemes
LayerNorm Scaling (LNS) [3 scales normalized output by inverse square root of layer depth with Equation (3). Since the scaling factor can be absorbed into the affine transformation inside LayerNorm, LNS is essentially Pre-LN with different initialization of LayerNorm weights. Thus, we apply GPAS to LNS in the same way as Pre + GPAS.
LNS: L = 21 + F(LN())/VI), fE [Attn, FFN]
(3)
Sandwich-LN: 4(41 = + (N(((NN()))) f (AAtt,,FN)
(4)
For Sandwich-LN [5], GPAS also applies after the residual connection with Equation (2) The only difference from Pre-LN is that the sub-layer forward has an additional LayerNorm as shown in Equation (4)
For Post-LN and DeepNorm [ [4], however, the scenarios are quite different. As shown in Equation 5 Post-LN breaks the residual connection by wrapping xι f(zi) with LayerNorm. DeepNorm inherits this architecture and introduces a scaling factor a to the shortcut 2ı and another scaling factor / to the sub-layer weights in where both and B are predefined and fixed during training. Empirically, we found applying GPAS after LayerNorm does not bring performance gain. Instead, we hypothesize that the scaling factor is not optimal since it does variation and training dynamics. Therefore, we apply GPAS to the scaled shortcut a 27, while keeping the input to attention and FFN unscaled. Ultimately, we found applying GPAS as Equation 7 and (8) boosts pretrain performance.
account for layerwise
Post-LN:
F141 LN(#) f(x)), EE {Aten, FFN].
(5)
Sa E {Attn, FFN]. DeepNorm + GPAS: x = xi - -SiLU(a)) ((((:o)sssmm))))) sg(mi),
DeepNorm:
T1+1 = LN(a I + fa(m)),
(6)
(7)
T++ = LN(a of + fa(m)), Sa E {Altn, FFN].
(8)
Mix-LN 12) combines Pre-LN and Post-LNN layers in the same model, so we apply GPAS differently its Pre- and Post-LN layers. For the Post-LN layers, we apply the same strategy Equation 7 For the Pre-LN layers, we use Equation (2)
4 Experiments and Results
4.1 Experiment Setup
Model architectures.
Based on our proposed architectural modifications in Section 3 we conduct
extensive experiments verify their effectiveness. Following (28) and (2)]. we employ LLaMA- based model architectures for implementing Post-LN []] DeepNorm [1]. Pre-LN (30], Sandwich- LN [3 Mix-LN [12] and LNS [.] All models share the same attention and FFN architectures as well as normalization layers except for normalization scheme. Specifically, all baseline architectures
(from Post-LN LNS) utilize RMSNorm [17. LLaMA attention and LLaMA MLP
32 with
SwiGLU activation [33] Moreover, they share the same initialization except for DeepNorm and LNS, which add additional scaling certain parameters. Scaled Embed [34] is applied stabilize pretraining. We then apply GPAS to these baseline architectures according to Section 3 and refer to the modified architectures as DeepNorm+ + GPAS, Pre GPAS, Sandwich + GPAS, Mix GPAS, and LNS + GPAS, respectively. Notice that we did not conduct experiment on Post GPAS, since we found DeepNorm to be a better starting point a preliminary study, and utilized our compute budget on DeepNorm GPAS instead.
Pretraining. We perform pretraining experiments across all architectures and at five model scales: 71M, 130M, 250M, 350M, and 1B. For the larger 7B configuration, we only pretrain Pre-LN and Pre+ + GPAS due to limited computational resources. Following (28), we adopt the Adam (35) opti- mizer and train on the C4 dataset [7]. We tokenize the pretraining corpus with T5 tokenizer (36

---

since it was trained on the C4 dataset. For models with GPAS, we initialize all learnable gates or O. When a; = GPAS does not scale the activations or alter the gradients, meaning the very first forward and backward passes during training are identical to those of the baseline model. We did not explore alternative initialization schemes for oq- We also use the same gate value for both attention and FFN sub-layers within the same layer.
carried out on NVIDIA H800 GPUs. Models below 350M can be trained in under to 8 hours, while each 1B model took 2 to 3 days of training. More detailed configurations are listed in Table
All experiments
Table I: Pretraining configurations for models of different sizes.
Model Size
Learning Rate
Warmup Steps
Training Steps
Batch Size
Train Tokens
Eval Tokens
71M
1 10-3
IK
10K
512
IB
10M
130M
1 10-3
2K
20K
512
2B
10M
250M
10-3
4K
40K
512
4B
350M
5 10
6K
60K
512
6B
10M
IB
5xx10" 10-1
10K
100K
512
10B
10M
Supervised finetuning. To determine whether the improvements observed during pretraining persist in subsequent training stages, we perform SFT on the pretrained models with 1B parameters from Section 4.2 and then evaluate their performance common reasoning benchmarks. Following [12]. we finetune the models on the Commonsensel 70K dataset [38] and evaluate the models on seven downstream tasks. The learnable gates aq frozen during SFT to avoid disturbing pretrained knowledge. We use a learning rate of 3 x 10-4 and train for 4 epochs using LISA 33). As for
evaluation, adopt the widely used LM Evaluation Harness 40)
4.2 Pretrain Results
Table 2: Perplexity (L) of pretrained models on evaluation set. improvement over the base method.
Numbers in parentheses show
Method
71M
130M
250M
350M
IB
Post-LN 1
33,80
26.50
1351,58
21.19
1406.66
DeepNorm [4
35.49
26.78
22.20
21.76
1400.39
DeepNorm GPAS
34.78 (-0.71)
26.62 (-0.16)
21.89 (-0.31)
21.29 (-0.47)
16.01 (-1384)
Pre-LN 20
33.98
26.61
21.54
20.71
16.53
Pre + GPAS
33.38 (-0.60)
26.25 (-0.36)
21.34 (-0.20)
19.77 (-0.94)
16.11 (-0.42)
Sandwich-LN [15
32.28
25.31
20.43
20.20
16.26
Sandwich + GPAS
31.44 (-0.84)
24.86 (-0.45)
20.38 (-0.05)
19.45 (-0.75)
15.85 (-0.41)
Mix-LN [12
33.88
26.29
21.52
20.73
15.87
Mix + GPAS
33,26 (-0.62)
26.03 (40.26)
21.43 (-0.09)
19.82 (-0.91)
15.38 (-0.49)
LNS [13
34.58
25.91
20.59
20.35
15.61
LNS GPAS
32.68 (-1.90)
24.95 (-0.96)
19.89
(-0.70)
19.38 (-0.97)
14.87 (-0.74)
Table 2 summarizes the pretraining results across all model sizes and architectures. Results for 7B- parameter Pre-LN and Pre GPAS are provided Appendix B We have the following observations: 0 Consistent Gains from GPAS. For nearly all architectures and model sizes, applying GPAS leads to lower perplexities (numbers in parentheses), confirming that preserving gradient magnitudes while adjusting activation scales is beneficial. @ Best Baselines at Different Scales. At smaller scale (71M), Sandwich-LN is the strongest baseline (32.28), and adding GPAS further reduces perplexity to 31.44. At larger scale (IB), LNS already outperforms other baselines (15.61), but GPAS still offers a notable improvement ( (14.88). range from moderate (== 0.7 drop in perplexity) to quite substantial (e.g., -1.90 for LNS 7IM). These consistent gains highlight that GPAS can be seamlessly integrated with various normalization schemes and model sizes, providing an effective and lightweight solution to boost convergence and performance in LLM pertaining.
Magnitude of Improvement. The improvements (in parentheses)
For detailed analysis how GPAS enhances performance, please refer to Section 5
5

---

Pro-LN
5.0
ProLN
Pre - GPAS
Pre GPAS
4.5
3.5
3.0
2.5
40K
J00K
Traiiing Step
Training Stop
Pretrain loss
Evaluation loss
Figure 2: Pretrain and evaluation loss of Pre-LN and Pre+ GPAS, 1B parameters
Table 3: 0-shot performance (1) on various benchmarks (IB ----marmettrr models). All numbers are
accuracy in %. Bold: higher among GPAS and baseline method; Bold italic: highest among all.
Method
MMLU
BoolQ
PIQA
SIQA
HellaSwag
WinoG
ARC-e
ARC-e
OBQA
Average
Post-LN
22.95
37.83
52.77
34,03
26.20
48.15
27.36
19.37
11.40
31.12
DeepNorm
22.95
37.83
52.77
34.08
26.20)
51.14
27.31
19.37
11.40
31.45
DcepNorm GPAS
26.46
62.11
69.53
46,93
34.37
52.09
49.24
22.61
20.40
42.64
Pre-LN
25.96
50.34
68.66
44.27
32.39
51.14
49.37
21.33
17.60
40.12
Pre + GPAS
26.68
59.79
69.31
46.52
33.64
52.49
49.79
22,70
22.00
42.55
Sandwich-LN
27.42
61.77
67.63
44.68
32.76
50.67
23.12
21.40
Sandwich GPAS
27.29
61.90
69.15
45.29
34,61
50.36
51,39
23.46
22.20
42.85
Mix-LN
26.24
61.93
68.
45.50
33.09
52.25
48.78
24.40
20.80
42,40
Mix GPAS
26.23
61.99
69.59
45.60
33,51
53,51
50,34
22.35
22.40
42.83
LNS
26.62
62.02
69.48
45.39
34.76
51.38
50.88
23.29
19,80
42.63
LNS+GPAS
27.78
61.56
71.00
47.49
36,19
51.22
$2.37
25.51
24.40
44.19
4.3 SFT Results
We present the finetuning results in Table and notice the following observations: Gains from GPAS Across Architectures. When GPAS is applied (bottom half), every architecture sees a mea- surable boost in average accuracy. For example, Pre GPAS achieved a notable 2.43% increase in average accuracy. DeepNorm + GPAS achieves the most dramatic improvement since vanilla Deep- Norm diverged during pretraining, while incorporating GPAS successfully avoided this catastrophic behavior. This suggests that DeepNorm's 's fixed residual scaling may not be optimal yet for network stability, and GPAS automatically learns to increase stability and convergence speed. Task-Specific Improvements. Breaking down the gains by individual datasets shows that the enhancements are typically consistent across tasks but vary in magnitude. For instance, Pre + GPAS's performance on BoolQ reached 9.45% increase in accuracy, OBQA also has notable gain of 4.4%. Meanwhile, LNS + GPAS attains improvements both OBQA (+4.60) and HellaSwag (+1.43), pushing its overall average 44.19 19% (+1.56).
Overall, by adaptively controlling the scale of the activations while preserving gradient magnitudes, GPAS consistently unlocks additional gains in downstream tasks.
5
Analysis
In this section, we provide thorough analysis of GPAS by examining its training dynamics and model properties. We base our analysis primarily on Pre-LN and its GPAS-augmented version with IB parameters. We also provide detailed ablation study in Appendix A
5.1 Learned Layerwise Gate Values
GPAS introduces learnable gate ay for each layer. We visualize the learned values of aj in Figure 3 All models are the IB-parameter models trained in Section 4.2 Generally, Pre-LN and Sandwich-L.N layers tend to learn positive gates, while Post-LN layers tend to learn negative gates. Notably, for the Mix + GPAS model, the first Post-LN layers learned negative gates, and the remaining Pre-LN layers learned positive values.

---

Across all model variants, the first layer tend to learn negative gate values. This means the network scales up the activation in the first layer. We hypothesize that this is due the low variance of the initial word embeddings, and the network automatically learns to scale up that embedding to match the variance of subsequent layers. This hypothesis is also supported by Figure [bb] where the input activation to the first layer (which is the initial word embedding) has much lower variance than the remaining layers.
Throughout training, the gates exhibit the most variation during the first 10-20% of steps. After that, they tend to remain relatively stable. One likely reason for this behavior the use of half- precision training with BFloat16, which may suppress small updates to the gate values due to limited numerical precision.
0.15
0.2
Gute Value
Gate Valuee
16
Index
0.05
0.0
0.00
Layer
-0.1
-0.05
-0.
GDAS
DecNNoe GBAS
-1.10
20
20K
40K
60K
80K
100K
Layer Index
Training Step
(a) Learned gate values au for different models
(b) Activated gate values SaLU((:) across training
steps of Pre + GPAS
Figure 3: Learned layerwise gate values
5.2 Activation Variance Comparison
We compare the activation variance of Pre-LN and Pre GPAS, across all pretraining steps and layers. Experiments are conducted on 1B parameter models. We only plot every 2 layers including the first and last layers for brevity. Figure 4 shows that Pre-LN has an exponential increase in activation variance across layers, while Pre + GPAS offers near 50% decrease in highest activation variance, with more uniform and compact activation variance distribution across layers.
500
250
Activation Vanunee
Activatioo Variunee
400
200
16
16
apyer Tndes
300
150
12
12
200
La9p
100
0
20X
40K
60K
80K
100K
20K
40K
(0X
80K
100K
Training Step
Training Step
(a) Pre-LN
(b) Pre GPAS
Figure 4: Layerwise activation variance with and without GPAS
5.3 Gradient Norm Comparison
We verify whether GPAS preserves gradient while scaling down activations. Figure shows the layerwise gradient norms across across training steps. Models are the IB parameter variants we trained in Section 4.2 It's evident that Pre GPAS has a much larger gradient compared to its baseline method. Most layers in Pre-LN have gradient norm around 0.05 -0.1, while Pre GPAS's are scaled to 0.05 0.5. Although the first layer's gradient of Pre + GPAS seemed overly aggressive, the model was able to adjust to that gradient scale and achieve faster convergence.
However, we did find that the gradient spike around step 10K in Figure 5b disturbed the pretraining process. The gate values at step 10K went through an aggressive fluctuation (see Figure (3)). which caused the model to adjust to that change at steps 10K 10K-30K. This is also observed in the loss curves

---

in Figure 2 where both pretrain and evaluation loss of Pre GPAS temporarily went above that of Pre-LN around step 10K. One simple way fix this issue to use gradient clipping on the gate values or. However, we found that, while stablizing training, applying an additional gradient clipping the gate values slightly degrades final performance (PPI. 16. 16.11-116.21 This suggests that the gate values might need an update scheme different from the main parameters. For our main experiments, we apply gradient clipping of 1.0 to all parameters except for the gates. For other GPAS-augmented baselines, we did not encounter such fluctuation in gate values which disturbs training.
1.5
20
Gradieni Norm
Gradient Norm
-16
2
16
Layer Index
1.0-
Layer
0.0
20K
40K
S0K
100K
20K
40K
60K
80K
100K
Trinning Step
Training Step
(a) Pre-LN
(h) Pre + GPAS
Figure 5: Layerwise gradient norm with and without GPAS
5.4 Weight Norm Comparison
Since Pre + GPAS scales up activations in the early layers according to Figure 4b we suspect that its early layer weights also have larger norms. Specifically, we compare the norm of weights in the attention and FFN modules. Figure shows that Pre GPAS has larger weights in almost all layers compared Pre-LN, especially in the early layers where activation variances are relatively small. This effectively scales up the output variance of attention and FFN modules, such that early layers and deeper layers share similar activation variance. We believe that the combination larger weights and stable activation variance across layers enables the network to tolerate larger gradients without crashing, which in turn speeds up convergence.
2.0
ratio=1
1.8
ratio gate_proj down_proj
SLproj
k_oroj
1.6
1.6
prof
o_proj
1.4
1.4
1.0
15
20
Layer Index
Layer Index
(a) Attention weight norm ratio
(b) FFN weight norm ratio
Figure 6: Attention and FFN weight norm ratios of Pre + GPAS over Pre-LN.
5.5 Layer Importance Comparison
We compare the layerwise importance of models with and without GPAS. The importance of a given layer measured as the performance drop after removing that layer. We take the finetuned models in Section 4.3 as baseline, and measure the difference in average score across those benchmarks listed Table after removing each one of the layers.
Figure 7 shows that Pre GPAS increases layer importance for most and especially the deeper layers, while some layers in vanilla Pre-LN
harmful to performance. This indicates that GPAS enables more efficient utilization of model parameters, while Pre-LN's exponential variance growth limited the learning capacity of deeper layers.
We also show layer importance LNS + GPAS in Figure 7b LNS mitigates variance growth in Pre-LN by downscaling LayerNorm output with square root of layer depth. In this case, GPAS still

---

amplifies the importance of each layer by small but noticeable amount, leading to a higher overall performance.
0.0
&Avenge Score
2.5
Avengge Score
2.5
-3.0
-7.5
-7.5
10.0
Pre-NN
10.0
LNS
GPAS
-12.5
-12.5
16
20
Removed Layer Index
Removed Layer Index
(a) Pre-LN vs. Pre GPAS
(b) LNS vs. LNS + GPAS
Figure 7: Layer importance measured by performance drop after removal.
5.6
Gradient Analysis
Based on our experimental findings, we provide a brief and preliminary analysis to the activation variance and gradient profiles of Pre-LN and Pre GPAS. By the result of [13 Theorem 1 and Lemma 1, and the work of (34]. we have the following result for Pre-LN transformers:
Lemma 1 (Variance and Gradient Growth in Pre-LN), For Pre-LN Transformer with L layers, using Equations ( and let We be the model parameter matrix at layer Assume that for all layers, me Tee and We are mutually independent and follow Gaussian distributions with mean zero. Let 0gg 8m
be the output of the L-th layer, and consider the gradient norm
Let the upper bound for this
2
norm denoted by UP(.). Then it satisfies:
dyL dr
) < O(L),
M < UP
When add the GPAS the Pre-LN transformers, have the following theorem
Theorem 1 (Variance and Gradient Growth in Pre-LN with GPAS). Using Equations ① and (2) and under the same assumptions as Lemma assume that each ar is bounded and varies slowly across layers. Let L(a) and (a) denote the layerwise lower and upper bounds of log(1 SiLU(ar)), respectively. Lei a denote the variance of dt. Then the upper bound of the gradient norm, denoted as
UP
8UL 02
). satisfies the following inequality:
{
1 σ
a I (L00)/2 11
}
Byy Oun
))
expl (M(())
minfL,or
exp
< UP
S exp
(10)
The detailed description as well as the complete proof, are provided in Appendix D From Theorem [ By comparing Lemma (before GPAS) with Theorem | (after GPAS), For the lower bound of the gradient norm upper estimate UP(
dul
). the term L(a) effectively acts as a compensatory factor
that offsets the influence a, thereby accelerating the growth of the bound beyond a constant rate. This leads strictly non-constant lower bound, which in turn ensures that the gradient magnitude retains sufficient variability across layers. particular, this precludes the possibility vanishing gradients and encourages meaningful gradient flow during backpropagation.
In contrast, for the upper bound, the term ((a) serves as multiplicative scaling factor that exponentially suppresses the bound. This results in a significantly lower gradient norm ceiling compared to the ((L) upper bound of standard Pre-LN Transformers. Consequently, the proposed modification mitigates gradient explosion and reduces the occurrence of large loss spikes, thereby enhancing training stability. Notably, the improved upper bound grows more slowly than L,, further dampening excessive gradient amplification and promoting smoother optimization dynamics.

---

Conclusion
We proposed Gradient- Activation Scaling, a simple method that mitigates the exponential growth of activation variance in Pre-LN Transformers, thereby accelerating convergence and enhanc- ing parameter efficiency. Beyond improving Pre-LN Transformers, GPAS also proves be a viable plugin for alternative normalization schemes such as DeepNorm, Sandwich-LN, Mix-LN, and LNS.
Limitations. While GPAS showed consistent gains in pretraining performance, our experiments are limited to 1B-parameter models due to computational constraints. Additionally, the current use of the SILU activation may still permit unstable gate updates that disrupt training dynamics- although gradient clipping partially mitigates this issue. Nevertheless, the learnable gates introduce uncertainty during pretraining, and a predefined gate schedule with theoretical guarantees might be more preferable, especially at larger scales. A deeper understanding of how GPAS affects training stability and convergence remains an open question. Finally, GPAS intended for training LLMs from scratch, and applying it models pretrained without GPAS will likely be suboptimal.
Broader Impact. GPAS offers a lightweight solution to improve the convergence speed and parameter efficiency of LLM pretraining. This may contribute more accessible and sustainable development of foundation models, with potential downstream benefits across applications such as education and scientific research.
10

---

Acknowledgments
This work was partially supported by grant from the Research Grants Council of the Hong Kong Special Administrative Region, China (Project Reference Number: A0E/E--001224-N)
References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NewrIPS, 2017.
[2] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama
3 herd of models. arXiv preprint arXiv:240722178; 2024.
[3] An Yang, Baosong Yang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Zhou, Chengpeng Li, Chengyuan Li, Dayiheng Liu, al. Qwen2 technical report. arXiv preprint arXiv:2407.10671 2024.
[4] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12945 2025.
[5] Alexci Baevski and Michael Auli. Adaptive input representations for neural language modeling. ICLR, 2019.
[6] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 2019.
[7] Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. ACL, 2019.
[8] Lu Yin, You Wu, Zhenyu Zhang, Cheng- Yu Hsieh, Yaging Wang, Yiling Jia, Mykola Pech- enizkiy, Yi Liang, Zhangyang Wang, and Shiwei Liu. Outlier weighed layerwise sparsity (owl): A missing secret sauce for pruning iims to high sparsity. ICML, 2024.
[9] Andrey Gromov, Kushal Tirumala, Hassan Shapourian, Paolo Glorioso, and Daniel A Roberts. The unreasonable ineffectiveness of the deeper layers. arXiv preprint arXiv:2403.1788. 2024.
[10] Xin Men, Mingyu Xu, Qingyu Zhang, Bingning Wang, Hongyu Lin, Yaojie Lu, Xianpei Han, and Weipeng Chen. Shortgpt: Layers large language models are more redundant than you expect. arXiv preprint 1xXiv:2403303853 2024.
[11] Jimmy Lei Ba. Layer normalization. arXiv preprint 1rXiv:1607.064500 2016.
[12] Pengxiang Li, Lu Yin, and Shiwei Liu. Mix-ln: Unleashing the power of deeper layers by combining pre-In and post-In. arXiv preprint arXiv:2412. 2024.
[13] Wenfang Sun, Xinyuan Song, Pengxiang Li, Lu Yin, Yefeng Zheng, and Shiwei Liu, The curse of depth in large language models. arXiv preprint arXiv:2502.05795 2025.
[14] Hongyu Wang, Shuming Ma, Li Dong, Shaohan Huang. Dongdong Zhang, and Furu Wei. Deepnet: Scaling transformers to 1,000 layers. TPAMI, 2024.
[15] 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. NewrIPS, 34:19822-1988 2021.
[16] Yuxin Wu and Kaiming He. Group normalization. In ECCV, pages 2018.
[17] Biao Zhang and Rico Sennrich. Root mean square layer normalization. NewrIPS, 32, 2019.
[18] Jingjing Xu, Xu Sun, Zhiyuan Zhang, Guangxiang Zhao, and Junyang Lin. Understanding and improving layer normalization. arXiv preprint
2019.

---

[19] Jiachen Zhu, Xinlei Chen, Kaiming He, Yann LeCun, and Zhuang Liu. Transformers without normalization, 2025.
[20] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization the transformer architecture. In ICML, pages 10524-10533. PMLR, 2020.
[21] Sho Takase, Shun Kiyono, Sosuke Kobayashi, and Jun Suzuki. B2t connection: Serving stability and performance in deep transformers. ACL, 2023.
[22] Jeonghoon Kim, Byeongchan Lee, Cheonbok Park, Yeontaek Oh, Beomjun Kim, Taehwan Yoo, Seongjin Shin, Dongyoon Han, Jinwoo Shin, and Kang Min Yoo. Peri-In: Revisiting normalization layer in the transformer architecture. arXiv preprint arXiv:2502.02732, 2025.
[23] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. arXiv preprint arXiv: 221. 2022.
[24] Joshua Ainslie, James Lee-Theee Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13244 2023.
[25] Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li. Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specialization mixture-of-erte ex-eeerss language models. arXiv preprint Mxii:2240100600 2024.
[26] DeepSeek-AI, Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang. Bo Liu, Chenggang Zhao, Chengqi Dengr, al. Deepseek-v2: A strong, economical, and efficient mixture-of-expprts language model. arXiv preprint arXiv:2405.0444
2024.
[27] Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions, 2017.
[28] Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: High- rank training through low-rank updates. In ICLR, 2023.
[29] Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Meforr-efficient Ilm training by gradient low-rank projection. ICML, 2024.
[30] Zihang Dai, Zhilin Yang, Yiming Yang. Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. ACL, 2019.
[31] Xinyu Gong, Wuyang Chen, Tianlong Chen, and Zhangyang Wang. Sandwich batch normal- ization: A drop-in replacement for feature distribution heterogeneity. In Proceedings of the IEEE/CVF Winter Conference Applications of Computer Vision, pages 2494 2504, 2022.
[32] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie- Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 2023.
[33] Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202,, 2020.
[34] Sho Takase, Shun Kiyono, Sosuke Kobayashi, and Jun Suzuki. Spike no more: Stabilizing the pre-training of large language models. arXiv preprint (xXiv:231211690) 2023.
[35] Diederik Kingma. Adam: A method for stochastic optimization. ICI.R, 2015.
[36] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lec, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-67 2020.
[37] Jesse Dodge, Maarten Sap, Ana Marasovié, William Agnew, Gabriel Ilharco, Dirk Groeneveld, Margaret Mitchell, and Matt Gardner. Documenting large webtext corpora: case study on the colossal clean crawled corpus, 2021.
12

---

[38] Zhiqiang Hu, Lei Wang. Yihuai Lan, Wanyu Xu, Ee-Peng Lim, Lidong Bing, Xing Xu, Soujanya Poria, and Roy Ka- Wei Lee. Amm-adapeers: An adapter family for parameter-efficien fine-tuning of large language models. EMNLP, 2023.
[39] Rui Pan, Xiang Liu, Shizhe Diao, Renjie Pi, Jipeng Zhang, Chi Han, and Tong Zhang. Lisa: Layerwise importance sampling for memory-efficient large language model finettunigg 2024.
[40] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024.
[41] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. Llama 2: Open foundation and finettuedd chat models, 2023.
[42] E. T. Whittaker and G. N. Watson. A Course of Modern Analysis. Cambridge Mathematical Library. Cambridge University Press, Cambridge, 4 edition, 1996.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
We proposed GPAS which acelerates convergence and boosts performance of L.LM pretraining, extensive experiments in Sections 4 and confirms the utility of our method.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discussed limitations our work Sections and addressing drawbacks such limited model scale and potential training instability issues.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification:
We provide proofs of our theoretical results (Section 5.6) in Section
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provided details of our training approach in Sections and covering model architecture, hyperparameter, additional techniques, etc. We also provide code to replicate our experiments in the supplementary files of this submission.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
15

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: In Section 4 we provided citations the open-source training framework and dataset that we used with our experiments, along with training configurations in Table We also provide code replicate our experiments in the supplementary files of this submission.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://niis.co cc/
mublic/guider poodeeSubmissiooolllee
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips wttpccccccwccpcdeeeeeeeeeeepeeee CodeSobmisiiionooooeooooo ierobbmmssioooolol.mcoo
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: We provide details of our training settings in Sections We also provide code to replicate our experiments in the supplementary files of this submission.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Running multiple pretraining experiments to produce error bars is very costly and consuminn,,, and out our computational budget. However, we do run our proposed method across various baselines to show that method consistently improves all baseline methods (Table and
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
16

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CL, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provided details on compute resources in Section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification:
We made sure follow the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss potential impacts of our work in Section 6
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
17

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI).
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: Our models are trained the well established C4 dataset, and there's 's little safety concern since the generative ability of our pretrained models are limited by model and data scale.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We provide license information of assets used our paper in Section
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
18

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our work does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
19

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification:
We only use LLMs for writing at the sentence level.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy pettts://neurre http///rrpp..ccoooffeeeeeeee//// Cc/Conferrecces/2225//LL for what should or should be described.
20

---

A
Ablation Study
We perform series of ablation studies on GPAS elucidate the rationale behind its architectural design. In particular, we investigate the impact of four key factors: the Activation function, the Application location of GPAS, the Necessity of stop gradient operator, and the choice between Learnable vs. predefined gate value. All ablations are conducted on 5S0M-parameter models, chosen because this scale typically serves as a reliable proxy for larger models while remaining feasible within our computational budget. We also limit the baseline architecture to Pre-LN only, and leave explorations of other architectures to future research.
Activation function. GPAS uses the SILU activation by default. To enable more customizable control over activated gate values, we generalize GPAS to Equation (1I). where Act(+) can be any activation function. We then examine how different activation functions affect performance, using a 350M- p Pre-LN model with GPAS. Table4]shows that Identity and Tanh achieve slightly lower perplexities than SiLU in the 350M setting. However, this advantage does not persist at larger scale: the IB setting, Identity and SiLU achieve perplexities of 16.49 and 16.25, respectively. We hypothesize that these activations, particularly Identity, permit overly aggressive gate updates during training: by contrast, SiLU imposes a smoother gradient profile, constraining updates to a more stable range.
D1+1 = 2/41 - Act((x) - sg(() (g(2141).
(11)
Table Ablations on activation function (left) and GPAS insertion position (right). Based on 350M Pre-LN. SiLU (8 = 8) refers to scaled SiLU: x (fz).
Activation
Perplexity
GPAS Position
Perplexity
Identity
20.08 (-1.27)
ReLU
21.17 (-0.18)
After sub-layer (default)
20.35 (-1.00)
LeakyReLU
21.17 (-0.18)
Before sub-layer
20.73 (-0.62)
Tanh
20.09 (-1.26)
After LayerNorm
21.33 (-0.02)
SiLU (8- 8)
20.26 (-1.09)
After Attn / FFN
21.23 (-0.12)
SILU (default)
20.35 (-1.00)
No GPAS
21.35
No GPAS
21.35
Where to apply GPAS. Since GPAS can, in principle, be integrated at any point where intermediate activations arise, we investigate the impact of inserting at different locations within the transformer
block. In Table 4 we compare four variants against a 350M-paramettr Pre-LN baseline: (1)
our default setting, which applies GPAS right after the sub-layer (residual module output); (2) applying GPAS before the sub-layer; (3) placing it immediately after the LayerNorm; and (4) after each Attn/FFN module. While each insertion strategy outperforms the no-GPAS baseline, the default approach (inserting GPAS after the sub-layer) offers the largest perplexity reduction (-1.00), indicating that modulating the combined residual and module output is the most effective choice.
Necessity of stop gradient operator. We conduct a control experiment to show the necessity of the stop gradient operator in GPAS. For the control experiment, we use Equation 12 to replace Equation
x1+1 = 2|ا SiLU(ar) 2/+11
(12)
We found that the stop gradient operator sg(-) is crucial for preserving gradients. The gradient norm of GPAS without sg(-) looks very similar to that of Pre-LN in Figure Although it still manages to scale down the activation variance by around 50% across all layers, GPAS without sg(.) does not offer meaningful improvement in perplexity compared to Pre-LN (see Table This highlights the importance of the stop gradient operator to prevent gradient vanishing issue associated with gradient downscaling.
Learnable vs. predefined gate value. GPAS adopts learnable gates by default. We investigate whether predefined gates can also be used, which could potentially offer more predictability and consistency. To determine reasonable set of predefined gate values, we extract the gates from
21

---

Table Ablations on stop-gradiens usage (left) and gating strategy (right) GPAS.
Method
Perplexity
Method
Perplexity
sg(
20.35 (-1.00)
Learnable gate
20.35 (-1.00)
w/o sg(-)
21.34 (-0.01)
Predefined gate
22.46 (+1.11)
No GPAS
21.35
No GPAS
21.35
pretrained Pre GPAS model, and those values initialize the gates for the control experiment, where the gate values will be fixed during training. Results are shown in Table 5 We found that fixing the gates in this way led to substantial drop in performance, especially in early stages when GPAS with learnable gates converges much quicker. This result confirms the importance of learnable gates to account for training dynamics. Another potential way of using predefined gate value is to introduce a warmup stage for the gates, which we leave to future work.
B Pretrain Results on 7B-Parameter Models
To further verify the effectiveness of GPAS on larger scale models, we perform pretraining exper- iments on Pre-LN and Pre GPAS with 7B parameters. follow 41 and use learning rate of 3 x 10 with 10K warmup steps and cosine decay. Batch size is set to 2048 and scheduled train for 150K steps on 60B tokens. We use gradient clipping of 0.01 on gate parameters a and 1.0 on other parameters to stabilize training. Due to constraints computational resources, the models up to 40K steps with 16B tokens. After 40K training steps, Pre-LN and Pre + GPAS reached evaluation perplexity of 15.27 and 13.82, respectively. As shown in Figure 8b Pre + GPAS achieves a much faster convergence than vanilla Pre-LN.
only train
3.75
Pre-LN Pre GPAS
3.25
Pre-UN
GBAS
2.75
5X
10K
15K
20K
25K
30K
35K
40K
5K
10K
15K
20K
25K
30K
40K
Training Stop
Training Step
(a) Pretrain loss
(b) Evaluation loss
Figure 8: Pretrain and evaluation loss of Pre-LN and Pre GPAS, 7B parameters
22

---

C
Gradient Preservation with GPAS
Following Section 3.1
the gradient 0L/Oz, is:
naive scaling: 0c
0C
L II
0m++1
dr1
d.l
dr
1=1
ac
[0x1+1
0r'41
II
0xi4
dzi
0c
0xl+1
L-
004+1
d
dxi
0x(+3
0c
0al+11
8t,
dr
Omn
Oc
ac
0ul+1
0m
GPAS:
dr1
Ox
0x(41
ac
7.
0al+1
I
1
Ox
1=1
0c
0zl+1
dx
8x
1=1
D Proof of Theorem of GPAS
Proof. By equation
and have the following:
y 2t-1 FFN LN(2)))
Ve
(13)
x/ = σe + Attn LN(I))).
VE
Pre-LN: 2/41 = 21 ((L(zz)), fe {Attn, FFN]
(14)
GPAS: #141 = 21111 SiT.U((r) sg(1(41)
(15)
Following the variance analysis in 13 both the FFN and Attn modules contribute equally to variance accumulation. For simplicity, we have:
22 H = -2, (1+ ) (Sin-)))(a))).
(16)
The variance with regard to ozi
(1- SiLU(ar))
2 21
1 1 oze
(17)
For the parameter &e we have
(1 (1+c-ar) a0
1 - -it) =1
(18)
a-ac
1+e -
1-1
log 02 log 22 +
Σ
log 1 1 1 + log (I-SLLU(a)) +C,
(19)
x1
1=1
23

---

where C' is an (unspecified) constant.
Using the original definition for SiLU we have
1 (iLU(a)))) log
1 t -4a)
)
logg( te ----aaa -ae) -logl (1+e
(20)
1 1
ltee (1te-aa- 14e-0e - which is function of og. Thus the inequality becomes
Next, denote L(ar)= = log
log o 2 log o
Σ
1
I(ae)) C. +1
(21)
Oze
k=1
For fixed 22 log L. (C + [(ae)), if oz < 0, there is no lower bound for the variance. So it reach O. To establish upper bound for o
Σ log log
1+
(22)
Ozk
azk
km
The SiLU value of this term is less than or equal to zero (since typically 1 - SiLU(as) S 1) or can be bounded by appropriate constant M (az). For derivation, we denote
log - (iL((a)())) & ((s).
(23)
Putting these bounds together have
1 + M(aa) +
log o I log 2
(24)
Oak
Next, we analyze the gradient stability of GPAS. Following Equation (38) in [13, and applying the formalization techniques from 42]. derive the result under the consideration that stopgrad is used, thereby eliminating any gradient contributions at that point. Consequently, we obtain:
UP
OyL dx
1)
II
1 1+
_
1
A+
BBBB 0
B):
(25)
Ozz
Substitute our lower bound in Equation 21 on of into the above product. Notice that if
Σ ( + I(ar)) + C.
1-1
22 2 22 exp
Si
with
S =
(26)
Jz
+1
k=1
Thus, we have the upper bounds on the inverses:
1
S expl
< 1
2 and
exp -Si
(27)
0
2
m1
Thus a valid lower bound for the UP norm is
dyL
A
Jnn +1 + ((as))
+ +c]
UP
1+
exp
am
2
1=1
k=1
(28)
B
1
M(ave) +c]})
exp
a2
Oza
+1
Similar to above, assuming that for each layer and based Equation 24
we have:

---

L-l Σ
S o0 exp
(29)
UsN
So we have:
1
1-1
- M
2
exp
Σ
(30)
Jz:
Jai
2
Hence the UP product is upper bounded by
UP(
OUL 0z1
Σ -
1-1
+-=+Mmmm M(ae)) +(0+
1
+
exp
2
Oze
2=1
k=
(31)
B
{-1
M(ar) + +c)31/
exp
o2
Oek
k=1
D.1 Upper bound of UP())
Assume that the M (ax) is bounded and will not change so much. Inserting this into (31 thus, in the case of constant a and indeeenndeeen M (a), we obtain the explicit bound
UP(1 (1O Dri 1,)
JyL
S exp{
{ A exp{ - (M(())(]}
1 1
2x
1 1-exp
20
(32)
exp
expo - (M((M())))))))))))))) + +c]]
+ o
1 - exp
It is easy to verify:
1 exp
I-1
20
=O(in((L,o))
(33)
[-1 1
- exp
-20
20
and similarly
1 exp
[--1
a
0 (nin(L,o))
(34)
- exp{
Plug Equations
33) and
into the bound. We deduce
UP(
dyL
1b)
0x
+ (=p(-(M(()) }
A
B exp _ (M(() + (m)om(mn(). + oz
< exp
A exp{ -/4 (Maa) + [0) min(L,o)
+0
_ exp ((M(a) + +0) minthhot)}
exp
expp{ - 1 (M())) minn(
exp
o
(35)
25

---

D.2 Lower bound for UP(-)
Assume that the L(ax) bounded and will change much. Inserting this into 28 thus, in the case of constant a and inyee-indeeendent M (a), we can obtain the explicit bound.
Assume with D = L(a:), For each term in the product (with index 1), define
-+I
B F(I) = 1 exp - - 1) σ exx(-((-1)D0))
A
1
(36)
Taking the logarithm of the whole product, get
2 Σ exp - (6DD+0) +0D+ oe exp( -000))+0((p---)CC)))
(37)
s=0
Each of the sums over is geometric series. For instance,
L-2 Σ Cexp(-2 ==0
L-2 - + 01) = C/2 Σ
=c- (L-1)D/2
(e (--/2)²
(38)
e-2/2
D/2
and similarly for the second term. Notice that when D/2 the whole sum (even as L 00) is e bounded by a constant. Exponentiating, we can write:
dyL UP()
1,) 2 exp
A
C/2
T-e-n +om++oom--
(39)
dz,
- 1-e-D12
0²
similar to the above process, we can get:
f0
cl t + +1)}
UP
Oyr. Ox
1b) 2 exp o
(40)
E License
We provide licenses of assets used in our paper, including model, dataset and code.
C4 dataset: This is the dataset used in pretraining experiments. It's released under the Open Data Commons License Attribution family License
Commonsense 170K dataset: This is the dataset we used for supervised fine-tuning, released License
26


Reviews:
[{"review_id": "5RUTEBEmke", "reviewer": "Reviewer_odJN", "summary": "The authors propose a simple modification of residual connections in transformers that aims to adjust the scale of intermediate activations without impacting the scale of gradients. They test the new GPAS method by pretraining and finetuning language models. The results show that the method improves performance and decreases the variance of intermediate activations.", "questions": "- Is there any negative impact of training speed? I would expect not, but it is always better to ask. I think it would be better to clearly mention this in the paper.\n- Since you decided to use a relatively established Commonsense170K dataset/benchmark for evaluation, why not reporting the performance on all datasets in the benchmark?\n- Is it correct to the SFT evaluation as zero-shot when the models are directly trained on training samples from the evaluated datasets?\n- In the experiments with predefined gates, what is the impact of this setting on the activation and gradient variance?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n- First of all, I really like the simplicity of this method. Researchers often try to impress each other by proposing complicated contraptions, but this is not the case; it tries to solve a complicated problem with a simple fix.\n- The authors have tested their method not only on the most popular pre-layernorm architecture but also on four other architectures. The improvement due to their modification seems to be consistent across them.\n- The paper is well-written and was easy to understand and follow,\n- The empirical analysis that demonstrates the effects of GPAS is extensive and clear.\n- The paper also contains detailed ablations studies that experimentally confirm the various choices of the GPAS modification. \n\n**Weaknesses**\n- The evaluation of the downstream performance could be improved. For example, the performance of all models on MMLU and WinoGrande is equivalent random guessing (all MMLU results are actually worse than the 25% random baseline). Thus, the advertised improvement on these benchmarks is most likely just random noise. Overall, the improvements are consistent, but small, so adding some confidence information would help to make the evaluation more reliable.\n- The mathematical justification in 5.5 does not seem to add much value. If I understand it correctly, it merely shows that the new parameterization has the potential of stabilizing the gradient norms, not that it actually achieves that when the parameters are learned. That would need a much more complicated analysis of the training dynamics.\n- The language models used for the empirical evaluation are all under-trained, at least according to the \"Chinchilla scaling laws\" [1]. I definitely do not think that they are under-trained to a degree that would invalidate the results, but it is still a concern. Modern language models are trained way past the point of \"optimal data size\", so it is important to test how the proposed method behaves in this regime. Note that that does not necessarily have to be done using a large and expensive language model.\n- Similarly, other slightly concerning aspect of the experimental setup is that the hyperparameters were not optimized specifically for each evaluated method. For example, could it be the case that GPAS works better just because of the chosen learning rate but models without GPAS would be equally good when trained with a slightly lower learning rate? I understand that doing a hyperparameter search for every experiment is prohibitively expensive, but it could again be done separately in a small-scale setting.\n\n---\n\n[1] [Training Compute-Optimal Large Language Models](https://arxiv.org/abs/2203.15556)", "was_revised": true, "final_justification": "I'm quite confident with keeping the score given to this paper the same, I believe that 5 (accept) is a fair assessment for the contributions of this paper.", "rebuttal": "We sincerely appreciate your constructive feedbacks and insightful comments! As for your concerns, we will address them one by one as follows.\n\n>W1. The evaluation of the downstream performance could be improved. For example, the performance of all models on MMLU and WinoGrande is equivalent random guessing (all MMLU results are actually worse than the 25% random baseline). Thus, the advertised improvement on these benchmarks is most likely just random noise. \n\nThank you for noticing. We believe the MMLU score was caused by using an older version of `lm-eval`. After updating to the latest version, here are the re-evaluated results:\n\n|Method|MMLU|\n|-|-|\n|Post-LN|22.95|\n|DeepNorm|22.95|\n|DeepNorm+GPAS|26.46|\n|Pre-LN|25.95|\n|Pre+GPAS|26.68|\n|Sandwich-LN|27.29|\n|Sandwich+GPAS|27.41|\n|Mix-LN|26.24|\n|Mix+GPAS|26.23|\n|LNS|26.62|\n|LNS+GPAS|27.81|\n\nThis is more aligned with the reported results of Mix-LN [2] and LNS [3]. \nWe summarize the reasons for near random score on MMLU and WinoG as follows:\n- MMLU is out-of-distribution relative to our pretraining corpus. Our models are trained on the C4 dataset, which is generic web corpus. This dataset is good for language modeling and learning everyday knowledge, but not ideal for knowledge/reasoning intensive tasks.\n- Our models have not undergone instruction finetuning (except for Commonsense170K SFT), which further limits their performance on tasks requiring instruction following or academic reasoning.\n\nNonetheless, this does not undermine our core claims. On Commonsense Reasoning benchmarks—which are more closely aligned with the C4 domain—GPAS still showed consistent and meaningful improvements.\n\n>Overall, the improvements are consistent, but small, so adding some confidence information would help to make the evaluation more reliable.\n\nWe fully agree with this point. We will add confidence information in our updated manuscript.\n  \n>W2. The mathematical justification in 5.5 does not seem to add much value. If I understand it correctly, it merely shows that the new parameterization has the potential of stabilizing the gradient norms, not that it actually achieves that when the parameters are learned. That would need a much more complicated analysis of the training dynamics.\n\nWe admit that the mathematical justification in Section 5.5 is preliminary. Our goal was to offer theoretical insights into how GPAS modulates gradient and activation variance. By leveraging the results from Curse of Depth [3], Theorem 1 and Lemma 1 in our manuscript already provide evidence that GPAS improves gradient propagation—even without an explicit training‐dynamics derivation. A full treatment of the training dynamics would indeed require a far more detailed analysis, which is part of our ongoing research efforts and will be systematically explored in future publications. \n\n>W3. The language models used for the empirical evaluation are all under-trained, at least according to the \"Chinchilla scaling laws\" \\[1\\]. I definitely do not think that they are under-trained to a degree that would invalidate the results, but it is still a concern. Modern language models are trained way past the point of \"optimal data size\", so it is important to test how the proposed method behaves in this regime. Note that that does not necessarily have to be done using a large and expensive language model.\n\nFirst we want to clarify that our pre-training setting closely follows the setting that are widely used in previous works [1,2,3]. We also agree that the models are under-trained by today's standards. An observation that alleviates this concern, is that the training loss gap between GPAS models and baseline models tend to stabilize after the initial 30% steps, and remain consistent gap for the rest of training. Based on this observation, it is likely that, when continued to trained on more data, the improvements of GPAS over baselines will be similar to the under-trained scenario.\n\n>W4. Similarly, other slightly concerning aspect of the experimental setup is that the hyperparameters were not optimized specifically for each evaluated method. For example, could it be the case that GPAS works better just because of the chosen learning rate but models without GPAS would be equally good when trained with a slightly lower learning rate? I understand that doing a hyperparameter search for every experiment is prohibitively expensive, but it could again be done separately in a small-scale setting.\n\n- We clarify that we followed the learning rate settings of previous works (Mix-LN and LayerNorm-Scaling) strictly and didn't specifically tune learning rate for our approach. We believe that directly using their optimal learning rate doesn't benefit our approach, but instead shows the good generalization of our approach. \n- Although we did not conduct systematic learning rate search for all model scales and architectures, we did perform a small search for the 350M models. We found that learning rate of 1e-3 tend to cause spikes in early training stages for Pre-LN, and causes the loss to diverge for 3 of the 6 baseline architectures (Post-LN, DeepNorm, Sandwich-LN). Therefore we adopted a learning rate of 5e-4 for this setting. \n\n>Q1. Is there any negative impact of training speed? I would expect not, but it is always better to ask. I think it would be better to clearly mention this in the paper.\n\nBased on our training framework and hardware, the training throughput after adding the GPAS becomes marginal as the model size scales up. The throughput percentage relative to baseline is around 97% for 71M to 350M models, and 99% for 1B to 7B models. Moreover, we believe that implementing the GPAS operation as a fused kernel will further decrease overhead, which our current implementation does not. We will add this observation to our revised version.\n\n>Q2. Since you decided to use a relatively established Commonsense170K dataset/benchmark for evaluation, why not reporting the performance on all datasets in the benchmark?\n\nWe chose the subset to align with the results of previous works (Mix-LN and LayerNorm Scaling), since they also reported these subtasks only. Another minor reason is to fit the table within the page width without using a too-small font, which could impact readability. For full disclosure, the complete scores:\n\n(**bold**: higher score between GPAS and base model; ***bold italic***: highest among all)\n|Model|MMLU|BoolQ|PIQA|SIQA|HellaSwag|WinoG|ARC-e|ARC-c|OBQA|AVG|\n|-|-|-|-|-|-|-|-|-|-|-|\n|Post-LN|22.95|37.83|52.88|32.91|26.24|49.49|27.44|19.11|11.60|31.16|\n|DeepNorm|22.95|37.83|52.61|32.91|26.20|49.64|27.23|19.28|11.60|31.14|\n|DeepNorm+GPAS|**26.46**|**53.82**|**67.95**|**33.42**|**31.43**|**50.67**|**47.22**|**22.78**|**21.20**|**39.44**|\n|Pre-LN|25.95|50.37|68.72|**32.91**|32.29|50.99|49.33|21.50|17.80|38.87|\n|Pre+GPAS|**26.68**|**59.76**|**68.99**|32.45|**33.72**|**52.25**|**49.87**|**22.01**|**22.20**|**40.88**|\n|Sandwich-LN|27.29|61.74|67.63|**33.62**|32.70|**51.07**|47.43|23.29|21.20|40.66|\n|Sandwich+GPAS|**27.41**|**61.96**|**69.10**|33.27|**34.62**|50.36|**50.97**|23.29|**22.20**|**41.46**|\n|Mix-LN|**26.24**|61.90|68.82|**33.37**|33.11|***52.88***|48.91|**24.06**|20.80|41.12|\n|Mix+GPAS|26.23|**61.93**|**70.18**|32.96|**33.52**|***52.88***|**49.79**|22.95|**21.00**|**41.27**|\n|LNS|26.62|***62.08***|69.37|***34.34***|34.68|**52.09**|51.18|23.38|20.00|41.53|\n|LNS+GPAS|***27.81***|61.59|***70.95***|32.14|***36.09***|52.01|***52.40***|***25.77***|***24.40***|***42.57***|\n\n>Q3. Is it correct to the SFT evaluation as zero-shot when the models are directly trained on training samples from the evaluated datasets?\n\nThe SFT stage is mainly for learning the instruction following capabilities in order to perform the benchmarks properly. We agree that our benchmarks are not ideally zero-shot by using in-domain training data, except for MMLU which does not have a training set. We will clarify that in the updated version.\n\n>Q4. In the experiments with predefined gates, what is the impact of this setting on the activation and gradient variance?\n\nPredefined gates introduces even slightly smaller activation variance and slightly larger gradient, but falls way short in perplexity. Results are shown in the following table:\n\n|Method|Pre-LN|Pre+GPAS|Pre+GPAS (predefined gates)|\n|-|-|-|-|\n|Layerwise Variance|0~120|0~60|0~40|\n|Layerwise Grad Norm|0.05~0.07 | 0.05~0.3|0.05~0.35|\n|Perplexity|21.35|20.35|22.46|\n\n[1] Zhao, J., Zhang, Z., Chen, B., Wang, Z., Anandkumar, A. and Tian, Y., 2024. Galore: Memory-efficient llm training by gradient low-rank projection. ICML 2024.\n\n[2] Li, P., Yin, L. and Liu, S., 2024. Mix-ln: Unleashing the power of deeper layers by combining pre-ln and post-ln. ICLR 2025.\n\n[3] Sun, W., Song, X., Li, P., Yin, L., Zheng, Y. and Liu, S., 2025. The curse of depth in large language models. arXiv preprint arXiv:2502.05795."}, {"review_id": "4IWddUGUXk", "reviewer": "Reviewer_MBu7", "summary": "This work builds on the observations that increasing the depth of models has come with diminishing returns by proposing an amendment to existing normalisation schemes. This modification, GPAS, aims at rescaling the variance of activations without affecting the scale of the gradients. The authors evaluate adding their technique to several normalisation schemes empirically, and explore the effect of GPAS on the variance of activations and the norm of the gradients, along with other ablations on the choices made in the design of GPAS. The authors also propose a theoretical analysis of how GPAS affects the norm of the gradients.", "questions": "**Questions**\n1. “Moreover, by scaling $x’_{l+1}$ , $\\alpha_l$ controls the mixing ratio between residual [...] effectively balancing information from these 2 variables” I am not sure I fully understand this sentence. How does $\\alpha_l$ intervene in the mixing ratio of the residual and the output of Attn or the FFN? My understanding of eq. 2 is that $\\alpha_l$ affects both the residual and the output of attention/FFN equally since it’s applied after they are added together (in the form of a scaling factor $1-\\text{SiLU}(\\alpha_l)$).\n2. p3: did you experiment with applying GPAS to the entire argument of LN in eq 7 and/or replacing $x_l$ with $x’_l$ in the argument of $f$ in the same equation? Since you mention experimenting to find how to best apply GPAS in the post-LN / DeepNorm cases, negative results would also be interesting to readers of the paper experimenting with architectural changes.\n3. “We believe that the combination of larger weights and stable activation variance across layers enables the network to tolerate larger gradients without crashing, which in turn speeds up convergence. “ This conclusion about GPAS not crashing would need to be explained. If anything, the previous section showed that GPAS made training less stable.\n4. Lemma 1: there are multiple weights involved at layer l: $W_Q, W_K, W_V, W_O$ and fully connected layers weights. There is also the LN weight. So what is $W_l$? The concatenation of all these weights?\n\n**Suggestions**\n1. p4: “We did not report Post + GPAS [...]” I suggest adding that information to the Tab 2 caption too so readers may get that information at a glance\n2. Tab 3: In my experience, the MMLU and WinoG results are all too close to random chance to not be within variance of said random chance. For example, when pretraining, MMLU tends to fluctuate within a few percent of 25% until it “takes off” and starts making consistent progress. As far as I am concerned this doesn’t weaken the claim about GPAS based on other evals, so it’s not an issue per se, but I want to draw attention to it in case you may want to allude to that in your discussion to anticipate readers pointing that out.\n3. Tab 4: “scaled SiLU” is commonly referred to as “SwiGLU”, I recommend using that term.\n4. p3: did you use any kind of weight decay? If so, please add it in the paper.\n5. Theorem 1: I suggest avoiding the simultaneous use of $L$ as the layer count and as a lower bound function for clarity\n6. More of a suggestion, but I heavily recommend configuring your bibliography-style to one supporting hyperlinks for your references. For example, “arXiv preprint arXiv:2407.21783” in the first reference being clickable would enhance readability.\n\n**Typos**\n1. “in LLM pertaining” (end of sec 4.1) should be “pretraining”\n2. Tab 3: For WinoG, the 2nd best performance is Pre+GPAS (52.25), not LNS (52.09)\n3. Lemma 1 would benefit from rephrasing. For example, the first sentence needs rephrasing and is potentially missing a reference to equation 2, or “Let the upper bound for this 217 norm denoted by UP($\\cdot$)” needs to be “be denoted”", "limitations": "Yes.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Summary of strengths**\n1. The proposed modification is simple and is computationally cheap, and aims to solve the important issue of suboptimal learning in deep layers of LLMs\n2. The authors discuss design ablations and negative results, which will help the community build upon this work\n3. The authors’ experiments show that at a small scale (of models and training flops), GPAS helps and the claim that GPAS helps normalising the variance’s behaviour across layers is supported experimentally. The experimental work is easy to follow.\n\n**Summary of weaknesses**\n\nGenerally, the paper has 3 kinds of weaknesses that would be essential to fix before I would consider it ready for publication:\n1. Claims: Issues with clarity of claims about how GPAS works or performs (notably, being invalidated by experiments, e.g. re: stability or not affecting gradients).\n2. Instability & gradient clipping: I am glad that the authors raise the point of stability and added Fig 4. A key aspect of training foundation models is balancing the speed of learning with stability, with the latter being increasingly harder to achieve as the model or token budget is scaled up. However, GPAS introduces instability, even at a very small scale (<= 1B params with <= 10B tokens). As the authors point out, they did not use gradient clipping, except for a model that could not recover from instability at all otherwise, even though it is a default when training LLMs. In my opinion, and you may anticipate this sentiment to be shared by many readers, it is vital that the authors reevaluate their method across the paper with gradient clipping, especially because the experiments do not include models larger than 1B parameters – which are typically stable to train – and already see stability issues. Since you mention requiring 2-3 days per 1B model, I think you could focus on Pre and Pre+GPAS so you can update Fig 4, and retrain + reevaluate the models after the rebuttal period. Again, as the authors claim in the motivation section, experimenting with foundation models is very costly, so I would be reluctant to use compute on a method introducing instabilities – in other words, you cannot leave a fix as a hypothetical.\n3. Theoretical analysis: the theoretical analysis needs significant attention due to phrasing issues and/or incorrect claims. Moreover, it needs to be more stand-alone, since at times, it seems necessary to read reference [4] (the curse of depth paper) to understand some arguments (which are stated in [4] but not so clearly in the present work).\n\n**More feedback**\n1. Minor, but missing citation of *The Unreasonable Ineffectiveness of the Deeper Layers* by Gromov et al. from early 2024 when motivating your work, which was one of the first works noticing the depth issue.\n2. p2: Giving a little bit more info about Mix-LN is important; specifically explaining how it is not just using both Pre-LN and Post-LN everywhere in the architecture (which would predate Mix-LN, Gemma 2 for example introduced that before).\n3. p3: $\\beta$ is introduced in the text but does not appear anywhere in the DeepNorm equations. A potential solution is to subscript $f$ by $\\beta$ and point to something more specific in the appendix. This is important because it’s not clear what happens to $\\beta$ with GPAS. Since the text only mentions adjusting $\\alpha$, I assume $\\beta$ is left untouched (i.e. keeps being used).\n4. p3: Eq 6 and 7 show how GPAS is used with DeepNorm. While I assume it corresponds to $\\alpha=1$ in eq. 7, I believe it would be important to either explicitly mention it near eq. 7, or to have an equation for PostLN+GPAS to avoid any confusion (and clarify in the text whether the sentence “Empirically, we found applying GPAS after LayerNorm does not bring performance gain” applies solely to DeepNorm or also to Post-LN).\n5. p4: “Details on how 121 we augment each baseline are discussed in Section 3” Section 3 does not mention LNS and LNS+GPAS. It is worth expanding on it in section 3 by giving the equations as done for the other baselines (especially since LNS+GPAS is arguably the best performing baseline in Tab 2).\n6. Lemma 1 / theorem 1: how can $x_l$ and $x’_l$ be independent when $x_l$ is a function of $x’_l$ such that $x’_l$ determines the value of $x_l$? You might suppose it for a proof that would make use of eq 2, but this supposition already rules out applying the theoretical results to GPAS since these variables are quite dependent. Reading from the curse of depth paper, their $x_l$ and $x’_l$ clearly are not “mutually independent” and they do not assume non-zero covariance, so I assume the phrasing of your assumption needs to be clarified. That being said, their work uses quite an unrealistic assumption that the output of the attention softmax is equal to $1/n I$. It could be justified by the exploding variance, but I don’t see them mentioning it (and it would change the analysis slightly since that would not be true for a finite and potentially large number of layers).\n7. Eq 8: $M$ is never introduced.\n8. “This leads to a strictly non-constant lower bound, which in turn ensures that the gradient magnitude retains sufficient variability across layers. In particular, this precludes the possibility of vanishing gradients and encourages meaningful gradient flow during backpropagation.” I may be missing something so correct me if I’m wrong, but Theorem 1’s lower bound is a lower bound on the upper bound of the gradient norm. It is not a lower bound on the gradient norm itself, so it might very well be 0 without violating theorem 1.\n9. “ Consequently, the proposed modification mitigates gradient explosion and reduces the occurrence of large loss spikes, thereby enhancing training stability” again, I do not understand the claim as this was not the case empirically, nor does the theory prove this in real world settings due to the assumptions being made as discussed in previous points.\n10. The claim about GPAS being an intervention that leaves the gradients intact, that appears to be the main claim as it is underlined on the first page and is mentioned in the abstract, is not correct as shown empirically (and as the authors want to show in the analysis since based on [4], it’s a way to show the intervention on the variance corrects undesirable behaviour observed with the norm of the gradient with respect to the first layer’s parameters). You may want to argue that unlike direct interventions on the gradients, you target a cause of problematic gradients and not just the symptoms.\n11. Theorem 1: should the $\\alpha$ argument of $M$ and $L$ be indexed by $l$ since $M, L$ are bounds on $1-\\text{SiLU}(\\alpha_l)$?\n12. Theorem 1, as opposed to the curse of depth paper, suggests that the upper bound will not scale as the number of layers increases. But the curse of depth papers argues that this is undesirable as it means that additional layers will act as the identity mapping, or in other words, that deep layers become ineffective. I tend to agree with that desiderata – by the chain rule, a model with very large $L$ should have (sufficiently) large gradients for every layer, and hence the norm of the gradient with respect to the first layer’s parameters should be increasingly large with increasing L, unless a method targeting the gradients is used (here, the authors claim GPAS “leav[es] the backward gradients intact”).", "was_revised": true, "final_justification": "**Post-rebuttal**\nThe authors will provide clarifications and fixes accordingly. I count on them to rerun experiments with gradient clipping for the camera ready version, which should not change the results significantly but is conceptually important to reflect realistic setup and address instabilities observed. I also count on them to remove claims that their method improves stability. I think it's great that they indicated the negative result for readers. I increased my score from 3 to 5.", "rebuttal": "We sincerely appreciate your very detailed review and constructive feedbacks! Due to word limit, we kindly acknowledge your following suggestions together, and conduct significant proofreading and polishing accordingly:\n- W3 (Weakness 3)\n- MF1 (More feedback 1), MF2, MF4\n- S1 (Suggestion 1), S3, S5, S6\n- Typos\n\nWe will now answer your remaining concerns.\n\n>W1. Further clarify how GPAS works or performs.\n\n- **How GPAS works**. GPAS addresses activation explosion in Pre-LN while preserving gradients. In LNS and DeepNorm, it enhances static scaling with a learnable one that adapts to training dynamics. For Sandwich-LN, it dynamically scales the shortcut branch for better blending. In Mix-LN, it enhances both Pre- and Post-LN layers.\n- **Stability**. Stability is only mildly affected in the 1B Pre-LN run due to a single warmup spike, which gradient clipping resolves (see W2). On 1B DeepNorm, GPAS prevents loss divergence, and on 7B Pre-LN (see W2), it improves convergence without instability. The 7B results will be added in Appendix A.\n- **Gradients**. The claim that GPAS \"does not affect gradients\" refers to Eq.2—its Jacobian is identity, so it does not scale backpropagated gradients, though it scales down the forward activations. We will clarify this in Section 3. The resulting increase in layerwise gradients arises because:\n    - The \"naive\" version (Eq.10, without stopgrad) empirically does not change layerwise gradient norm compared to the baseline, but reduces activations similar to Eq.2.\n    - Standard GPAS (Eq.2), avoids gradient downscaling of Eq.10—thus amplifying layerwise gradients relative to naive scaling.\n\n>W2. Instability & gradient clipping.\n\nWe agree gradient clipping is generally recommended for LLM training. However, we initially excluded it for the following reasons:\n- Follow settings of prior works (Mix-LN, LNS).\n- Stress-test GPAS under harsher conditions.\n- Gradient clipping has limited impact on small models and does not affect GPAS gains, PPL shown below:\n\n|Method|71M Pre-LN|71M Pre+GPAS|350M Pre-LN|350M Pre+GPAS|\n|-|-|-|-|-|\n|no gradclip|34.23|33.49|21.35|20.35|\n|gradclip (1.0)|33.99|33.33|20.92|19.96|\n\nWe also ran 7B pretrain with gradient clipping, confirming that GPAS scales well and remains stable. Although we couldn’t rerun 1B due to crowded server, we tested GPAS on 7B Pre-LN before the rebuttal using LR 3e-4, clipping 1.0, and 40K steps on 16B tokens. At step 40K, Pre-LN and Pre+GPAS achieve perplexities of 15.27 and 13.82, respectively. The full eval loss is shown below:\n\n|Training Step|5K|10K|15K|20K|25K|30K|35K|40K|\n|-|-|-|-|-|-|-|-|-|\n|Pre-LN (7B)|3.854|3.191|2.982|2.886|2.826|2.784|2.752|2.726|\n|Pre+GPAS (7B)|3.844|3.078|2.87|2.777|2.72|2.681|2.651|2.626|\n\n>MF3. $\\beta$ should appear in the DeepNorm equations as subscript.\n\nYes, $\\beta$ is left untouched for DeepNorm+GPAS. We've added $\\beta$ as a subscript in our updated manuscript.\n\n>MF5. Section 3 does not mention LNS and LNS+GPAS.\n\nThanks for pointing out. We'll give separate equations in Section 3. The LNS forward pass is $x_{l+1}'= x_l+f(\\mathrm{LN}(x_l)/\\sqrt{l})$, and LNS+GPAS adds Equation 2 same as Pre-LN.\n \n>MF6. Lem/Thm.1: how can $x_l$ and $x’_l$ be independent...\n\nThank you for pointing this out. We did not mean $x_l$ and $x'_l$ are mutually independent—only that both are independent of the weight matrix $W_l$. We will clarify this.\n\nYou are right that the Softmax output is not uniform—our statement was a simplification. However, in our analysis we focus on the asymptotic analysis of variance of the Softmax. This involves a summation over N terms (N denotes the number of key-value pairs), which are assumed to be independent. When N is large, by the theory of concentration of measure in high dimensions, the variance converges to that of the uniform case. Thus, the non-uniformity does not affect the result. This is also supported by Lem A.1 in the updated Curse of Depth paper.\n\n>MF7. Eq 8: $M$ is never introduced.\n\nYou are right that $M$ was never defined. In our notation, $M$ denotes an arbitrary non-zero constant. Equivalently, this means that the stated upper bound actually admits a strictly positive constant lower bound $M>0$. We will correct this in the later version.\n\n>MF8. Thm.1’s lower bound...might very well be 0 without violating Thm.1.\n\nWe believe there's a misunderstanding, and we offer the following clarification. Our result establishes a nontrivial range for the upper bound on the gradient norm; it does not assert a direct lower bound on the gradient norm itself. Indeed, the actual gradient norm can become zero after $L$ layers due to cancellations or “ups and downs” in the network. So the lower bound must be zero. However, characterizing how the upper bound behaves is still important: it shows that no matter how deep the network grows, the worst-case gradient magnitude remains confined within a fixed range and does not blow up or vanish exponentially. Moreover, no matter how large $L$ becomes, the upper bound remains a constant—giving rise to the “curse of depth” phenomenon in which deeper layers become entirely ineffective. In our paper, Thm1 can avoid this because the upper bound must be larger than a constant $M$. Hence, characterizing the behavior of this upper bound remains meaningful. Finally, note that loss spikes are directly linked to this gradient upper bound (see the discussion in the Curse of Depth paper).\n\n>MF9. \"Consequently...thereby enhancing training stability\" again, I do not understand the claim...\n\nWe want to clarify that our theoretical analysis is more of a statement about the potential of GPAS to enhance stability, and did not take training dynamics into account (mainly $\\alpha_l$ being learnable). A full treatment of training dynamics would indeed require a far more detailed analysis, which is part of our ongoing research efforts.\n\n>MF10. Claim about GPAS leaving gradients intact seems incorrect empirically.\n\nAs we have clarified in W1, \"leaving the gradient intact\" is with respect to the naive scaling case.\n\n>MF11. Thm.1: should the $\\alpha$ argument of $M$ and $L$ be indexed by $l$?\n\nThanks for the question. In our notation, $M(\\alpha)$ and $L(\\alpha)$ are uniform bounds that apply to all $\\alpha_l$ for $l \\leq L$. Since these bounds hold simultaneously across all layers, they are not indexed by $l$. We will clarify this point in the revised manuscript.\n\n>MF12. Thm.1's upper bound seems to conflict with the curse of depth paper.\n\nYour point is right, but our results do not conflict with theirs. The Curse of Depth paper shows that the gradient-norm upper bound can remain a fixed constant, which is undesirable. In Lemma 1, We do not say our gradient norm is a constant, but the upper bound of gradient norm must be larger than a constant. In contrast, we tighten the range of that upper bound: instead of a constant lower bound, we show that the lower bound itself grows with depth. At the same time, the upper bound’s growth rate is reduced, so it does not explode too rapidly. As a result, the overall band of possible gradient norms is narrower, which means that, with GPAS, we achieve better performance by avoiding both unbounded growth and constancy of the upper bound.\n\n>Q1. How does $\\alpha_l$ affect mixing ratio?\n\nThanks for the question. While GPAS is applied after adding the outputs of Attn/FFN and the shortcut within the same sub-layer, across adjacent layers it effectively scales only the shortcut branch—since the LayerNorm before Attn/FFN removes any scaling from the previous layer. In this sense, $\\alpha_l$ modulates the mixing ratio between the shortcut and the Attn/FFN outputs. We will clarify this interpretation in the revised manuscript.\n\n>Q2. Clarify DeepNorm+GPAS negative trials.\n\nWe experimented with applying GPAS to the entire argument of Equation 5, specifically to the $x_{l+1}$ in Equation 5. We did not experiment with combining both approaches. The negative trials are not included due to page limit. We are actively reproducing negative results to be put in the appendix of the revised manuscript.\n\n>Q3. Explain conclusion about GPAS not crashing.\n\nThank you for raising this point. First we clarify that GPAS does not make training less stable—standard gradient clipping fully resolves the observed issues (see W1 & W2).\nBy \"tolerate larger gradients\" we mean GPAS layers operate with larger gradients than baselines, especially in early layers. \"Larger weights\" are less sensitive to gradient noise, making them more robust. \"Stable activations\" refer to the consistent activation magnitudes across layers, as shown in Fig 3b.\n\n>Q4. Lem 1: what's $W_l$?\n\nTo clarify, by $W_l$ we refer to a standard weight matrix (exluding LN weights). In our setting, all weight matrices are treated uniformly as i.i.d. with zero mean and variance \n. Therefore, $W_l$ in this context refers to any such weight matrix e.g. $W_Q, W_K, W_V, W_O$. We will clarify this point in the revised manuscript.\n\n>S2. Clarify MMLU and WinoG results.\n\nThank you for pointing this out. The original MMLU score was due to using an outdated version of `lm-eval`. After updating, the corrected results are:\n\n|Method|MMLU|\n|-|-|\n|Post-LN|22.95|\n|DeepNorm|22.95|\n|DeepNorm+GPAS|26.46|\n|Pre-LN|25.95|\n|Pre+GPAS|26.68|\n|Sandwich-LN|27.29|\n|Sandwich+GPAS|27.41|\n|Mix-LN|26.24|\n|Mix+GPAS|26.23|\n|LNS|26.62|\n|LNS+GPAS|27.81|\n\nThis is more aligned with the reported results of Mix-LN [1] and LNS [2]. \nWe believe the low MMLU and WinoG scores were due to:\n- MMLU being out-of-distribution for our C4 pretraining corpus, which favors everyday knowledge over academic reasoning.\n- Lack of instruction finetuning, which limits performance on instruction-based benchmarks.\n\nThat said, our core claims remain intact—on commonsense reasoning tasks closer to C4, GPAS still yields consistent gains.\n\n>S4. Weight decay usage.\n\nFollowing Mix-LN and LNS, we did not use weight decay for the main experiments.\n\n[1] Mix-LN. 2025.\n[2] LNS. 2025."}, {"review_id": "5cGTbAVoWD", "reviewer": "Reviewer_XJJw", "summary": "This paper studies the exponential activation variance problem in Pre-LN Transformers. The authors propose GPAS, a simple yet effective method that scales down activations in the forward pass while maintaining gradient magnitudes intact. Experiments conducted across various model sizes (71M to 1B parameters) and architectures (Pre-LN, Sandwich-LN, Post-LN, DeepNorm, Mix-LN, and LNS) demonstrate that GPAS consistently enhances both training stability and model performance.", "questions": "- I feel that learning \\alpha in GPAS is not straightforward. How do you set your hyperparameters for learning \\alpha?\n- What is the intuition behind GPAS resulting in larger gradient norms across layers?\n- We can control the activation variance of each layer to be at a similar magnitude at initialization. Specifically, this can be achieved by adjusting the initialization standard deviation of weight matrices and using modified residual connections. For instance, some works use: x^{l+1} = 1/L f(x^{l}) + (1- 1/L) x^{l}.\n  - Would GPAS still be effective in this setting?\n  - If it remains effective, could you explain why?\n- I am confused about why \"exponential growth in activation variance across layers\" leads to deeper layers contributing less to learning. I don't clearly see the direct connection between activation variance and the abstract notion of \"contribution,\" and I think merely observing training/validation loss insufficient.\n  - Which experiments in your paper explicitly support this claim, or could you provide further justification?\n  - How do you quantify or characterize the contribution of individual layers (especially deeper layers) to the overall learning process of the model?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The proposed GPAS method is both simple to implement and broadly applicable to various transformer architectures (pre-LN, sandwich-LN, etc.)\n- The performance gain is large. The ablation study is insightful.\n\n\nWeaknesses:\n- Although acknowledged by the authors, the large gradient norm led by GPAS might cause instability and present issues in some training scenarios. In such settings, a gradient clipping algorithm might be needed.\n- The model scale up to 1B is good, but the data size is relatively small. I suggest you conduct your experiments with at least a 30B token budget.\n- In fact, the gain of GPAS is a bit bigger than I expected. To truly show the effectiveness of the GPAS technique, experiments on more datasets are needed. I strongly recommend the authors to try your GPAS on recently proposed datasets such as DCLM or fineweb, and more tokens.", "was_revised": true, "final_justification": "Based on the authors' reply, I did gain some intuition about language model pretraining. However, my concerns have not been fully addressed. Therefore, I have decided to increase my score from 3 to 4.", "rebuttal": "We sincerely appreciate your detailed review and constructive feedbacks. We will address your concerns one by one.\n\n>W1. large gradient norm led by GPAS might cause instability...\n\nThank you for mentioning this point. We would like to clarify that it is the single gradient spike in warmup stage that led to instability (which only happened in one of the runs), not the large gradient norm itself. Nevetheless, as validated by our 7B experiments, simply applying gradient clipping fixes this potential issue, which we will detail below. Moreover, for the 1B DeepNorm experiment, applying GPAS prevents loss divergence, which shows GPAS introduces stability.\n\nWe further tested GPAS on 7B Pre-LN before the rebuttal period. We used a learning rate of 3e-4, gradient clipping of 1.0, and trained for 40K steps on 16B tokens, before we stopped the training due to budget constraint. The evaluation perplexities at step 40K are 15.27 and 13.82 for Pre-LN and Pre+GPAS, which shows more significant gains from GPAS than the 1B setting. The evaluation loss across training steps are shown in the following table:\n\n|Training Step|5K|10K|15K|20K|25K|30K|35K|40K|\n|-|-|-|-|-|-|-|-|-|\n|Pre-LN (7B)|3.854|3.191|2.982|2.886|2.826|2.784|2.752|2.726|\n|Pre+GPAS (7B)|3.844|3.078|2.87|2.777|2.72|2.681|2.651|2.626|\n\n>W2. The model scale up to 1B is good, but the data size is relatively small. I suggest you conduct your experiments with at least a 30B token budget.\n\nWe agree that the data scale is relatively small by today's standards. We had to compromise due to limited time and compute budget. However, this limitation does not invalidate our results. Moreover, we further scale up the experiments to 7B parameters and found consistent gains from our proposed method (please see our response to W1). Finally, the loss gap between GPAS and baseline models tend to stabilize after initial stage of training. This means if we extend to train on more tokens, the gains of GPAS will likely remain consistent.\n\n>W3. recommend to try GPAS on recently proposed datasets and more tokens.\n\nWe understand that training on more diverse datasets as well as more tokens will further validate our proposed method. Although we cannot rerun our experiments on more datasets due to budget constraint, we believe that this limitation is minor and does not affect our claims.\n\n>Q1. How do you set your hyperparameters for learning \\\\alpha?\n\nThank you for raising this question. We set the initial value of $\\alpha_l$ to be zero, which corresponds to not scaling activations. There are no other hyperparameters regarding the GPAS operation. As for how GPAS learns $\\alpha_l$ during training, we provided a brief explanation in Section 3. We elaborate on this point as follows:\n- Within the same sub-layer, GPAS scales activations after adding Attn/FFN output and the shortcut.\n- When looked across adjacent layers, GPAS can be viewed as scaling the shortcut branch only, since the LayerNorm before Attn/FFN cancels out scaling from the previous layer.\n- By learning to scale the shortcut branch, GPAS is effectively balancing information from Attn/FFN output and the shortcut itself. We believe this is one of the reasons GPAS learns $\\alpha_l$ automatically. This intuition also corresponds to the reason why large activation variance is harmful—it makes Attn/FFN outputs overpowered by the shortcut, essentially ignoring their contributions.\n \n>Q2. Intuition behind GPAS resulting in larger gradient norms across layers?\n\nThank you for pointing this out. We've added a brief explanation in the method section, as well as in our theoretical analysis. From an experimental point of view, GPAS introduces larger gradients for the following two reasons:\n- When we apply GPAS without using stop gradient operator (Equation 10), the layer-wise gradients are of the same magnitude as the baseline model, while the activations are scaled down to a similar degree as standard GPAS with stopgrad. This means that **naively scaling down activations does not change layerwise gradient magnitude**, compared to baseline model with no scaling.\n- Unlike the naive scaling above, GPAS does not scale down backward gradients while scaling down forward activations. This means that during backpropagation, **GPAS gradients are free from the scaling factor aggregated across layers**. Although each layer only introduces a downscaling factor of 0.9~1.0, these factors aggregate across layers, and eliminating them leads to larger gradients of GPAS.\n \n>Q3.\n> - Does GPAS work with scaled initialization?\n> - Does GPAS work with  $x^{l+1} = 1/L f(x^{l}) + (1- 1/L) x^{l}$?\n\nThank you for raising this question. We will answer your questions one by one.\n- Whether GPAS works with scaled initialization:\n    - We tested Mitchell initialization with Pre-LN (proposed in OLMo [1]), and apply GPAS using Equation 2. **The final evaluation perplexities are listed in the following table**. \"Default\" refers to Pre-LN with default init, where each weight init uses the same std range:\n\n    |Method|Default|Mitchell|Mitchell+GPAS|\n    |------|------|------|------|\n    |71M|34.23|34.01|32.95 (-1.06)|\n    |350M|21.35|20.49|19.79 (-0.70)|\n\n    - Why GPAS works with Mitchell initialization: After inspecting training statistics, we found that while Mitchell init introduces lower layerwise activation variance (0 ~ 80) compared to default init (0 ~ 120), the layerwise variance is still exponentially growing. Mitchell+GPAS steers layerwise variance to be more evenly distributed (0 ~ 50), as opposed to growing monotonically with layer index.\n    - LayerNorm-Scaling is essentially a scaled initialization of Pre-LN, specifically it initializes LayerNorm weights as $1/\\sqrt l$ instead of $1$ ($l$ being the layer index). According to our main experiments (Table 2), LNS+GPAS achieves similar improvement in perplexity compared to Pre-LN.\n    - DeepNorm also uses scaled initialization in Attn/FFN weights, and GPAS provides notable performance gain.\n\n- Whether GPAS works with $x^{l+1} = 1/L f(x^{l}) + (1- 1/L) x^{l}$:\n\n    - We'll name this method \"**ScaleRes**\" for the sake of discussion, and assume $L$ is the total number of layers. We'll apply GPAS to $x^{l+1}$ the same way as Equation 2, and name it as \"**ScaleRes+GPAS**\". The table below shows result on 350M models. It seems ScaleRes does not perform as well as Pre-LN. Nevertheless, applying GPAS still offers a 0.84 decrease in evaluation perplexity.\n\n    |Method|Pre-LN|ScaleRes|ScaleRes+GPAS|\n    |---|---|---|---|\n    |Perplexity|21.35|23.22|22.38 (-0.84)|\n\n    - Why GPAS works in this setting: we believe that scaling $f(x_l)$ by $1/L$ is too aggressive, making signals from Attn/FFN even more overpowered by the shortcut branch. In this case, GPAS learns to scale down the activations and effectively scales down the shortcut branch, making Attn/FFN outputs less overpowered, and increases their contribution to learning.\n\n>Q4. Explain the connection between activation variance and the abstract notion of \"contribution\"\n>- experimental support or further justification?\n>- quantification for layerwise contribution?\n\nWe apologize for not being more elaborate when motivating our method. We've polished our writing in Section 1 and 2 to better motivate our work. For the question that you raised, we will answer them one by one.\n- **Connection between \"exponential growth\" and \"deeper layers contributing less\"**: In Pre-LN, since the input to Attn/FFN is normalized, the output of Attn/FFN is usually stable in scale (at least during initial training). However, the shortcut branch is not normalized and its variance accumulates and overpowers the signal from Attn/FFN. Thus Attn/FFN would have to learn to scale up their outputs to make meaningful transformation to the shortcut. Ultimately, this causes Attn/FFN's output to also grow with layer depth, despite their input being normalized. This means that **part of the compute is simply dedicated to learning to \"scale up\" the normalized inputs**, leading to inferior learning efficiency. The learning outcome is also inferior, because Attn/FFN still tend to have much smaller outputs than the shortcut branch, thus **their signal is still overpowered**, limiting their ability to transform hidden states.\n- **Justification for the above claim**: In early stages of our study, we quantified and visualized the above phenomenon by calculating layerwise Jacobian. We found that **deeper layers' Jacobian tend to be close to identity matrix**, which means deeper layers hardly transforms hidden states. But due to the page limit, and that previous works (Mix-LN, LNS, ...) already provide comprehensive analysis of this phenomenon, we did not include our observations in the paper. We will elaborate on this phenomenon in Section 1 to better motivate our work in the updated manuscript.\n- **Quantification for layerwise contribution**: We've added a \"Layer Importance Analysis\" subsection to Section 5. We quantify layer importance of layer $l$ as the **difference in average benchmark score after removing that layer**, which is an established practice adopted by previous works (Mix-LN, LNS, ...). This is implemented via directly feeding the output of layer $l-1$ to the input of layer $l+1$, essentially ignoring the Attn/FFN outputs of layer $l$. The results show that deeper layers of Pre+GPAS still contribute as much as previous layers, while some layers of vanilla Pre-LN are even harmful to performance. The performance delta of removing certain deep layers (1B model, 24 layers in total) are shown in the following table. A ***positive delta*** means benchmark score increases when removing the corresponding layer, which is undesirable.\n\n|Layer Index|17|18|19|20|21|22|23|24|\n|-|-|-|-|-|-|-|-|-|\n|Pre-LN|***0.42***|***0.19***|***0.02***|***0.44***|-0.69|***0.6***|***0.37***|-0.59|\n|Pre+GPAS|-1.22|-0.86|-1.19|-0.22|-0.64|-0.87|-1.99|-2.31|\n\n[1] Team OLMo et al, 2025. 2 OLMo 2 Furious. arXiv preprint arXiv:2501.00656"}, {"review_id": "3fwlFlJZWX", "reviewer": "Reviewer_aEp8", "summary": "The paper proposes Gradient-Preserving Activation Scaling (GPAS), an adaptation to Pre-LayerNorm used in LLMs. Pre-LN has an exponential growth in activation variance across layers, which affects the ability of deeper layers to learn effectively. GPAS scales activations, but keeps their gradients unchanged to avoid the vanishing gradient problem.", "questions": "### Questions:\n* Did you consider the original LLaMA model in your comparisons (e.g., Table 2 and 3)?\n* Would you expect higher gains in performance in larger/deeper model ?\n\n### Typos:\n* L184: across across \n* L247,248: correct to \"while,..they fail\"", "limitations": "yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths: \n* The paper addresses an important problem with potential high impact and is well-written \n* The experiments include several layer normalization schemes, and the evaluation shows consistent gains \n* The evaluation considers  pre-training performance (perplexity) and performance on down-stream tasks \n\n\n### Weaknesses: \n* It is not clear if the gains in performance are significant and whether they generalize to larger architectures (also mentioned by the authors in Limitations)", "was_revised": true, "final_justification": "I think the paper addresses an important problem, is well-written and shows good empirical results. I raised minor concerns regarding the generalization to larger model, and these were partially addressed in the rebuttal. I believe it is impossible to evaluate on very large models, as this requires huge computational costs. That said, I must admit that this is not my area of expertise and it is possible that I missed something.", "rebuttal": "Thank you for your constructive feedbacks and thoughtful comments! We have proofread and polished our writing following your suggestions. For your concerns, we will answer them one by one.\n\n>W1. It is not clear if the gains in performance are significant and whether they generalize to larger architectures (also mentioned by the authors in Limitations)\n\nThank you for mentioning this limitation. We further experimented on 7B parameter models for Pre-LN and Pre+GPAS before the rebuttal period, and GPAS showed even more gains in perplexity compared to smaller scales. We used a learning rate of 3e-4, gradient clipping of 1.0, and trained for 40K steps on 16B tokens, before we stopped the training due to budget constraint. The evaluation perplexities at step 40K are 15.27 and 13.82 for Pre-LN and Pre+GPAS, respectively. The evaluation loss across training steps are shown in the following table:\n\n|Training Step|5K|10K|15K|20K|25K|30K|35K|40K|\n|-|-|-|-|-|-|-|-|-|\n|Pre-LN (7B)|3.854|3.191|2.982|2.886|2.826|2.784|2.752|2.726|\n|Pre+GPAS (7B)|3.844|3.078|2.87|2.777|2.72|2.681|2.651|2.626|\n\n>Q1. Did you consider the original LLaMA model in your comparisons (e.g., Table 2 and 3)?\n\nThank you for mentioning this. Our main experiments are based on the original LLaMA architecture, specifically its Attn/FFN implementations. We did not compare with the official LLaMA model checkpoints released by Meta [1], since their scale of parameter and data is different from our setting. We have polished our writing of experiment setup in Section 4, so that readers will be more clear about the architectures and checkpoints used.\n\n>Q2. Would you expect higher gains in performance in larger/deeper model ?\n\nBased on our experiments with 7B parameter models, the gains in perplexity are indeed more significant than the 1B setting. When trained for 40K steps, Pre-LN and Pre+GPAS reached perplexities of 15.27 and 13.82, a 1.45 gain in perplexity for GPAS, which is higher than the 0.67 gain in 1B Pre-LN.\n\n[1] Hugo T., Louis M., Kevin S. et al, 2023. Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv preprint arXiv: 2307.09288"}]

Author rebuttal:

Meta-review: This paper introduces a method for faster convergence speed. The method is scaling down the intermediate activations while keeping their gradients unchanged. The method is a simple technique that can be used in combination with existing approaches and many transformers variants. The paper is well-written and the experimental results show the effectiveness of GPAS clearly.
Decision: Accept (poster)