Title: On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling
Venue: NeurIPS 2025
Accepted: yes
Keywords: Width Scaling

Feature Learning

Large Learning Rates

Maximal Update Parameterization

Standard Parameterization

Initialization

Scaling Theory

Tensor Programs
Abstract: Scaling limits, such as infinite-width limits, serve as promising theoretical tools to study large-scale models. However, it is widely believed that existing infinite-width theory does not faithfully explain the behavior of practical networks, especially those trained in *standard parameterization* (SP) meaning He initialization with a global learning rate. For instance, existing theory for SP predicts instability at large learning rates and vanishing feature learning at stable ones. In practice, however, optimal learning rates decay slower than theoretically predicted and networks exhibit both stable training and non-trivial feature learning, even at very large widths. Here, we show that this discrepancy is not fully explained by finite-width phenomena.

Instead, we find a resolution through a finer-grained analysis of the regime previously considered unstable and therefore uninteresting. In particular, we show that, under the cross-entropy (CE) loss, the unstable regime comprises two distinct sub-regimes: a catastrophically unstable regime and a more benign controlled divergence regime, where logits diverge but gradients and activations remain stable. Moreover, under large learning rates at the edge of the controlled divergence regime, there exists a well-defined infinite width limit where features continue to evolve in all the hidden layers. In experiments across optimizers, architectures, and data modalities, we validate that neural networks operate in this controlled divergence regime under CE loss but not under MSE loss. Our empirical evidence suggests that width-scaling considerations are surprisingly useful for predicting empirically maximal stable learning rate exponents which provide useful guidance on optimal learning rate exponents. Finally, our analysis clarifies the effectiveness and limitations of recently proposed layerwise learning rate scalings for standard initialization.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the discrepancy between theoretically predicted optimal (stable) learning rates (LRs) and empirically observed ones as model scale increases under standard parameterization (SP). The authors particularly highlight the role of cross-entropy (CE) loss, which in contrast to MSE loss allows stable training under large LRs even with divergent logits. Interestingly, as shown in the paper, neither long-training dynamics nor the catapult mechanism arguments are able to explain the stability of large LRs in SP. Instead, thanks to the properties of the CE loss derivative, training dynamics remains stable with LR exponent up to $-1/2$, which aligns well with practice. Experiments confirm the theoretical predictions and help further elucidate the intuition behind different choices of architectural elements, optimizers, and parameterizations.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n\nThis paper provides a solid foundation for the empirically observed LR scaling laws, bridging the gap between theory and practice and enabling better understanding and more accurate prediction of optimal LR values at scale.\nThe idea to particularly consider the role of the loss function seems novel and interesting, which is a solid contribution of this work.\nIt is indeed counterintuitive to anticipate stable training dynamics and even feature learning when model outputs (logits) enter the divergence regime.\nThe authors did a great job of explaining the practical benefits of cross-entropy loss, and I specifically like the part demonstrating that MSE loss can possess similar benign properties under $\\mu$P, which makes it both theoretically attractive and practically usable.\n\nThe authors clearly position their work in relation to prior art, explicitly focusing on missing parts of the puzzle and suggesting possible explanations.\n\nThe theoretical results are supported by extensive practical experiments and detailed discussions.\n\nWriting is clear and consistent. \nThe paper is well-written, although it is still relatively hard to follow by a nonspecialist in deep learning scaling theory.\n\n# Weaknesses\n\nSince I am not an expert in deep learning scaling theory, I cannot fully assess the possible technical limitations of this work.\nFor me, the main weakness of this work is the unclear practical implications of the proposed theory, which has a descriptive rather than predictive nature.\nI would suggest that the authors pay more attention to this aspect in a future revision.\n\nQUESTIONS:\nPlease refer to Weaknesses.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates why standard parameterization (SP) with He initialization remains effective at large learning rates, despite infinite-width theory predicting instability. The authors demonstrate that cross-entropy loss fundamentally changes the stability landscape compared to MSE loss, enabling a regime where logits diverge but training remains stable. They provide theoretical analysis and validate their predictions empirically across small tasks on MLPs and Transformers.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n- The paper addresses a fundamental puzzle in deep learning theory - why standard parameterization works despite theoretical predictions of instability. \n- The key insight about CE loss enabling a controlled divergence regime is novel and provides a principled explanation for empirical observations.\n- Empirical validation across multiple architectures (MLPs, Transformers), optimizers (SGD, Adam), and datasets strengthens the theoretical claims.\n- The connection between width-scaling theory and practical learning rate selection cloud be valuable for practitioners.\n\n### Weaknesses\n\n- I found the paper a bit hard to read and missed a clear \"golden thread\".\n- Proposition 1 (line 208) is marked as \"informal\" but is central to the paper's claims. The formal statement in Appendix C.3 should be in the main text.\n- Numerical precision issues when logits diverge are mentioned only briefly (line 354). This is a critical practical concern that undermines the applicability of the controlled divergence regime.\n- Section 5.2 on SP-full-align feels disconnected from the main narrative and the claim about \"breaking learning rate transfer on image datasets\" (line 319) needs more thorough investigation.\n- While the paper explains why SP works, it doesn't lead to better training methods or actionable insights for practitioners. The contribution is primarily explanatory rather than enabling new capabilities.\n- In general, I found the scope of the paper a bit unclear and it lacks a strong empirical validation that the width-scaling considerations provide good predictions of the maximal stable learning rate and that the maximal stable learning rate often dominates the optimal learning rate.\n\nQUESTIONS:\n- How does numerical precision affect the controlled divergence regime in practice?\n- What are the practical implications beyond explaining why SP works?\n- What happens at very large widths or training times?\n\nLIMITATIONS:\nThe authors briefly mention numerical considerations but do not adequately address how floating point limitations affect the practical applicability of their theoretical results.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors try to answer why standard parameterization remains stable and effective at large learning rates, despite the theoretical predictions. While the catapult effect and edge of stability theory does not fully explain this phenomenon, the authors find a reason from the choice of cross-entropy (CE) loss as a loss function. The authors insist that CE loss allows stable training under large learning rates even when the output logits diverge.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1.\tThe paper is easy to follow even for those who are not familiar with the subject. The informal explanations of the propositions and theorems as well as the proof sketches highly enhance the readability and help the readers understand the intuition behind the theories.\n2.\tThe paper has strong theoretical contributions. The focus on the cross-entropy loss seems very interesting to me.\n3.\tThe experimental results seem to support the theoretical results as well. The experiments seem to be extensive and include examples on large models like GPT.\n\nWeaknesses\n\n1.\tI am not very much fond of the idea of inducing the maximal stable learning rate from an optimal learning rate. The authors briefly explain that Proposition 2 suggests that the idea is true since it is the only setting under which feature learning is preserved at large width in all hidden layers. I only partially agree with the authors’ opinion, and I think authors need to provide a more rigorous explanation to confirm this. The result is practically useful, though.\n2.\tIn Page 30, it says that the authors train a single epoch to prevent confounding from multi-epoch overfitting effects. I personally cannot find a particular reason why this is necessary, since feature learning often occurs after a certain period of training time, especially when the data is complex and grokking occurs.\n\nQUESTIONS:\n1.\tIs inducing the maximal stable learning rate from an optimal learning rate a frequently exercised custom or wisdom in this area?\n2.\tPage 30 mentions that the authors train a single epoch to prevent confounding from multi-epoch overfitting effects. Could you explain more on why this is necessary for the experiments?\n3.\tOne question for my sheer curiosity. It is known that a cross-entropy loss leads to a max-margin solution. What would be the relationship between the learning rate and a converged solution? Beyond a maximal stable learning rate, would the parameter not converge to a max-margin solution?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}