[Paper Review] Adafactor: Adaptive Learning Rates with Sublinear Memory Cost

Noam Shazeer, Mitchell Stern|arXiv (Cornell University)|Apr 11, 2018

Stochastic Gradient Optimization Techniques10 references163 citations

TL;DR

Adafactor introduces a memory-efficient adaptive optimizer by factoring second-moment estimates for matrix parameters into per-row and per-column sums, enabling sublinear memory usage with competitive performance to Adam on Transformer training; it also proposes update clipping and decay-rate schedules to stabilize training and introduces relative step sizes (Adafactor) for scale-aware updates.

ABSTRACT

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.

Motivation & Objective

• Motivate memory constraints in adaptive gradient methods as model sizes grow.
• Propose a factored second-moment estimator to reduce per-parameter memory from O(nm) to O(n+m) for matrix parameters.
• Identify instability issues when second-moment decays are too slow and provide remedies (update clipping and adaptive decay).
• Demonstrate that Adafactor achieves comparable results to Adam on Transformer training with substantially less memory.
• Extend optimization with relative step sizes that scale updates to parameter magnitudes.

Proposed method

• Introduce a factored second-moment representation V ≈ RS where R ∈ R^{n×k} and S ∈ R^{k×m} with k ≪ n,m for matrix-valued parameters.
• Derive an analytic solution for the rank-1 (k=1) case ensuring compatibility with exponential smoothing for row/column sums (V1_m1_n^T V / 1_n^T V 1_m).
• Implement Adam with factored second moments using per-row and per-column accumulators (R_t and C_t) and normalize to form actoredullV_t = (R_t C_t)/(1_n^T R_t).
• Propose update clipping to cap large unscaled updates when RMS(U_t) exceeds threshold d.
• Propose increasing decay schedules for the second moment (ast, ollow Reddi et al.) and alternative schedules to stabilize training.
• Define Adafactor as a relative-step optimizer where the actual step alpha_t is computed from parameter-scale RMS and a relative step rho_t, with per-parameter updates U_t = G_t / sqrt(V_hat_t) and optional clipping.

Experimental results

Research questions

• RQ1 Can a factored (per-row/per-column) second-moment estimate match the performance of full second-moment accumulators in adaptive optimizers?
• RQ2 Does reducing memory via factorization affect convergence and model quality on large-scale tasks like Transformer training?
• RQ3 What stability issues arise when decoupling momentum and using adaptive learning rates, and how can update clipping and decay-rate schedules mitigate them?
• RQ4 Do relative step sizes that scale with parameter magnitudes improve robustness across differently scaled parameter initializations?
• RQ5 How do the proposed remedies (update clipping, increasing decay, relative step sizes) interact in practice on a modern neural MT task?

Key findings

• Factored second-moment estimation reduces memory from O(nm) to O(n+m) for matrix parameters while delivering comparable Transformer BLEU scores to full-accumulator Adam.
• Removing momentum can destabilize training, but update clipping and appropriate decay schedules restore stability.
• Update clipping with threshold d improved stability in non-warmup setups; with d=1 it significantly mitigated instability (d=2 showed no improvement).
• Relative step sizes combined with Adafactor maintained competitive performance and were robust across different embedding parameter scales.
• An increasing decay schedule for the second moment (e.g., 1 - t^{-c}) yielded stable and convergent results for certain c values, notably c=0.5 (and related variants) in combination with clipping.
• Using a rank-1 or factorized representation and the proposed scaling allowed training of Transformer models with sublinear auxiliary storage while achieving BLEU scores close to Adam-based baselines.