cs.LGOct 9, 2026

On the Depth Transfer of Neural Optimizers via L-Smoothness

Authors: Maximilian Wong, Chenyu Xu, Jiajin Li, Yiping Lu

Organizations: University of Chicago · Peking University · Sauder School of Business University of British Columbia · Beijing International Center for Mathematical Research Peking University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Scaling deep learning models efficiently requires predictable performance transfer across compute regimes. Establishing reliable hyperparameter transfer allows optimal training dynamics and structural choices to be discovered on low-cost proxy models, avoiding catastrophic failures during massive-scale training. Existing depth-transfer analyses are architecture-specific and heuristic: hyperparameter transfer is only justified through toy blocks, and both rely on infinite-depth dynamics, a specific initialization, and feature-learning criteria. Building on LL-smoothness framework of Xu et al. (2026), we understand hyperparameter transfer as how do the Lipschitz and LL-smoothness constants of the loss scale with the depth KK? We show that inverse-depth residual branch scaling α=Θ(1/K)α=Θ(1/K) is sufficient to make both constants depth-independent at fixed width, uniformly over the bounded parameter domain and for arbitrary p→qp\to q geometries. Because the bounds depend on each residual branch only through its Lipschitz and smoothness constants, the argument applies to general residual networks (including branches of block depth ≥2\geq 2 such as transformer and ResNet blocks) without assumptions on initialization or correlations between layer weights. Since the smoothness constant controls the curvature penalty in the descent inequality, our bounds motivate a learning-rate scale without an additional depth-dependent curvature correction for updates bounded in the prescribed norm. This provides a geometric explanation of inverse-depth residual scaling and its application to the operator-norm geometries associated with AdamW, Muon, and MOGA. Learning-rate sweeps on GPT-2 (depths 2--128) and CIFAR-100 ResNets (8--32 residual blocks) confirm that Muon and MOGA exhibit stable learning-rate transfer under these prescriptions, with a common optimal rate across a 64-fold increase in depth.

Figures & tables

Figure 1: Final-epoch training loss against learning rate for residual MLPs of varying depth trained on CIFAR-10 with AdamW, Muon, and MOGA. Each branch contains two linear layers, and width is fixed across depths. Under standard unscaled residual branches (left), the best learning rate varies with depth. Our inverse-depth scaling, ( α=Kbase/K\alpha=K_{\mathrm{base}}/K ) with ( Kbase=2K_{\mathrm{base}}=2 ) (right), show that optimal learning rate is aligned across depths.
Figure 2: Learning-rate transfer across GPT-2 depths at fixed width. Panels show MOGA with p∈{1.5,2,2.5}p\in\{1.5,2,2.5\} , Muon, and AdamW. Colors and markers indicate transformer depth. Stars mark the lowest training loss among the tested learning rates at each depth.
Figure 3: Learning-rate transfer on CIFAR-100 with depth-scaled ResNets. Panels show MOGA with p∈{1.5,2,2.5}p\in\{1.5,2,2.5\} . Curves report final-epoch training cross-entropy for ResNet-18, ResNet-34, and ResNet-66, corresponding to K∈{8,16,32}K\in\{8,16,32\} residual blocks. Stars indicate the lowest loss among the tested learning rates for each model.
Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Learning-rate transfer across GPT-2 depths at fixed width. Panels show MOGA with p∈{1.5,2,2.5}p\in\{1.5,2,2.5\} , Muon, and AdamW. Colors and markers indicate transformer depth. Stars mark the lowest validation loss among the tested learning rates at each depth.
Figure 5: Learning-rate transfer across ResNet depths. Final validation loss is plotted against the peak learning rate, with each curve corresponding to a different depth. Colors and markers indicate ResNet depth. Stars mark the lowest validation loss among the tested learning rates at each depth.
More in this topic

Deep Learning Optimization

Deep learning optimization develops algorithms and training strategies that help neural networks converge, generalize, and use compute efficiently.

See all 205 papers

Explore similar work

CardsList
Mar 10, 2026

On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfer

A central question in modern deep learning is how to design optimizers whose behavior remains stable as the network width ww increases. We address this question by interpreting several widely used neural-network optimizers, including AdamW and Muon, as instances of steepest descent under matrix operator norms. This perspective links optimizer geometry with the Lipschitz structure of the network forward map, and enables width-independent control of both Lipschitz and smoothness constants. However, steepest-descent rules induced by standard p→qp \to q operator norms lack layerwise composability and therefore cannot provide width-independent bounds in deep architectures. We overcome this limitation by introducing a family of mean-normalized operator norms, denoted \pmean→\qmean\pmean \to \qmean, that admit layerwise composability, yield width-independent smoothness bounds, and give rise to practical optimizers such as rescaled AdamW, row normalization, and column normalization. The resulting learning rate width-aware scaling rules recover μμP scaling~\cite{yang2021tensor} as a special case and provide a principled mechanism for cross-width learning-rate transfer across a broad class of optimizers. We further show that Muon can suffer an O(w)\mathcal{O}(\sqrt{w}) worst-case growth in the smoothness constant, whereas a new family of row-normalized optimizers we propose achieves width-independent smoothness guarantees. Based on the observations, we propose MOGA (Matrix Operator Geometry Aware), a width-aware optimizer based only on row/column-wise normalization that enables stable learning-rate transfer across model widths. Large-scale pre-training on GPT-2 and LLaMA shows that MOGA, especially with row normalization, is competitive with Muon while being notably faster in large-token and low-loss regimes.

Ruihan Xu, Jiajin Li, Yiping Lu

University of Chicago · Sauder School of Business, University of British Columbia · Industrial Engineering & Management Science, Northwestern University

Neural Network OptimizationHyperparameter TransferMatrix Optimization
AskPDF
Jun 16, 2026

Loop-μμP: Residual Scaling and Hyperparameter Transfer in Looped Transformers

Looped Transformers have drawn growing attention because looping improves multi-step reasoning and increases the executed depth of a model without adding parameters. They achieve this by reusing a stack of LL layers for NN loops. However, increasing the loop count NN not only destabilizes training but also shifts the best hyperparameters, so they must be searched again for every loop count. Inspired by Depth-μμP (Yang et al., 2024), which multiplies each residual branch by L−1/2L^{-1/2} for depth LL so that a learning rate tuned on a shallow untied network transfers to a deep one, we study how the residual branches and the learning rate should scale with the loop count NN. We find that loops cannot simply be treated as extra depth, because the loops share weights. At initialization, distinct layers add nearly uncorrelated increments to the residual stream, so their sum grows as L\sqrt{L}, whereas a shared layer adds positively correlated increments at every loop, so their sum grows linearly in NN. For a single layer looped NN times, we show in the infinite-width limit that multiplying the residual branch by 1/N1/N and keeping the learning rate independent of NN is the only parameterization with both a stable forward pass and maximal residual-stream updates. We then generalize the analysis to a stack of LL distinct layers looped NN times to obtain a parameterization we call Loop-μμP, in which the loop factor 1/N1/N combines with the depth factor L−1/2L^{-1/2}. Empirically, we find in language modeling that the base learning rate tuned on the smallest model stays best under Loop-μμP for loop counts up to 88 and executed depths up to 384384 layers. These results make Loop-μμP a practical recipe for scaling looped Transformers.

Shaowen Wang, Bingrui Li, Tansheng Zhu +4

Tsinghua University · ByteDance Seed · M-A-P

Residual LearningHyperparameter TransferRecurrent TransformersNeural Network Training DynamicsTransformer Training Dynamics
AskPDF
May 20, 2026

Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate

Hyperparameter transfer allows extrapolating optimal optimization hyperparameters from small to large scales, making it critical for training large language models (LLMs). This is done either by fitting a scaling law to the hyperparameters or by a judicious choice of parameterization, such as Maximal Update (μμP), that renders optimal hyperparameters approximately scale invariant. In this paper, we first develop a framework to quantify hyperparameter transfer through three metrics: (1) the quality of the scaling law fit, (2) the robustness to extrapolation errors, and (3) the asymptotic loss penalty due to choice of parameterization. Next, we investigate through a comprehensive series of ablations why μμP appears to offer high-quality learning rate transfer relative to standard parameterization (SP), as existing theory is inadequate. We find that the overwhelming benefit of μμP relative to SP when training with AdamW arises simply from maximizing the learning rate of the embedding layer. In SP, the embedding layer learning rate acts as a bottleneck that induces training instabilities; increasing it by a factor of width to match μμP dramatically smooths out training while improving hyperparameter transfer. We also find that weight decay improves the scaling law fits, while, in the fixed token-per-parameter setting, it hurts the robustness of the extrapolation.

Dayal Singh Kalra, Maissam Barkeshli

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

Language Model Scaling LawsHyperparameter TransferLLM Training
AskPDF