Analysis of the AdaFactor Optimizer (with Open Source Implementation)

Generalized KL Divergence

In SGD, all parameters share a single scalar learning rate. In Adam, every parameter has its own learning rate \alpha_t / \left(\sqrt{\hat{v}_t} + \epsilon\right). We know that through fine-tuning the learning rate, SGD can actually achieve good results. This suggests that the fact that "every parameter has its own learning rate" is not particularly critical, or in other words, "fine-tuning every single parameter’s individual learning rate" is not vital.

This inspires us to replace \hat{v}_t with an approximation that uses fewer parameters. For an approximation with fewer parameters, low-rank decomposition naturally comes to mind. For an m \times n matrix C, we hope to find an m \times k matrix A and a k \times n matrix B such that: AB \approx C When k is sufficiently small, the total number of parameters in A and B is less than that of C. To save as much as possible, AdaFactor directly sets k=1, seeking \{a_i\}_{i=1}^m and \{b_j\}_{j=1}^n such that: a_i b_j \approx c_{i,j} Since we are approximating, we need a metric. A common metric is the Euclidean distance: \sum_{i,j} (a_i b_j - c_{i,j})^2 However, under this distance, a_i, b_j do not have an analytical solution. Furthermore, during the optimization process, c_{i,j} (i.e., \hat{v}_t) is non-negative, but a_i b_j optimized via the above objective cannot be guaranteed to be non-negative, which might disrupt the optimization process.

The authors of the original paper cleverly switched to a different metric that allows for an analytical solution for a_i, b_j. Specifically, they used the "Generalized KL Divergence," also known as "I-Divergence" (I-Divergence), which takes the form: 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} This metric originates from the inequality x \log x \geq x - 1 (\forall x > 0), where equality holds if and only if x=1. By substituting x = p/q (p, q > 0) and multiplying both sides by q, we get: p \log \frac{p}{q} - p + q \geq 0 Equality holds if and only if p=q. If p, q have multiple components, summing the results for all components yields the metric in Eq. [eq:i-div].

Clearly, Generalized KL Divergence is a natural extension of the KL divergence for probabilities, but it does not require c_{i,j} and a_i b_j to be normalized; it only requires them to be non-negative, which perfectly fits the AdaFactor scenario. Ingeniously, under this objective, there is an analytical solution: a_i = \sum_{j}c_{i,j},\quad b_j = \frac{\sum_{i}c_{i,j}}{\sum_{i,j}c_{i,j}}\label{eq:aibj} This analytical solution is quite intuitive: sum the rows and columns respectively, multiply them, and then divide by the total sum.

Derivation Process

Taking the partial derivatives of Eq. [eq:i-div] and setting them to zero: \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. Rearranging gives: \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. Note that if (a_i, b_j) is an optimal solution, then (\lambda a_i, b_j/\lambda) is also one. Simply put, if all a_i are multiplied by a constant and all b_j are divided by the same constant, a_i b_j remains unchanged. Thus, we can arbitrarily specify \sum_i a_i or \sum_j b_j, as they are merely scaling scalars. Without loss of generality, if we specify \sum_j b_j = 1, we obtain Eq. [eq:aibj].

Intuitive Understanding

We can also understand the result in Eq. [eq:aibj] from another perspective. Since c_{i,j} is non-negative, we can normalize it to have the properties of a probability distribution, i.e., \hat{c}_{i,j} = \frac{c_{i,j}}{\sum_{i,j}c_{i,j}}. We then attempt the decomposition \hat{c}_{i,j} \approx \hat{a}_i \hat{b}_j. Since \hat{c}_{i,j} now acts as a bivariate joint probability distribution, \hat{a}_i and \hat{b}_j correspond to their marginal distributions: \hat{a}_i = \sum_j \hat{c}_{i,j} = \frac{\sum_{j}c_{i,j}}{\sum_{i,j} c_{i,j}},\quad \hat{b}_j = \sum_i \hat{c}_{i,j} = \frac{\sum_{i}c_{i,j}}{\sum_{i,j}c_{i,j}} To get from \hat{c}_{i,j} back to c_{i,j}, we need to multiply by \sum_{i,j}c_{i,j}. We can multiply this into either \hat{a}_i or \hat{b}_j. Without loss of generality, assuming we multiply it into \hat{a}_i, we obtain Eq. [eq:aibj].

AdaFactor Prototype

With the result in Eq. [eq:aibj], we can construct a more memory-efficient optimizer, which is the prototype of AdaFactor. Simply put, when the parameter \theta is a standard 1D vector, the optimization process remains unchanged. However, when \theta is an m \times n matrix, the calculated gradient g_t is also a matrix, and thus g_t^2 is a matrix. We then perform low-rank decomposition on g_t^2 and maintain two sets of buffer variables v^{(r)}_t \in \mathbb{R}^m and v^{(c)}_t \in \mathbb{R}^n, which are moving averages of the low-rank decomposition results. Finally, v^{(r)}_t and v^{(c)}_t are used together to adjust the learning rate: \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_{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_{i}\left(g_{i,j;t}^2+\epsilon\right)\\ &v_{i,j;t} = v^{(r)}_{i;t} v^{(c)}_{j;t} / \sum_{j}v^{(c)}_{j;t}\\ &\hat{v}_t = v_t / \left(1 - \beta_2^t\right)\\ &\theta_t = \theta_{t-1} - \alpha_t g_t / \sqrt{\hat{v}_t} \end{aligned}\right. (Adding \epsilon to g_t^2 rather than \hat{v}_t is the form chosen by AdaFactor, not a mistake by the author.)

Equivalent Form

To recognize this, let’s rewrite the update process for \hat{v}_t in Adam: \begin{aligned} \hat{v}_t &= v_t / \left(1 - \beta_2^t\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} So if we set \hat{\beta}_{2,t} = \beta_2\frac{1 - \beta_2^{t-1}}{1 - \beta_2^t}, the update formula becomes: \hat{v}_t = \hat{\beta}_{2,t}\hat{v}_{t-1} + \left(1 - \hat{\beta}_{2,t}\right)g_t^2 The question is whether this \hat{\beta}_{2,t} is reasonable. The answer is likely not. When t=1, \hat{\beta}_{2,t}=0, and \hat{v}_t is simply g_t^2, meaning the real-time gradient is used to correct the learning rate, which is the maximum correction strength. As t \to \infty, \hat{\beta}_{2,t} \to \beta_2. Here, v_t is a weighted average of the accumulated squared gradients and the current squared gradient. Since \beta_2 < 1, the weight of the current gradient 1 - \beta_2 is non-zero. This may lead to training instability because, in the later stages of training, gradients become smaller and training itself tends to stabilize; the significance of correcting the learning rate diminishes. Therefore, the correction strength should decrease, and as t \to \infty, the learning rate should ideally become constant (equivalent to decaying into SGD). This requires \hat{\beta}_{2,t} \to 1 as t \to \infty.

New Decay Strategy

To achieve this, AdaFactor adopts the following decay strategy: \hat{\beta}_{2,t} = 1 - \frac{1}{t^c} \label{eq:beta2} It satisfies \hat{\beta}_{2,1}=0 and \lim_{t\to\infty} \hat{\beta}_{2,t}=1. However, not just any c is suitable; it must be 0 < c < 1. While c > 0 is easy to understand, why c < 1? The original paper contains an analysis of this, which readers can explore, but I find the derivation in the paper somewhat obscure, so I will provide my own interpretation.

First, for \hat{v}_t, a simple idea is to use the average of all squared gradients: \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 This is equivalent to setting \hat{\beta}_{2,t} = 1 - \frac{1}{t}. The downside of this scheme is that every step’s gradient is weighted equally, which is counter-intuitive because older gradients should generally be less important. Thus, the weight of the historical part should be appropriately reduced. When c < 1, 1 - \frac{1}{t^c} < 1 - \frac{1}{t}. Therefore, a concise solution is to take c < 1 in Eq. [eq:beta2]. AdaFactor’s default c is 0.8.