Making "Alchemy" More Scientific (IX): Classic Adaptive Gradient Algorithms

By 苏剑林 | September 05, 2026

The previous eight articles in this series centered around SGD and its learning rate. Starting from this article, we will officially enter the world of adaptive gradient algorithms. It can be said that all modern adaptive gradient algorithms share a common origin: the 2011 classic work "Adaptive Subgradient Methods for Online Learning and Stochastic Optimization"—the famous AdaGrad paper.

This paper generalizes the classic convergence conclusions of SGD to a general preconditioned matrix form, subsequently deriving that the optimal preconditioned matrix is precisely the second moment of the gradients (to the power of −1/2), and then applies a diagonal approximation to obtain AdaGrad. While AdaGrad itself might not be the most popular today, the subsequent works it inspired, such as RMSProp and Adam, are the undisputed mainstream optimizers. Furthermore, newer optimizers like Shampoo and KL-Shampoo can still be considered its descendants.

In this article, let us revisit this classic masterpiece together.

Difficulty Analysis

In previous articles, we briefly discussed the idea of adaptive learning rates, such as in "Making Alchemy More Scientific (V): Fine-tuning Learning Rates Based on Gradients", but that was still within the SGD framework, tuning a scalar learning rate. This time, we consider a general update rule:

Where $\boldsymbol{H}_t$ is a positive definite symmetric matrix, and $\boldsymbol{H}_t^{-1}$ is usually called the "preconditioner." We can view it as a matrix-valued learning rate; clearly, when $\boldsymbol{H}_t=\boldsymbol{I}$, it reduces to SGD.

Why expand to a matrix learning rate? From an input-output perspective, we want to input a gradient vector and output an update vector. The most general linear transformation between vectors is matrix multiplication, so a matrix learning rate is the ultimate generalization. As for the restriction of positive definiteness, it ensures that $\boldsymbol{g}^{\top}\boldsymbol{H}_t^{-1}\boldsymbol{g}\geq 0$ for any $\boldsymbol{g}$, so that $-\boldsymbol{H}_t^{-1}\boldsymbol{g}$ remains a direction that can decrease the loss.

In this section, let's briefly analyze the difficulties encountered if we directly apply the proof process of SGD to a general $\boldsymbol{H}_t$. The starting point for all previous proofs was the following identity:

We then separate the term $(\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$ and relate it to the loss value via the convexity assumption:

However, when we consider the general update rule \eqref{eq:H-g}, if we use the same identity, we get $(\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot \boldsymbol{H}_t^{-1}\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$. Due to the presence of $\boldsymbol{H}_t^{-1}$, even if we assume the convexity of the loss function, we cannot relate this term to the loss function—the right side of the convexity inequality \eqref{leq:convex} must be a "clean" $(\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$; sandwiching $\boldsymbol{H}_t^{-1}$ in the middle prevents its application.

Generalized Identity

Therefore, we must generalize this identity simultaneously. The strategy is to use the identity \(\newcommand{tr}{\mathop{\text{tr}}}\boldsymbol{x}\cdot \boldsymbol{y} = \boldsymbol{x}^{\top}\boldsymbol{y} = \tr(\boldsymbol{y}\boldsymbol{x}^{\top})\). Note that $\boldsymbol{y}\boldsymbol{x}^{\top}$ is an outer product of vectors, resulting in a matrix. This means we can first establish a matrix-valued identity at the outer product level and then take the $\tr$ on both sides at the appropriate moment to recover the scalar form. In the matrix-valued identity, $\boldsymbol{H}_t^{-1}$ can be moved out smoothly.

Specifically, we have (for simplicity, we denote $\boldsymbol{g}(\boldsymbol{x}_t,\boldsymbol{\theta}_t)$ as $\boldsymbol{g}_t$):

Multiplying by $\boldsymbol{H}_t$ on the right and rearranging gives:

Taking the $\tr$ on both sides and repeatedly using $\tr(\boldsymbol{A}\boldsymbol{B}) = \tr(\boldsymbol{B}\boldsymbol{A})$, we get:

This is the identity we desire, where $\Vert\boldsymbol{x}\Vert_{\boldsymbol{H}} \triangleq \sqrt{\boldsymbol{x}^{\top}\boldsymbol{H}\boldsymbol{x}}$ is called the "Mahalanobis norm." The Euclidean norm is $\Vert\boldsymbol{x}\Vert_{\boldsymbol{I}}$, which we denote as $\Vert\boldsymbol{x}\Vert$. Now the left side restores $(\boldsymbol{\theta}_t- \boldsymbol{\varphi})\cdot\boldsymbol{g}_t$, allowing us to connect it to the convexity inequality \eqref{leq:convex}, at the cost of changing the distance metric on the right side from the Euclidean norm to a time-varying Mahalanobis norm.

Classic Scaling

The next steps are largely identical to those in "Making Alchemy More Scientific (I): Average Loss Convergence of SGD" and "Making Alchemy More Scientific (II): Generalizing Conclusions to Unbounded Domains". We use inequality \eqref{leq:convex} to relate the left side to the loss, and then sum over $t=1,2,\cdots,T$ to get:

However, because the $\Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\varphi}\Vert_{\boldsymbol{H}_t}^2$ term in the first summation on the right is not $\Vert\boldsymbol{\theta}_{t+1} - \boldsymbol{\varphi}\Vert_{\boldsymbol{H}_{t+1}}^2$, the sum does not telescopically cancel. Here we use a classic trick from Part I:

To continue scaling, we need to add an assumption to $\boldsymbol{H}_t$: for any $t$, $\boldsymbol{H}_t - \boldsymbol{H}_{t-1}$ is positive semi-definite, denoted as $\boldsymbol{H}_t \succeq \boldsymbol{H}_{t-1}$. This corresponds to the previous non-increasing learning rate assumption. Furthermore, for any positive semi-definite matrix $\boldsymbol{H}$, it holds that $\boldsymbol{H} \preceq \tr(\boldsymbol{H})\boldsymbol{I}$, thus $\Vert\boldsymbol{x}\Vert_{\boldsymbol{H}}\leq \sqrt{\tr(\boldsymbol{H})}\Vert\boldsymbol{x}\Vert$. This scaling is quite loose, but sufficient for qualitative conclusions.

Using the above assumption and scaling, and dropping the non-positive $-\Vert\boldsymbol{\theta}_{T+1} - \boldsymbol{\varphi}\Vert_{\boldsymbol{H}_T}^2$, we have:

Here $R = \max(\Vert\boldsymbol{\theta}_1 - \boldsymbol{\varphi}\Vert,\cdots,\Vert\boldsymbol{\theta}_T - \boldsymbol{\varphi}\Vert)$ is the common upper bound for all $\Vert\boldsymbol{\theta}_t - \boldsymbol{\varphi}\Vert$. If there is doubt about this step, we can simply assume, as in Part I, that optimization occurs within a bounded domain, so all $\Vert\boldsymbol{\theta}_t - \boldsymbol{\varphi}\Vert$ naturally have a common upper bound $R$.

The Result Emerges

Substituting the scaling from the previous section back into \eqref{leq:mid-1} and choosing $\boldsymbol{\varphi}$ to be the theoretical optimum $\boldsymbol{\theta}^*$, we get:

For simplicity, consider the case where $\eta_t$ is a constant $\eta$:

This is the general convergence conclusion for $\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \eta \boldsymbol{H}_t^{-1}\boldsymbol{g}_t$. Now, thinking in reverse: for a good optimizer, the right side of the above equation should be as small as possible. Therefore, minimizing the upper bound might yield a better optimizer, consistent with the logic in Part V. To this end, we first use the assumption $\boldsymbol{H}_1\preceq \boldsymbol{H}_2 \preceq \cdots\preceq \boldsymbol{H}_T$ to get $\Vert\boldsymbol{g}_t\Vert_{\boldsymbol{H}_t^{-1}}^2\geq \Vert\boldsymbol{g}_t\Vert_{\boldsymbol{H}_T^{-1}}^2$, thus:

That is, the upper bound is minimized when all $\boldsymbol{H}_t$ take the same $\boldsymbol{H}$. Next, we rewrite the right side of the above equation as:

Where $\boldsymbol{\Sigma}=\sum_{t=1}^T \boldsymbol{g}_t\boldsymbol{g}_t^{\top}$ is the (unnormalized) second moment matrix. It is not hard to see that $\eta$ is a redundant parameter; the truly independent variable is $\boldsymbol{H}/\eta$. This tells us that $\eta$ can be any positive number. For simplicity, we take $\eta=R$, which makes $\frac{R^2}{2T\eta}=\frac{\eta}{2T}$. Thus, minimizing the upper bound is equivalent to:

Interestingly, this is essentially the optimization objective of PSGD!

Solving the Objective

To solve the objective above, we again utilize the cyclic property of the trace to rewrite the two terms respectively as:

Then, using the arithmetic-geometric mean inequality:

Next, using the Cauchy-Schwarz inequality:

The Cauchy-Schwarz equality holds when there exists a $\lambda\geq 0$ such that $\boldsymbol{H}^{1/2}=\lambda\boldsymbol{\Sigma}^{1/2}\boldsymbol{H}^{-1/2}$, i.e., $\boldsymbol{H}=\lambda\boldsymbol{\Sigma}^{1/2}$. The arithmetic-geometric mean equality holds when $\Vert\boldsymbol{H}^{1/2}\Vert_F = \Vert\boldsymbol{\Sigma}^{1/2}\boldsymbol{H}^{-1/2}\Vert_F$. Substituting this further determines $\lambda=1$, so the optimal solution is:

This is the most classic result of the AdaGrad paper. It can be seen as the "ancestor" of adaptive learning rate optimizers based on second moments. This shows that adjusting learning rates based on second moments is not a random heuristic trick, but a natural product of the "minimize convergence upper bound" principle. When $\boldsymbol{\Sigma}$ is not invertible, we can add $\kappa \boldsymbol{I}$ to it; this is an implementation detail not expanded upon in the theoretical derivation.

Intuitively, $\boldsymbol{H}^{-1}\boldsymbol{g}_t=\boldsymbol{\Sigma}^{-1/2}\boldsymbol{g}_t$ performs a "whitening" operation on the gradient, making the updates more isotropic, which is consistent with the philosophy behind modern Muon optimizers.

Various Variants

However, although $\boldsymbol{H}^* = \boldsymbol{\Sigma}^{1/2}$ is theoretically optimal, it violates causality because it assumes we know the gradients at all time steps from the start. We are no strangers to this scenario; we encountered it in Part V. The solution is to truncate the sum to the current moment:

It is worth noting that $\boldsymbol{H}_t$ defined this way still satisfies $\boldsymbol{H}_t\succeq\boldsymbol{H}_{t-1}$, so the convergence conclusion up to \eqref{leq:const-lr-ada} still holds. However, keeping the full matrix $\boldsymbol{H}_t$ is extremely expensive in terms of both storage and computation and is not realistic. A naive solution is to consider a diagonal approximation, taking only the diagonal part of $\boldsymbol{\Sigma}_t$ and then its square root:

This is what we call the AdaGrad optimizer, where $\odot$ represents the Hadamard product, and vector division is also element-wise division in the Hadamard sense. Note that the learning rates for each component in AdaGrad are monotonically decreasing, which is usually considered an undesirable property in practice (it's easy to "get stuck" in late-stage training). Thus, RMSProp replaces the sum with an exponential moving average (EMA):

Note that while the practical effect of switching to EMA is often better, it can no longer guarantee $\boldsymbol{H}_t\succeq\boldsymbol{H}_{t-1}$; this is where theory and practice begin to "part ways." If we then add momentum, we get the classic of classics—the Adam optimizer (Bias Correction is omitted for simplicity):

Before Muon, Adam was the most mainstream optimizer, bar none. Even with Muon, Adam remains a critical component of mainstream optimizers (for Embeddings, LM Heads, etc.). RMSProp is typically used for training generative models, especially GANs, because adding momentum can lead to Mode Collapse in such models.

In addition to these classic optimizers, newer optimizers like PSGD, Shampoo, KL-Shampoo, and AdaBK are essentially just different structural approximations of $\boldsymbol{\Sigma}$ or $\boldsymbol{\Sigma}^{-1/2}$ targeted at matrix multiplication parameters. In other words, to this day, the preconditioned matrix $\boldsymbol{\Sigma}^{-1/2}$ derived by AdaGrad remains one of the most important guiding principles for optimizers.

Article Summary

This article reviewed the foundational work of adaptive gradient algorithms—AdaGrad—and followed the roadmap of "convergence analysis → upper bound minimization → optimal preconditioned matrix" to reproduce its complete derivation. It can be said that AdaGrad's proposal to "adjust learning rates based on the second moment of gradients" is the underlying principle shared by almost all subsequent adaptive learning rate optimizers.