From Spectral Norm Gradients to a New Kind of Weight Decay
In the post Appreciating the Muon Optimizer: A Genuine Leap from Vectors to Matrices], we introduced a new optimizer called "Muon," one perspective on which is that it performs steepest gradient descent under spectral norm regularization — a view that seems to reveal a more fundamental optimization direction for matrix parameters. As we know, for matrix parameters we often also add weight decay, which can be understood as the gradient of the squared $F$ norm. So from Muon's point of view, could constructing a new form of weight decay from the gradient of the squared spectral norm achieve better results?
This raises the question: what does the gradient (or derivative) of the spectral norm actually look like? And what would the new weight decay designed from it look like? Let's dig into these questions.
Background Recap
The spectral norm, also known as the "$2$ norm," is one of the most commonly used matrix norms. Compared to the simpler $F$ norm (Frobenius norm), it tends to reveal signals more intrinsically tied to matrix multiplication — this is because it is defined in terms of matrix multiplication itself: for a matrix parameter $\boldsymbol{W}\in\mathbb{R}^{n\times m}$, its spectral norm is defined as
\begin{equation}\Vert\boldsymbol{W}\Vert_2 \triangleq \max_{\Vert\boldsymbol{x}\Vert=1} \Vert\boldsymbol{W}\boldsymbol{x}\Vert\end{equation}
Here $\boldsymbol{x}\in\mathbb{R}^m$ is a column vector, and the right-hand side $\Vert\Vert$ is the (Euclidean) norm of the vector. From another angle, the spectral norm is the smallest constant $C$ such that the following inequality holds for all $\forall \boldsymbol{x}\in\mathbb{R}^m$:
\begin{equation}\Vert\boldsymbol{W}\boldsymbol{x}\Vert \leq C\Vert\boldsymbol{x}\Vert\end{equation}
It is not hard to show that when $C$ takes the $F$ norm $\Vert W\Vert_F$, the above inequality also always holds, so we can write $\Vert \boldsymbol{W}\Vert_2\leq \Vert \boldsymbol{W}\Vert_F$ (because $\Vert \boldsymbol{W}\Vert_F$ is just one of the constants $C$ that makes the inequality hold, while $\Vert \boldsymbol{W}\Vert_2$ is the smallest such $C$). This result also shows that if we want to control the magnitude of the output, using the spectral norm as a regularizer is more precise than using the $F$ norm.
Back in Lipschitz Constraints in Deep Learning: Generalization and Generative Models], written six years ago, we already discussed the spectral norm, in two application scenarios: first, WGAN explicitly imposes a Lipschitz constraint on the discriminator, one implementation of which is spectral norm-based normalization; second, some works have shown that using the spectral norm as a regularizer yields better performance than using the $F$ norm as a regularizer.
Deriving the Gradient
Now let's get to the main topic and try to derive the gradient of the spectral norm, $\nabla_{\boldsymbol{W}} \Vert\boldsymbol{W}\Vert_2$. We know that the spectral norm is numerically equal to the largest singular value, as we proved in the section "Matrix Norms]" of The Road to Low-Rank Approximation (II): SVD]. This means that if $\boldsymbol{W}$ has SVD $\sum\limits_{i=1}^{\min(n,m)}\sigma_i \boldsymbol{u}_i\boldsymbol{v}_i^{\top}$, then
\begin{equation}\Vert\boldsymbol{W}\Vert_2 = \sigma_1 = \boldsymbol{u}_1^{\top}\boldsymbol{W}\boldsymbol{v}_1\end{equation}
where $\sigma_1 \geq \sigma_2 \geq \cdots \geq \sigma_{\min(n,m)} \geq 0$ is a singular value of $\boldsymbol{W}$. Taking the differential of both sides, we get
\begin{equation}d\Vert\boldsymbol{W}\Vert_2 = d\boldsymbol{u}_1^{\top}\boldsymbol{W}\boldsymbol{v}_1 + \boldsymbol{u}_1^{\top}d\boldsymbol{W}\boldsymbol{v}_1 + \boldsymbol{u}_1^{\top}\boldsymbol{W}d\boldsymbol{v}_1\end{equation}
Note that
\begin{equation}d\boldsymbol{u}_1^{\top}\boldsymbol{W}\boldsymbol{v}_1 = d\boldsymbol{u}_1^{\top}\sum_{i=1}^{\min(n,m)}\sigma_i \boldsymbol{u}_i\boldsymbol{v}_i^{\top}\boldsymbol{v}_1 = d\boldsymbol{u}_1^{\top}\sigma_1 \boldsymbol{u}_1 = \frac{1}{2}\sigma_1 d(\Vert\boldsymbol{u}_1\Vert^2)=0\end{equation}
and similarly $\boldsymbol{u}_1^{\top}\boldsymbol{W}d\boldsymbol{v}_1=0$, so
\begin{equation}d\Vert\boldsymbol{W}\Vert_2 = \boldsymbol{u}_1^{\top}d\boldsymbol{W}\boldsymbol{v}_1 = \text{Tr}((\boldsymbol{u}_1 \boldsymbol{v}_1^{\top})^{\top} d\boldsymbol{W}) \quad\Rightarrow\quad \nabla_{\boldsymbol{W}}\Vert\boldsymbol{W}\Vert_2 = \boldsymbol{u}_1 \boldsymbol{v}_1^{\top}\end{equation}
Note that a key condition for this proof is $\sigma_1 > \sigma_2$, because if $\sigma_1=\sigma_2$, then $\Vert\boldsymbol{W}\Vert_2$ can be expressed either as $\boldsymbol{u}_1^{\top}\boldsymbol{W}\boldsymbol{v}_1$ or as $\boldsymbol{u}_2^{\top}\boldsymbol{W}\boldsymbol{v}_2$, and the gradients derived by the same method would be $\boldsymbol{u}_1 \boldsymbol{v}_1^{\top}$ and $\boldsymbol{u}_2 \boldsymbol{v}_2^{\top}$ respectively. The non-uniqueness of the result means the gradient does not exist. Of course, from a practical standpoint, the probability of two numbers being exactly equal is vanishingly small, so we can ignore this issue.
(Note: the derivation above follows a Stack Exchange answer], but that answer does not prove $d\boldsymbol{u}_1^{\top}\boldsymbol{W}\boldsymbol{v}_1=0$ and $\boldsymbol{u}_1^{\top}\boldsymbol{W}d\boldsymbol{v}_1=0$ — the author has filled in that part here.)
Weight Decay
Based on this result and the chain rule, we have
\begin{equation}\nabla_{\boldsymbol{W}}\left(\frac{1}{2}\Vert\boldsymbol{W}\Vert_2^2\right) = \Vert\boldsymbol{W}\Vert_2\nabla_{\boldsymbol{W}}\Vert\boldsymbol{W}\Vert_2 = \sigma_1 \boldsymbol{u}_1 \boldsymbol{v}_1^{\top}\label{eq:grad-2-2}\end{equation}
Compare this with the result under the $F$ norm:
\begin{equation}\nabla_{\boldsymbol{W}}\left(\frac{1}{2}\Vert\boldsymbol{W}\Vert_F^2\right) = \boldsymbol{W} = \sum_{i=1}^{\min(n,m)}\sigma_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top}\end{equation}
Placed side by side like this, the picture becomes clear: the weight decay derived from the squared $F$ norm as a regularizer penalizes all singular values simultaneously, whereas the weight decay corresponding to the squared spectral norm only penalizes the largest singular value. If our goal is to shrink the magnitude of the output, shrinking only the largest singular value is the "just right" approach — shrinking all singular values might achieve a similar effect, but it can also compromise the expressive power of the parameters.
By the "Eckart-Young-Mirsky theorem]," the rightmost expression in Eq. $\eqref{eq:grad-2-2}$ has another meaning: it is the "optimal rank-1 approximation" of the matrix $\boldsymbol{W}$. In other words, spectral weight decay changes the operation of subtracting the matrix itself at each step into subtracting its optimal rank-1 approximation at each step — a weaker penalty, but in some sense one that "strikes more directly at the essence."
We can call this "Spectral Weight Decay."
Numerical Computation
For practical purposes, the crucial question now is: how do we compute $\sigma_1 \boldsymbol{u}_1 \boldsymbol{v}_1^{\top}$? SVD is of course the simplest and most direct approach, but it also has the highest computational cost, so we need to find a more efficient way.
Without loss of generality, let $n\geq m$. First, note that
\begin{equation}\sigma_1 \boldsymbol{u}_1 \boldsymbol{v}_1^{\top} = \sum_{i=1}^m\sigma_i \boldsymbol{u}_i \boldsymbol{v}_i^{\top} \boldsymbol{v}_1 \boldsymbol{v}_1^{\top} = \boldsymbol{W}\boldsymbol{v}_1 \boldsymbol{v}_1^{\top}\end{equation}
This shows that computing $\sigma_1 \boldsymbol{u}_1 \boldsymbol{v}_1^{\top}$ only requires knowing $\boldsymbol{v}_1$, and by our discussion in The Road to Low-Rank Approximation (II): SVD], $\boldsymbol{v}_1$ is actually the eigenvector corresponding to the largest eigenvalue of the matrix $\boldsymbol{W}^{\top}\boldsymbol{W}$. This way, we've turned the problem of computing the SVD of a general matrix $\boldsymbol{W}$ into an eigendecomposition of the real symmetric matrix $\boldsymbol{W}^{\top}\boldsymbol{W}$ — already a reduction in complexity, since eigendecomposition is usually noticeably faster than SVD.
If that still feels too slow, we can bring out the principle underlying many eigendecomposition algorithms — "Power Iteration]":
When $\sigma_1 > \sigma_2$, the iteration
\begin{equation}\boldsymbol{x}_{t+1} = \frac{\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{x}_t}{\Vert\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{x}_t\Vert}\end{equation}
converges to $\boldsymbol{v}_1$ at a rate of $(\sigma_2/\sigma_1)^{2t}$.
Each step of power iteration only requires two "matrix-vector" multiplications, with complexity $\mathcal{O}(nm)$; the total complexity for $t$ iterations is $\mathcal{O}(tnm)$, which is quite favorable. The downside is that convergence is slow when $\sigma_1,\sigma_2$ are close to each other. But in practice power iteration tends to perform even better than the theory might suggest — many early works got decent results with just a single iteration, because when $\sigma_1,\sigma_2$ are close, the two eigenvectors are, to some degree, interchangeable, and even when power iteration hasn't fully converged, what it produces is essentially an average of the two eigenvectors, which is often good enough.
Proof of Convergence
In this section we complete the proof of power iteration. It's not hard to see that power iteration can be equivalently written as
\begin{equation}\lim_{t\to\infty} \frac{(\boldsymbol{W}^{\top}\boldsymbol{W})^t \boldsymbol{x}_0}{\Vert(\boldsymbol{W}^{\top}\boldsymbol{W})^t \boldsymbol{x}_0\Vert} = \boldsymbol{v}_1\end{equation}
To prove this limit, we start from $\boldsymbol{W}=\sum\limits_{i=1}^m\sigma_i \boldsymbol{u}_i\boldsymbol{v}_i^{\top}$ and substitute to compute
\begin{equation}\boldsymbol{W}^{\top}\boldsymbol{W} = \sum_{i=1}^m\sigma_i^2 \boldsymbol{v}_i\boldsymbol{v}_i^{\top},\qquad(\boldsymbol{W}^{\top}\boldsymbol{W})^t = \sum_{i=1}^m\sigma_i^{2t} \boldsymbol{v}_i\boldsymbol{v}_i^{\top}\end{equation}
Since $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_m$ forms an orthonormal basis of $\mathbb{R}^m$, $\boldsymbol{x}_0$ can be written as $\sum\limits_{j=1}^m c_j \boldsymbol{v}_j$, so we have
\begin{equation}(\boldsymbol{W}^{\top}\boldsymbol{W})^t \boldsymbol{x}_0 = \sum_{i=1}^m\sigma_i^{2t} \boldsymbol{v}_i\boldsymbol{v}_i^{\top}\sum_{j=1}^m c_j \boldsymbol{v}_j = \sum_{i=1}^m\sum_{j=1}^m c_j\sigma_i^{2t} \boldsymbol{v}_i\underbrace{\boldsymbol{v}_i^{\top} \boldsymbol{v}_j}_{=\delta_{i,j}} = \sum_{i=1}^m c_i\sigma_i^{2t} \boldsymbol{v}_i\end{equation}
and
\begin{equation}\Vert(\boldsymbol{W}^{\top}\boldsymbol{W})^t \boldsymbol{x}_0\Vert = \left\Vert \sum_{i=1}^m c_i\sigma_i^{2t} \boldsymbol{v}_i\right\Vert = \sqrt{\sum_{i=1}^m c_i^2\sigma_i^{4t}}\end{equation}
Due to random initialization, the probability of $c_1=0$ is extremely small, so we can assume $c_1\neq 0$, and then
\begin{equation}\frac{(\boldsymbol{W}^{\top}\boldsymbol{W})^t \boldsymbol{x}_0}{\Vert(\boldsymbol{W}^{\top}\boldsymbol{W})^t \boldsymbol{x}_0\Vert} = \frac{\sum\limits_{i=1}^m c_i\sigma_i^{2t} \boldsymbol{v}_i}{\sqrt{\sum\limits_{i=1}^m c_i^2\sigma_i^{4t}}} = \frac{\boldsymbol{v}_1 + \sum\limits_{i=2}^m (c_i/c_1)(\sigma_i/\sigma_1)^{2t} \boldsymbol{v}_i}{\sqrt{1 + \sum\limits_{i=2}^m (c_i/c_1)^2(\sigma_i/\sigma_1)^{4t}}}\end{equation}
When $\sigma_1 > \sigma_2$, all the $\sigma_i/\sigma_1(i\geq 2)$ are less than 1, so as $t\to \infty$ the corresponding terms all go to zero, and the final limit is $\boldsymbol{v}_1$.
Related Work
The earliest paper to propose spectral norm regularization is probably the 2017 work Spectral Norm Regularization for Improving the Generalizability of Deep Learning], which compared weight decay, adversarial training, spectral norm regularization, and other methods, finding that spectral norm regularization gave the best generalization performance.
That paper's approach was not, as in this post, to derive $\nabla_{\boldsymbol{W}}\Vert\boldsymbol{W}\Vert_2^2 = 2\sigma_1\boldsymbol{u}_1 \boldsymbol{v}_1^{\top}$ directly, but rather to estimate $\Vert\boldsymbol{W}\Vert_2$ via power iteration and then add $\Vert\boldsymbol{W}\Vert_2^2$ as a weighted term into the loss function, letting the optimizer compute the gradient itself. This is somewhat less efficient, and it also can't easily be decoupled from the optimizer in the form of weight decay. The approach in this post is comparatively more flexible, letting us — just as with AdamW — keep weight decay separate from the optimization of the main loss.
Of course, from today's LLM perspective, the biggest problem with those early experiments was that their scale was simply too small to be very convincing. Still, given that the Muon optimizer has already paved the way with the spectral norm, I think it's worth revisiting and experimenting with spectral weight decay. That said, whether it's the $F$ norm or the spectral norm form of weight decay, these "generalization-oriented" techniques often carry an element of luck — so it's best to keep expectations moderate.
My own preliminary experiments on language models suggest a slight improvement at the loss level (hopefully not wishful thinking — though at worst, no degradation was observed either). The experimental setup used power iteration to compute an approximation of $\boldsymbol{v}_1$ (initialized as an all-ones vector, iterated 10 times), and then changed the original weight decay $-\lambda \boldsymbol{W}$ into $-\lambda \boldsymbol{W}\boldsymbol{v}_1\boldsymbol{v}_1^{\top}$, leaving the value of $\lambda$ unchanged.
Summary
In this post, we derived the gradient of the spectral norm, used it to construct a new form of weight decay, and shared some thoughts on it.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.