A Brief Analysis of the AdaFactor Optimizer (with Open-Source Implementation)
Since pretrained models such as GPT and BERT became popular, one clear trend has been that models keep getting bigger, because larger models combined with more thorough pretraining tend to top the leaderboards more effectively. However, while ideals can be boundless, reality is often cramped: sometimes a model is so large that even if you own a GPU with huge memory, or even a TPU, you can still feel utterly hopeless. For instance, the largest version of GPT2 has 1.5 billion parameters, and the largest T5 model even reaches 11 billion parameters. At that scale, even on a TPU cluster you simply can't push the batch size very high.
At this point, we usually turn to the optimization process itself — for example, using mixed-precision training (in TensorFlow you can also use a new floating-point format called bfloat16), which saves GPU memory while also speeding up training; or using an optimizer that is more memory-efficient, such as RMSProp, which uses less memory than Adam. This post introduces AdaFactor, a new optimizer proposed by Google, first presented in the paper Adafactor: Adaptive Learning Rates with Sublinear Memory Cost. AdaFactor has the property of adaptive learning rates, but is even more memory-efficient than RMSProp, and it specifically addresses a number of shortcomings of Adam.
Adam
Let's first review the update process of the commonly used Adam optimizer. Let $t$ denote the iteration step, $\alpha_t$ the current learning rate, $L(\theta)$ the loss function, $\theta$ the parameters being optimized, and $\epsilon$ a small positive number that prevents overflow. Then the update process of Adam is
\begin{equation}\left\{\begin{aligned}&g_t = \nabla_{\theta} L(\theta_{t-1})\\ &m_t = \beta_1 m_{t-1} + \left(1 - \beta_1\right) g_t\\ &v_t = \beta_2 v_{t-1} + \left(1 - \beta_2\right) g_t^2\\ &\hat{m}_t = m_t\left/\left(1 - \beta_1^t\right)\right.\\ &\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.\\ &\theta_t = \theta_{t-1} - \alpha_t \hat{m}_t\left/\left(\sqrt{\hat{v}_t} + \epsilon\right)\right. \end{aligned}\right.\end{equation}
To save GPU memory, we first need to know where the memory is being spent. First of all, both the computational cost and the bulk of the memory usage clearly go to $\nabla_{\theta} L(\theta_{t-1})$ — that is, computing the gradient is quite resource-intensive, which is also why "although ALBERT has far fewer parameters than BERT, its training speed isn't noticeably faster." Besides that, the memory consumption mainly comes from $m,v$: we need to maintain two sets of cache variables to compute the moving averages of the first two moments of the gradient (namely $m$ and $v$), which are used to compute the parameter updates. Each of these two sets of variables is as large as the trainable parameters themselves, so for models with a large parameter count, the memory consumed by these two cache sets is substantial.
AdaFactor
In this section we'll introduce the AdaFactor optimizer in relatively fine detail, involving quite a few formulas and derivations. Readers who only want a rough overview can feel free to skip some of the mathematical content.
Dropping momentum
As we know, CV models often rely on "SGD + momentum" to reach their best performance, and adaptive-learning-rate optimizers usually don't train to the best results there. But for NLP models, the situation is somewhat reversed — adaptive learning rates seem to matter more, and it's rare to hear of tuning NLP models purely with SGD. So, as a first step toward saving memory, we can drop the momentum term in Adam, which removes one set of cache parameters and naturally saves memory:
\begin{equation}\left\{\begin{aligned}&g_t = \nabla_{\theta} L(\theta_{t-1})\\ &v_t = \beta_2 v_{t-1} + \left(1 - \beta_2\right) g_t^2\\ &\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.\\ &\theta_t = \theta_{t-1} - \alpha_t g_t\left/\sqrt{\hat{v}_t + \epsilon}\right. \end{aligned}\right.\end{equation}
This is essentially a variant of RMSProp, with the extra step of $\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.$ compared to RMSProp.
Low-rank decomposition
After dropping $m$, the cache variables are immediately cut in half, but AdaFactor isn't satisfied yet — it wants to keep the adaptive learning rate feature while further compressing the parameter count of the cache variable $v$. This time, it makes use of low-rank matrix decomposition.
Generalized KL divergence
In SGD, all parameters share a single scalar learning rate; in Adam, each parameter has its own learning rate $\alpha_t\left/\sqrt{\hat{v}_t + \epsilon}\right.$. We know that with careful tuning of the learning rate, SGD can actually achieve quite good results too, which suggests that "every parameter having its own learning rate" isn't all that critical — or put another way, "finely tuning each parameter's own learning rate" isn't especially important.
This suggests that it might be enough to replace $\hat{v}_t$ with some approximation that has fewer parameters. And the natural candidate for "an approximation with fewer parameters" is low-rank decomposition. For a $m\times n$ matrix $C$, we want to find a $m\times k$ matrix $A$ and a $k\times n$ matrix $B$ such that
\begin{equation}AB \approx C\end{equation}
When $k$ is small enough, the total number of parameters in $A$ and $B$ is smaller than the parameter count of $C$. To save as much as possible, AdaFactor goes all the way and sets $k=1$, i.e., it looks for $\{a_i\}_{i=1}^m$ and $\{b_j\}_{j=1}^n$ such that
\begin{equation}a_i b_j \approx c_{i,j}\end{equation}
Since we're approximating something, we need a metric to measure the approximation quality. The most obvious candidate is Euclidean distance:
\begin{equation}\sum_{i,j} (a_i b_j - c_{i,j})^2\end{equation}
But under this distance, $a_i,b_j$ has no closed-form solution. Moreover, during optimization $c_{i,j}$ (i.e., $\hat{v}_t$) is non-negative, while the $a_i b_j$ obtained by optimizing the objective above cannot be guaranteed to be non-negative, which could easily disrupt the optimization process.
The authors of the original paper cleverly switched to a different metric, one under which $a_i,b_j$ does have a closed-form solution. Specifically, they used the "generalized KL divergence," also known as the "I-divergence" (I-Divergence), which takes the form:
\begin{equation}l = \sum_{i,j} c_{i,j}\log \frac{c_{i,j}}{a_i b_j} - c_{i,j} + a_i b_j \label{eq:i-div}\end{equation}
This metric originates from the inequality $x\log x\geq x - 1(\forall x > 0)$, with equality holding if and only if $x=1$. So substituting $x = p / q\,(p,q > 0)$ and multiplying both sides by $q$, we get
\begin{equation}p\log \frac{p}{q} - p + q \geq 0\end{equation}
with equality holding if and only if $p=q$. If $p,q$ has multiple components, we simply sum the results over all components, which gives us the metric $\eqref{eq:i-div}$.
Clearly, the generalized KL divergence is a natural extension of the KL divergence between probability distributions, except that it doesn't require $c_{i,j}$ and $a_i b_j$ to be normalized — it only requires them to be non-negative, which fits the AdaFactor scenario perfectly. And, elegantly, this setting paired with this objective happens to have a closed-form solution:
\begin{equation}a_i = \sum\limits_{j}c_{i,j},\quad b_j = \frac{\sum\limits_{i}c_{i,j}}{\sum\limits_{i,j}c_{i,j}}\label{eq:aibj}\end{equation}
This closed-form solution is actually quite intuitive: it's simply the row sums and column sums multiplied together, then divided by the grand total.
Derivation
Taking the partial derivative directly with respect to $\eqref{eq:i-div}$ and setting it to zero, we get
\begin{equation}\left\{\begin{aligned} &\frac{\partial l}{\partial a_i}=\sum_j -\frac{c_{i,j}}{a_i} + b_j = 0\\ &\frac{\partial l}{\partial b_j}=\sum_i -\frac{c_{i,j}}{b_j} + a_i = 0 \end{aligned}\right.\end{equation}
which rearranges to
\begin{equation}\left\{\begin{aligned} &a_i \sum_{j} b_j = \sum_j c_{i,j}\\ &b_j \sum_{i} a_i = \sum_i c_{i,j} \end{aligned}\right.\end{equation}
Notice that if $(a_i,b_j)$ is one optimal solution, then $(\lambda a_i,b_j/\lambda)$ is also optimal — in other words, if all the $a_i$'s are multiplied by a constant and all the $b_j$'s are divided by the same constant, $a_i b_j$ remains unchanged. So we're free to fix either $\sum\limits_{i} a_i$ or $\sum\limits_{j} b_j$, since they're just a scaling factor. Without loss of generality, let's fix $\sum\limits_{j} b_j=1$, and then we solve for $\eqref{eq:aibj}$.
Intuitive understanding
We can also understand the result $\eqref{eq:aibj}$ from another angle. Since $c_{i,j}$ is non-negative, we can normalize it so that it behaves like a probability distribution, i.e., $\hat{c}_{i,j}=\frac{c_{i,j}}{\sum\limits_{i,j}c_{i,j}}$, and then try to perform the factorization $\hat{c}_{i,j}\approx \hat{a}_i \hat{b}_j$. Since $\hat{c}_{i,j}$ now acts as a joint probability distribution over two variables, $\hat{a}_i,\hat{b}_j$ correspond to their marginal distributions, i.e.,
\begin{equation}\hat{a}_i = \sum_j \hat{c}_{i,j} = \frac{\sum\limits_{j}c_{i,j}}{\sum\limits_{i,j} c_{i,j}},\quad \hat{b}_j = \sum_i \hat{c}_{i,j} = \frac{\sum\limits_{i}c_{i,j}}{\sum\limits_{i,j}c_{i,j}}\end{equation}
Now, going from $\hat{c}_{i,j}$ to $c_{i,j}$ requires multiplying by an extra factor of $\sum\limits_{i,j}c_{i,j}$, which we can fold into either $\hat{a}_i$ or $\hat{b}_j$. Without loss of generality, suppose we fold it into $\hat{a}_i$, giving us $\eqref{eq:aibj}$.
The prototype of AdaFactor
With the result $\eqref{eq:aibj}$ in hand, we can now use it to build a more memory-efficient optimizer — this is the prototype of AdaFactor. In short, when a parameter $\theta$ is an ordinary one-dimensional vector, the optimization process stays the same; but when $\theta$ is a $m\times n$ matrix, the computed gradient $g_t$ is also a matrix, and hence $g_t^2$ is a matrix too. In this case we apply low-rank decomposition to $g_t^2$, then maintain two sets of cache variables $v^{(r)}_t\in \mathbb{R}^m,v^{(c)}_t\in\mathbb{R}^n$, each taking the moving average of the decomposed results, and finally use $v^{(r)}_t,v^{(c)}_t$ jointly to adjust the learning rate:
\begin{equation}\left\{\begin{aligned}&g_{i,j;t} = \nabla_{\theta} L(\theta_{i,j;t-1})\\ &v^{(r)}_{i;t} = \beta_2 v^{(r)}_{t-1;i} + \left(1 - \beta_2\right) \sum\limits_{j}\left(g_{i,j;t}^2+\epsilon\right)\\ &v^{(c)}_{j;t} = \beta_2 v^{(c)}_{t-1;j} + \left(1 - \beta_2\right) \sum\limits_{i}\left(g_{i,j;t}^2+\epsilon\right)\\ &v_{i,j;t} = v^{(r)}_{i;t} v^{(c)}_{j;t}\left/\sum\limits_{j}v^{(c)}_{j;t}\right.\\ &\hat{v}_t = v_t\left/\left(1 - \beta_2^t\right)\right.\\ &\theta_t = \theta_{t-1} - \alpha_t g_t\left/\sqrt{\hat{v}_t}\right. \end{aligned}\right.\end{equation}
(Adding $\epsilon$ to $g_t^2$ rather than to $\hat{v}_t$ is a design choice made by AdaFactor itself — don't blame the author for it!)
The moving-average weight
In Adam, as well as in the AdaFactor prototype above, the moving-average weight $\beta_2$ is always held constant. AdaFactor argues that this isn't well justified, and proposes a new strategy.
An equivalent form
To see this, let's rewrite Adam's update process for $\hat{v}_t$:
\begin{equation}\begin{aligned} \hat{v}_t =& v_t\left/\left(1 - \beta_2^t\right)\right.\\ =&\frac{\beta_2 v_{t-1} + (1-\beta_2) g_t^2}{1 - \beta_2^t}\\ =&\frac{\beta_2 \hat{v}_{t-1}\left(1 - \beta_2^{t-1}\right) + (1-\beta_2) g_t^2}{1 - \beta_2^t}\\ =&\beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}\hat{v}_{t-1} + \left(1 - \beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}\right)g_t^2 \end{aligned}\end{equation}
So, if we set $\hat{\beta}_{2,t}=\beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}$, the update formula becomes
\begin{equation}\hat{v}_t =\hat{\beta}_{2,t}\hat{v}_{t-1} + \left(1 - \hat{\beta}_{2,t}\right)g_t^2\end{equation}
The question is: is this choice of $\hat{\beta}_{2,t}$ reasonable enough? The answer is probably not quite. When $t=1$, $\hat{\beta}_{2,t}=0$, and in that case $\hat{v}_t$ is exactly $g_t^2$ — that is, the current gradient alone is used to correct the learning rate, giving the strongest possible correction. When $t\to\infty$, $\hat{\beta}_{2,t}\to \beta_2$, and in that case $v_t$ is a weighted average of the accumulated squared gradients and the current squared gradient. Since $\beta_2 < 1$, this means the weight $1 - \beta_2$ on the current gradient never reaches 0, which can lead to instability in training, because as training proceeds gradients shrink and training itself becomes more stable, so the correction applied to the learning rate should also shrink; ideally, $t\to\infty$ should hold, meaning that as $t\to\infty$, the learning rate correction should become negligible, i.e., $\hat{\beta}_{2,t}\to 1$, essentially degenerating into a fixed, constant learning rate (equivalent to SGD).
The new decay strategy
To achieve this, AdaFactor adopts the following decay strategy:
\begin{equation}\hat{\beta}_{2,t} =1 - \frac{1}{t^c}\label{eq:beta2}\end{equation}
which satisfies $\hat{\beta}_{2,1}=0,\lim\limits_{t\to\infty} \hat{\beta}_{2,t}=1$. But even so, not every choice of $c$ works — we must also have $0 < c <1$. $c > 0$ is easy to understand, but why do we need $c < 1$? The original paper includes an analysis of this, which readers can look up, but the author finds the original derivation somewhat obscure, so here's my own take on it.
First, for $\hat{v}_t$, the most obvious approach one might think of is a simple average of all squared gradients so far, i.e.:
\begin{equation}\hat{v}_t = \frac{1}{t}\sum_{i=1}^t g_i^2=\frac{t-1}{t}\hat{v}_{t-1} + \frac{1}{t}g_t^2\end{equation}
which is equivalent to setting $\hat{\beta}_{2,t} =1 - \frac{1}{t}$. The one drawback of this scheme is that every step's gradient gets equal weight, which is counterintuitive, since intuitively more distant gradients should matter less. So we should reduce the weight of the historical part somewhat, and since $c < 1$, $1 - \frac{1}{t^c} < 1 - \frac{1}{t}$, a simple approach is to take $c < 1$ in equation $\eqref{eq:beta2}$. AdaFactor's default value of $c$ is $0.8$.
Layer-wise adaptivity
Finally, we can further correct the update magnitude according to the norm of the parameters — this idea comes from the LAMB optimizer, which was also introduced in an earlier post A Simple Introduction to Six Derived Optimizers, with Implementation. Put simply, it normalizes the final update, then multiplies it by the parameter norm — in other words, no matter how you tweak things, only the direction of the final update matters, while its magnitude is jointly determined by the parameter's own norm and the preset learning rate, keeping the relative degree of change consistent across all layers and all parameters.
The complete AdaFactor algorithm
With that, we can finally write out the complete update process of AdaFactor:
\begin{equation}\left\{\begin{aligned}&g_{i,j;t} = \nabla_{\theta} L(\theta_{i,j;t-1})\\ &\hat{\beta}_{2,t} =1 - t^{-c}\\ &v^{(r)}_{i;t} = \hat{\beta}_{2,t} v^{(r)}_{t-1;i} + \left(1 - \hat{\beta}_{2,t}\right) \sum\limits_{j}\left(g_{i,j;t}^2+\epsilon_1\right)\\ &v^{(c)}_{j;t} = \hat{\beta}_{2,t} v^{(c)}_{t-1;j} + \left(1 - \hat{\beta}_{2,t}\right) \sum\limits_{i}\left(g_{i,j;t}^2+\epsilon_1\right)\\ &\hat{v}_{i,j;t} = v^{(r)}_{i;t} v^{(c)}_{j;t}\left/\sum\limits_{j}v^{(c)}_{j;t}\right.\\ &u_t = g_t\left/\sqrt{\hat{v}_t}\right.\\ &\hat{u}_t = u_t \left/\max\left(1, \left. RMS(u_t)\right/d\right)\right.\times \max\left(\epsilon_2, RMS(\theta_{t-1})\right)\\ &\theta_t = \theta_{t-1} - \alpha_t \hat{u}_t \end{aligned}\right.\end{equation}
Here $RMS(x)=\sqrt{\frac{1}{n}\sum\limits_{i=1}^n x_i^2}$ is a variant of the norm, and the step $\max\left(1, \left. RMS(u_t)\right/d\right)$ amounts to a clipping operation, i.e., normalization is only applied when $RMS(u_t) > d$. The default parameters in the original paper are
$$\begin{array}{c|c} \hline \epsilon_1 & 10^{-30}\\ \hline \epsilon_2 & 10^{-3}\\ \hline d & 1\\ \hline \hat{\beta}_{2,t} & 1 - t^{-0.8}\\ \hline \end{array}$$
If the parameter is an ordinary one-dimensional vector rather than a matrix, then $\hat{v}_t$ just uses the plain update formula $\hat{v}_t = \hat{\beta}_{2,t} v_{t-1} + \left(1 - \hat{\beta}_{2,t}\right) \left(g_t^2+\epsilon_1\right)$. Additionally, the paper proposes that if no learning rate is supplied, $a_t = \min\left(10^{-2},\frac{1}{\sqrt{t}}\right)$ can be used as a default learning rate, but from reading the source code, the author found this default is rarely used in practice — you basically still need to supply your own learning rate.
Open-source implementation
For everyone's convenience, the author has open-sourced their own implementation of AdaFactor:
GitHub repository: https://github.com/bojone/adafactor
The open-source release includes both a pure-Keras version and a tf.keras version; usage is the same as any ordinary Keras optimizer, and the tf.keras version can also be used as a plain TensorFlow optimizer. The implementation is based on the mesh_tensorflow source code, for which many thanks. The optimizer has also been built into bert4keras for easy use.
One thing worth noting: when using AdaFactor, it's best to use a fairly large batch_size, because the low-rank decomposition itself introduces error, and if the batch_size is too small, the gradient estimate itself also carries substantial error — the combination of the two might prevent the optimization process from converging at all. For pretraining, batch_size is typically already large, which is why quite a few pretrained models have started using AdaFactor; for ordinary downstream tasks, it's also worth trying AdaFactor, though it may take some extra experimentation ("alchemy") to beat plain vanilla Adam. Also, one more reminder: when using AdaFactor, set the learning rate on the larger side — around $10^{-3}$ is generally a good scale, even during the fine-tuning stage.
Summary
This post introduced Google's AdaFactor optimizer, an optimizer designed to reduce memory usage, which also specifically analyzes and addresses a number of shortcomings of Adam. The author believes that the analysis AdaFactor carries out with respect to Adam is quite a classic piece of work, well worth studying carefully, and it serves as an invaluable case study for readers interested in optimization research.
Of course, there's no method that's guaranteed to work — there's only
however good the method, making it actually work in practice still takes careful "alchemy."
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.