YellowFin and the Art of Momentum Tuning

TL;DR

This work tackles the substantial burden of hyperparameter tuning in deep learning by examining momentum SGD and demonstrating that a single learning rate and momentum can rival adaptive optimizers like Adam. It introduces YellowFin, an automatic tuner that uses gradient-derived curvature and variance statistics to set both $\alpha$ and $\mu$ on the fly, with an optional closed-loop variant for asynchronous training that tracks total momentum $\mu_T$. The authors establish robustness properties of the momentum operator via generalized curvature and a robust region, then validate YellowFin across ResNets and LSTMs, achieving up to 3.28x synchronous and 2.69x asynchronous speedups over Adam while maintaining stability. The approach avoids manual tuning, provides practical measurement tools, and reports strong empirical gains across multiple tasks, supporting broader adoption of momentum-based optimization with automatic tuning.

Abstract

Hyperparameter tuning is one of the most time-consuming workloads in deep learning. State-of-the-art optimizers, such as AdaGrad, RMSProp and Adam, reduce this labor by adaptively tuning an individual learning rate for each variable. Recently researchers have shown renewed interest in simpler methods like momentum SGD as they may yield better test metrics. Motivated by this trend, we ask: can simple adaptive methods based on SGD perform as well or better? We revisit the momentum SGD algorithm and show that hand-tuning a single learning rate and momentum makes it competitive with Adam. We then analyze its robustness to learning rate misspecification and objective curvature variation. Based on these insights, we design YellowFin, an automatic tuner for momentum and learning rate in SGD. YellowFin optionally uses a negative-feedback loop to compensate for the momentum dynamics in asynchronous settings on the fly. We empirically show that YellowFin can converge in fewer iterations than Adam on ResNets and LSTMs for image recognition, language modeling and constituency parsing, with a speedup of up to 3.28x in synchronous and up to 2.69x in asynchronous settings.

Figures (15)

• Figure 1: YellowFin in comparison to Adam on a ResNet (CIFAR100, cf. Section \ref{['sec:experiments']}) in synchronous and asynchronous settings.
• Figure 1: German-English translation validation metrics using convolutional seq-to-seq model.
• Figure 2: Spectral radius of momentum operator on scalar quadratic for varying $\alpha$.
• Figure 3: (a) Non-convex toy example; (b) linear convergence rate achieved empirically on the example in (a) tuned according to \ref{['eqn:noiseless_tuning_rule']}; (c,d) LSTM on MNIST: as momentum increases from $0.9$ to $0.99$, the global learning rate and momentum falls in robust regions of more model variables. The convergence behavior (shown in grey) of these variables follow the robust rate $\sqrt{\mu}$ (shown in red).
• Figure 4: When running YellowFin, total momentum $\hat{\mu}_t$ equals algorithmic value in synchronous settings (left); $\hat{\mu}_t$ is greater than algorithmic value on 16 asynchronous workers (middle). Closed-loop YellowFin automatically lowers algorithmic momentum and brings total momentum to match the target value (right). Red dots are total momentum estimates, $\hat{\mu}_T$, at each iteration. The solid red line is a running average of $\hat{\mu}_T$.

Theorems & Definitions (14)

• Definition 1: Convergence rate
• Definition 2: Generalized curvature
• Lemma 3: Robustness of the momentum operator
• Definition 4: Generalized condition number
• Lemma 5
• Lemma 6
• Lemma 7
• proof
• Lemma 8
• proof