By 苏剑林 | August 23, 2026
An optimizer that uses momentum as a state variable typically follows this 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}
Differences between optimizers are mainly reflected in the update function $\phi$, such as SGDM, SignSGD, Muon, etc. New explorations generally revolve around $\phi$ because the first equation for momentum is so simple that people don't think there's much room for modification.
However, the protagonist of this article is precisely momentum. We will provide a new perspective on the momentum mechanism: momentum can be viewed not only as an average of gradients but also as the solution to an online regression problem. Starting from this point, we can naturally derive some recent works.
Basic Concepts
Similar to Muon, this article primarily considers matrix parameters of linear layers: assume a linear layer $Y=XW$, where $X \in \mathbb{R}^{b \times d_{in}}$ is the input, $W \in \mathbb{R}^{d_{in} \times d_{out}}$ is the weight, and $Y \in \mathbb{R}^{b \times d_{out}}$ is the output. Let the loss function be denoted as $\mathcal{L}(Y) = \mathcal{L}(XW)$, then
\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, according to the idea in "Why Do We Prefer Isotropy? An Understanding Based on Steepest Descent", we believe that parameters are essentially by-products of the model, and changes at the feature level of the model are most relevant to model performance. Ideally, we hope 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 $Y$ is not a variable that can be modified at will; what we can directly modify is only $W$. Therefore, we must find a way to indirectly achieve this effect by modifying $W$.
Regression Objective
How do we achieve this indirectly? Suppose the final update rule is $W \leftarrow W - \eta \boldsymbol{\Phi}$, then $Y \leftarrow Y - \eta \boldsymbol{X}\boldsymbol{\Phi}$. We want it to be as close as possible to the effect of Eq. (4), which means we want $\boldsymbol{X}\boldsymbol{\Phi} \approx \frac{\partial \mathcal{L}}{\partial \boldsymbol{Y}}$. Thus, 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. In fact, this is just a linear regression problem, which can be solved directly as:
\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 input data. To synthesize contributions from different batches, we consider applying EMA to both $\boldsymbol{X}^{\top}\boldsymbol{X}$ and $\boldsymbol{X}^{\top}\frac{\partial \mathcal{L}}{\partial \boldsymbol{Y}}$, which leads to 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 Perspectives
Now we know that by using the corrected gradient $(\boldsymbol{X}^{\top}\boldsymbol{X} + \lambda \boldsymbol{I})^{-1}\boldsymbol{G}$, we can achieve feature-level gradient descent. We can treat this as a "more reliable gradient." Similarly, the corresponding $\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t$ can be viewed as a "more reliable momentum."
From this perspective, we can try to replace the momentum in various momentum-based optimizers with $\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t$, such as 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 the Newton-Muon optimizer. We say "roughly" because the original Newton-Muon actually uses a scheme of correction followed by EMA (i.e., applying EMA to $(\boldsymbol{X}^{\top}\boldsymbol{X} + \lambda \boldsymbol{I})^{-1}\boldsymbol{G}$ to serve as momentum), which in principle saves one set of state variables. As for matrix inversion, one can refer to "Efficient Calculation of Matrix r-th Roots and Inverse r-th Roots".
If the input is isotropic, we can expect $\boldsymbol{Z}_t = \sigma^2 \boldsymbol{I}$. In this case, Newton-Muon degenerates into Muon. In other words, Muon can be understood as feature gradient descent under the assumption of isotropy, which brings us back to the conclusion of the article "Why Do We Prefer Isotropy? An Understanding Based on Steepest Descent". Furthermore, when $\lambda \to \infty$, Newton-Muon also degenerates into Muon, so we can adjust the degree of approximation to Muon by tuning $\lambda$.
Incremental Update
A more universal and playable approach is: do not seek an analytical solution, but directly use gradient descent to optimize the objective \eqref{eq:obj}! First, we can find the gradient of Eq. \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 for this "inner layer" optimization. Continuing the idea from the previous section, if we say $\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t$ is a more reliable momentum, then the above equation is a "more reliable momentum update rule" derived based on the idea of gradient descent! This is DeltaMomentum, which introduces the Delta Rule into momentum. Compared to old momentum, it is analogous to the difference between GDN and Vanilla Linear Attention, while $\boldsymbol{Z}_t^{-1}\boldsymbol{M}_t$ corresponds to MesaNet.
Of course, to make DeltaMomentum work, some details must be noted, mainly concerning the normalization of $\boldsymbol{X}$ and the choice of $\gamma$; please refer to the original paper. The author does not entirely agree with some treatments in the original paper and suggests readers refer to it with discretion. Regardless, DeltaMomentum provides a new direction for the momentum mechanism and does not require expensive operations like inversion, making it worth reflecting on.
Related Work
In fact, recently, there have been more and more explorations into building preconditioners based on inputs. In addition to Newton-Muon and DeltaMomentum, Xie Tian previously published another slightly different result in "Steepest Descent in Feature Space". Using our notation, it is 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 in the form of the $-1/2$ power both inside and outside the $\msign$. The derivation principle is steepest descent under spectral norm constraints in the feature space, which is consistent with the central idea of this article.
There may be other similar works, but I honestly can't recall them. Readers are welcome to supplement in the comments. From the experimental results, some attempts have indeed yielded positive results; for example, Newton-Muon performs better than Muon on Speedrun (refer to here). So overall, this direction has a certain degree of merit.
Extended Reflections
However, this input-oriented preconditioner design also has some unsatisfactory aspects.
The first is at the implementation level: it requires us to record the self-correlation matrix $\boldsymbol{X}_t^{\top}\boldsymbol{X}_t$ during the forward calculation and then pass it to the optimizer. This breaks the independence of the optimizer—the optimizer is no longer a "black box" that can work just by receiving gradients; it needs to be coupled with model definition details. This is hardly elegant, not to mention the extra communication and calculation overhead introduced.
The second is at the theoretical level: building a preconditioner based on actual inputs may limit the optimizer's exploration to the subspace spanned by the inputs, leading to insufficient exploration in other directions and thus poor performance. A simple mitigation strategy is to increase $\lambda$ to reduce the preconditioning strength—since $\lambda \to \infty$ is equivalent to no preconditioning, choosing an appropriate $\lambda$ always offers a chance to tune the performance, but this also introduces an extra hyperparameter that needs careful adjustment.
In general, this direction is still in its nascent stage—the prototype has emerged, but engineering and theoretical issues have not been fully resolved. There are no fewer problems than conclusions.
Article Summary
Starting from the idea of making "parameter gradient descent" approximate "feature gradient descent," this article reinterprets momentum as the solution to an online regression problem, thereby introducing some related work. Interestingly, this path is extremely consistent with the evolution of linear attention from Vanilla to DeltaNet and then to MesaNet; the two seem to have much to learn from each other.