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

TL;DR

This paper introduces a theoretically grounded adaptive warm-up scheduling method for norm-constrained optimizers, improving large language model training by automatically tuning warm-up duration without extra hyperparameter tuning.

Contribution

It develops a new adaptive warm-up schedule based on a generalized smoothness assumption, with proven convergence guarantees and practical implementation for large-scale language models.

Findings

Adaptive warm-up outperforms manual schedules in LLaMA training.

The method automatically adjusts warm-up duration without extra hyperparameter tuning.

Theoretical analysis supports the natural emergence of warm-up and decay phases.

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