Why Do We Need Warm-up? A Theoretical Perspective
Foivos Alimisis ⋅ Rustem Islamov ⋅ Aurelien Lucchi
Abstract
Learning rate warm-up -- increasing the learning rate at the beginning of training -- has become a ubiquitous heuristic in modern deep learning, yet its theoretical foundations remain poorly understood. In this work, we provide a principled explanation for why warm-up improves training. We rely on a generalization of the $(L_0, L_1)$-smoothness condition, which bounds local curvature as a linear function of the loss suboptimality and exhibits desirable closure properties. We show -- both theoretically and empirically -- that this condition is satisfied by common neural architectures and accurately captures the curvature of the optimization landscape early in training. Adapting the learning rate in response to this curvature condition naturally induces a warm-up–like schedule, and we show that this choice yields provably faster convergence guarantees than using a fixed learning rate. Experiments on language and vision models show that the resulting one-parameter warm-up schedule can match tuned linear warm-up and improve over no warm-up.
Lay Summary
The goal of our paper is to identify a smoothness-type condition to explain the need of learning rate warm-up when training neural networks, and especially large language models. We propose a generalization of the popular $(L_0,L_1)$-condition (which we call $(H_0,H_1)$). This condition is satisfied by general classes of neural networks, while it is amenable to optimization analysis. Under this condition, we prove that, in realistic function classes, gradient descent with learning rate warm-up performs better than with a fixed learning rate in terms of iteration complexity. Under extra assumptions, we extend this result also to stochastic gradient descent. We empirically validate the tightness of the $(H_0,H_1)$-condition in the initial phase of training and show that the derived warm-up schedule is practically competitive with other well-tuned warm-up strategies.
Successful Page Load