Where Does Warm-Up Come From? Adaptive Scheduling for Norm-Constrained Optimizers

TL;DR

This work provides a theoretical basis for learning-rate warm-up in norm-constrained optimizers by introducing a generalized smoothness model where local curvature scales with the suboptimality gap. It derives an adaptive warm-up schedule that naturally transitions to decay, and implements a practical LR scheduler that determines warm-up duration automatically using only standard hyperparameters. Empirically, the adaptive warm-up matches or outperforms hand-tuned schedules across Muon, Lion, and normSGD on large-language-model pretraining. The approach demonstrates robustness to target-loss choice and scales across model sizes, batch regimes, and optimizer geometries, offering a principled path to hyperparameter-free warm-up in practice.

Abstract

We study adaptive learning rate scheduling for norm-constrained optimizers (e.g., Muon and Lion). We introduce a generalized smoothness assumption under which local curvature decreases with the suboptimality gap and empirically verify that this behavior holds along optimization trajectories. Under this assumption, we establish convergence guarantees under an appropriate choice of learning rate, for which warm-up followed by decay arises naturally from the proof rather than being imposed heuristically. Building on this theory, we develop a practical learning rate scheduler that relies only on standard hyperparameters and adapts the warm-up duration automatically at the beginning of training. We evaluate this method on large language model pretraining with LLaMA architectures and show that our adaptive warm-up selection consistently outperforms or at least matches the best manually tuned warm-up schedules across all considered setups, without additional hyperparameter search. Our source code is available at https://github.com/brain-lab-research/llm-baselines/tree/warmup

Figures (6)

• Figure 1: Empirical smoothness ratio $\mathcal{K}^t$ versus suboptimality gap $\Delta^t$ for Lion on large-scale pretraining. The trajectory is well-fitted by a quadratic dependence, indicating $K_\rho > 0$. All hyperparameter setup is provided in Appendix \ref{['appendix:hyper_ass']}.
• Figure 3: Validation loss on LLaMA 124M model with $\mathrm{bs} = 32$ (top) and $256$ (bottom) as a function of manually selected warmup length (solid) vs. adaptive warmup (dashed). Across optimizers, the adaptive method outperforms or at least matches the best manually tuned value without any grid search.
• Figure 4: Example of learning rate schedule produced by Algorithm \ref{['alg:lr_scheduler']} for the Lion optimizer on FineWeb.
• Figure 5: Validation loss on LLaMA 210M model with $\mathrm{bs} = 32$ (top) and $256$ (bottom) as a function of manually selected warmup length (solid) vs. adaptive warmup (dashed). Across optimizers, the adaptive method outperforms or at least matches the best manually tuned value without search.
• Figure 6: Ablation over the target loss $f^\star$ for LLaMA 124M pretraining with $\mathrm{bs}=32$ using the Muon optimizer. The adaptive scheduler remains stable across a broad range of $f^\star$ values when $f^\star$ is chosen within a reasonable neighborhood of the optimal loss.

Theorems & Definitions (12)

• theorem 1
• theorem 2
• theorem 3
• Lemma 4
• proof
• Lemma 5
• proof
• corollary 1
• proof
• proof