Steepest Descent on Manifolds: 5. Dual Gradient Descent

流形上的最速下降：5. 对偶梯度下降 苏剑林 · 2025-11-03 · Mathematics matrix optimizer muon constraint steepest descent

In the previous four posts, we solved several concrete steepest-descent problems with equality constraints imposed on the parameters. For the problems in the third and fourth posts, no closed-form solution could be found, so I proposed corresponding fixed-point iteration methods. Among these, the third post, Steepest Descent on Manifolds: 3. Muon + Stiefel, studied "Muon on the Stiefel manifold," a problem originally posed by Jeremy Bernstein in his article Orthogonal manifold.

For this problem, Jeremy Bernstein ultimately gave his own solution as well, which I call "Dual Gradient Descent." It's also well worth studying.

Basic Concepts

Jeremy Bernstein's solution was eventually published in Thinking Machines Lab's blog post Modular Manifolds, the lab's second blog post, where it is called "Dual Ascent." Here, however, in keeping with the style of the previous four posts, I'll call it "Dual Gradient Descent." more

In fact, dual gradient descent can be regarded as a natural consequence of the method of Lagrange multipliers. However, a rigorous discussion of Lagrange multipliers is actually rather involved—for instance, it requires invoking the Minimax theorem. So in this series, to avoid such complications, we've adopted a derivation approach based on "undetermined coefficients," which makes dual gradient descent feel somewhat less natural. But that's fine—we can still work our way to it following our usual line of reasoning, even if it takes a bit more space.

Let's first recall the various notations. $\boldsymbol{W}\in\mathbb{R}^{n\times m}$ is a matrix parameter, and without loss of generality we assume $n\geq m$; $\boldsymbol{G}\in\mathbb{R}^{n\times m}$ is its gradient. $\Vert\boldsymbol{G}\Vert_2$ denotes the spectral norm of matrix $\boldsymbol{G}$, equal to its largest singular value; $\Vert\boldsymbol{G}\Vert_*$ denotes the nuclear norm of matrix $\boldsymbol{G}$, equal to the sum of all its singular values. In particular, according to the conclusion of The Derivative of the SVD, we have

\begin{equation}\nabla_{\boldsymbol{G}}\Vert\boldsymbol{G}\Vert_* = \sum_i \nabla_{\boldsymbol{G}} \sigma_i = \sum_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top} = \boldsymbol{U}\boldsymbol{V}^{\top} = \newcommand{msign}{\mathop{\text{msign}}}\msign(\boldsymbol{G}) \label{eq:nuclear-grad}\end{equation}

where $\boldsymbol{G}=\sum_i \sigma_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top} = \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$ is the SVD of $\boldsymbol{G}$. In other words, the gradient of the nuclear norm is exactly the $\msign$ operator, which is an important foundation for the derivation that follows.

Problem Description

We'll continue introducing dual gradient descent by following our earlier line of reasoning, so this section first restates the problem and the existing results.

In Steepest Descent on Manifolds: 3. Muon + Stiefel, the problem we needed to solve was

\begin{equation}\newcommand{tr}{\mathop{\text{tr}}}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1,\,\, \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I},\,\,\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{0} \label{eq:muon-stiefel}\end{equation}

The solution is $\boldsymbol{\Phi} = \msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})$, where $\boldsymbol{X}\in\mathbb{R}^{m\times m}$ is an undetermined symmetric matrix such that $\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} = \boldsymbol{0}$.

In Steepest Descent on Manifolds: 4. Muon + Spectral Sphere, the problem we needed to solve was

\begin{equation}\max_{\boldsymbol{\Phi}} \tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) \qquad \text{s.t.}\qquad \Vert\boldsymbol{\Phi}\Vert_2 = 1,\,\, \tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0 \label{eq:muon-spectral}\end{equation}

The answer is $\boldsymbol{\Phi} = \msign(\boldsymbol{G} + \lambda\boldsymbol{\Theta})$, where $\lambda$ is an undetermined coefficient such that $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0$.

As we can see, in both cases our final task reduces to finding undetermined coefficients that satisfy some additionally introduced equality constraint—which is essentially solving a (system of) nonlinear equation(s). Dual gradient descent turns this equation-solving problem into the minimization of some objective function, which can then be solved by gradient descent.

The Dual Objective

The key to this transformation is the gradient identity for the nuclear norm, $\eqref{eq:nuclear-grad}$. For simplicity, let's first look at the "Muon + Spectral Sphere" problem $\eqref{eq:muon-spectral}$, where the undetermined coefficient is just a scalar, making things easier to observe. It's not hard to verify that

\begin{equation}\nabla_{\lambda} \Vert\boldsymbol{G} + \lambda\boldsymbol{\Theta}\Vert_* = \tr(\boldsymbol{\Theta}^{\top}\msign(\boldsymbol{G} + \lambda\boldsymbol{\Theta})) = \tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})\end{equation}

This means that solving the equation $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0$ is equivalent to finding a point where the gradient of $\Vert\boldsymbol{G} + \lambda\boldsymbol{\Theta}\Vert_*$ vanishes, which could be a (local) minimum or maximum. Since $\Vert\boldsymbol{G} + \lambda\boldsymbol{\Theta}\Vert_*$ clearly has no maximum, we instead look for its minimum:

\begin{equation}\lambda^* = \newcommand{argmin}{\mathop{\text{argmin}}}\argmin_{\lambda} \Vert\boldsymbol{G} + \lambda\boldsymbol{\Theta}\Vert_*\label{eq:muon-spectral-obj}\end{equation}

Let's walk through the steps here again:

1. Our goal is to solve the equation $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0$; finding any one solution will do.
2. $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})$ is exactly the gradient of $\Vert\boldsymbol{G} + \lambda\boldsymbol{\Theta}\Vert_*$ with respect to $\lambda$.
3. This turns the problem into finding (local) minimum/maximum points, since the gradient vanishes there.
4. It's easy to see there's no maximum, so we can only look for the minimum.

Gradient Descent

Once the objective $\eqref{eq:muon-spectral-obj}$ is fixed, we can solve it via gradient descent, where the gradient is readily available, namely $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})$. The gradient descent update is then

\begin{equation}\lambda \quad \leftarrow\quad \lambda - \eta \tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})\end{equation}

Of course, we could also apply a $\newcommand{sign}{\mathop{\text{sign}}}\sign$ to $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})$, turning it into SignSGD—these kinds of variations are all fair game. In terms of the iteration scheme, dual gradient descent is much simpler than the fixed-point iterations we proposed earlier. However, in many cases, dual gradient descent requires far more iteration steps to converge, and may need careful learning-rate tuning or the introduction of momentum before convergence can be achieved.

So, as far as solving the equation $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0$ goes, dual gradient descent isn't a particularly ideal approach. But our ultimate goal isn't actually to solve the equation $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0$—it's to compute $\boldsymbol{\Phi}$ as the optimization direction for the model. Model training is itself already an iterative process, so we can cache the historical $\lambda$ and adopt an approximate strategy that updates $\lambda$ synchronously with the model parameters:

\begin{equation}\boldsymbol{\Phi} = \msign(\boldsymbol{G} + \lambda\boldsymbol{\Theta}), \quad \boldsymbol{W}\leftarrow\boldsymbol{W}- \eta_1 \boldsymbol{\Phi},\quad \lambda \leftarrow\lambda - \eta_2 \tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})\end{equation}

In this way, each training step only requires one extra, nearly free computation of $\lambda - \eta_2 \tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})$, giving an approximate realization of the original objective $\eqref{eq:muon-spectral}$. Formally, this amounts to a kind of adaptive weight decay for Muon.

On the Stiefel Manifold

Having discussed the relatively simpler "Muon + Spectral Sphere" case, let's turn to "Muon + Stiefel," i.e., the objective $\eqref{eq:muon-stiefel}$. Here the undetermined matrix $\boldsymbol{X}$ is subject to the constraint $\boldsymbol{X}=\boldsymbol{X}^{\top}$, which we remove by setting $\boldsymbol{X}=\boldsymbol{\Lambda}+\boldsymbol{\Lambda}^{\top}$, where $\boldsymbol{\Lambda}\in\mathbb{R}^{m\times m}$ is an arbitrary matrix. We can then show that

\begin{equation}\nabla_{\boldsymbol{\Lambda}}\Vert\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}\Vert_* = \boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W} \end{equation}

Here $\boldsymbol{\Phi} = \msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X})$. So, solving the system of equations $\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W}=\boldsymbol{0}$ can likewise be turned into finding the minimum of the function $\Vert\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}\Vert_*$, and then solved with gradient descent:

\begin{equation}\boldsymbol{\Lambda} \quad\leftarrow\quad \boldsymbol{\Lambda} - \eta(\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W}) \end{equation}

Since $\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W}$ is necessarily symmetric, it's also feasible to directly use $\boldsymbol{X} \leftarrow\boldsymbol{X} - \eta(\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W})$. Combining this with $\boldsymbol{W}$ to update them synchronously, we obtain

\begin{equation}\boldsymbol{\Phi} = \msign(\boldsymbol{G} + \boldsymbol{W}\boldsymbol{X}), \quad \boldsymbol{W}\leftarrow\boldsymbol{W}- \eta_1 \boldsymbol{\Phi},\quad \boldsymbol{X} \leftarrow\boldsymbol{X} - \eta_2(\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W})\end{equation}

This gives us an approximation of the objective $\eqref{eq:muon-stiefel}$, where the extra $\boldsymbol{X} - \eta_2(\boldsymbol{W}^{\top}\boldsymbol{\Phi}+\boldsymbol{\Phi}^{\top}\boldsymbol{W})$ computed at each step is likewise nearly free.

Lagrange Multipliers

In both examples, the equations we needed to solve happened to equal the gradient of some nuclear-norm objective. Is this merely a coincidence? Certainly not—as we mentioned in the "Basic Concepts" section, this is a natural consequence of the method of Lagrange multipliers, and this section will expand on that point.

For ease of understanding, let's again use the relatively simple objective $\eqref{eq:muon-spectral}$ as an example. It can be equivalently written as

\begin{equation}\max_{\Vert\boldsymbol{\Phi}\Vert_2\leq 1} \min_{\lambda\in\mathbb{R}}\tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) + \lambda\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})\end{equation}

To understand this transformation, we only need to notice that the above expression must satisfy $\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi})=0$; otherwise, the step $\min$ could always be pushed to negative infinity, forcing the final $\max$ result to also be negative infinity. As for changing $\Vert\boldsymbol{\Phi}\Vert_2 = 1$ to $\Vert\boldsymbol{\Phi}\Vert_2\leq 1$, this doesn't change the maximum value obtained (since the maximum is always attained on the boundary), but it does make the feasible region of $\boldsymbol{\Phi}$ into a convex set.

With this equivalent form in hand, we can use the Minimax theorem to swap the order of $\min$ and $\max$:

\begin{equation}\begin{aligned} &\,\max_{\Vert\boldsymbol{\Phi}\Vert_2\leq 1} \min_{\lambda\in\mathbb{R}}\tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) + \lambda\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi}) \\ =&\, \min_{\lambda\in\mathbb{R}}\max_{\Vert\boldsymbol{\Phi}\Vert_2\leq 1}\tr(\boldsymbol{G}^{\top}\boldsymbol{\Phi}) + \lambda\tr(\boldsymbol{\Theta}^{\top} \boldsymbol{\Phi}) \\ =&\, \min_{\lambda\in\mathbb{R}} \Vert\boldsymbol{G} + \lambda \boldsymbol{\Theta}\Vert_* \end{aligned}\end{equation}

Here, the step of taking $\max$ over $\Vert\boldsymbol{\Phi}\Vert_2\leq 1$ is exactly the basic result from the Muon derivation, so computing $\max$ first poses no difficulty. In this way we obtain the dual objective $\Vert\boldsymbol{G} + \lambda \boldsymbol{\Theta}\Vert_*$ of the original problem $\eqref{eq:muon-spectral}$.

Some readers might wonder: doesn't this method of Lagrange multipliers look different from what I learned? That's because here the method of Lagrange multipliers has been generalized to arbitrary convex sets, and we've rigorously discussed the interchangeability of $\min,\max$, in order to guarantee the final result is what we want. The version of Lagrange multipliers typically taught is just a heuristic procedure for solving constrained optimization problems in $\mathbb{R}^n$, without much discussion of the theoretical guarantees involved.

Summary

In this post, we introduced the idea of using dual gradient descent to find the direction of steepest descent on a manifold, which is also the method used a while back in Thinking Machines Lab's blog post Modular Manifolds to solve for Muon on the Stiefel manifold.