How to Scale Second-Order Optimization

Thank you for your thoughtful feedback. Inspired by your comments, we provide additional results below, as well as several clarifications.

On the novelty of this work

We would like to clarify that the novelty of this work goes beyond identifying that second-order optimizers have different Chinchilla scaling laws. As you pointed out, we derived hyperparameter scaling rules for a number of second-order optimizers that prior works did not consider. The only existing work studying µP scaling for second-order optimizers is [1], which only addresses width scaling and only covers K-FAC and Shampoo. By contrast, we presented a different and more general approach, avoiding the push-through identity used in [1], which does not apply when $\xi’\neq\xi$ in Equation 8 of our paper. Our approach allows us to derive the scaling rules for a much wider range of optimizers, including Muon, Shampoo with blocking, Shampoo with grafting, and SOAP, and easily handles both width and depth scaling. Our results have high practical relevance as 1) these optimizers have been shown to significantly outperform K-FAC and vanilla Shampoo considered in [1], and 2) joint width and depth scaling is more efficient than width-only scaling at scale as demonstrated in [2,3].

Testing our scaling rules in larger-scale, full vocabulary experiments

While Figure 1 (left) shows that the scaling rules had a relatively small impact on transformers, the lack of consistent feature learning under SP in Figure (right) and the growing gap between SP and our scaling rules in Figure 2 (left) strongly suggest that the benefit of correct scaling rules will be more pronounced at larger scales. To demonstrate that, we ran experiments with even larger models using the full vocabulary version of higher quality Fineweb dataset, showing that the distinction between SP and our scaling rules indeed become more pronounced at larger scales.

Compute-optimal scaling (c.f. Figure 3)

Furthermore, we conduct full vocabulary experiments on the Fineweb dataset in the compute-optimal setting. Due to time constraints, we only compare Adam with Muon, but will include additional optimizers in the final paper. In the following tables, we find that our observations that second-order optimizers 1) reduce Chinchilla-optimal tokens per parameter by about 2x, and 2) lead to about 1.5x - 2x compute efficiency gain relative to Adam continue to hold in this more standard experiment setup. Here, we also extended the model size up to 200M parameters. The optimal token per parameter is estimated using the same fit procedure in Section 4.1.

Please also refer to our response to Reviewer xRrA subtitled "Larger-scale, full vocabulary experiments" to see our experiments regarding stable optimal learning rate.

Transformer’s robustness to learning rate

Our statement that the transformer architecture is overall highly robust to suboptimal learning rates is directly based on empirical evidence in Figure 1 (left), where the impact of the correct scaling rules is relatively small, as you also pointed out. Furthermore, our experiments in Figure 6 Appendix G show that MLPs are much less robust to suboptimal learning rates, unlike transformers, even with second-order optimizers, directly showing that the transformer architecture indeed leads to improved robustness.

Despite being derived from MLPs, our scaling rules apply to transformers as the underlying principles of µP do not depend on specific architecture details and have been shown to apply to transformers both in theory and in practice [3]. We note that it is standard practice in the µP literature to perform theoretical derivations with MLPs for simplicity [3,4,5]. We will clarify this in the next version. Finally, the consistent feature learning strength across scales in Figure 1 (right) provides a direct empirical verification that our scaling rules are correct.

Significance of Random Initialization

We follow prior works on µP in assuming the weights are randomly initialized. This choice both reflects practice and allows one to prove existence of well-defined infinite limits using the central limit theorem and the law of large numbers [4,5,6] which underlie our derivation for the scale of the gradients, activations, and their updates in Section 3.1. We note that recent work shows that the random initialization assumption is not necessary, and can be replaced by conditions on the spectral norms of the weights at initialization [6].

Why asymptotic arguments apply

While the aspect ratio $A$ does not grow to infinity, asymptotic arguments in µP only require the width (or depth in the case of depth-scaling) to grow to infinity in order to appeal to the central limit theorem and the law of large numbers [4,5].

Choice of target Loss in Figure 3 To highlight the efficiency multipliers across different target losses, we extend results in Figure 3 into a more detailed table. Naive refers to using the same data scaling as is optimal for Adam, while ours indicate using corrected scaling we suggested that reduce the number of tokens per parameter for Muon and Shampoo.

Wall-clock speed ups relative to Adam on Openwebtext:

Wall-clock speed ups relative to Adam on Chess:

Exponents and Constants of Figure 3

In Figure 3, we fit a power law for loss

$L(C) = \left( \frac{C}{C_0} \right) ^{-b} + L_0,$

where $C$ is the amount of training compute. We fix $L_0$ across the three optimizers so that the values of the other parameters ($C_0$ and $b$) are meaningfully comparable. We then show the fitted coefficients:

Openwebtext:

Chess:

We see that only the constant factor improved significantly (by ~2x) while the exponent changed less than 10%. This suggests that second-order optimizers primarily change the constant but not the scaling exponent. We will discuss this important finding in the updated version.

Why our batch size scaling experiments are valuable

Our results on batch size scaling are valuable as they reveal that existing theoretical models for batch size scaling based on SDEs do not allow effective hyperparameter transfer in practice. In the paper, we fully acknowledge that a better theoretical understanding of batch size scaling is needed. We believe our thorough evaluation and honest presentation should be viewed as a strength rather than a weakness, as acknowledged by Reviewer K9rj.

Improving paper organization

Thank you for your suggestions in improving the organization of our paper. We separately submitted the main paper and the appendix following the NeurIPS submission guidelines, which prevented the link from working properly. Per your suggestion, we will include the figures referenced in Section 3.4 to the main text using the additional page in the camera-ready revision to improve readability. We deferred those figures to the appendix due to space constraints. We apologize for the formatting error that led to an extra section in the appendix (section E) and will correct it in the revision. No content is intended for section E.

References:

1. Ishikawa and Karakida. "On the Parameterization of Second-Order Optimization Effective towards the Infinite Width." International Conference on Learning Representations (ICLR), 2024
2. Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361, 2020.
3. Dey, Nolan, et al. "Don't be lazy: CompleteP enables compute-efficient deep transformers." arXiv preprint arXiv:2505.01618, 2025.
4. Yang et al. "Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer." Advances in Neural Information Processing Systems (NeurIPS), 2021.
5. Yang and Littwin. "Tensor programs ivb: Adaptive optimization in the infinite-width limit." International Conference on Learning Representations, 2023
6. Yang, Greg, James B. Simon, and Jeremy Bernstein. "A spectral condition for feature learning." arXiv preprint arXiv:2310.17813, 2023.
7. Hoffmann, Jordan, et al. "Training compute-optimal large language models." Advances in Neural Information Processing Systems (NeurIPS), 2022.

Thank you for your supportive review and constructive feedback! We appreciate that you recognize the comprehensiveness of our scaling rules analysis and its demonstrated practical benefits. We respond to your comments and questions below.

Theory for batch size scaling

Thank you for recognizing our honesty in stating the limitations of our batch size scaling analysis as a strength. We believe improving our understanding of batch size scaling, both theoretically and empirically, is an exciting next step, and hope our work highlights its importance.

Overhead of Muon and Shampoo

You are correct that we did not report the overhead of second-order optimizers in terms of FLOPs. This was due to the complexity involved in quantifying the exact number of operations required by routines such as eigendecomposition and matrix roots, which are highly dependent on their implementations. Furthermore, these operations tend to be more sequential in nature and less compute-bound compared to matrix multiplication in the forward and backward pass, making FLOPs no longer an accurate indicator of real runtime overhead. Therefore, we instead directly report wall-clock time in Figure 4 for a more realistic estimate of efficiency gains. Figure 4 shows that 1) the preconditioner overhead indeed dramatically reduces the wall-clock speedup for Shampoo as model size increases if batch size stays constant, and 2) by scaling batch size linearly with model width, both Muon and Shampoo consistency outperform Adam in runtime. To see the concrete runtime savings, please refer to our response subtitled "Faster wall-clock time" to Reviewer zik3.

We therefore conclude that second-order optimizers can consistently outperform Adam provided we sufficiently co-scale batch size with model size. While we agree first-order methods such as Adam do not require large batch size in principle, practitioners often train at the critical batch size (the maximum batch size that does not significantly increase the loss for a given token budget) purely for the benefit of parallelization [1,2]. We will clarify these takeaways in the update paper.

Response to Questions

• $\xi$ refers to the sample that we take the gradient step on to update the weights, while $\xi’$ refers to an arbitrary point in the input space. As defined in Line 166, $x$ is the shorthand for $x(\xi)$ and $x’$ is the shorthand for $x(\xi’)$, where $x$ denotes the features for a sample in a particular hidden layer. The arbitrariness of $\xi’$ means the µP conditions should be satisfied for any point in the domain. We will make this clearer in the next version.

• As described in Appendix F, the base model configuration is depth 3 and width 128. We did not experiment with other values. We chose these small values to demonstrate optimal learning rate is stable across a wide range of widths and depths.

• In the plot in question, the circular markers represent individual data points, while the dashed line corresponds to the linear regression line in log-log space (power-law fit). We’ll revise the figure caption to make this clearer in the final version.

References:

1. Chowdhery, A. et al. “PaLM: Scaling Language Modeling with Pathways.” Journal of Machine Learning Research, 2023
2. Meta AI Llama Team. “The Llama 3 Herd of Models.” arXiv preprint arXiv:2407.21783, 2024

Thank you for your detailed review and thoughtful suggestions. Inspired by your comments, we provide a range of new experimental results and clarifications.

On applicability to the transformer architecture

We would like to clarify that all experiments in the main text are conducted on transformers that include standard components such as RMSNorms, embeddings, and softmax attention. Only the additional experiments in Appendix G use MLPs. We also note that it is standard practice in the µP literature to perform theoretical derivations with MLPs for simplicity [1,2,3], as the underlying principles do not depend on specific architecture details and have been shown to apply to transformers both in theory and in practice [1]. We will clarify this point in the updated paper.

Analysis beyond the first gradient step

While we only explicitly analyzed the µP conditions for the first gradient step, it is a well-established result in the µP literature that the same argument straightforwardly extends to any step $t$ during training via induction, so long as $t$ is constant with respect to model size [1,2,3]. Indeed, as illustrated in the right panel of Figure 1, the difference in RMS (Root Mean Square) of the feature updates remains notably more consistent under µP scaling than standard parameterization (SP) throughout the training process up to $10^4$ steps.

Overhead of preconditioners at scale

As shown in Section 4.3, it’s crucial to scale the batch size appropriately, ideally proportional to the model width (or $P^{1/2}$ where $P$ is the number of parameters), as the model size increases to so that the overhead of second-order optimizers like Muon and Shampoo is only ever a constant fraction of the training compute. In order for this scaling not to negate the benefits of preconditioning, we need the critical batch size to scale at least as fast as $P^{1/2}$. Fortunately, our empirical result in Figure 4 (left) suggests that the critical batch size indeed scales at a rate close to this. Under this scaling, we found both Muon and Shampoo consistently outperform Adam in runtime (Figure 4, right plot). Beyond these encouraging results, whether second-order optimizers truly can outperform Adam at much larger model sizes must be ultimately determined empirically, which we believe is an exciting and important future direction.

Larger-scale, full vocabulary experiments

We appreciate your concerns regarding the use of character-level tokenization and experiment scale. To address this, we conducted both the µP and scaling experiments using a full vocabulary tokenizer on the Fineweb dataset with larger models.

While we regret that we cannot include images of the results here, we provide the key outcomes in tabular format below. Notably, the conclusions from the character-level experiments continue to hold under this more realistic setup, reinforcing the robustness and generality of our findings despite the initial small-scale configuration.

Stable optimal learning rate (c.f. Figure 1)

Based on the setup in Appendix E, this configuration differs by using the full Fineweb vocabulary, processing a total of 100 million tokens, including a 10 million token warm-up phase. The largest model in this setting has 3 transformer blocks with width up to 2048 (280M parameters in total). While only the result of Muon is presented here, we will include all optimizers in the camera-ready version. At this scale, the performance gap between µP and SP becomes more evident.

Width scaling for Muon: our results show that µP maintains a consistent optimal learning rate, reflected as a vertical stripe at zero, while SP demonstrates a decreasing trend in its stable learning rate range.

SP:

µP:

As in Figure 1, we visualize changes in the RMS of the hidden features (∆h). Additionally, we include the RMS change of the logits (∆f). From the table, we show that ∆h and ∆f only remain stable across different widths under µP, highlighting scale-dependent inconsistencies and further underscoring the limitations of SP.

Features and logits change for Muon:

Depth scaling for Muon: similar to the conclusion of width scaling, the optimal learning rate starts to shift as the model gets deeper.

SP:

Depth-Scaled:

Compute-optimal scaling (c.f. Figure 3)

Furthermore, we conduct full vocabulary experiments on the Fineweb dataset in the compute-optimal setting. Due to time constraints, we only compare Adam with Muon, but will include additional optimizers in the final paper. In the following tables, we find that our observations that second-order optimizers 1) reduce Chinchilla-optimal tokens per parameter by about 2x, and 2) lead to about 1.5x - 2x compute efficiency gain relative to Adam continue to hold in this more standard experiment setup. Here, we also extended the model size up to 200M parameters. The optimal token per parameter is estimated using the same fit procedure in Section 4.1.

Paper organization Your point regarding the focus of the presentation is well-taken. To improve readability, we have deferred the experimental details to the appendix so the main text can focus on a complete, self-contained guide for practitioners interested in scaling second-order optimizers with µP. Due to space constraints, we deferred the detailed derivation for Muon scaling rules based on its update rule in Appendix B. Similarly, we apply and extend the SDE formulation by Malladi et al. for batch size scaling in Appendix D and compare its performance against our alternative, empirically determined scaling rule in Figure 2. We will use the additional page in the camera-ready version to clarify these details and findings.

Thanks again for your review. We believe the content of this response, which required a significant effort, will strengthen the paper. We will additionally correct all formatting and typos. We would really appreciate it if you would consider raising your score in light of our response.

References

1. Yang et al. "Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer." Advances in Neural Information Processing Systems (NeurIPS), 2021.
2. Yang and Littwin. "Tensor programs ivb: Adaptive optimization in the infinite-width limit." International Conference on Learning Representations, 2023
3. Everett et al. "Scaling exponents across parameterizations and optimizers." arXiv preprint arXiv:2407.05872, 2024
4. Liu, Jingyuan, et al. "Muon is scalable for LLM training." arXiv preprint arXiv:2502.16982, 2025.

We hope our additional experiments and clarifications have helped address your concerns. Please let us know if have further questions regarding the paper or our rebuttal. We are happy to continue the discussion.

Thank you for the supportive review! Inspired by your comments, we provide additional results and clarifications here and will include them in the updated paper. We hope you will consider raising your score in light of our response.

Faster wall-clock time

We would like to clarify that we did show in Figure 3 that Muon and Shampoo are significantly more efficient compared to Adam in terms of faster wall-clock time for a given target loss, provided that batch size is scaled together with model size. Based on the results in Figure 3, we estimate the following relative wall-clock time saving relative to Adam for different final losses achieved by Adam:

Wall-clock time savings for Openwebtext dataset:

Wall-clock time savings for Chess dataset:

As you pointed out, efficient implementation of second-order optimizers for distributed training, such as distributing the preconditioning operation across devices, is crucial for maximizing their performance gains in practice. As we did not implement distributed versions of second-order optimizers and performed redundant computations across devices, we expect their efficiency gains in wall-clock times can be made even larger than what we estimated, making the case for second-order optimizers even more favorable. Regarding stale preconditioners, we update the preconditioner every 10 steps for Shampoo variants. We will include these details in the updated paper.

Extending µP to cover hyperparameters specific to 2nd order optimizers

We did consider how to scale hyperparameters specific to second-order optimizers with respect to width and depth, including the damping factor in all optimizers, and the learning rate and damping factor for grafting. We applied these scaling rules in our experiments and listed them in Table 1 with derivations in Appendix B. In Appendix B, we also derived how the learning rate should scale with block size under blocking.

Regarding how to scale the number of Newton-Schultz iterations, from the µP perspective, keeping the number of iterations constant is sufficient for a well-defined infinite width limit. However, this is not the only choice, as scaling the number of iterations to infinity also leads to a well-defined limit (converging to the true matrix sign via SVD), though it leads to a much higher overhead. For blocking, both keeping the block size constant and equal to the full matrix size produce well-defined infinite-width limits (Adam-like limit vs full Shampoo limit), and it is unclear which limit is preferable a priori. We agree that exploring scaling rules for these hyperparameters is an interesting future direction, and will include this discussion in the updated paper.

Efficiency considerations

We use techniques such as blocking and periodic updates to both speed up our experiments as well as to demonstrate that our proposed scaling rules work for practical instantiations of second-order optimizers. We view this as a strength of our paper rather than a weakness, as we ultimately care about scaling rules appropriate for practical instantiations of second-order optimizers rather than their idealized implementations, which generally require different scaling rules.

Theoretical explanation for constant momentum

We believe better theoretically understanding how to optimally scale momentum is an exciting and important future direction. Our current hypothesis for why it’s better to keep the momentum constant rather than following the SDE scaling rules is that there is a tension between keeping the first vs second moment of momentum buffers invariant to batch size. Specifically, as shown in [1], a large value of $\beta$ is necessary for keeping the variance of the momentum buffers low regardless of the batch size, while a decreasing value of $\beta$ is required for keeping its first moment (mean) invariant as we scale batch size. We will elaborate on this discussion in the revision.

References:

1. Malladi et al. "On the SDEs and scaling rules for adaptive gradient algorithms." Advances in Neural Information Processing Systems (Neurips), 2022

Dear reviewer, we wanted to check in to see if the additional results and clarifications have helped addressed your concerns. We would really appreciate it if you consider raising your score. Many thanks for your supportive review.

We really appreciate your supportive review and thoughtful suggestions!

Hyperparameter tuning details

We described our hyperparameter tuning procedure in Appendix F. For each optimizer, we determine the base learning rate by choosing its optimal value on a small model of 3 transformer blocks and an embedding dimension of 128 (Figure 1) and scaling to larger models via our scaling rules. For momentum and batch size, our results in Section 3.4 suggest $\beta_1=0.9$ and $\beta_2=0.95$ is near optimal for a wide range of batch sizes. We thus used these values for all other experiments unless stated otherwise.

Whether SP specifically fails for second-order methods

SP fails not only for second-order methods, but for Adam as well, as demonstrated extensively in prior works (e.g. [1]). We provide the following additional experiment to demonstrate that SP Adam indeed has a shifting optimal learning rate while µP Adam has a stable optimal learning rate:

Loss under SP

Loss under µP

Chinchilla scaling rule with SP

We do not believe finding Chinchilla scaling rules for SP would yield useful insights. Under SP, larger models eventually diverge after the first gradient step [1], trivializing the resulting scaling laws. In the Chinchilla paper [2], they avoided this issue by manually adjusting learning rates across model sizes, presumably tuned via additional experiments, suggesting that directly using SP would lead to worse performance.

Novelty

We do believe our work has important novelty. In addition to deriving scaling rules for many second-order optimizers that prior works did not consider (Muon, Grafted Shampoo, SOAP etc.) and demonstrating their empirical effectiveness, our work contains at least two other novel and impactful contributions: 1) a simple procedure for deriving µP and depth scaling for a general second-order optimizer; and 2) the observation that second-order optimizers significantly change the Chinchila scaling laws and that language models should generally be trained with 1.5x~2x fewer tokens per parameter if using Muon or Shampoo for compute-efficient scaling. We believe these findings, while ultimately simple to state, are key to fully realizing the benefits of second-order optimizers.

Experiment scale

We would like to clarify that our experiments are scaled up to 5e18 FLOPs, as shown in Figure 3, much larger than 1e9 FLOPs that you suggested. We acknowledge that validating our results with even larger-scale experiments, such as with billion-parameter models, is an exciting future direction.

Seed

Thank you for being cautious about uncertainties. We will add error bars to the figures in the camera-ready version for clarity. In our GPT-2 setup, averaging over multiple seeds won't change the conclusions as the seed-to-seed standard deviation in the final loss is small, compared to the difference across learning rates and optimizers. To demonstrate this, here we show the mean final losses and the standard deviations in the brackets for five different random seeds.

Losses for the base width model:

Precision

We agree that numerical precision is particularly important for second-order optimizers, especially for matrix inversion and eigendecompositions. Exploring scaling rules for second-order optimizers that minimize numerical errors in these key operations, such as via unit-scaled µP [3], is an exciting future direction.

References:

1. Yang, Greg, et al. "Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer." Advances in Neural Information Processing Systems (NeurIPS), 2021.
2. Hoffmann, Jordan, et al. "Training compute-optimal large language models." Advances in Neural Information Processing Systems (NeurIPS), 2022.
3. Blake, Charlie, et al. "u-$\mu$ P: The Unit-Scaled Maximal Update Parametrization." International Conference on Machine Learning (ICML) Workshop, 2024.