What do near-optimal learning rate schedules look like?

This paper introduces a search procedure to identify near-optimal learning rate schedule shapes across various workloads, revealing that while warmup and decay are robust features, commonly used schedules are suboptimal and the ideal shape is significantly influenced by hyperparameters like weight decay.

The Gist

Imagine you are training a very smart, but slightly clumsy, robot to learn a new skill. The robot learns by taking steps, and the size of each step is determined by something called the learning rate.

If the robot takes steps that are too big, it stumbles and falls (the training fails). If the steps are too small, it moves so slowly it never finishes the task (the training is inefficient).

For years, scientists have known that the size of the steps matters most. But they've been arguing about the shape of the step-size plan. Should the robot start with tiny steps, get bigger, and then get smaller again? Should it start big and slowly shrink? Or should it just pick one size and stick with it?

This paper, titled "What do near-optimal learning rate schedules look like?", is like a massive experiment where the researchers built a robot that tries thousands of different "step-size plans" to see which one actually works best.

Here is the breakdown of their journey, using simple analogies:

1. The Problem: We're Guessing the Shape

Most people use a standard "recipe" for these step-size plans. The most popular recipe is called Cosine Decay. It looks like a smooth hill: you start small (warmup), go up to a peak, and then slowly slide down the other side (decay).

But the researchers asked: "Is this hill the best shape? Or is there a weird, jagged, or wavy shape that would make the robot learn even faster?"

2. The Experiment: The "Shape Hunter"

To find the answer, they didn't just guess. They created a Search Procedure. Think of this as a super-fast robot that runs thousands of simulations on three different "training grounds":

• Linear Regression: A simple math problem (like predicting house prices based on size).
• CIFAR-10: A task where a computer learns to recognize pictures of cats, dogs, and cars.
• WikiText-103: A task where a computer learns to write sentences like a human (language modeling).

They tested many different "families" of shapes:

• The Constant: A flat line (same step size forever).
• The Cosine: The classic smooth hill.
• The Spline: A flexible rubber band that can be bent into almost any shape.
• The Smooth Non-Monotonic: A wild card that can go up and down however it wants, with no rules.

3. The Big Discoveries

A. The "Base Rate" is the King

The most important finding is that the size of the steps matters way more than the shape of the plan.

• Analogy: Imagine you are driving a car. It doesn't matter if you have a fancy GPS route (the shape); if you are driving at 10 mph when you should be doing 60, you aren't going to get there fast.
• Takeaway: Before you worry about making a fancy schedule, you must tune the "Base Learning Rate" (the peak speed). If you get that wrong, no amount of fancy shaping will save you.

B. The "Warmup" and "Cool Down" are Real

Even when they let the "wild card" robot (the Smooth Non-Monotonic family) try to be weird and skip the rules, it still figured out that it needed a Warmup and a Cool Down.

• The Warmup: Start with tiny steps to get your bearings, then speed up.
• The Cool Down: As you get close to the finish line, slow down so you don't overshoot the target.
• Why? It turns out these aren't just habits; they are fundamental physics of how these robots learn. Even if you didn't tell the robot to do it, the best-performing robots chose to do it.

C. Simple vs. Complex Worlds

Here is where it gets interesting. The "best" shape depends on what the robot is learning:

• In the Simple World (Linear Regression): The best plan was actually no warmup at all! It was a flat line of big steps, followed by a sudden, sharp drop to zero at the very end. It's like sprinting the whole race and then slamming on the brakes.
• In the Complex World (Images and Language): The best plan was a gentle hill. You need the warmup to stabilize the chaotic learning process, and a slow, gentle decay to fine-tune the details.
• Lesson: What works for simple math doesn't necessarily work for complex AI. Don't copy-paste rules from simple problems to complex ones.

D. The "Weight Decay" Connection

They also found that the learning rate schedule is deeply connected to another setting called Weight Decay (a way to prevent the robot from memorizing the training data too perfectly).

• Analogy: If you increase the "friction" (weight decay) on the robot's wheels, the best step-size plan changes. It turns out that if you have high friction, you should keep your steps big for longer before slowing down.

4. The Conclusion: What Should You Do?

The researchers concluded that while we can find "near-perfect" shapes using their search method, the gains are often small.

1. Don't obsess over the shape: The difference between a perfect curve and a standard Cosine curve is often tiny.
2. Do obsess over the Base Rate: Tuning the peak speed is 100x more important than tweaking the curve.
3. Stick to the basics: A standard Warmup + Cosine Decay is a very strong, reliable strategy that works almost everywhere.
4. Beware of the "Simple" trap: Don't assume a strategy that works on a simple math problem will work on a complex language model.

In a nutshell:
The paper is like a master chef testing thousands of recipes. They found that while there are some "secret sauce" shapes that are slightly better, the most important thing is to get the main ingredient (the base learning rate) right. And, surprisingly, even if you let the chef invent a completely new recipe from scratch, they will almost always end up adding a "starter" (warmup) and a "finisher" (decay) because those are the only ways to cook the dish perfectly.

Technical Summary

Here is a detailed technical summary of the paper "What do near-optimal learning rate schedules look like?" by Naganuma et al.

1. Problem Statement

Despite the consensus that learning rate (LR) schedules generally benefit from an initial warmup phase and a final decay phase, there is no rigorous understanding of the optimal shape of these schedules for specific workloads.

• Current Practice: Researchers typically tune a handful of parameters (warmup duration, peak LR, decay start) within fixed functional forms (e.g., linear, cosine, inverse-square-root).
• The Gap: It is unclear if these standard shapes are optimal or if the optimal shape is highly workload-dependent. Furthermore, the interplay between schedule shape and other optimizer hyperparameters (like weight decay) is not well characterized.
• Goal: To systematically search for near-optimal schedule shapes within parameterized families and determine how these shapes relate to workload characteristics and other hyperparameters.

2. Methodology

2.1. Schedule Families

The authors defined several parameterized schedule families (functions mapping normalized time $t/T \in [0,1]$ to LR multipliers $\in [0,1]$):

• Standard/Popular: Constant (con), Cosine (cos-std), Generalized Cosine (cos-gen with tunable exponent), Square-root Decay (sqrt), Generalized Rex (rex).
• Flexible/Interpolated: Two-Point Spline (tps), Two-Point Linear (tpl). These use spline interpolation with two control points to define the decay profile.
• Unconstrained: Smooth Non-Monotonic (snm). A fully general two-control-point spline that does not enforce warmup or monotonic decay, allowing the search to "discover" these features if beneficial.

2.2. Workloads

Experiments were conducted on three distinct tasks to ensure generality:

1. Linear Regression: A synthetic problem with prescribed covariance. This serves as a "ground truth" test case where the theoretical optimal schedule can be computed analytically.
2. Image Classification: A small Convolutional Neural Network (CNN) on CIFAR-10.
3. Language Modeling: An 8M parameter Transformer on WikiText-103.

2.3. Search Procedure

To find near-optimal shapes, the authors employed a two-stage search strategy:

1. Search Step: Randomly sampled thousands of schedule shapes (3,600 for CIFAR-10, 600 for WikiText-103) from each family. For each shape, they swept over 16 base learning rates. Performance was evaluated using a median of training losses across multiple random seeds (10 for CIFAR-10, 5 for WikiText-103) to ensure robustness against initialization variance.
2. Evaluation Step: The top $k$ schedules (ranked by median score) were re-evaluated with 100 seeds (10 initializations $\times$ 10 data orderings) to obtain precise performance metrics and confidence intervals.

• Regime: Experiments were run in an "optimization-limited" regime (insufficient steps to reach saturation), ensuring that differences in schedule shapes resulted in measurable performance gaps.

3. Key Contributions

1. Ground Truth Benchmark: Provided the first known optimal schedule for linear regression trained with SGD, validating the search procedure against a theoretical optimum.
2. Discovery of Robust Features: Demonstrated that warmup and monotonic decay are fundamental features of near-optimal schedules for deep neural networks, even when searching families (like snm) that do not enforce them.
3. Workload Dependence: Showed that optimal schedules for deep learning (non-convex) differ significantly from those for linear regression (convex).
4. Hyperparameter Interactions: Quantified how the optimal schedule shape depends on other optimizer hyperparameters, specifically finding a strong dependency on weight decay.
5. Search Validation: Provided evidence that random search is sufficient to find near-optimal shapes for flexible families (like tps) but struggles with highly unconstrained families (snm) due to the sparsity of "good" regions in the parameter space.

4. Key Results

4.1. Linear Regression (Ground Truth)

• Optimal Shape: The theoretically optimal schedule has no warmup and features a flat, high learning rate for most of training followed by a sharp decay at the end.
• Search Performance: The random search recovered schedules with similar qualitative features (little warmup, sharp decay) but failed to match the exact theoretical optimum, particularly for the snm family.
• Insight: The optimal strategy involves using a large LR to make progress in small eigendirections early on, then decaying to converge large eigenmodes.

4.2. Deep Learning Workloads (CIFAR-10 & WikiText-103)

• Universal Warmup & Decay: Unlike linear regression, warmup is crucial for deep learning workloads. The snm family, despite having no built-in warmup constraint, "discovered" schedules with significant warmup (10–30% of training time) and monotonic decay.
• Base LR Dominance: The base learning rate is the most critical factor for success. Once a schedule includes warmup and decay, the specific shape matters less than the magnitude of the base LR.
• Flexible Families: More flexible families (tps, cos-gen, sqrt) consistently outperformed standard cos-std and con schedules, though the gains were small but statistically significant (e.g., reducing CIFAR-10 error from 0.092 to ~0.063).
• Stability: Schedules found via search were stable across different random seeds, whereas standard schedules sometimes failed to converge or were suboptimal.

4.3. Impact of Other Hyperparameters

• Weight Decay ($\lambda_{WD}$): This had the strongest effect on schedule shape. Increasing weight decay favored schedules that decay later (keeping the LR high for longer).
• Momentum ($\beta_1, \beta_2$): Variations in $\beta_1$ and $\beta_2$ had minor effects on the optimal shape, though higher $\beta_1$ slightly favored later decay.
• Training Horizon: Increasing the total number of steps favored gentler decay curves, while the warmup fraction remained stable relative to the total horizon.

4.4. Search Efficiency

• The snm family was difficult to optimize via random search because the "good" region (warmup + decay) is a small fraction of its total parameter space.
• Families like tps and cos-gen were well-searched; the top 3 schedules within these families were nearly identical, suggesting the search had converged to the local optimum of the family.

5. Significance and Implications

• Validation of Heuristics: The study provides empirical proof that the "warmup + decay" heuristic is not just a convention but a fundamental requirement for optimizing deep neural networks, distinct from convex optimization problems.
• Practical Guidance:
  • Researchers should prioritize tuning the base learning rate over fine-tuning complex schedule shapes.
  • If tuning schedule shapes, flexible families (like Two-Point Spline) offer marginal but real gains over standard Cosine decay.
  • Weight decay and learning rate schedules are coupled; changing one requires re-tuning the other.
• Future Directions: The paper suggests that automatic learning rate selectors could be built by predicting optimal schedule shapes based on simple training metrics (loss trajectory, gradient norms). It also highlights the need for better search algorithms (e.g., Bayesian optimization) for highly flexible, unconstrained schedule families.

In conclusion, this work represents the most comprehensive analysis to date on near-optimal learning rate schedules, bridging the gap between theoretical convex optimization results and the practical realities of training deep non-convex models.