How to More Scientifically Estimate the Spectral Norm of a Matrix
The spectral norm is a core concept in matrix analysis, and it also plays an important role in deep learning — from the Lipschitz constraints required in the WGAN era, to the training stability of today's LLMs, to the increasingly popular Muon optimizer, all of these are closely tied to the spectral norm of matrix parameters. Therefore, how to estimate the spectral norm efficiently and accurately is becoming ever more important, and deserves careful study.
As is well known, Power Iteration is the standard method for estimating the spectral norm, but there is still plenty of room for improvement. This post will briefly go over some ideas for estimating the spectral norm, including improving the convergence rate of power iteration, and how to estimate a strict upper bound on the spectral norm, among other things.
Spectral Norm
The definition of the spectral norm is
\begin{equation}\Vert\boldsymbol{W}\Vert_2 = \max_{\Vert\boldsymbol{x}\Vert_2=1} \Vert \boldsymbol{W}\boldsymbol{x}\Vert_2\end{equation}more
where $\boldsymbol{W}\in\mathbb{R}^{n\times m}, \boldsymbol{x}\in\mathbb{R}^m$, and without loss of generality, let $n\geq m$. From this definition, the spectral norm represents the "expansion rate" of a linear layer — after passing through the linear layer, the norm of the output vector can be at most $\Vert\boldsymbol{W}\Vert_2$ times that of the input vector. The spectral norm also equals the largest singular value of the matrix (see the proof in The Road to Low-Rank Approximation (II): SVD]); we won't dwell on these basics too much here.
The first time I encountered the spectral norm was back when GANs were popular. WGAN] burst onto the scene, emphasizing the necessity for the discriminator to satisfy a Lipschitz constraint. The Lipschitz constant of a linear layer $\boldsymbol{W}\boldsymbol{x}$ is precisely the spectral norm of the matrix $\boldsymbol{W}$, which is how techniques like spectral normalization and spectral regularization came about. Interested readers can refer to Lipschitz Constraints in Deep Learning: Generalization and Generative Models].
In recent years, the increasingly popular Muon] optimizer is also closely connected to the spectral norm — it is usually viewed as steepest descent under the spectral norm. Meanwhile, in Beyond MuP: 4. Sticking to Parameter Stability], we also proposed constraining the spectral norm of parameters to ensure training stability. All of this creates a demand for accurately computing the spectral norm.
Power Iteration
Computing the spectral norm via SVD is the most direct approach, but it's clearly too expensive. We usually use power iteration instead:
\begin{equation}\boldsymbol{v}^{(t)} = \frac{\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{v}^{(t-1)}}{\Vert\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{v}^{(t-1)}\Vert_2}\end{equation}
which converges to the right principal singular vector $\boldsymbol{v}_1$ at a rate of $(\sigma_2/\sigma_1)^{2t}$. After $T$ steps of iteration, we have
\begin{equation}\sigma_1 \approx \Vert\boldsymbol{W}\boldsymbol{v}^{(t)}\Vert_2\end{equation}
For a proof, see From Spectral Norm Gradients to Thoughts on a New Style of Weight Decay]. In practice, if we only care about the spectral norm and not about the singular vectors, the convergence rate is often more optimistic than $(\sigma_2/\sigma_1)^{2t}$. In addition, the result of power iteration can fluctuate depending on the initial value $\boldsymbol{v}^{(0)}$, so we might consider running $k$ parallel copies and then taking the maximum.
If we simultaneously compute $k$ copies, and at each step replace the L2 normalization with an orthogonalization of these $k$ vectors, i.e.,
\begin{equation}\newcommand{QR}{\mathop{\text{QR}}}\boldsymbol{V}^{(t)} = \QR(\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{V}^{(t-1)})\end{equation}
then the result will converge to the top $k$ right singular vectors, from which we can further obtain the top $k$ singular values. The underlying principle can be found in Muon Implementation via Streaming Power Iteration: 4. Principles]. We won't expand on this extension here though, and will instead keep our focus on estimating the spectral norm — i.e., the largest singular value.
Computing the Gradient
First, here is a reference implementation of power iteration based on Jax:
import jax, jax.lax as lax
import jax.numpy as jnp
@jax.jit(static_argnums=(1,))
def l2_normalize(x, axis=-2):
return x / jnp.linalg.vector_norm(x, axis=axis, keepdims=True)
@jax.jit(static_argnums=(1, 2, 3))
def spec_norm_v1(w, T=10, k=1, key=42):
v_shape = w.shape[:-2] + w.shape[-1:] + (k,)
v_init = jax.random.normal(jax.random.PRNGKey(key), v_shape)
v_step = lambda i, v: l2_normalize(w.mT @ (w @ v))
v = lax.fori_loop(0, T, v_step, v_init)
return jnp.linalg.vector_norm(w @ v, axis=-2).max(axis=-1)
Clearly, the complexity of power iteration is $\mathcal{O}(mnTk)$, which is already quite reasonable in terms of complexity alone. If our purpose is just monitoring, then the code above is already sufficient. But if we're using it for spectral normalization or spectral regularization, we need the gradient of the spectral norm, and following the implementation above would require backpropagating step by step through the power iteration loop, which is computationally expensive.
In fact, in From Spectral Norm Gradients to Thoughts on a New Style of Weight Decay] we already derived the gradient of the spectral norm, $\nabla_{\boldsymbol{W}}\sigma_1 = \boldsymbol{u}_1 \boldsymbol{v}_1^{\top}$, so we can simply define its gradient directly as $\nabla_{\boldsymbol{W}}\sigma_1 = \boldsymbol{u}_1 \boldsymbol{v}_1^{\top}$, reducing the cost of backpropagation. Beyond that, we can also make use of $\sigma_1 = \boldsymbol{u}_1^{\top}\boldsymbol{W}\boldsymbol{v}_1$, directly outputting in the forward pass
\begin{equation}\sigma_1 = \color{skyblue}{[}\boldsymbol{u}_1^{\top}\color{skyblue}{]_{sg}}\boldsymbol{W}\color{skyblue}{[}\boldsymbol{v}_1\color{skyblue}{]_{sg}}\end{equation}
where $\color{skyblue}{[\cdot]_{sg}}$ is the stop-gradient operator, so that during automatic differentiation, we only differentiate with respect to $\boldsymbol{W}$, and the result is exactly $\boldsymbol{u}_1 \boldsymbol{v}_1^{\top}$, avoiding backpropagation through the internal power iteration of $\boldsymbol{u}_1,\boldsymbol{v}_1$.
Don't Waste Anything
Through $T$ steps of power iteration, we obtain a sequence of vectors $\boldsymbol{v}^{(1)},\boldsymbol{v}^{(2)},\cdots,\boldsymbol{v}^{(T)}$, but in the end we only use $\boldsymbol{v}^{(T)}$ to estimate $\sigma_1$, which seems rather "wasteful." So people have thought about making use of all of them, to form a better estimate of the spectral norm.
Barring anything unusual, $\boldsymbol{v}^{(1)},\boldsymbol{v}^{(2)},\cdots,\boldsymbol{v}^{(T)}$ are all mutually distinct, but they get closer and closer to $\boldsymbol{v}_1$. We can think of them as "$\boldsymbol{v}_1$ plus some kind of noise," and if we combine them together for "denoising," we may indeed be able to obtain a more accurate result. Specifically, we consider constructing a better approximation to $\boldsymbol{v}_1$ via a linear combination of $\boldsymbol{v}^{(1)},\boldsymbol{v}^{(2)},\cdots,\boldsymbol{v}^{(T)}$, and the subspace spanned by this set of vectors is called the "Krylov subspace."
For simplicity, we first perform one orthogonalization to obtain an equivalent orthonormal basis $\boldsymbol{Q} = [\boldsymbol{q}_1,\boldsymbol{q}_2,\cdots,\boldsymbol{q}_T]\in\mathbb{R}^{m\times T}$. Next, we want to find coefficients $\boldsymbol{x}=[x_1,x_2,\cdots,x_T]^{\top}\in\mathbb{R}^T$ such that the vector $\sum_{i=1}^T x_i \boldsymbol{q}_i = \boldsymbol{Q}\boldsymbol{x}$ satisfies, as closely as possible,
\begin{equation}\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{Q}\boldsymbol{x}\approx \sigma_1^2\boldsymbol{Q}\boldsymbol{x}\end{equation}
Multiplying both sides by $\boldsymbol{Q}^{\top}$ gives $\boldsymbol{Q}^{\top}\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{Q}\boldsymbol{x}\approx \sigma_1^2\boldsymbol{x}$, which shows that $\sigma_1^2,\boldsymbol{x}$ are respectively an eigenvalue and eigenvector of $(\boldsymbol{W}\boldsymbol{Q})^{\top}\boldsymbol{W}\boldsymbol{Q}$. This means we only need to perform an eigenvalue decomposition on it, take its largest eigenvalue, and take the square root, to obtain a better estimate. Since this is only a $T\times T$ matrix, when $T$ is relatively small, eigenvalue decomposition is not costly at all, so we can simply call an off-the-shelf function.
This idea is essentially a simplified version of the "Lanczos algorithm" used in modern SVD solvers — real SVD implementations involve more elaborate handling. This acceleration idea also has something in common with Randomized SVD: randomized SVD typically projects using a Gaussian random matrix, whereas here we project using the orthonormal basis of the Krylov subspace.
Acceleration Technique
Here is reference code for power iteration incorporating Krylov subspace acceleration:
@jax.jit(static_argnums=(1, 2))
def spec_norm_v2(w, T=10, key=42):
v_shape = w.shape[:-2] + w.shape[-1:] + (1,)
v_init = jax.random.normal(jax.random.PRNGKey(key), v_shape)
v_step = lambda v, _: (l2_normalize(w.mT @ (w @ v)),) * 2
v = lax.scan(v_step, v_init, length=T)[1].swapaxes(0, -1)[0]
v = jnp.linalg.qr(v)[0]
return jnp.linalg.eigvalsh((u := w @ v).mT @ u)[..., -1]**0.5
Empirical results show that subspace acceleration is highly effective: at $T=5$, its accuracy already exceeds that of vanilla power iteration at $T=10$, and it's faster too. At $T=10$, roughly two-thirds of the time is spent on power iteration and one-third on QR decomposition, while the time spent on eigenvalue decomposition is negligible.
Readers who have read the "Streaming Power Iteration]" series might think of using the Cholesky QR method introduced there to speed up the QR decomposition. Unfortunately, since the results of power iteration become increasingly collinear, the condition number of the matrix to be QR-decomposed keeps getting worse — which is exactly the regime where Cholesky QR is "helpless," with a very high failure rate. So it's essentially unusable for this acceleration.
A more effective acceleration technique is to use only the results from the last few steps of power iteration — for example, the last three, $\boldsymbol{v}^{(T-2)},\boldsymbol{v}^{(T-1)},\boldsymbol{v}^{(T)}$ — to build the Krylov subspace acceleration. This already captures most of the benefit, while reducing the computational cost of the QR and eigenvalue decompositions, thereby achieving a speedup.
Estimating the Upper Bound
Strictly speaking, power iteration and its accelerated version only give a lower bound on the spectral norm, which is admittedly sufficient in many cases. But if in some scenario we absolutely need an upper bound — for instance, if we want to use spectral normalization to guarantee that some matrix's norm is strictly less than 1 — then we need to look for another approach.
One basic option here is to compute the Schatten norm], which is indeed an upper bound on the spectral norm. Let the SVD of the matrix $\boldsymbol{W}$ be $\boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$, with singular values $\sigma_1\geq\sigma_2\geq\cdots\geq\sigma_m$. Then its Schatten-$p$ norm is defined as
\begin{equation}\Vert\boldsymbol{W}\Vert_{S,p} = \sqrt[\uproot{10}p]{\sum_{i=1}^m \sigma_i^p} \quad\geq\quad \sigma_1\end{equation}
The larger $p$ is, the closer this gets to the exact value. When $p$ is even, computing the Schatten-$p$ norm is relatively feasible, mainly by exploiting the following identity:
\begin{gather}\newcommand{tr}{\mathop{\text{tr}}}\Vert\boldsymbol{W}\Vert_{S,2k} = \sqrt[\uproot{10}2k]{\sum_{i=1}^m \sigma_i^{2k}} = \sqrt[\uproot{3}2k]{\tr(\boldsymbol{V}\boldsymbol{\Sigma}^{2k}\boldsymbol{V}^{\top})} = \sqrt[\uproot{3}2k]{\tr((\boldsymbol{W}^{\top}\boldsymbol{W})^k)} \\ \Vert\boldsymbol{W}\Vert_{S,4k} = \sqrt[\uproot{10}4k]{\sum_{i=1}^m \sigma_i^{4k}} = \sqrt[\uproot{3}2k]{\Vert\boldsymbol{V}\boldsymbol{\Sigma}^{2k}\boldsymbol{V}^{\top}\Vert_{S,2}} = \sqrt[\uproot{3}2k]{\Vert(\boldsymbol{W}^{\top}\boldsymbol{W})^k\Vert_{S,2}}\label{eq:S-4k}\end{gather}
It's easy to see that the Schatten-$2$ norm is actually just the Frobenius norm, which can be computed by summing the squares of all elements and taking the square root. So the two identities above show that as long as we compute $(\boldsymbol{W}^{\top}\boldsymbol{W})^k$, we can obtain $\Vert\boldsymbol{W}\Vert_{S,2k}$ and $\Vert\boldsymbol{W}\Vert_{S,4k}$ at relatively low cost.
The first step, $\boldsymbol{W}^{\top}\boldsymbol{W}$, has complexity $\mathcal{O}(nm^2)$, and each subsequent squaring step has complexity $\mathcal{O}(m^3)$, so the complexity of computing $\mathcal{O}(nm^2 + Tm^3)$ works out to $(\boldsymbol{W}^{\top}\boldsymbol{W})^{2^T}$. In theory, this complexity far exceeds that of power iteration — in fact, when $n,m$ is fairly large, even the first step $\boldsymbol{W}^{\top}\boldsymbol{W}$ alone already exceeds the cost of power iteration. So if all we want is an estimate of the spectral norm without requiring an upper bound, power iteration should still be the preferred choice.
That said, matrix multiplication is generally highly parallelizable, so when $n,m$ is not too large, or $n\gg m$, the Schatten-$p$ norm doesn't necessarily create a computational bottleneck, and in that case we can consider using it. Or alternatively, when we absolutely need a strict upper bound, this seems to be the most direct route available.
Preserving Numerical Stability
To compute $\Vert\boldsymbol{W}\Vert_{S,2^{T+2}}$, a naive implementation would start from $\boldsymbol{M} = \boldsymbol{W}^{\top}\boldsymbol{W}$, repeat $\boldsymbol{M}\leftarrow \boldsymbol{M}^2$ a total of $T$ times to get $(\boldsymbol{W}^{\top}\boldsymbol{W})^{2^T}$, and then substitute into the formula $\eqref{eq:S-4k}$. But since this is a super-exponential operation, it will quickly explode to NaN or collapse to zero.
In general, $\Vert\boldsymbol{W}\Vert_{S,2^{T+2}}$ itself won't blow up numerically — the problem lies in explicitly computing $(\boldsymbol{W}^{\top}\boldsymbol{W})^{2^T}$. The solution is to renormalize after each power step 【$\boldsymbol{M}\leftarrow \boldsymbol{M}^2/\tr(\boldsymbol{M}^2)$】, which prevents numerical explosion and also ensures the scaling stays compact enough to avoid collapsing to zero. At the same time, we accumulate the normalization factor from each step in the log domain, for use in the final computation of $\Vert\boldsymbol{W}\Vert_{S,2^{T+2}}$.
The computational procedure is summarized as follows:
$$\begin{array}{|l|} \hline \text{compute}\Vert\boldsymbol{W}\Vert_{S,2^{T+2}}\text{as upper bound of spectral norm} \\[4pt] \hline \begin{array}{ll} 1: & \text{Initialize }\log S = \log\tr(\boldsymbol{W}^{\top}\boldsymbol{W}), \boldsymbol{M}=\frac{\boldsymbol{W}^{\top}\boldsymbol{W}}{ \tr(\boldsymbol{W}^{\top}\boldsymbol{W})}\\ 2: & \textbf{For }t=1,2,\cdots,T\textbf{ do } \\ 3: & \qquad \log S\leftarrow 2\log S + \log \tr(\boldsymbol{M}^2) \\ 4: & \qquad \boldsymbol{M} \leftarrow \frac{\boldsymbol{M}^2}{ \tr(\boldsymbol{M}^2)} \\ 5: & \text{Output } \exp\left(\frac{\log S + \log\Vert\boldsymbol{M}\Vert_F}{2^{T+1}}\right) \end{array} \\ \hline \end{array}$$
A simple reference implementation:
@jax.jit
def tr(w):
return w.trace(axis1=-1, axis2=-2)[..., None, None]
@jax.jit(static_argnums=(1,))
def spec_norm_v3(w, T=5):
m = (m := w.mT @ w) / (s := tr(m))
ms_step = lambda i, ms: ((m := ms[0] @ ms[0]) / (s := tr(m)), 2 * ms[1] + jnp.log(s))
m, logs = lax.fori_loop(0, T, ms_step, (m, jnp.log(s)))
logf = 0.5 * jnp.log((m**2).sum(axis=[-1, -2], keepdims=True))
return jnp.exp((logs + logf) / 2**(T + 1))
Higher-Order Moments
To compute $\Vert\boldsymbol{W}\Vert_{S,2^{T+2}}$, we need to compute $\boldsymbol{W}^{\top}\boldsymbol{W},\cdots,(\boldsymbol{W}^{\top}\boldsymbol{W})^{2^{T-1}},(\boldsymbol{W}^{\top}\boldsymbol{W})^{2^T}$ by some means, which means we can simultaneously obtain $\Vert\boldsymbol{W}\Vert_{S,2},\cdots,\Vert\boldsymbol{W}\Vert_{S,2^{T+1}},\Vert\boldsymbol{W}\Vert_{S,2^{T+2}}$, but in the end only $\Vert\boldsymbol{W}\Vert_{S,2^{T+2}}$ is used — again, this seems somewhat "wasteful."
Is there a way, just as with the Krylov subspace method for power iteration, to make use of all these results to improve estimation accuracy? Indeed there is! Fast Tight Spectral-Norm Bounds] offers an idea for improving the estimate via nonlinear programming. Let's start with a simple case: suppose we've obtained the 2nd and 4th moments of the singular values,
\begin{equation}\sum_{i=1}^m \sigma_i^2 = S_2, \qquad \sum_{i=1}^m \sigma_i^4 = S_4\end{equation}
Then, according to the results of the previous two sections, $S_4^{1/4}$ is closer to the spectral norm than $S_2^{1/2}$, so we return $S_4^{1/4}$ as the upper bound on the spectral norm, discarding $S_2$. But is $S_2$ really useless? Consider: if $m=2$, then we'd have exactly two equations and two unknowns, and in theory we could solve for $\sigma_1,\sigma_2$ exactly! When $m > 2$, exact solving isn't possible, but we can still narrow down the range of $\sigma_1$. Specifically, we have
\begin{equation}\frac{S_2 - \sigma_1^2}{m-1} = \frac{1}{m-1}\sum_{i=2}^m \sigma_i^2 \leq \sqrt{\frac{1}{m-1}\sum_{i=2}^m \sigma_i^4} = \sqrt{\frac{S_4 - \sigma_1^4}{m-1}}\end{equation}
i.e., $(S_2 - \sigma_1^2)^2\leq (m-1)(S_4 - \sigma_1^4)$, which is essentially a quadratic inequality in one variable, and is easy to solve to get
\begin{equation}\sigma_1 \leq \sqrt{\frac{S_2+\sqrt{(m-1)(mS_4-S_2^2)}}{m}}\label{eq:s2-s4}\end{equation}
This gives us an upper bound on $\sigma_1$ expressed in terms of $S_2$ and $S_4$, which is a better estimate than $S_4^{1/4}$. Fast Tight Spectral-Norm Bounds] further generalizes this to a form that simultaneously makes use of $S_2,S_4,S_6,S_8$, which is more compact but also more involved — interested readers can consult the original paper. Using higher-order moments to improve the estimate is theoretically feasible, but in practice it often requires solving complex systems of nonlinear equations, which limits its practical value.
Limitations
In theory, the result $\eqref{eq:s2-s4}$ can also be generalized to using arbitrary $S_{2k}$ and $S_{4k}$ to obtain a more accurate upper bound:
\begin{equation}\sigma_1 \leq \sqrt[\uproot{10}2k]{\frac{S_{2k}+\sqrt{(m-1)(mS_{4k}-S_{2k}^2)}}{m}}\end{equation}
However, when $k$ is fairly large, this result has very limited practical significance, because it requires explicitly computing $S_{2k}$ and $S_{4k}$, which will explode or collapse when $k$ is fairly large — not a realistic option. One possible improvement is to factor out $S_{4k}^{1/4k}$:
\begin{equation}\sqrt[\uproot{10}2k]{\frac{S_{2k}+\sqrt{(m-1)(mS_{4k}-S_{2k}^2)}}{m}} = S_{4k}^{1/4k}\cdot\sqrt[\uproot{10}2k]{\frac{S_{2k}/S_{4k}^{1/2}+\sqrt{(m-1)(m-S_{2k}^2/S_{4k})}}{m}}\end{equation}
and then let $S_{2k}/S_{4k}^{1/2}=e^{\epsilon}$, expanding approximately in the log domain:
\begin{equation}\sqrt[\uproot{10}2k]{\frac{S_{2k}/S_{4k}^{1/2}+\sqrt{(m-1)(m-S_{2k}^2/S_{4k})}}{m}} \approx 1 - \frac{\epsilon^2}{4k(m-1)}\end{equation}
However, our goal is to obtain an upper bound on the spectral norm, and how to preserve the upper-bound property while doing this approximate expansion seems fairly complicated. On the other hand, if we've already computed a fairly large $k$, then $S_{4k}^{1/4k}$ itself is already quite accurate, so using programming-based ideas to further improve precision doesn't add much value.
The End
This post has briefly summarized several approaches to estimating the spectral norm. In practical applications, if we only need an approximate estimate for monitoring purposes, power iteration and its subspace-accelerated variant are usually sufficient; if we absolutely need a strict upper bound, we can consider computing the Schatten norm. The "don't waste anything" philosophy underlying each of these two approaches — the Krylov subspace method exploiting the iteration history, and nonlinear programming exploiting information from higher-order moments — also offers valuable inspiration for algorithm design in general.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.