Technical Summary: Why Do We Need Warm-up? A Theoretical Perspective

Problem Statement

Learning rate (LR) warm-up, the practice of linearly increasing the learning rate at the start of training, is a ubiquitous heuristic in modern deep learning, particularly for Transformers. While empirical evidence suggests it mitigates training instabilities, reduces gradient variance, and improves robustness to peak LR choices, a rigorous theoretical justification for its necessity has remained elusive. Existing convergence analyses for stochastic gradient descent (SGD) primarily focus on the learning rate decay phase, and prior theoretical frameworks (such as $(L_0, L_1)$-smoothness) often fail to explain the initial warm-up phase, sometimes even prescribing decreasing step-sizes when gradients are observed to increase early in training.

Methodology

The $(H_0, H_1)$-Smoothness Condition

The authors introduce a generalized smoothness condition, termed $(H_0, H_1)$-smoothness, to characterize the loss landscape during the early stages of training. A function $f$ is $(H_0, H_1)$-smooth if its local curvature (spectral norm of the Hessian) is bounded by an affine function of the loss suboptimality:
$$\|\nabla^2 f(w)\|_2 \leq H_0 + H_1(f(w) - f^*)$$
where $f^*$ is the global minimum.

This condition differs from the standard $(L_0, L_1)$-smoothness, which bounds curvature by the gradient norm. The authors argue that dependence on the gradient norm leads to counterintuitive LR scheduling implications during early training, whereas dependence on the loss suboptimality (which decays monotonically) provides a more realistic model for the initial phase.

Theoretical Justification

The paper provides theoretical evidence that common neural architectures satisfy this condition:

1. Deep Linear Networks: Under weak or strong balancedness assumptions, deep linear networks with MSE loss satisfy $(H_0, H_1)$-smoothness, particularly in the initial phase where weight norms are small and loss is large.
2. Non-Linear Networks: The condition holds for 2-layer networks with cross-entropy loss and $L_2$ regularization, as well as deep networks with Leaky-ReLU activations under specific regularity conditions.
3. Transformers: An informal result demonstrates that in-context learning loss functions for simple transformer architectures satisfy this condition under mild regularity assumptions on input/output distributions.

Convergence Analysis

The authors analyze Gradient Descent (GD) and SGD under the $(H_0, H_1)$ assumption:

• Fixed Step-size: They derive lower complexity bounds showing that a fixed step-size must be chosen conservatively to account for the worst-case early curvature, leading to slow convergence.
• Adaptive Warm-up: They propose an adaptive step-size schedule derived from the smoothness condition: $\eta_k \propto \frac{1}{H_0 + H_1(f(w_k) - f^*)}$. Since the loss $f(w_k)$ decreases during training, the step-size naturally increases, mimicking a warm-up schedule.
• Convergence Guarantees: Under additional structural assumptions (Aiming condition, Polyak-Łojasiewicz condition, or convexity), the authors prove that GD with this adaptive warm-up schedule achieves faster convergence rates than fixed-step GD, particularly when the initial loss is large.

Key Results

Empirical Verification

The authors empirically verify the $(H_0, H_1)$-smoothness condition on:

• Language Models: Transformers of varying sizes (70M, 160M, 410M) trained on the FineWeb dataset.
• Vision Models: ResNet50 and ViT-Tiny trained on ImageNet-32.

Using a small constant learning rate and gradient clipping to probe the landscape near initialization, they observe that the estimated local smoothness decays approximately linearly with the training loss during the initial phase. This validates the affine relationship proposed in the $(H_0, H_1)$ condition.

Performance Evaluation

The authors compare three schedules: no warm-up, tuned linear warm-up, and the proposed $(H_0, H_1)$-driven warm-up (a one-parameter schedule based on the loss trajectory).

• Language Modeling: The $(H_0, H_1)$ warm-up schedule matches the performance of tuned linear warm-up and significantly outperforms training without warm-up.
• Image Classification: Similarly, the proposed schedule matches linear warm-up performance on ViT-Tiny and improves upon the no-warm-up baseline.
• Initialization Effects: Experiments suggest that initialization schemes with depth-scaling (common in Transformers) result in higher initial sharpness, necessitating a warm-up phase to stabilize training.

Significance and Claims

The paper claims to provide the first principled theoretical explanation for learning rate warm-up based on the geometry of the loss landscape. By introducing $(H_0, H_1)$-smoothness, the authors:

1. Bridge Theory and Practice: They demonstrate that the curvature of neural network loss landscapes in the early training phase is better captured by a function of loss suboptimality rather than gradient norm.
2. Justify Warm-up: They show that an adaptive step-size responding to this curvature naturally induces a warm-up-like schedule, which is provably faster than fixed-step methods under standard non-convex assumptions (Aiming, PL).
3. Practical Utility: The resulting one-parameter schedule is competitive with standard linear warm-up, offering a theoretically grounded alternative that reduces the need for extensive hyperparameter tuning of warm-up lengths.

The authors remain modest, noting that their theoretical guarantees apply specifically to the early sharpness-reduction phase and that the analysis is currently limited to GD and SGD, with future work needed to extend these results to practical optimizers like Adam and to generalize beyond Euclidean geometry.