Noise-Adaptive Layerwise Learning Rates: Accelerating Geometry-Aware Optimization for Deep Neural Network Training (2510.14009v1)

Abstract: Geometry-aware optimization algorithms, such as Muon, have achieved remarkable success in training deep neural networks (DNNs). These methods leverage the underlying geometry of DNNs by selecting appropriate norms for different layers and updating parameters via norm-constrained linear minimization oracles (LMOs). However, even within a group of layers associated with the same norm, the local curvature can be heterogeneous across layers and vary dynamically over the course of training. For example, recent work shows that sharpness varies substantially across transformer layers and throughout training, yet standard geometry-aware optimizers impose fixed learning rates to layers within the same group, which may be inefficient for DNN training. In this paper, we introduce a noise-adaptive layerwise learning rate scheme on top of geometry-aware optimization algorithms and substantially accelerate DNN training compared to methods that use fixed learning rates within each group. Our method estimates gradient variance in the dual norm induced by the chosen LMO on the fly, and uses it to assign time-varying noise-adaptive layerwise learning rates within each group. We provide a theoretical analysis showing that our algorithm achieves a sharp convergence rate. Empirical results on transformer architectures such as LLaMA and GPT demonstrate that our approach achieves faster convergence than state-of-the-art optimizers.

• The paper presents a noise-adaptive optimizer (LANTON) that estimates per-layer gradient noise and scales learning rates to achieve faster convergence.
• It employs LMO-based updates combined with dynamic noise estimation to improve sample efficiency and robustness in transformer architectures.
• Empirical results confirm LANTON's superior performance across varying base learning rates and training scales on both small- and large-scale models.

Noise-Adaptive Layerwise Learning Rates for Geometry-Aware Deep Network Optimization

Motivation: Heterogeneous Gradient Noise in Deep Networks

The paper addresses a critical inefficiency in geometry-aware optimizers for deep neural networks (DNNs): the use of fixed learning rates within parameter groups, despite substantial heterogeneity in gradient noise across layers and over the course of training. Empirical evidence demonstrates that, even within groups of layers sharing the same norm or matrix shape, the stochastic gradient noise can vary by orders of magnitude and evolve dynamically during optimization.

Figure 1: The stochastic gradient noise is heterogeneous across groups and layers in transformers. The first subfigure shows that average gradient noise in hidden layers varies across parameter groups defined by matrix shape and evolves over training. The last three subfigures illustrate that, within each layer group, the gradient noise varies substantially across layers.

This heterogeneity is especially pronounced in transformer architectures, where query-key (QK), value-output (VO), and MLP layers exhibit distinct and time-varying noise profiles. Standard geometry-aware optimizers such as Muon, D-Muon, and Scion, which leverage norm-constrained linear minimization oracles (LMOs) for matrix parameters, do not adapt learning rates at the layer level, leading to suboptimal convergence and sample efficiency.

LANTON: Noise-Adaptive Layerwise Learning Rate Scaling

The paper introduces LANTON (LAyer-wise Noise-adaptive learning raTe scaling with Operator Norms), a geometry-aware optimization algorithm that dynamically estimates the gradient variance in the dual norm induced by the LMO for each layer and uses this to adapt the learning rate per layer throughout training. The key algorithmic steps are:

1. Per-layer Gradient and Momentum Computation: For each layer $\ell$, compute the stochastic gradient $G_t^\ell$ and maintain a momentum buffer $B_t^\ell$.
2. LMO-based Update Direction: Obtain the update direction $O_t^\ell = \mathrm{LMO}(B_t^\ell)$, where the LMO and norm are chosen according to the parameter type (e.g., RMS$\to$RMS for hidden layers, $1\to\infty$ for embeddings).
3. Noise Estimation: Maintain a moving average $H_t^\ell$ of the squared dual-norm difference between consecutive gradients (or two independent stochastic gradients for theoretical analysis), serving as an estimator of the local gradient noise.
4. Adaptive Learning Rate Scaling: Compute a scaling factor $$\alpha_t^\ell = \alpha / \sqrt{\alpha^2 + H_t^\ell}$$, and set the effective learning rate for layer $\ell$ as $$\eta_t^\ell = \eta_t \sqrt{\alpha_t^\ell / \alpha_t^m}$$, where $\alpha_t^m$ is the maximum scaling factor within the group.
5. Parameter Update: Update $X_{t+1}^\ell = X_t^\ell - \eta_t^\ell O_t^\ell$.

This approach ensures that layers with higher estimated noise receive smaller learning rates, while layers with lower noise can exploit larger steps, improving both convergence speed and sample efficiency.

Theoretical Analysis: Improved Convergence under Layerwise Noise

The paper provides a non-asymptotic convergence analysis for LANTON under a layerwise $L$-smoothness assumption and almost-surely bounded (above and below) layerwise gradient noise in the dual norm. The main result establishes that LANTON achieves a convergence rate of

$$\tilde{O}\left(\frac{1}{\sqrt{T}} + \sqrt{\frac{\sum_\ell \bar{\sigma}_\ell}{T^{1/4}}}\right)$$

where $\bar{\sigma}_\ell$ is an upper bound on the gradient noise for layer $\ell$. This rate improves upon prior geometry-aware optimizers that use a uniform noise bound $\bar{\sigma}_{\max}$, as $$\sqrt{\sum_\ell \bar{\sigma}_\ell} \ll \sum_\ell \bar{\sigma}_{\max}$$ in practice. The analysis leverages high-probability martingale concentration and norm equivalence arguments to control the adaptive learning rate scaling and the bias introduced by momentum.

Empirical Results: Accelerated Training and Robustness

Small-Scale LLM Pretraining

On GPT2-small and GPT2-medium models trained from scratch on OpenWebText-100k, LANTON consistently outperforms AdamW, Muon, MARS, SCION, and D-Muon in both training and validation loss, with faster convergence from the earliest iterations.

Figure 2: Training/validation loss on Openwebtext-100k datasets.

Large-Scale Transformer Pretraining

For LLaMA-1.1B on C4 and LLaMA-0.5B on Minipile, LANTON achieves lower training and validation loss than all baselines, including block-wise and layer-wise learning rate methods (LAMB, BW-AdamW). Notably, LANTON adapts to the evolving noise landscape and outperforms block-wise tuned optimizers that use fixed ratios.

Figure 3: Training/validation loss on C4 and Minipile datasets.

Sample Efficiency and Token Budget

LANTON demonstrates higher sample efficiency: to reach a comparable loss, D-Muon requires approximately $1.5\times$ more training tokens than LANTON.

Figure 4: (a) LANTON outperforms layer-wise/block-wise learning rate methods on C4. (b) LANTON achieves higher sample efficiency than D-Muon.

Robustness to Base Learning Rate

LANTON exhibits robustness to the choice of base learning rate, maintaining superior convergence across a range of $\eta_{\max}$ values.

Figure 5: LANTON is robust to the choices of base learning rates.

Wall-Clock Efficiency

LANTON achieves faster reduction in loss per wall-clock time compared to D-Muon and other baselines, indicating practical efficiency gains in real-world training scenarios.

Figure 6: Training and validation loss vs. wall-clock time.

Implementation Considerations

• Norm and LMO Selection: The method requires careful assignment of operator norms and corresponding LMOs for each parameter group. For hidden layers, the RMS$\to$RMS norm and nuclear norm dual are used, with the LMO implemented via Newton-Schulz iterations. For embedding/LM head layers, the $1\to\infty$ norm and signum LMO are used.
• Noise Estimation: In practice, the noise estimator uses the difference between consecutive gradients to avoid doubling the batch size, which is a practical compromise with negligible empirical impact.
• Integration with Existing Optimizers: LANTON can be implemented as a wrapper around D-Muon or other geometry-aware optimizers, requiring only the addition of per-layer noise tracking and adaptive scaling.
• Computational Overhead: The additional cost is dominated by the per-layer noise estimation and LMO computation, which is amortized in large-scale distributed training.

Implications and Future Directions

The results demonstrate that geometry-aware optimization can be substantially improved by incorporating noise-adaptive layerwise learning rates, especially in architectures with pronounced heterogeneity in gradient noise. This approach reduces the need for manual learning rate tuning and block-wise heuristics, and is theoretically justified to yield better noise dependence in convergence rates.

Potential future directions include:

• Extending the approach to even larger models and distributed training regimes.
• Investigating the interaction between noise-adaptive scaling and other forms of adaptivity (e.g., adaptive moment estimation, gradient clipping).
• Exploring the impact of noise-adaptive geometry-aware optimization on generalization and downstream transfer.

Conclusion

LANTON provides a principled and empirically validated solution to the inefficiency of fixed learning rates in geometry-aware deep network optimization. By dynamically adapting learning rates at the layer level based on dual-norm gradient noise, it achieves faster convergence, higher sample efficiency, and robustness to hyperparameter choices. Theoretical analysis confirms improved noise dependence, and extensive experiments on transformer architectures substantiate the practical benefits. This work establishes noise-adaptive layerwise scaling as a key component for scalable and effective training of deep neural networks.