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

TL;DR

LANTON introduces a noise-adaptive layerwise learning-rate scheme on top of geometry-aware optimization to address heterogeneous and evolving gradient noise across layers. By tracking gradient variance in the dual norm induced by the chosen LMO and rescaling per-layer updates accordingly, it achieves improved convergence rates $\tilde{O}(1/\sqrt{T} + \sqrt{\sum_{\ell}\bar{\sigma}_{\ell}}/T^{1/4})$ and faster transformer training. The method maintains compatibility with existing geometry-aware frameworks (Muon/D-Muon) while applying distinct LMOs per layer type (e.g., UV$^{\top}$ for hidden layers, Signum for embeddings/LM heads, RMS normalization for LayerNorm). Empirically, LANTON demonstrates ~1.5x speedups over D-Muon in transformer pretraining and improved sample efficiency on GPT2 and LLaMA setups, with modest overhead and robust performance across learning-rate choices.

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.

Figures (6)

• 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.
• Figure 2: Training/validation loss on Openwebtext-100k datasets.
• Figure 3: Training/validation loss on C4 and Minipile datasets.
• Figure 4: Training/validation loss on C4 datasets. (a) Comparison with algorithms using layer-wise/block-wise learning rates. (b) LANTON shows higher sample efficiency than D-Muon.
• Figure 5: LANTON is robust to the choices of base learning rates.

Theorems & Definitions (17)

• Theorem 5.3
• Lemma 5.3
• Lemma 5.3
• proof : Proof sketch of \ref{['thm:muon']}
• Lemma A.1: Azuma-Hoeffding inequality
• Lemma A.2: liu2023nearlyoptimal
• proof : Proof of \ref{['lem:MDS']}
• Lemma A.3: Equivalence of norms
• Remark A.4
• Lemma B.0