A New Understanding of Momentum: Approximating Gradient Descent at the Feature Level

动量的新理解：逼近特征层面的梯度下降 苏剑林 · 2026-08-23 · Machine Learning linear gradient optimizer muon

An optimizer with momentum as its state variable typically has the following basic form:

\begin{equation}\begin{aligned} \boldsymbol{M}_t =&\, \beta \boldsymbol{M}_{t-1} + (1 - \beta) \boldsymbol{G}_t \\[4pt] \boldsymbol{W}_t =&\, \phi(\boldsymbol{W}_{t-1}, \boldsymbol{M}_t, \boldsymbol{G}_t, t) \end{aligned}\end{equation}

Different optimizers differ mainly in the update function $\phi$, as in SGDM, SignSGD, Muon, etc., and new explorations have largely revolved around $\phi$, since the momentum term in the first equation is so simple that nobody thinks there's much room left to modify it.

But momentum is exactly what this post is about. We're going to give the momentum mechanism a new interpretation: momentum can be seen not only as an average of gradients, but also as the solution to an online regression problem. Starting from this idea, we can naturally arrive at some recent works. more

Basic Concepts

As with Muon, we mainly consider the matrix parameters of a linear layer here: suppose we have a linear layer $\boldsymbol{Y}=\boldsymbol{X}\boldsymbol{W}$, where $\boldsymbol{X}\in\mathbb{R}^{b\times d_{in}}$ is the input, $\boldsymbol{W}\in\mathbb{R}^{d_{in}\times d_{out}}$ is the weight, and $\boldsymbol{Y}\in\mathbb{R}^{b\times d_{out}}$ is the output. Let the loss function be denoted $\mathcal{L}(\boldsymbol{Y})=\mathcal{L}(\boldsymbol{X}\boldsymbol{W})$, so that

\begin{equation}\boldsymbol{G} = \frac{\partial \mathcal{L}}{\partial\boldsymbol{W}} = \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}

The simplest optimizer is gradient descent (at the parameter level)

\begin{equation}\boldsymbol{W}\quad\leftarrow\quad \boldsymbol{W} - \eta\frac{\partial \mathcal{L}}{\partial\boldsymbol{W}} = \boldsymbol{W} - \eta \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}

However, following the idea in Why Do We Favor Isotropy? A Perspective from Steepest Descent, we hold that parameters are essentially just a byproduct of the model — it's the change at the level of model features that is most closely tied to model performance. Ideally, we would like to achieve gradient descent at the feature level:

\begin{equation}\boldsymbol{Y}\quad\leftarrow\quad \boldsymbol{Y} - \eta\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\label{eq:gd-y}\end{equation}

The problem is that $\boldsymbol{Y}$ is not a variable we can modify at will; the only thing we can directly modify is $\boldsymbol{W}$. So we have to find a way to achieve this effect indirectly, by modifying $\boldsymbol{W}$.

The Regression Objective

How do we achieve this indirectly? Let the final update rule be $\boldsymbol{W}\leftarrow\boldsymbol{W} - \eta\boldsymbol{\Phi}$, so that $\boldsymbol{Y}\leftarrow\boldsymbol{Y}-\eta \boldsymbol{X}\boldsymbol{\Phi}$. We want this to get as close as possible to the effect of equation $\eqref{eq:gd-y}$, i.e., we want $\boldsymbol{X}\boldsymbol{\Phi}\approx\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}$. So we consider minimizing

\begin{equation}\min_{\boldsymbol{\Phi}} \frac{1}{2}\left\Vert\boldsymbol{X}\boldsymbol{\Phi} - \frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\right\Vert_F^2 + \frac{\lambda}{2}\Vert\boldsymbol{\Phi}\Vert_F^2 \label{eq:obj}\end{equation}

where $\lambda > 0$ is a regularization coefficient. This is, in fact, nothing more than a linear regression problem, and it can be solved directly to give

\begin{equation}\boldsymbol{\Phi}^* = (\boldsymbol{X}^{\top}\boldsymbol{X} + \lambda \boldsymbol{I})^{-1} \boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}\end{equation}

Notice that $\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}$ is exactly $\boldsymbol{G}=\frac{\partial \mathcal{L}}{\partial\boldsymbol{W}}$, and $(\boldsymbol{X}^{\top}\boldsymbol{X} + \lambda \boldsymbol{I})^{-1}$ is a preconditioner based on the input data. To combine contributions across different batches, we consider applying an EMA to both $\boldsymbol{X}^{\top}\boldsymbol{X}$ and $\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial\boldsymbol{Y}}$, which gives us an SGDM variant:

\begin{equation}\begin{aligned} \boldsymbol{M}_t =&\, \beta \boldsymbol{M}_{t-1} + (1 - \beta) \boldsymbol{G}_t \\[4pt] \boldsymbol{Z}_t =&\, \beta \boldsymbol{Z}_{t-1} + (1 - \beta) (\boldsymbol{X}_t^{\top}\boldsymbol{X}_t + \lambda \boldsymbol{I}) \\[4pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta \boldsymbol{Z}_t^{-1}\boldsymbol{M}_t \end{aligned}\end{equation}

Switching Perspective

We now know that, through the corrected gradient $(\boldsymbol{X}^{\top}\boldsymbol{X} + \lambda \boldsymbol{I})^{-1}\boldsymbol{G}$, we can achieve gradient descent at the feature level. So we can regard this as a kind of "more reliable gradient," and correspondingly, the associated $\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t$ can be regarded as a kind of "more reliable momentum."

From this perspective, we can try replacing the momentum in various momentum-based optimizers with $\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t$ — for instance, in Muon:

\begin{equation}\begin{aligned}\newcommand{msign}{\mathop{\text{msign}}} \boldsymbol{M}_t =&\, \beta \boldsymbol{M}_{t-1} + (1 - \beta) \boldsymbol{G}_t \\[4pt] \boldsymbol{Z}_t =&\, \beta \boldsymbol{Z}_{t-1} + (1 - \beta) (\boldsymbol{X}_t^{\top}\boldsymbol{X}_t + \lambda \boldsymbol{I}) \\[4pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta \msign(\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t) \end{aligned}\end{equation}

This is, roughly speaking, the Newton-Muon optimizer. I say "roughly" because the original Newton-Muon actually uses a "correct-then-EMA" scheme (i.e., it takes the EMA of $(\boldsymbol{X}^{\top}\boldsymbol{X} + \lambda \boldsymbol{I})^{-1}\boldsymbol{G}$ as the momentum), which in principle can save a set of state variables. As for the matrix inversion, you can refer to Efficient Computation of Matrix r-th Roots and Inverse r-th Roots.

If the input is isotropic, then we can expect $\boldsymbol{Z}_t=\sigma^2\boldsymbol{I}$, in which case Newton-Muon degenerates to Muon. In other words, Muon can be understood as feature-level gradient descent under the isotropy assumption — which again brings us back to the conclusion of Why Do We Favor Isotropy? A Perspective from Steepest Descent. Moreover, when $\lambda\to\infty$, Newton-Muon also degenerates to Muon, so we can tune $\lambda$ to control how closely it approximates Muon.

Incremental Updates

A more general and more "playable" approach is: instead of trying to solve for the analytical solution, just directly optimize the objective $\eqref{eq:obj}$ via gradient descent! First, we can compute the gradient of equation $\eqref{eq:obj}$ with respect to $\boldsymbol{\Phi}$:

\begin{equation}(\boldsymbol{X}^{\top}\boldsymbol{X}+\lambda\boldsymbol{I})\boldsymbol{\Phi}-\boldsymbol{X}^{\top}\frac{\partial\mathcal{L}}{\partial\boldsymbol{Y}} = (\boldsymbol{X}^{\top}\boldsymbol{X}+\lambda\boldsymbol{I})\boldsymbol{\Phi}-\boldsymbol{G}\end{equation}

If we use gradient descent to update $\boldsymbol{\Phi}$, we get

\begin{equation}\begin{aligned} \boldsymbol{\Phi}_t =&\, \boldsymbol{\Phi}_{t-1} - \gamma[(\boldsymbol{X}_t^{\top}\boldsymbol{X}_t+\lambda\boldsymbol{I})\boldsymbol{\Phi}_{t-1}-\boldsymbol{G}_t] \\[4pt] =&\, (1-\gamma\lambda)\boldsymbol{\Phi}_{t-1} + \gamma[\boldsymbol{G}_t - (\boldsymbol{X}_t^{\top}\boldsymbol{X}_t)\boldsymbol{\Phi}_{t-1}] \end{aligned}\end{equation}

where $\gamma > 0$ is the learning rate of this "inner" optimization. Continuing the idea from the previous section — if $\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t$ is some kind of more reliable momentum, then the above equation is a "more reliable momentum update rule" derived from the idea of gradient descent! This is exactly DeltaMomentum, which introduces the Delta Rule into momentum. Compared with old-style momentum, this is analogous to the difference between GDN and Vanilla Linear Attention, while $\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t$ is analogous to MesaNet.

Of course, to actually get DeltaMomentum running, there are some details to pay attention to, mainly regarding the normalization of $\boldsymbol{X}$ and the choice of $\gamma$ — please refer to the original paper for these. I'm not entirely convinced by some of the choices made in the original paper either, so I'd suggest reading it with a critical eye. Regardless, DeltaMomentum offers a new direction for exploring the momentum mechanism, one that doesn't even require expensive operations like matrix inversion — well worth savoring.

Related Work

In fact, over the past while, there has been a growing number of explorations into building preconditioners from the input. Besides Newton-Muon and DeltaMomentum, Xie Tian previously published a somewhat different result in Steepest Descent in Feature Space. Using our notation here, it's roughly:

\begin{equation}\begin{aligned} \boldsymbol{M}_t =&\, \beta \boldsymbol{M}_{t-1} + (1 - \beta) \boldsymbol{G}_t \\[4pt] \boldsymbol{Z}_t =&\, \beta \boldsymbol{Z}_{t-1} + (1 - \beta) (\boldsymbol{X}_t^{\top}\boldsymbol{X}_t + \lambda \boldsymbol{I}) \\[4pt] \boldsymbol{W}_t =&\, \boldsymbol{W}_{t-1} - \eta \boldsymbol{Z}_t^{-1/2} \msign(\boldsymbol{Z}_t^{-1/2}\boldsymbol{M}_t) \end{aligned}\end{equation}

That is, $\boldsymbol{Z}_t$ acts on the inside and outside of $\msign$ respectively, each raised to the power $-1/2$. The derivation is based precisely on steepest descent under a spectral-norm constraint at the feature level, which is consistent with the central idea of this post.

There may be other similar works, but I honestly can't think of any more — readers are welcome to add to the list in the comments (note: one reader has already pointed out DoPr). Looking at the experimental results, some of these attempts have indeed achieved positive results — for example, Newton-Muon outperforms Muon in the Speedrun setting (see here) — so overall, this direction does seem to have some merit.

Further Thoughts

That said, this style of input-oriented preconditioner design also has some unsatisfying aspects.

The first is at the implementation level: it requires us to record the autocorrelation matrix $\boldsymbol{X}_t^{\top}\boldsymbol{X}_t$ during the forward pass and then pass it into the optimizer. This breaks the independence of the optimizer — the optimizer is no longer a "black box" that can just take the gradient and do its job; it needs to be coupled with details at the level of the model definition. This is hard to call elegant from an implementation standpoint, not to mention the extra communication and computation overhead it introduces.

The second is at the theoretical level: constructing the preconditioner from the actual input may confine the optimizer's exploration to the subspace spanned by the inputs, leaving other directions under-explored, which could hurt performance. A relatively simple mitigation is to increase $\lambda$ to weaken the preconditioning effect — since at $\lambda\to\infty$ it's equivalent to having no preconditioning at all, there is always a chance of tuning performance by choosing an appropriate $\lambda$. But this also adds one more hyperparameter that needs careful tuning.

Overall, this direction is still in its early, budding stage — the basic shape has emerged, but neither the engineering nor the theoretical issues have been fully resolved. There are no fewer open questions than there are conclusions.

Summary

Starting from the idea of making "parameter-level gradient descent" approximate "feature-level gradient descent," this post reinterprets momentum as the solution to an online regression problem, which naturally leads into some related work. Interestingly, this line of thinking closely parallels the evolution of linear attention from Vanilla, to DeltaNet, to MesaNet — the two seem to have quite a lot to learn from each other.