PaperlayerView exact revision on arXiv

Source-linked AI summary

Lookahead Optimizer: k steps forward, 1 step back

Michael R. Zhang, James Lucas, Geoffrey Hinton, Jimmy Ba

arXiv:1907.08610v2cs.LGcs.NEstat.ML

TL;DR

Deep-network optimization commonly relies on SGD variants, but adaptive and accelerated improvements still motivate more stable, less tuning-sensitive methods. Lookahead addresses this by alternating k fast-weight updates from another optimizer with a slow-weight synchronization step, and reports improved convergence across multiple tasks with minimal overhead.

  • Problem

    Existing SGD improvements use adaptive learning rates or acceleration, while neural-network training still requires methods that are stable and less sensitive to hyperparameters.

  • Method

    Lookahead updates fast weights k times with any standard optimizer, then updates slow weights toward the final fast weights and resets the fast weights.

  • Results

    Lookahead improves convergence across image classification, language modeling, and neural machine translation, often improving generalization and remaining robust to hyperparameter changes.

  • Takeaways & Limitations

    Lookahead can be combined with standard optimizers such as SGD or Adam to improve convergence with minimal computational overhead.

  • Takeaways & Limitations

    The theoretical analysis assumes diagonal curvature and noise covariance structures, including a non-trivial co-diagonalizable noise assumption.

Abstract

from arXiv · show

The vast majority of successful deep neural networks are trained using variants of stochastic gradient descent (SGD) algorithms. Recent attempts to improve SGD can be broadly categorized into two approaches: (1) adaptive learning rate schemes, such as AdaGrad and Adam, and (2) accelerated schemes, such as heavy-ball and Nesterov momentum. In this paper, we propose a new optimization algorithm, Lookahead, that is orthogonal to these previous approaches and iteratively updates two sets of weights. Intuitively, the algorithm chooses a search direction by looking ahead at the sequence of fast weights generated by another optimizer. We show that Lookahead improves the learning stability and lowers the variance of its inner optimizer with negligible computation and memory cost. We empirically demonstrate Lookahead can significantly improve the performance of SGD and Adam, even with their default hyperparameter settings on ImageNet, CIFAR-10/100, neural machine translation, and Penn Treebank.

1 Introduction

Lookahead is proposed as an optimizer orthogonal to adaptive learning-rate and accelerated SGD approaches. It combines an inner optimizer’s fast-weight updates with slower synchronized updates and improves convergence across several deep-learning settings.

  • Motivation: Lookahead complements adaptive learning-rate and accelerated optimization schemes rather than replacing either approach.Recent SGD improvements include AdaGrad and Adam, plus heavy-ball and Nesterov momentum.
  • Contribution: Lookahead updates fast weights k times with a standard optimizer, then moves slow weights toward the final fast weights.The method can use inner optimizers such as SGD or Adam.
  • Empirical evaluation: Lookahead improves convergence over its inner optimizer across CIFAR, ImageNet, Penn Treebank, and WMT 2014 English-to-German tasks.The experiments use image classifiers, LSTM language models, and Transformer-based translation models.
  • Empirical evaluation: The method is reported to be robust to changes in the inner optimizer, synchronization period, and slow-weights learning rate.It often also improves generalization performance while remaining robust to hyperparameter changes.

2 Method

Lookahead maintains fast and slow weights: an inner optimizer explores for k steps, then slow weights interpolate toward the resulting fast weights before resetting the fast weights. The design smooths optimization trajectories, reduces variance, and uses fixed α in the main experiments because it offers similar convergence benefits with lower overhead.

  • Algorithm: Lookahead applies any standard optimizer A to fast weights for k updates, then linearly interpolates slow weights toward the final fast weights.After synchronization, the fast weights are reset to the current slow-weights value.
  • Optimization behavior: Fast weights can make progress along low-curvature directions while slow-weight interpolation smooths oscillations in high-curvature directions.The paper connects this combination with improved learning and reduced variance.
  • Optimization behavior: In a ResNet-32 CIFAR-100 visualization, Lookahead’s slow-weight path moves toward an improved test-accuracy region that SGD leaves unexplored after 20 updates.The fast weights explore around the minimum while the slow update pushes the trajectory toward that region.
  • Algorithm: The slow-weight trajectory is an exponential moving average of final fast weights, emphasizing recent proposals while retaining influence from earlier ones.This averaging behavior is independent of the inner optimizer.
  • Computational cost: The optimizer adds constant amortized overhead from parameter copying and arithmetic, while maintaining one additional parameter copy.The operation count is described as O((k+1)/k) times that of the inner optimizer.
  • Selecting α: A quadratic-loss analysis provides a principled slow-weight step size α based on curvature approximations.The practical adaptive estimate uses an empirical Fisher approximation and clips α between αlow and 1 for stability.
  • Selecting α: Fixed α offers similar convergence benefits to adaptive α and tends to generalize better in practice.The fixed choice also avoids estimating the empirical Fisher, reducing memory and computational cost when the inner optimizer lacks that estimate.

3 Convergence Analysis

Lookahead is analyzed on noisy and deterministic quadratic models, where its slow-weight interpolation reduces steady-state variance and can accelerate convergence, especially during momentum-induced oscillations.

  • 3.1 Noisy quadratic analysis: Lookahead is analyzed on a noisy quadratic model as a proxy for understanding convergence guarantees relevant to neural network optimization.The analysis compares Lookahead with SGD under equal learning rates and simulated hyperparameter settings.
  • 3.1 Noisy quadratic analysis: Under 0 < γ < 2/L, SGD and Lookahead converge in expectation to zero, while their variances approach analytically characterized fixed points.Here L = max_i a_i.
  • 3.1 Noisy quadratic analysis: For α ∈ (0, 1), Lookahead has a strictly smaller variance fixed point than its SGD inner optimizer at the same learning rate.The reduction follows because the first product term in the Lookahead variance fixed point is smaller than 1.
  • 3.1 Noisy quadratic analysis: After 1000 updates, fixed Lookahead hyperparameters can reach SGD’s steady-state risk more quickly, and Lookahead outperforms SGD across the simulated α values.Figure 3 compares expected losses at matched final risk, using k = 5 and overlaid α values in (0, 1].
  • 3.2 Deterministic quadratic convergence: Lookahead significantly improves classical momentum’s convergence rate in the under-damped regime, where oscillations occur, but is slightly worse in the over-damped regime.The deterministic analysis varies the condition number under the optimal learning rate while fixing the momentum coefficient.

4 Related work

Lookahead performs parameter averaging during optimization, distinguishing it from prior averaging, meta-learning, variance-reduction, and extrapolation methods that use related inner–outer structures.

  • Parameter averaging: Unlike SWA and other parameter-averaging approaches focused mainly on end-of-training parameters, Lookahead averages parameters during training to accelerate convergence.The paper also reports lower validation error for Lookahead on additional CIFAR-100 architectures.
  • Related optimization methods: Reptile also uses inner and outer optimization loops, but samples separate tasks to learn parameters that initialize future tasks, unlike Lookahead’s optimization setting.Lookahead does not sample new tasks for each outer loop.
  • Related optimization methods: Katyusha pulls parameters toward checkpoints during inner-loop steps and uses an SVRG correction, whereas Lookahead pulls back only at inner-loop completion without that correction.The paper notes that the correction could nevertheless be incorporated into Lookahead.
  • Related optimization methods: Anderson acceleration and related extrapolation methods retain all inner-loop iterates and compute linear combinations, creating memory and combination-selection challenges absent from Lookahead’s simpler scheme.The related method of Scieur et al. requires on the order of k times more memory than Lookahead.
  • Empirical comparison: Figure 5 compares CIFAR-100 training loss and CIFAR ResNet-18 validation accuracy across optimizers, with Lookahead and Polyak wrapped around SGD.Other optimizers receive learning-rate and weight-decay grid searches.

5 Experiments

Lookahead was evaluated across image classification, language modeling, and neural machine translation, consistently improving convergence and often generalization or robustness over inner optimizers. On ImageNet, it reached strong top-1 accuracy in fewer epochs, while language-modeling and translation experiments showed favorable convergence and hyperparameter behavior.

  • CIFAR-10 and CIFAR-100: Lookahead achieved significantly faster convergence on CIFAR experiments despite using a learning-rate schedule optimized for the inner optimizer.The authors note that this matters for anytime performance when datasets, hyperparameters, and schedules are not well calibrated.
  • ImageNet: 75% single-crop top-1 accuracy was reached on ImageNet in 50 epochs, and 75.5% in 60 epochs with aggressive learning-rate decay.The learning rate was decayed by a factor of 10 at epochs 30, 48, and 58.
  • ImageNet: 77% single-crop top-1 accuracy was reached on ResNet-152 in 49 epochs, and 77.96% in 60 epochs.This matched the reported 77% result while scaling Lookahead to ResNet-152 with an aggressive decay schedule.
  • Language modeling: Lookahead with Adam achieved the fastest convergence and best training, validation, and test perplexity on Penn Treebank.SGD took around 700 epochs to converge and did not match Adam's final performance, while ASGD improved on Adam but did not match Lookahead.
  • Neural machine translation: Lookahead sped up early Transformer training relative to Adam and later AdaFactor, while all methods converged to similar training loss.On WMT 2014 English-to-German, Lookahead also enabled learning rates of 0.02, 0.04, and 0.06 to reach similar final losses.
  • Inner and outer loop evaluation: Fast weights could substantially degrade task performance within an inner loop, whereas the slow-weight step restored test accuracy and outer-loop variance.This behavior was observed when plotting test accuracy for every update during epoch 65.

6 Conclusion

Lookahead combines with standard optimization methods by using fast-weight trajectories to update slow weights. The paper links this design to variance reduction, improved convergence, and strong results across deep-learning benchmarks.

  • Conclusion: Lookahead can be combined with any standard optimization method and computes updates by looking ahead at fast weights generated by another optimizer.The algorithm maintains fast and slow weights and uses the fast-weight sequence to guide updates.
  • Conclusion: The paper presents variance reduction and strong empirical results across multiple deep-learning benchmark datasets and architectures.These results support Lookahead's use as an optimizer wrapper rather than a replacement for a specific inner optimizer.

A Noisy quadratic analysis

The noisy quadratic analysis characterizes SGD and Lookahead through their expectation and variance dynamics. Lookahead can reduce steady-state loss through variance reduction, although its expectation term may converge more slowly under the same learning rate.

  • Dynamics: Lookahead’s noisy-quadratic dynamics are analyzed through its slow-weight trajectories and recursively computed variance and covariance terms.The analysis derives the slow-weight variance from the variance recurrence and the covariance between slow and fast weights.
  • Steady-state analysis: Under a contraction learning-rate condition, SGD and Lookahead converge in expectation to zero and their variances approach fixed points.The contraction argument uses Banach’s fixed point theorem before solving the variance fixed points.
  • Loss trade-off: For the same learning rate, Lookahead achieves smaller loss when its greater variance reduction outweighs its slower expectation-term convergence.The expectation comparison is between 1 − α + α(I − γA)^k for Lookahead and (I − γA)^k for SGD.
  • Finite-time convergence: In Figure 11, SGD can outperform Lookahead after only a few updates for some small α values, but other α choices outperform SGD even at that horizon.The short-horizon difference is attributed to SGD making faster progress on the expectation term.
  • Finite-horizon tuning: Lookahead dominates SGD in expected loss over all finite time horizons when each optimizer’s constant-through-time hyperparameters are tuned for that horizon.The comparison uses grid searches over SGD’s learning rate and Lookahead’s learning rate and α with k = 5 fixed.

B Deterministic quadratic convergence analysis

The deterministic quadratic analysis represents Lookahead and momentum updates as linear dynamical systems. Matrix products and eigenvalues then provide convergence bounds, while interpolation and adaptive slow-weight choices describe the method’s design.

  • Quadratic model: The deterministic analysis models the quadratic objective with optimum θ* = 0 and uses linear dynamics to study convergence.The quadratic formulation is introduced as a general case while placing the optimum at zero for simplicity.
  • Dynamical-system representation: Fast-weight iterates are stacked and transformed using matrices for Lookahead interpolation, inner momentum updates, and iterate realignment.L denotes interpolation, B denotes classical momentum, and T realigns fast-weight iterates.
  • Matrix structure: Each transition matrix shifts iterates through identity blocks, applies the inner update, and encodes Lookahead interpolation in its top-right structure.The matrix blocks distinguish the iterate shift, momentum update, and interpolation operation.
  • Convergence bounds: Eigenvalues of the matrix product bound convergence, with kth roots required because the system represents k inner updates followed by one slow-weight update.Standard linear-system solvers are used after forming the appropriate matrix product.
  • Slow-weight step size: The analysis also derives an optimal slow-weight step size by differentiating the quadratic loss along the interpolation direction.A further approximation uses the Fisher as a Hessian approximation and stochastic mini-batch gradients without additional computation.

C Experimental setup

The experimental setup evaluates Lookahead across image-classification architectures using established training implementations and tuned or official optimizer settings. CIFAR experiments search baseline and Lookahead hyperparameters, while ImageNet uses a standard 90-epoch ResNet protocol.

  • General protocol: The experiments use publicly available setups, multiple random seeds, and standard CIFAR image preprocessing, with results plotted as means and one-standard-deviation error bars.A reviewer noted that the CIFAR ResNet-18 implementation has wider channels and more parameters than the original.
  • CIFAR baselines: CIFAR baseline searches sweep SGD learning rates and weight decay, selecting learning rate 0.05 and weight decay 0.001; Adam is represented by AdamW.SGD uses momentum 0.9, and choices are based on final validation performance.
  • CIFAR Lookahead: For CIFAR Lookahead, the best reported setting uses inner SGD learning rate 0.1, α = 0.8, and k = 5, although choices are described as fairly robust.The search considers α ∈ {0.2, 0.5, 0.8} and k ∈ {5, 10}, and the reported version resets momentum.
  • ImageNet implementation: ImageNet experiments wrap Lookahead around the official PyTorch settings and additionally test an aggressive decay schedule on four Nvidia P100 GPUs with batch size 256.The aggressive schedule decays at epochs 30, 48, and 58, with weight decay 1e-4.

C.3 Language modeling

The language-modeling and translation experiments apply Lookahead to established LSTM and Transformer setups. The reported settings suggest robustness across k and α, while the translation models reach around 27 BLEU after 250k steps.

  • Penn Treebank setup: Penn Treebank experiments use a three-layer LSTM with 1150 hidden units per layer, 400-dimensional embeddings, and specified dropout rates.The setup follows the model and code provided by Merity et al.
  • Lookahead search: Lookahead language-model experiments search k ∈ {5, 10, 20} and α ∈ {0.2, 0.5, 0.8} after selecting competitive SGD and Adam baselines.Final performance is reported using parameters giving the best validation performance during training.
  • Hyperparameter robustness: For language modeling, α = 0.5 or 0.8 with k = 5 or 10 worked best, and the method was largely robust to different k and α choices.The passage notes that joint optimization with the underlying optimizer could yield better results.
  • Neural machine translation: After 250k steps, the translation models using Adam, AdaFactor, and Lookahead settings each reach around 27 BLEU on Newstest2014.The task uses a single TPU core with eight workers and minibatch size 2048 per worker.

D Additional Experiments

Lookahead’s convergence is robust to how the inner optimizer’s momentum state is handled, and faster training-loss convergence corresponds to better validation performance on CIFAR and ImageNet.

  • Maintaining, interpolating, or resetting momentum all improve Lookahead convergence versus SGD.The sensitivity study compares these three momentum-state treatments for SGD with heavy-ball momentum.
  • Faster convergence in training loss corresponds to better validation performance on CIFAR and ImageNet.The paper presents validation-accuracy curves throughout training for both datasets.

D.3 Comparison to Stochastic Weight Averaging

Lookahead and Stochastic Weight Averaging serve different purposes and can be complementary. In CIFAR-100 experiments, Lookahead improves performance during training and in the weight-averaged network, while avoiding SWA’s averaging-start decision.

  • SWA requires choosing when tail averaging begins, and starting too early or too late can harm performance.Figure 15 illustrates the sensitivity of test accuracy to the averaging start point.
  • Figure 16 reports CIFAR-100 test accuracy for SWA and Lookahead using Wide ResNet-28-10, including accuracy during training and for the SWA network.SWA begins at epoch 161 in this experiment.
  • Lookahead is applied from initialization and therefore does not require a tail-averaging start decision.SWA can perform better at intermediate training stages by averaging multiple models, whereas Lookahead targets variance reduction and final performance.
  • Lookahead wrapped around SGD achieves higher test accuracy throughout training and in the weight-averaged network than SGD used as SWA’s inner loop.This comparison uses three Wide ResNet-28-10 runs on CIFAR-100.
  • Figure 15 reports CIFAR-100 test accuracy for SWA and Lookahead using PreResNet-110 with α = 0.8 and k = 10 for Lookahead.The experiments follow the repository’s hyperparameter settings and use a higher-than-typical end-of-training learning rate.
Loading 1907.08610v2…