Title: GPAS: Accelerating Convergence of LLM Pretraining via Gradient-Preserving Activation Scaling
Venue: NeurIPS 2025
Accepted: yes
Keywords: Large Language Models

Pretraining

Normalization
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.


Peer reviews:
comments: 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.

STRENGTHS AND WEAKNESSES:
### Strengths: 
* The paper addresses an important problem with potential high impact and is well-written 
* The experiments include several layer normalization schemes, and the evaluation shows consistent gains 
* The evaluation considers  pre-training performance (perplexity) and performance on down-stream tasks 


### Weaknesses: 
* 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)

QUESTIONS:
### Questions:
* Did you consider the original LLaMA model in your comparisons (e.g., Table 2 and 3)?
* Would you expect higher gains in performance in larger/deeper model ?

### Typos:
* L184: across across 
* L247,248: correct to "while,..they fail"

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-06-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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”).\n\nQUESTIONS:\n**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”\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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)\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}