Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

Noam Shazeer; Mitchell Stern

Source-linked AI summary

Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

Noam Shazeer, Mitchell Stern

arXiv:1804.04235v1Apr 11, 2018cs.LGcs.AIstat.ML

TL;DR

Adaptive optimizers improve empirical training performance but require per-parameter state, creating a memory burden for large models. The paper factors matrix second moments and addresses instability from stale accumulators, achieving comparable results to Adam with much less auxiliary storage. It also scales updates relative to parameter magnitude.

Problem
Per-parameter second-moment estimators make adaptive optimization memory-intensive, especially as neural network models grow.
Method
The paper stores row- and column-sum moving averages for matrix parameters, adds instability remedies, drops momentum, and scales updates relative to parameter magnitude.
Results
The factored approach attains comparable convergence and stability results to Adam with sub-linear additional memory requirements.
Takeaways & Limitations
The approach supports training significantly larger models on the same memory-constrained hardware while retaining comparable quality to Adam.
Takeaways & Limitations
The proposed decay schedule must not increase too fast, because past gradients can then maintain a weight bounded away from zero throughout training.

Abstract

from arXiv · show

In several recently proposed stochastic optimization methods (e.g. RMSProp, Adam, Adadelta), parameter updates are scaled by the inverse square roots of exponential moving averages of squared past gradients. Maintaining these per-parameter second-moment estimators requires memory equal to the number of parameters. For the case of neural network weight matrices, we propose maintaining only the per-row and per-column sums of these moving averages, and estimating the per-parameter second moments based on these sums. We demonstrate empirically that this method produces similar results to the baseline. Secondly, we show that adaptive methods can produce larger-than-desired updates when the decay rate of the second moment accumulator is too slow. We propose update clipping and a gradually increasing decay rate scheme as remedies. Combining these methods and dropping momentum, we achieve comparable results to the published Adam regime in training the Transformer model on the WMT 2014 English-German machine translation task, while using very little auxiliary storage in the optimizer. Finally, we propose scaling the parameter updates based on the scale of the parameters themselves.

1. Introduction and Background

Adaptive optimizers can outperform SGD empirically but require substantial auxiliary memory. The paper proposes factored second-moment estimates, remedies for instability, and parameter-scale-based update scaling to reduce memory while retaining adaptivity.

Adam keeps two additional values per parameter, tripling memory requirements during training.
For an n × m matrix, factored second-moment storage reduces memory from O(nm) to O(n+m).The method tracks moving averages of row and column sums and reconstructs a low-rank accumulator approximation.
The factored approach achieves comparable performance to full accumulators on a large-scale machine translation task.
Removing momentum can cause training instability because second-moment accumulators become out of date.The paper proposes two remedies for this instability.
The paper scales update sizes relative to parameter scale because relative parameter change may be more relevant than absolute step size.

2. A Brief Review of Adam

Adam updates parameters using bias-corrected first- and second-moment estimates of stochastic gradients. Its step-size schedule may also use warmup followed by decay, which the paper finds yields more stable results in experiments.

Adam scales updates by the inverse square root of the bias-corrected second-moment estimate plus a regularization constant.
Adam maintains exponential running averages of the first and second gradient moments, then applies bias correction before updating parameters.
The optimizer receives a stochastic objective realization and gradient at each step, such as a minibatch loss and its gradient.
The experiments use a linear warmup followed by inverse square-root decay and find this schedule produces more stable results.

3. Factored Second Moment Estimation

The paper factors matrix second-moment accumulators into row and column statistics, reducing storage while preserving useful adaptive estimates. A Transformer experiment finds comparable results to Adam, although a column-mean simplification performs much worse.

Adaptive methods’ auxiliary storage becomes prohibitive as models grow, motivating a structure-exploiting low-memory alternative.
For a matrix parameter, standard Adam stores an exponential moving average of every squared gradient entry.
The rank-1 factorization has an analytic solution, whereas general rank-k factors require a nontrivial procedure such as alternating minimization.
The rank-1 projection exactly recovers rank-1 matrices and is expressed using row and column sums.
The factored method stores row- and column-sum moving averages, requiring space proportional to n + m rather than nm.The resulting approximation is compatible with exponential smoothing because row and column sums are linear functions of the accumulator.
The factored Transformer variants achieve results similar to baseline Adam, while the column-mean scheme performs much worse.The row-mean scheme is similar to baseline for this model.

4. No Momentum

The paper removes Adam’s first-moment accumulator by setting β1 = 0 and studies when momentum is necessary. Warmup preserves similar quality without momentum, whereas removing warmup makes training unstable.

Setting β1 = 0 turns off momentum and removes Adam’s first-moment accumulator.
BLEU = 23.6 vs. 23.4 with a warmup schedule shows similar model quality without and with momentum.
BLEU = 0.1 vs. 23.1 without warmup shows greater instability for the model without momentum.

5. A Problem with Adam: Out-of-Date Second Moment Estimator

The second-moment decay rate creates a trade-off: fast decay can impair convergence, while slow decay can make estimates stale and destabilize training. The resulting update errors are measurable through RMS(U_t), whose deviations from 1 indicate inaccurate scaling.

Fast decay (β2 = 0.9) has convergence problems, whereas slow decay (β2 = 0.999) has stability problems.
Slow decay bases the second-moment estimate on older gradients, which can produce smaller-than-desired or larger-than-desired updates as the model evolves.
RMS(U_t) measures the root-mean-square of the unscaled parameter updates, and values far from 1 signal that the second-moment estimator is inaccurate.
For a Transformer weight matrix, fast decay kept RMS(U_t) near 1, while slow decay caused substantial fluctuations; values above 1 indicate overly large updates.
Slow decay’s association with both oversized updates and instability supports, but does not prove, that oversized updates cause the instability.

6. Update Clipping

Update clipping scales down actual parameter updates when their RMS exceeds a threshold, addressing oversized updates more directly than gradient clipping. In fast-decay experiments without warmup, d = 1 substantially reduced instability, whereas d = 2 did not.

Update clipping scales down a weight-vector or matrix update whenever RMS(U_t) exceeds threshold d.
The actual parameter update is α_t times the clipped unscaled update.
6.1. Comparison to Gradient Clipping: Unlike gradient clipping, update clipping caps the norm of the actual adaptive update rather than only the gradient.
6.2. Experiments: d = 1 significantly ameliorated instability without learning-rate warmup, while d = 2 did not improve it.
6.2. Experiments: Update clipping did not significantly affect experiments with warmup, which had no instability problems.

7. Increasing Decay Parameter

The paper reformulates Adam’s bias correction as an increasing second-moment decay parameter and proposes schedules that combine early stability with later performance. The schedule must increase slowly enough for past-gradient contributions to vanish, and experiments show a t^-0.5 schedule can maintain stability and convergence.

7.2. Proposed Alternative: The proposed family of schedules uses a scalar c > 0 to control how quickly the second-moment decay parameter increases.
7.2. Proposed Alternative: The schedule starts with low ˆβ2t and increases toward 1, combining early stability with the later performance benefits of high decay.
7.2. Proposed Alternative: Past-gradient contributions vanish only when the schedule’s increase is sufficiently slow; the proposed condition requires c ≤ 1.If c > 1, past gradients retain weight bounded away from 0 throughout training.
7.2. Proposed Alternative: The schedule ˆβ2t = 1−t^-0.5 maintained stability and convergence, and with update clipping produced results similar to constant high β2 with clipping.

8. Relative Step Size

Relative step sizes define updates relative to parameter scale rather than with absolute step sizes. The paper combines this design with Adafactor and evaluates it under alternative embedding initialization and scaling schemes.

8. Relative Step Size: Relative step sizes are multiplied by the parameter scale, defined as the root-mean-square of the parameters with a lower bound ϵ2.The lower bound allows zero-initialized parameters to escape 0.
8. Relative Step Size: Adafactor combines relative step sizes with the paper’s other proposals, using recommended hyperparameters specified in Algorithm 6.
8.1. Experiments: The experiments compare Adam and Adafactor across the Transformer’s clever embedding scheme and two naive initialization-and-scaling schemes.

9. Experimental Setup

The experiments evaluate Adam, Adafactor, factored and unfactored variants, relative and absolute step sizes, and SGD on Transformer translation. Results indicate that sublinear-memory schemes can match Adam’s convergence and stability, while some decay and baseline choices are unstable.

9. Experimental Setup: The evaluation uses Transformer on WMT 2014 English-to-German translation with the latest Tensor2Tensor architecture.
9. Experimental Setup: Models were trained for 100,000 steps with batches containing approximately 4,096 input tokens and 4,096 target tokens.These batches were about eight times smaller than those used by Vaswani et al. (2017), speeding training to less than two hours per run on one Google TPU v2.
9. Experimental Setup: The experiments test linear-warmup inverse-square-root decay, flat-learning-rate schedules without warmup, relative step sizes, and SGD.
9.1. Results: The schemes with warmup mostly achieved very similar results, while fast decay of the second-moment estimator was significantly worse.
9.1. Results: Without warmup, the baseline became unstable, but momentum, fast decay, variable decay, and gradient clipping each relieved the instability.
9.1. Results: Rows (J) and (N) attained comparable convergence and stability to Adam with momentum while using sub-linear additional memory.
9.1. Results: SGD produced poorer and less stable results than Adam and depended strongly on the learning-rate choice.

10. Conclusion

On a popular machine translation task, the paper reports quality similar to Adam while using sublinear extra accumulator storage. It also introduces update clipping as a potentially broadly useful stabilization technique for adaptive methods.

10. Conclusion: Adafactor achieved similar quality to Adam on a popular machine translation task using a sublinear amount of extra accumulator space.
10. Conclusion: The reduced storage is intended to enable significantly larger models on the same memory-constrained hardware.
10. Conclusion: The paper introduces update clipping as a potentially more generally useful technique for stabilizing adaptive gradient methods.