Technical Summary: Theory of Optimal Learning Rate Schedules and Scaling Laws for a Random Feature Model

Problem Statement

Training deep learning models requires selecting hyperparameters such as learning rate (LR), batch size, and training horizon. While empirical "trial and error" is common, there is a lack of theoretical understanding regarding how optimal hyperparameters scale with the training horizon $T$. Existing strategies for transferring hyperparameters across model sizes often fail to transfer across different training horizons. Furthermore, while various LR schedules (e.g., linear decay, warmup-stable-decay, cosine annealing) are used in practice, the theoretical conditions dictating which schedule is optimal for a given task remain unclear. This paper addresses the need for a principled framework to derive optimal LR schedules that depend on the training horizon and the specific structure of the learning task.

Methodology

The authors employ a power-law random feature model trained with Stochastic Gradient Descent (SGD) to derive optimal hyperparameter schedules. The model assumes input data $x$ with a target function $y(x) = w^* \cdot \psi(x) + \sigma_0 z$, where features $\psi(x)$ exhibit power-law spectral decay ($\lambda_k \sim k^{-b}$) and teacher weights scale as $(w^*_k)^2 \lambda_k \sim k^{-a}$.

The core methodology involves optimal control theory:

1. Dynamics Formulation: The expected test loss $L_t$ is derived as a function of the learning rate schedule $\{\eta_t\}$ and batch size schedule $\{m_t\}$. The loss dynamics are governed by a linear recurrence relation for mode errors $c_{t,k}$.
2. Optimization Objective: The goal is to minimize the final test loss $L_T$ at a fixed horizon $T$. The authors formulate three optimization problems:
   • Optimal constant LR.
   • Optimal LR schedule (time-varying $\eta_t$).
   • Joint optimal LR and batch size schedule.
3. Solution Approaches:
   • Numerical: Using a single-shooting method (implemented in CasADi) to solve the discrete-time optimal control problem for finite $T$.
   • Analytical: Deriving continuous-time approximations for large $T$ using Noether's theorem and asymptotic analysis of the loss functional. This involves minimizing a functional of the integrated learning rate $\chi(t) = \int_t^T \eta(s) ds$.

Key Contributions and Results

1. Identification of Two Distinct Phases

The analysis reveals two distinct regimes for optimal learning rate schedules, determined by the relationship between the spectral exponents $a$ (task difficulty) and $b$ (feature complexity):

• The Easy Phase ($a > b$):

  • The optimal schedule follows a polynomial decay form: $\eta^*_T(t) \sim T^{b/a - 1} f(t/T)$, where $f(z) \propto (1-z)^{2b-1}$.
  • The learning rate gradually decays to balance bias and variance reduction throughout the training trajectory.
  • The loss scaling exponent is $\zeta_T = (a-1)/a$, matching the minimax-optimal rate for fixed batch sizes.

• The Hard Phase ($a < b$):

  • The optimal schedule resembles Warmup-Stable-Decay (WSD). It maintains a constant maximum learning rate ($\eta_{max}$) for the majority of training steps, followed by a rapid annealing phase over a vanishing fraction of the total time ($1 - t_s/T \sim T^{-(b-a)/(2b-1)}$).
  • This strategy prioritizes bias suppression via a maximal constant LR, compressing variance reduction into the final steps.
  • The loss scaling exponent improves to $\zeta_T = (a-1)/b$, outperforming constant LR strategies.

2. Superiority Over Benchmarks

The derived optimal schedules strictly outperform standard benchmarks:

• Constant LR: The optimal schedule achieves better scaling exponents in both phases.
• Power-Law Schedules: Even when optimizing the exponents of power-law schedules ($\eta \sim T^{-\xi} t^{-\delta}$), they asymptotically recover the same scaling behavior as the optimal constant LR baseline, failing to match the performance of the horizon-dependent optimal schedules derived via control theory.

3. Joint Optimization of Learning Rate and Batch Size

The authors extend the analysis to jointly optimize $\eta(t)$ and $m(t)$ under a fixed data budget $B_{tot}$:

• Easy Phase: A "batch ramp" (increasing batch size over time) combined with a constant maximal learning rate reduces wall-clock time. The loss scales as $T^{-(a-1)/b}$, improving the wall-clock exponent compared to fixed batch sizes, while the sample-budget exponent remains $(a-1)/a$.
• Hard Phase: The joint optimization yields the same scaling exponents as the fixed-batch optimum, with the batch size remaining bounded.

4. Optimal Momentum Schedules

For tasks with momentum, the authors optimize the momentum parameter $\beta(t)$ alongside $\eta(t)$.

• In the Easy Phase, the optimal momentum schedule yields similar scaling laws to optimal SGD.
• In the Hard Phase, optimizing the momentum schedule provides an improvement in the scaling exponent beyond what is achievable with SGD alone. The optimal momentum dynamics show an initial transient decrease followed by an increase as $t/T \to 1$.

5. Empirical Validation in Deep Networks

The theoretical phases are validated on deep learning benchmarks:

• ResNet on CIFAR-5M: Exhibits Hard Phase behavior. The optimal base learning rate for annealing schedules is approximately constant across horizons, consistent with the WSD profile.
• GPT-2 Style Transformers on C4: Exhibit Easy Phase behavior. The optimal base learning rate shifts (decreases) with the horizon $T$, even under annealing, consistent with the polynomial decay profile.

Significance and Claims

The paper claims to provide a theoretical foundation for understanding why learning rate schedules vary across different tasks and training horizons. By utilizing optimal control theory on a solvable random feature model, the authors demonstrate that:

1. Horizon Dependence: Optimal learning rates are not automatically transferable across training horizons; the schedule must be explicitly designed based on $T$.
2. Task Structure: The "difficulty" of a task (characterized by spectral exponents $a$ and $b$) dictates whether a task falls into the easy or hard phase, determining the optimal schedule shape (polynomial decay vs. WSD).
3. Efficiency: Horizon-dependent schedules and joint optimization of batch size can significantly improve scaling laws and reduce wall-clock training time compared to standard heuristics.

The authors acknowledge limitations, noting that their theory currently applies to linear models in the lazy/kernel regime and does not yet analytically derive the specific improved exponent for momentum in the hard phase or account for sophisticated optimizers like Adam or preconditioned methods. However, the work offers a principled explanation for the effectiveness of schedules like WSD and suggests that empirical tuning of base learning rates across horizons should account for the underlying spectral properties of the model and task.