cs.LGFeb 24, 2026

Extending μμP: Spectral Conditions for Feature Learning Across Optimizers

Authors: Akshita Gupta, Marieme Ngom, Sam Foreman, Venkatram Vishwanath

Organizations: Purdue University · Argonne National Laboratory

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Several variations of adaptive first-order and second-order optimization methods have been proposed to accelerate and scale the training of large language models. The performance of these optimization routines is highly sensitive to the choice of hyperparameters (HPs), which are computationally expensive to tune for large-scale models. Maximal update parameterization (μ(μP)) is a set of scaling rules which aims to make the optimal HPs independent of the model size, thereby allowing the HPs tuned on a smaller (computationally cheaper) model to be transferred to train a larger, target model. Despite promising results for SGD and Adam, deriving μμP for other optimizers is challenging because the underlying tensor programming approach is difficult to grasp. Building on recent work that introduced spectral conditions as an alternative to tensor programs, we propose a novel framework to derive μμP for a broader class of optimizers, including AdamW, ADOPT, LAMB, Sophia, Shampoo and Muon. We implement our μμP derivations on multiple benchmark models and demonstrate zero-shot learning rate transfer across increasing model width for the above optimizers. Further, we provide empirical insights into depth-scaling parameterization for these optimizers.

Figures & tables

Figure 1: μ\mu P for Sophia (trained on Llama2) - Coordinate check plots for the word embedding and output logits layers (left); Zero-shot learning rate transfer across increasing model width (middle); Decreasing training loss with increasing model width (right).
OptimizerΨ(⋅)\mathbf{\Psi(\cdot)}
AdamW / ADOPTm^(t)v^(t)+ϵ+λWl(t)\dfrac{\hat{\mathbf{m}}^{(t)}}{\sqrt{\hat{\mathbf{v}}^{(t)}}+\epsilon}+\lambda\mathbf{W}_{l}^{(t)}
Sophiaclip ⁣(m(t)max⁡{γh(t),ϵ},1)+λWl(t)\text{clip}\!\left(\dfrac{\mathbf{m}^{(t)}}{\max\{\gamma\mathbf{h}^{(t)},\epsilon\}},1\right)+\lambda\mathbf{W}_{l}^{(t)}
LAMBϕ(∣∣Wl(t)∣∣F)∣∣rl(t)+λWl(t)∣∣F(rl(t)+λWl(t))\dfrac{\phi(||\mathbf{W}_{l}^{(t)}||_{\text{F}})}{||\mathbf{r}_{l}^{(t)}+\lambda\mathbf{W}_{l}^{(t)}||_{\text{F}}}\left(\mathbf{r}_{l}^{(t)}+\lambda\mathbf{W}_{l}^{(t)}\right)
Shampoo(L(t))−1/4  ∇Wl(t)L  (R(t))−1/4(\mathbf{L}^{(t)})^{-1/4}\;\nabla_{\mathbf{W}_{l}^{(t)}}\mathcal{L}\;(\mathbf{R}^{(t)})^{-1/4}
Muonnlnl−1Ol(t)\sqrt{\frac{n_{l}}{n_{l-1}}}\mathbf{O}_{l}^{(t)}
Table 1: Values of Ψ(⋅)\Psi(\cdot) for different optimizers. Auxiliary variables are defined in Section 4 and Appendix A .
 Input WeightsOutput WeightsHidden Weights
Init. Var.1\;\color[rgb]{1,0,0}{(\tfrac{1}{n_{l-1}})}1\;\color[rgb]{1,0,0}{(\tfrac{1}{n^{2}_{l-1}})}1\;\color[rgb]{1,0,0}{(\tfrac{1}{n_{l-1}})}
Multiplier\tfrac{1}{\sqrt{n_{l-1}}}\;\color[rgb]{1,0,0}{(1)}\tfrac{1}{n_{l-1}}\;\color[rgb]{1,0,0}{(1)}\tfrac{1}{\sqrt{n_{l-1}}}\;\color[rgb]{1,0,0}{(1)}
AdamW / ADOPT1\;\color[rgb]{1,0,0}{(1)}\tfrac{1}{n_{l-1}}\;\color[rgb]{1,0,0}{(\tfrac{1}{n_{l-1}})}\tfrac{1}{n_{l-1}}\;\color[rgb]{1,0,0}{(\tfrac{1}{n_{l-1}})}
Sophia LR1\;\color[rgb]{1,0,0}{(-)}\tfrac{1}{n_{l-1}}\;\color[rgb]{1,0,0}{(-)}\tfrac{1}{n_{l-1}}\;\color[rgb]{1,0,0}{(-)}
LAMB LR1\;\color[rgb]{1,0,0}{(-)}1\;\color[rgb]{1,0,0}{(-)}1\;\color[rgb]{1,0,0}{(-)}
Shampoo LR\sqrt{n_{l}}\;\color[rgb]{1,0,0}{(-)}\tfrac{1}{\sqrt{n_{l-1}}}\;\color[rgb]{1,0,0}{(-)}\sqrt{\tfrac{n_{l}}{n_{l-1}}}\;\color[rgb]{1,0,0}{(-)}
Table 2: Comparison of μ\mu P from spectral conditions (black) vs. tensor programs ( Yang et al., 2021 , Table 3) (red).
Figure 2: (NanoGPT) Mean validation loss for increasing model width and different learning rates across four optimizers: ADOPT (top left), LAMB (top right), Sophia (bottom left), and Shampoo (bottom right). The plots demonstrate zero-shot learning rate transfer under μ\mu P (Table 2 ).
Figure 3: (Llama2) Validation loss for increasing model width and different learning rates across three optimizers: AdamW (left), ADOPT (middle), and LAMB (right). The plots demonstrate zero-shot learning rate transfer under μ\mu P (Table 2 ).
Figure 4: (Llama2 model) AdamW optimizer - Coordinate check plots under standard parameterization (top left) and under μ\mu P (bottom left) for the word embedding and output logits layers; Decreasing training loss with increasing model width under μ\mu P (right).
Appendix figures & tables30 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: μ\mu P for Muon (trained on Llama2) - Coordinate check plots for the word embedding and output logits layers (left); Zero-shot learning rate transfer across increasing model width (middle); Decreasing training loss with increasing model width (right).
ArchitectureNanoGPT Karpathy (2022)
Width128 (scaled to 2048)
Depth8
Number of heads2
Total parameters1.59 M (scaled to 403 M)
DatasetTiny Shakespeare
Vocab size65
Appendix
Table 3: Hyperparameter values and training settings to test μ\mu P on NanoGPT model.
ArchitectureNanoGPT Karpathy (2022)
Width256
Depth2 (scaled to 64)
Total parameters1.6 M (scaled to 50.56 M)
DatasetTiny Shakespeare
Vocab size65
Tokens per iteration8192
Appendix
Table 4: Hyperparameter values and training settings to test depth-scaling parameterization on NanoGPT model.
Figure 6: Coordinate check plots for AdamW under standard parameterization (top row), μ\mu P (middle row); depth scaling (bottom row) for NanoGPT model.
LR / Width12825651210242048
2×10−12\times 10^{-1}2.541111952.547703192.501325852.535593832.45719266
2×10−22\times 10^{-2}2.570098962.565837072.579006512.533859172.51431378
2×10−32\times 10^{-3}2.634747662.60228072.646793372.634496612.55710355
2×10−42\times 10^{-4}3.388270543.55441573.388969983.449416643.44561863
2×10−52\times 10^{-5}4.092213474.088714284.052577974.088373034.08405908
Appendix
Table 5: Mean validation loss for increasing model width and different learning rates for AdamW on NanoGPT model. The minimum loss for each width is highlighted in green.
LR / Depth248163264
2×10−12\times 10^{-1}2.535259172.551927652.535109442.503575562.512949632.53008548
5×10−25\times 10^{-2}2.527007982.494226772.503349862.294282362.451760292.36860998
2×10−22\times 10^{-2}2.556829772.521766662.565835632.304228622.455001122.5650301
2×10−32\times 10^{-3}2.597457812.630784752.602283162.615881362.640656632.65051214
2×10−42\times 10^{-4}3.413961253.416778333.554415543.458015043.432854893.47577778
2×10−52\times 10^{-5}4.092979594.059707964.088714284.081131464.067128344.10902596
Appendix
Table 6: Mean validation loss for increasing model depth and different learning rates for AdamW on NanoGPT model. The minimum loss for each depth is highlighted in green.
Figure 7: Coordinate check plots for ADOPT optimizer under SP (top row); μ\mu P (middle row); depth scaling (bottom row) for NanoGPT model.
LR / Width12825651210242048
2×10−12\times 10^{-1}2.551201342.546164042.541780792.55242962.54457998
7×10−27\times 10^{-2}2.485604762.443169752.370871232.507335342.50883015
2×10−22\times 10^{-2}2.431756972.588474512.570063752.543236972.53191725
2×10−32\times 10^{-3}2.630169312.60735522.656817442.661189562.55337548
2×10−42\times 10^{-4}3.5284043.490652323.490652323.427891333.43255997
2×10−52\times 10^{-5}4.091835984.088323754.05216984.088065944.08391444
Appendix
Table 7: Mean validation loss for increasing model width and different learning rates for ADOPT on NanoGPT model. The minimum loss for each width is highlighted in green.
LR / Depth248163264
2×10−12\times 10^{-1}2.561293682.514524382.547889872.514560782.522719222.55469418
9×10−29\times 10^{-2}2.486955722.474775632.531248012.481453022.506874722.54724765
2×10−22\times 10^{-2}2.567184132.504190292.588472762.444479542.549960692.52524622
2×10−32\times 10^{-3}2.679927982.629497132.60735522.604336182.617539882.6286815
2×10−42\times 10^{-4}3.410525963.465389573.567573943.478564423.436080223.56190586
2×10−52\times 10^{-5}4.092677594.059293914.088323754.080744434.066752594.10877307
Appendix
Table 8: Mean validation loss for increasing model depth and different learning rates for ADOPT on NanoGPT model. The minimum loss for each depth is highlighted in green.
Figure 8: Coordinate check plots for Sophia optimizer under SP (top row); μ\mu P (middle row); depth scaling (bottom row) for NanoGPT model.
LR / Width12825651210242048
2×10−12\times 10^{-1}3.09693982.571441172.568752612.625730362.57240287
2×10−22\times 10^{-2}2.274506092.278308472.316326382.533479051.98427689
2×10−32\times 10^{-3}2.54565972.614300572.55943022.548694852.65462987
2×10−42\times 10^{-4}3.354090133.546143693.360898023.358623823.36431138
2×10−52\times 10^{-5}4.087663814.088591264.060697564.088117124.08371623
Appendix
Table 9: Mean validation loss for increasing model width and different learning rates for Sophia on NanoGPT model. The minimum loss for each width is highlighted in green.
LR / Depth248163264
2×10−12\times 10^{-1}2.52135033.010813163.226491053.348552153.243104463.12229093
2×10−22\times 10^{-2}2.47170482.272322892.247361142.474757512.460612461.93401444
2×10−32\times 10^{-3}2.541031922.581362332.610355932.6106122.450684152.55488427
2×10−42\times 10^{-4}3.408877213.527654253.545875633.406694813.339977423.47574107
2×10−52\times 10^{-5}4.092673144.065767614.088591264.081404054.0665524.10874732
Appendix
Table 10: Mean validation loss for increasing model depth and different learning rates for Sophia on NanoGPT model. The minimum loss for each depth is highlighted in green.
Figure 9: Coordinate check plots for LAMB optimizer under SP (top row); μ\mu P (middle row); depth scaling (bottom row) for NanoGPT model.
LR / Width12825651210242048
2×10−12\times 10^{-1}3.33069152.919924742.756582342.847240922.84511503
2×10−22\times 10^{-2}2.274277692.553309442.532503452.506948952.51612274
2×10−32\times 10^{-3}2.467624192.427230282.475710552.491525492.46575729
2×10−42\times 10^{-4}3.696729743.709617143.668777783.23704293.37923479
2×10−52\times 10^{-5}4.169295314.16947544.16841034.16745794.16771809
Appendix
Table 11: Mean validation loss for increasing model width and different learning rates for LAMB on NanoGPT model. The minimum loss for each width is highlighted in green.
LR / Depth248163264
2×10−12\times 10^{-1}2.765341362.859497792.881156213.269327323.240937873.097018
2×10−22\times 10^{-2}2.508583072.511643892.553555012.339676622.483084442.11406271
7×10−37\times 10^{-3}2.451171722.466918152.502312342.456914352.486299362.45780365
2×10−32\times 10^{-3}2.504836242.542846842.427231232.432919032.432621722.42000318
2×10−42\times 10^{-4}3.64417063.793676063.709633433.573737383.614025753.42223287
2×10−52\times 10^{-5}4.169815064.16914864.16947544.169329334.168173954.16773876
Appendix
Table 12: Mean validation loss for increasing model depth and different learning rates for LAMB on NanoGPT model. The minimum loss for each depth is highlighted in green.
Figure 10: Coordinate check plots for Shampoo optimizer under SP (top row); μ\mu P (middle row); depth scaling (bottom row) for NanoGPT model.
LR / Width12825651210242048
1×10−21\times 10^{-2}2.644320653.008410063.267297113.395126824.17380921
9×10−39\times 10^{-3}2.66503312.895494543.207410653.453219183.41602135
5×10−35\times 10^{-3}2.631228052.676930433.302152793.322653533.36052688
3×10−33\times 10^{-3}2.673031572.851034013.371943873.469758433.49201838
1×10−31\times 10^{-3}2.905831652.979756283.610351173.572247353.72281067
Appendix
Table 13: Mean validation loss for increasing model width and different learning rates for Shampoo on NanoGPT model. The minimum loss for each width is highlighted in green.
LR / Depth248163264
3×10−23\times 10^{-2}2.834688192.946374813.38116053.273786233.325345833.31375853
1×10−21\times 10^{-2}2.639170892.63838142.668230143.22788083.248644353.20088768
7×10−37\times 10^{-3}2.641900222.610072532.739912273.128639383.209857783.37485345
5×10−35\times 10^{-3}2.777039452.722951572.727944612.936291223.254318083.37258538
3×10−33\times 10^{-3}2.71435422.973687892.853654863.320306623.279885373.40830247
Appendix
Table 14: Mean validation loss for increasing model depth and different learning rates for Shampoo on NanoGPT model. The minimum loss for each depth is highlighted in green.
Figure 11: Mean validation loss for increasing model depth and different learning rates for AdamW (left) and ADOPT (right) on NanoGPT model.
Figure 12: Mean validation loss for increasing model depth and different learning rates for LAMB (left) and Sophia (right) on NanoGPT model.
Figure 13: Mean validation loss for increasing model depth and different learning rates for Shampoo on NanoGPT model.
ArchitectureLlama 2
Width256 (scaled to 2048)
Depth16
Number of attention heads32
Total parameters154M (scaled to 1.38 B)
DatasetWikitext-103
Sequence length4096
Appendix
Table 15: Hyperparameter values and training settings to test μ\mu P on Llama2 model.
Figure 14: Coordinate check plots for AdamW optimizer under SP (top row); μ\mu P (middle row); depth scaling (bottom row) for Llama2 model.
LR / Width12825651210242048
5×10−15\times 10^{-1}4.554914.026763.812513.735733.79477
3×10−13\times 10^{-1}4.249783.902423.832523.894843.75046
1×10−11\times 10^{-1}4.486964.213144.052654.021013.95419
5×10−25\times 10^{-2}4.704214.43534.397534.341694.31635
1×10−11\times 10^{-1}5.577955.562845.561735.557715.55774
Appendix
Table 16: Validation loss for increasing model width and different learning rates for AdamW on Llama2 model. The minimum loss for each width is highlighted in green.
Figure 15: Coordinate check plots for ADOPT optimizer under SP (top row); μ\mu P (middle row); depth scaling (bottom row) for Llama2 model.
LR / Width12825651210242048
5×10−15\times 10^{-1}4.390334.020073.839323.777323.76814
3×10−13\times 10^{-1}4.117893.855363.725523.678023.66973
2×10−12\times 10^{-1}4.237653.879493.782423.800163.78846
1×10−11\times 10^{-1}4.323354.075973.99123.916543.95519
7×10−27\times 10^{-2}4.438194.225744.135654.068524.0683
5×10−25\times 10^{-2}4.641214.380964.315824.221864.21248
Appendix
Table 17: Validation loss for increasing model width and different learning rates for ADOPT on Llama2 model. The minimum loss for each width is highlighted in green.
Figure 16: Coordinate check plots for LAMB optimizer under SP (top row); μ\mu P (middle row); depth scaling (bottom row) for Llama2 model.
LR / Width12825651210242048
3×10−23\times 10^{-2}7.184526.350596.03846.529666.13429
1×10−21\times 10^{-2}5.588785.56385.560495.791746.01439
5×10−35\times 10^{-3}6.574766.604546.663986.980937.0471
1×10−31\times 10^{-3}10.2511210.2399810.2257510.2119910.19599
5×10−45\times 10^{-4}10.3299710.3277610.3239810.3206210.31677
Appendix
Table 18: Validation loss for increasing model width and different learning rates for LAMB on Llama2 model. The minimum loss for each width is highlighted in green.
Figure 17: Coordinate check plots for Sophia optimizer under SP (top row); μ\mu P (middle row); depth scaling (bottom row) for Llama2 model.
LR / Width12825651210242048
5×10−15\times 10^{-1}7.194036.995766.689926.603766.31375
3×10−13\times 10^{-1}6.176045.908265.806945.67385.71962
1×10−11\times 10^{-1}4.141223.836543.759263.674193.62891
7×10−27\times 10^{-2}4.427584.317024.057563.935613.94189
5×10−25\times 10^{-2}4.766324.510224.413584.344524.30914
3×10−23\times 10^{-2}4.823054.795924.730674.674734.74689
Appendix
Table 19: Validation loss for increasing model width and different learning rates for Sophia on Llama2 model. The minimum loss for each width is highlighted in green.
More in this topic

Hyperparameter Transfer

Hyperparameter transfer studies how to reuse training settings across model sizes, architectures, tasks, and training horizons.

See all 39 papers

Explore similar work

CardsList

  1. Spectral Condition for μμP under Width-Depth Scaling

    Feb 28, 2026Chenyu Zheng, Rongzhen Wang, Xinyu Zhang +1Deep Learning OptimizationHyperparameter Transfer

    Gaoling School of AI, Renmin University of China · ByteDance Seed

  2. Deriving Hyperparameter Scaling Laws via Modern Optimization Theory

    Mar 16, 2026Egor Shulgin, Dimitri von Rütte, Tianyue H. Zhang +3LMO MethodsScaling Laws

    King Abdullah University of Science and Technology (KAUST), Thuwal, Saudi Arabia · ETH Zürich, Zürich, Switzerland · Mila, Quebec AI Institute & Université de Montréal, Quebec, Canada +1

  3. Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate

    May 20, 2026Dayal Singh Kalra, Maissam BarkeshliLanguage Model Scaling LawsHyperparameter Transfer

    Department of Computer Science, University of Maryland, College Park · Department of Physics, University of Maryland, College Park · Joint Quantum Institute, University of Maryland, College Park +1