Exploring Landscapes for Better Minima along Valleys

AI summary

Overview

Research area: Deep learning optimization — specifically stochastic gradient-based optimizer design, loss-landscape geometry, and large-batch training.

Technical level: Advanced. The paper combines an intuitive geometric argument about "valleys" in the loss landscape with a formal convergence analysis (convex and non-convex cases) and a large-scale empirical study spanning computer vision, natural language processing, and reinforcement learning.

Scope in one sentence: The paper proposes a drop-in "adaptor" that modifies any gradient-based optimizer so it keeps exploring along flat valleys instead of stopping at the first local minimum, and demonstrates the resulting optimizer (ALTO) on large-batch training tasks.

What This Paper Is About

Standard gradient-based optimizers such as SGD and Adam stop searching once they are trapped by a local minimum, but local information (loss values and gradients) cannot guarantee that the point found is the lowest one or the one that generalizes best. The authors propose an add-on term "E" that lets an optimizer keep moving along landscape valleys — regions of low, nearly identical loss — in order to reach lower and flatter minima, which are usually associated with better generalization. They then apply the adapted optimizer to large-batch training, a setting where fewer parameter updates are made but nearly the same test accuracy is expected.

Key Contributions

1. A generic optimizer adaptor ("E"). A simple replacement of the gradient, g_k + α a_k → g_k with a_k an exponential moving average of the gradient difference g_k − g_{k−1}, which can be inserted into any gradient-based optimizer. The sign of α controls exploration (α < 0) versus exploitation (α > 0).

2. A geometric rationale for valley-following. The authors show that −∇‖∇f(θ_k)‖² = −2 H_k ḡ_k, and that this direction can be approximated cheaply as ḡ_k − ḡ_{k−1} ≈ H_k(θ_k − θ_{k−1}), so the Hessian never has to be computed directly. This gives an optimizer a way to escape sharp small-scale minima while remaining captured by large-scale ones.

3. Convergence proofs for both convex and non-convex settings. Theorem 1 covers the non-convex case (an O(ε²) bound on the average squared gradient norm) and Theorem 2 covers the convex case, giving a regret bound R(T) ≤ O(T^{0.5} G_∞ d^{1.5} D_∞²).

4. A large-batch optimizer, ALTO, with extensive evaluation. ALTO is the Lamb-adapted version. The authors report results across 17 CV and NLP tasks plus RL and LSTM experiments, and an ablation/hyperparameter study.

Main Findings

• Adapted optimizers explore valleys. In Figure 1, ESGD and EAdam continue along the valley on 2D test functions (the square of a cardioid and the Rosenbrock function), whereas SGD and Adam stop at the first minimum.

• Sign of α trades off flatness and speed. With α < 0 the optimizer converges slowly but tends to flat minima; with α > 0 it converges fast but tends to sharp minima. Lambs's setting α = 0 and ALTO's α = −5 are compared in Figure 3 on ResNet20 / CIFAR100.

• Large-batch accuracy gains. ALTO increases the test accuracy of the current state-of-the-art optimizer (Lamb) by an average of 2.5% across a variety of large-batch training tasks. It reports better accuracy in all 17 CV and NLP tasks in their comparison.

• Computation time saved. ALTO can save 29.68% of computation time on a typical CV task while reaching the same accuracy.

• Small-batch vs. large-batch. On ResNet20/CIFAR-10 at batch size 128, ALTO scores 91.24 versus SGD 91.85, Adam 89.88, AdamW 90.54, Lamb 90.89, and AdaBelief 91.12. At batch size 16384, ALTO scores 88.83 versus SGD 80.86, Adam 87.34, AdamW 82.29, Lamb 83.56, and AdaBelief 88.03.

• CIFAR-100 results. At batch size 128, ALTO reaches 65.74 versus SGD 64.93, Adam 64.35, AdamW 64.62, Lamb 61.29, AdaBelief 64.44. At batch size 16384, ALTO reaches 57.78 versus SGD 44.20, Adam 54.91, AdamW 52.95, Lamb 56.06, and AdaBelief 52.94.

• ImageNet at extreme batch sizes. On ResNet34/ImageNet with batch size 4086, ALTO reaches 70.83% versus SGD 49.35, Adam 54.96, AdamW 68.40, Lamb 70.34, and AdaBelief 70.18 — beating SGD at a much smaller batch size of 256 (70.64%). At batch size 256, ALTO scores 69.95, below SGD's 70.64.

• ResNet-50 / ImageNet across batch sizes (90 epochs). ALTO scores 77.22 (1K), 77.25 (2K), 77.35 (4K), 77.10 (8K), 76.87 (16K), and 76.70 (32K), outperforming Lamb's reported values of 77.06, 77.11, 76.92, 76.89, 76.66, and 76.42 at the same batch sizes. Adam, AdamW, and AdaBelief all score lower at every batch size.

• GPT-2 perplexity. ALTO achieves a test perplexity of 78.37 with batch size 4096, compared to Lamb's 83.13 (described as SOTA).

• Convergence guarantees are tighter than Lamb's in the stated comparison. Table 2 lists Lamb requiring T ≥ O(ε^{−4}), b ≥ O(ε^{−4}), η ≤ O(ε^{−2}), whereas ALTO requires T ≥ O(ε^{−2}), b ≥ O(ε^{−2}), and η ≤ O(1); the Lamb analysis in that table holds for only a small set of β values (0 or 1), while ALTO's covers a more general case.

• Timing caveat. Because of the extra acceleration term a, ALTO's per-epoch computation time is longer than Lamb's, but it reaches a given accuracy in less total training time.

• Gains grow with batch size. The authors state that the larger the batch size, the larger ALTO's advantage over other optimizers, and that in the 16384 batch case ALTO reaches the best accuracy of the compared optimizers using only half the number of epochs.

Methodology in Plain English

The authors start from two requirements that a persistent-exploration optimizer must satisfy at the same time: it must be caught by large-scale minima (so it follows the valley rather than wandering off) and it must escape small-scale sharp minima (so it keeps exploring). They observe that the gradient of the squared gradient norm, −∇‖∇f‖², behaves like an optimization direction plus a repulsion from sharp minima, and that the difference between it and the ordinary gradient is exactly the Hessian. Since computing the Hessian is unaffordable, they use the finite-difference identity ḡ_k − ḡ_{k−1} ≈ H_k(θ_k − θ_{k−1}) as a cheap surrogate, smoothed with an exponential moving average for stability.

The resulting update is a one-line substitution at the point where the gradient enters the optimizer: every gradient g_k is replaced with g_k + α a_k, with a_k = β₁ a_{k−1} + (1 − β₁)(g_k − g_{k−1}). This is applied to SGD (ESGD), Adam (EAdam), biased-Lamb (ALTO Vanilla), and Lamb (ALTO). ALTO additionally uses bias correction and layerwise regularization inherited from Lamb.

For theory, the authors fix standard assumptions (L-smoothness; unbiased, independent, variance-bounded stochastic gradients; bounded gradients; bounded parameters; a monotonicity assumption; convexity for the convex case) and prove convergence for the vanilla variant in both non-convex and convex settings, noting the biased variant differs only by a bounded factor.

For experiments, all runs used a single node with 4 NVIDIA 80GB A100 GPUs connected by PCI-E 3.0; multi-node runs were not performed. The authors tune 3–5 hyperparameters per task for every optimizer, tune only β₁ and the learning rate for ALTO while tuning everything for the baselines, and repeat each result three times, reporting the mean. They use α = 0.5, β₁ = 0.01 for batch sizes below 1K and α = −5, β₁ = 0.99 for batch sizes of 1K or above, with β₂ = 0.9, β₃ = 0.99, λ = 10⁻⁴, ε₁ = 10⁻⁶, ε₂ = 10⁻⁶, ε₃ = 10⁻¹⁰ unless stated otherwise.

Why This Matters

Impact on research. The paper reframes optimizer design around landscape geometry — deliberately staying in valleys and rejecting sharp minima — rather than only around the size of the update step. It also provides a convergence analysis under assumptions presented as weaker or more general than those used for Lamb, and it may open a new direction in optimization algorithm design.

Real-world applications.

• Large-batch pretraining of vision models, where the goal is to use many accelerators without losing test accuracy.
• Large language model pretraining (the paper reports GPT-2 results), where large effective batch sizes and low perplexity matter.
• Distributed data-parallel training pipelines that want to reach a target accuracy in less wall-clock time.
• Hyperparameter-limited production training, where the adaptor is a drop-in change to an existing optimizer.

Industry relevance. The adaptor requires no change to model architecture and only a small change to optimizer code, and the paper reports a 29.68% computation-time saving to reach a given accuracy on a typical CV task — a direct cost lever for teams training at scale. The authors note their code is available (linked as code and PyPI in the paper).

Future Directions

• Broadening the theory. The proofs are given for ALTO Vanilla (Algorithm 1), with the biased variant (Algorithm 2) stated to hold only up to a bounded factor; a full analysis of the biased, layerwise-regularized version under the same assumptions remains open.
• Reducing the per-epoch overhead. ALTO's per-epoch time exceeds Lamb's because of the extra a term, so lowering that cost or amortizing it would strengthen the practical case further.
• Extending beyond the tested settings. The authors did not run multi-node experiments due to limited hardware; whether the advantage holds in genuinely distributed, multi-node regimes is untested here.
• Characterizing when α > 0 helps. The paper reports that positive α sometimes gives marginal accuracy or speed benefits in small-batch training while consistently producing sharp minima — understanding when the exploration/exploitation trade-off should flip is left open.

Target Audience

This paper will be most useful to machine learning systems and optimization researchers, graduate students working on optimizer design or loss-landscape geometry, and practitioners who train large models with large batch sizes on many GPUs and need accuracy to hold up. Readers should be comfortable with stochastic optimization notation, Hessians, and convergence-rate statements; the empirical sections are more accessible than the theory sections.

Authors’ abstract

Finding lower and better-generalizing minima is crucial for deep learning. However, most existing optimizers stop searching the parameter space once they reach a local minimum. Given the complex geometric properties of the loss landscape, it is difficult to guarantee that such a point is the lowest or provides the best generalization. To address this, we propose an adaptor "E" for gradient-based optimizers. The adapted optimizer tends to continue exploring along landscape valleys (areas with low and nearly identical losses) in order to search for potentially better local minima even after reaching a local minimum. This approach increases the likelihood of finding a lower and flatter local minimum, which is often associated with better generalization. We also provide a proof of convergence for the adapted optimizers in both convex and non-convex scenarios for completeness. Finally, we demonstrate their effectiveness in an important but notoriously difficult training scenario, large-batch training, where Lamb is the benchmark optimizer. Our testing results show that the adapted Lamb, ALTO, increases the test accuracy (generalization) of the current state-of-the-art optimizer by an average of 2.5% across a variety of large-batch training tasks. This work potentially opens a new research direction in the design of optimization algorithms.