The Derivative of SVD

SVD的导数 苏剑林 · 2025-04-26 · Mathematics calculus analysis matrix SVD gradient

SVD (Singular Value Decomposition) is a common matrix factorization algorithm, and I believe many readers already have some familiarity with it — we also introduced it specifically in Road to Low-Rank Approximation (II): SVD. However, have readers ever considered that SVD can actually be differentiated? I myself was rather surprised when I first learned this result, because intuitively "decomposition" tends to be non-differentiable. But the fact is that SVD is indeed differentiable in general cases, which means that, in principle, we can embed SVD into a model and train it end-to-end with gradient-based optimizers.

The question is: since SVD is differentiable, what does its derivative actually look like? Next, following the reference Differentiating the Singular Value Decomposition, we will derive the formula for the derivative of SVD step by step.

Foundations for the Derivation

Suppose $\boldsymbol{W}$ is a full-rank $n\times n$ matrix, with all singular values pairwise distinct — this is the relatively easy case to discuss, and later we'll also discuss how far these conditions can be relaxed. Let the SVD of $\boldsymbol{W}$ be:

\begin{equation}\boldsymbol{W} = \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}\end{equation}more

So-called "differentiating SVD" essentially means finding the gradient, or differential, of $\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V}$ with respect to $\boldsymbol{W}$. To do this, let's first differentiate both sides of the above equation:

\begin{equation}d\boldsymbol{W} = (d\boldsymbol{U})\boldsymbol{\Sigma}\boldsymbol{V}^{\top} + \boldsymbol{U}(d\boldsymbol{\Sigma})\boldsymbol{V}^{\top} + \boldsymbol{U}\boldsymbol{\Sigma}(d\boldsymbol{V})^{\top}\end{equation}

Left-multiplying by $\boldsymbol{U}^{\top}$, right-multiplying by $\boldsymbol{V}$, and using $\boldsymbol{U}^{\top}\boldsymbol{U} = \boldsymbol{I}, \boldsymbol{V}^{\top}\boldsymbol{V} = \boldsymbol{I}$, we get

\begin{equation}\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V} = \boldsymbol{U}^{\top}(d\boldsymbol{U})\boldsymbol{\Sigma} + d\boldsymbol{\Sigma} + \boldsymbol{\Sigma}(d\boldsymbol{V})^{\top}\boldsymbol{V}\label{eq:core}\end{equation}

This will be the basis for the derivation that follows. Note that if we differentiate both sides of the identity $\boldsymbol{U}^{\top}\boldsymbol{U} = \boldsymbol{I}, \boldsymbol{V}^{\top}\boldsymbol{V} = \boldsymbol{I}$, we also obtain

\begin{equation}(d\boldsymbol{U})^{\top}\boldsymbol{U} + \boldsymbol{U}^{\top}(d\boldsymbol{U}) = \boldsymbol{0},\quad (d\boldsymbol{V})^{\top}\boldsymbol{V} + \boldsymbol{V}^{\top}(d\boldsymbol{V}) = \boldsymbol{0}\end{equation}

which shows that $\boldsymbol{U}^{\top}(d\boldsymbol{U})$ and $(d\boldsymbol{V})^{\top}\boldsymbol{V}$ are both antisymmetric matrices.

Singular Values

A characteristic feature of antisymmetric matrices is that their diagonal entries are all zero, while $\boldsymbol{\Sigma}$ is a diagonal matrix whose off-diagonal entries are all zero. This suggests that we may need to handle the diagonal and off-diagonal entries separately.

First, let's define the matrices $\boldsymbol{I}$ and $\bar{\boldsymbol{I}}$: $\boldsymbol{I}$ is simply the identity matrix, i.e., all diagonal entries are 1 and all off-diagonal entries are 0; $\bar{\boldsymbol{I}}$ is the complement of the identity matrix, i.e., all diagonal entries are 0 and all off-diagonal entries are 1. Using $\boldsymbol{I}$, $\bar{\boldsymbol{I}}$, and the Hadamard product $\otimes$, we can extract the diagonal and off-diagonal parts of equation $\eqref{eq:core}$ respectively:

\begin{align} \boldsymbol{I}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}) =&\, \boldsymbol{I}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{U})\boldsymbol{\Sigma} + d\boldsymbol{\Sigma} + \boldsymbol{\Sigma}(d\boldsymbol{V})^{\top}\boldsymbol{V}) = d\boldsymbol{\Sigma} \label{eq:value} \\[8pt] \bar{\boldsymbol{I}}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}) =&\, \bar{\boldsymbol{I}}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{U})\boldsymbol{\Sigma} + d\boldsymbol{\Sigma} + \boldsymbol{\Sigma}(d\boldsymbol{V})^{\top}\boldsymbol{V}) = \boldsymbol{U}^{\top}(d\boldsymbol{U})\boldsymbol{\Sigma} + \boldsymbol{\Sigma}(d\boldsymbol{V})^{\top}\boldsymbol{V}\label{eq:vector} \end{align}

Now let's look at equation $\eqref{eq:value}$, which can equivalently be written as

\begin{equation}d\sigma_i = \boldsymbol{u}_i^{\top}(d\boldsymbol{W})\boldsymbol{v}_i\label{eq:d-sigma}\end{equation}

This is the differential of the $i$-th singular value $\sigma_i$, where $\boldsymbol{u}_i,\boldsymbol{v}_i$ are respectively the $i$-th columns of $\boldsymbol{U},\boldsymbol{V}$. The spectral-norm gradient discussed in From the Spectral Norm Gradient to a New Kind of Weight Decay is in fact just the special case here where $i=1$.

Singular Vectors

Next let's look at equation $\eqref{eq:vector}$:

\begin{equation}\bar{\boldsymbol{I}}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}) = \boldsymbol{U}^{\top}(d\boldsymbol{U})\boldsymbol{\Sigma} + \boldsymbol{\Sigma}(d\boldsymbol{V})^{\top}\boldsymbol{V}\label{eq:vector-1}\end{equation}

Taking the transpose,

\begin{equation}\begin{aligned} \bar{\boldsymbol{I}}\otimes(\boldsymbol{V}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}) =&\, \boldsymbol{\Sigma}(d\boldsymbol{U})^{\top}\boldsymbol{U} + \boldsymbol{V}^{\top}(d\boldsymbol{V})\boldsymbol{\Sigma} \\[6pt] =&\, -\boldsymbol{\Sigma}\boldsymbol{U}^{\top}(d\boldsymbol{U}) - (d\boldsymbol{V})^{\top}\boldsymbol{V}\boldsymbol{\Sigma} \end{aligned}\label{eq:vector-2}\end{equation}

the second equality uses the fact that "$\boldsymbol{U}^{\top}(d\boldsymbol{U})$ and $(d\boldsymbol{V})^{\top}\boldsymbol{V}$ are both antisymmetric matrices." Equations $\eqref{eq:vector-1}$ and $\eqref{eq:vector-2}$ form a system of linear equations in $d\boldsymbol{U},d\boldsymbol{V}$, from which we want to solve for $d\boldsymbol{U},d\boldsymbol{V}$.

The approach is ordinary elimination. First, from $\eqref{eq:vector-1}\times\boldsymbol{\Sigma} + \boldsymbol{\Sigma}\times\eqref{eq:vector-2}$ we obtain

\begin{equation}\bar{\boldsymbol{I}}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}\boldsymbol{\Sigma} + \boldsymbol{\Sigma}\boldsymbol{V}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}) = \boldsymbol{U}^{\top}(d\boldsymbol{U})\boldsymbol{\Sigma}^2 - \boldsymbol{\Sigma}^2\boldsymbol{U}^{\top}(d\boldsymbol{U})\end{equation}

Here we've used the fact that the diagonal matrix $\boldsymbol{\Sigma}$ satisfies $\boldsymbol{\Sigma}(\bar{\boldsymbol{I}}\otimes \boldsymbol{M}) = \bar{\boldsymbol{I}}\otimes (\boldsymbol{\Sigma}\boldsymbol{M})$, as well as $(\bar{\boldsymbol{I}}\otimes \boldsymbol{M})\boldsymbol{\Sigma} = \bar{\boldsymbol{I}}\otimes (\boldsymbol{M}\boldsymbol{\Sigma})$. We know that left- (right-) multiplying by a diagonal matrix is equivalent to multiplying each row (column) of a matrix by the corresponding diagonal entry, so if we define the matrix $\boldsymbol{E}$, where $\boldsymbol{E}_{i,j} = \sigma_j^2 - \sigma_i^2$, then $\boldsymbol{U}^{\top}(d\boldsymbol{U})\boldsymbol{\Sigma}^2 - \boldsymbol{\Sigma}^2\boldsymbol{U}^{\top}(d\boldsymbol{U}) = \boldsymbol{E}\otimes (\boldsymbol{U}^{\top}(d\boldsymbol{U}))$, and the above equation can be written as

\begin{equation}\bar{\boldsymbol{I}}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}\boldsymbol{\Sigma} + \boldsymbol{\Sigma}\boldsymbol{V}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}) = \boldsymbol{E}\otimes (\boldsymbol{U}^{\top}(d\boldsymbol{U}))\label{eq:dU-0}\end{equation}

from which we can then solve

\begin{equation}d\boldsymbol{U} = \boldsymbol{U}(\boldsymbol{F}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}\boldsymbol{\Sigma} + \boldsymbol{\Sigma}\boldsymbol{V}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}))\label{eq:dU}\end{equation}

Similarly, from $\boldsymbol{\Sigma}\times \eqref{eq:vector-1} + \eqref{eq:vector-2}\times \boldsymbol{\Sigma}$ we solve:

\begin{equation}d\boldsymbol{V} = \boldsymbol{V}(\boldsymbol{F}\otimes(\boldsymbol{V}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}\boldsymbol{\Sigma} + \boldsymbol{\Sigma}\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}))\label{eq:dV}\end{equation}

Equation $\eqref{eq:dU},\eqref{eq:dV}$ is thus the differential of the eigenvectors [singular vectors]. Here

\begin{equation}\boldsymbol{F}_{i,j} = \left\{\begin{aligned} &\, 1/(\sigma_j^2 - \sigma_i^2), &\, i\neq j \\ &\, 0, &\, i = j \end{aligned}\right.\end{equation}

Gradient (I)

Now that we have the differential, how do we get the gradient? This turns out to be somewhat troublesome — not technically troublesome, but troublesome in terms of notation. For instance, $\boldsymbol{W},\boldsymbol{U}$ is an $n\times n$ matrix, so the full gradient of $\boldsymbol{U}$ with respect to $\boldsymbol{W}$ is a rank-4 tensor of shape $n\times n\times n\times n$, and higher-order tensors are something most people, myself included, aren't very comfortable with.

To sidestep the trouble of higher-order tensors, we have two options. First, from a programming perspective, there's really no need to derive the explicit form of the gradient at all — we can just write an equivalent forward expression based on the differential result, and let the framework's automatic differentiation handle the rest. For example, from equation $\eqref{eq:d-sigma}$ we can conclude that the gradient of $\sigma_i$ equals the gradient of $\newcommand{\sg}[1]{\color{skyblue}{#1}} \sg{\boldsymbol{u}_i}^{\top} \boldsymbol{W} \sg{\boldsymbol{v}_i}$, i.e.,

\begin{equation}\nabla_{\boldsymbol{W}} \sigma_i = \nabla_{\boldsymbol{W}} (\sg{\boldsymbol{u}_i}^{\top} \boldsymbol{W} \sg{\boldsymbol{v}_i})\end{equation}

Here we change the color of the symbol to $\sg{\blacksquare}$ to represent the stop_gradient operator, to avoid making the formula too cluttered. It happens that $\sigma_i$ also equals $\boldsymbol{u}_i^{\top}\boldsymbol{W}\boldsymbol{v}_i$, so we only need to replace every occurrence of $\sigma_i$ in the code with $\sg{\boldsymbol{u}_i}^{\top} \boldsymbol{W} \sg{\boldsymbol{v}_i}$, and we'll automatically get the correct gradient. More generally, we have

\begin{equation}\nabla_{\boldsymbol{W}} \boldsymbol{\Sigma} = \nabla_{\boldsymbol{W}} (\boldsymbol{I}\otimes(\sg{\boldsymbol{U}}^{\top} \boldsymbol{W} \sg{\boldsymbol{V}}))\end{equation}

that is, replace every $\boldsymbol{\Sigma}$ with $\boldsymbol{I}\otimes(\sg{\boldsymbol{U}}^{\top} \boldsymbol{W} \sg{\boldsymbol{V}})$.

Similarly, from equation $\eqref{eq:dU}$ we know that

\begin{equation}\nabla_{\boldsymbol{W}}\boldsymbol{U} = \nabla_{\boldsymbol{W}}(\sg{\boldsymbol{U}}(\sg{\boldsymbol{F}}\otimes(\sg{\boldsymbol{U}}^{\top}\boldsymbol{W}\sg{\boldsymbol{V}\boldsymbol{\Sigma}} + \sg{\boldsymbol{\Sigma}\boldsymbol{V}}^{\top}\boldsymbol{W}^{\top}\sg{\boldsymbol{U}})))\end{equation}

We can verify that $\boldsymbol{U}(\boldsymbol{F}\otimes(\boldsymbol{U}^{\top}\boldsymbol{W}\boldsymbol{V}\boldsymbol{\Sigma} + \boldsymbol{\Sigma}\boldsymbol{V}^{\top}\boldsymbol{W}^{\top}\boldsymbol{U}))$ is exactly the zero matrix, so we only need to replace every occurrence of $\boldsymbol{U}$ in the code with

\begin{equation}\boldsymbol{U} \quad \to \quad \sg{\boldsymbol{U}} + \sg{\boldsymbol{U}}(\sg{\boldsymbol{F}}\otimes(\sg{\boldsymbol{U}}^{\top}\boldsymbol{W}\sg{\boldsymbol{V}\boldsymbol{\Sigma}} + \sg{\boldsymbol{\Sigma}\boldsymbol{V}}^{\top}\boldsymbol{W}^{\top}\sg{\boldsymbol{U}}))\end{equation}

and that will preserve the correct forward result while also giving the correct gradient. Based on the same principle, the replacement format for $\boldsymbol{V}$ is

\begin{equation}\boldsymbol{V} \quad \to \quad \sg{\boldsymbol{V}} + \sg{\boldsymbol{V}}(\sg{\boldsymbol{F}}\otimes(\sg{\boldsymbol{V}}^{\top}\boldsymbol{W}^{\top}\sg{\boldsymbol{U}\boldsymbol{\Sigma}} + \sg{\boldsymbol{\Sigma}\boldsymbol{U}}^{\top}\boldsymbol{W}\sg{\boldsymbol{V}}))\end{equation}

Gradient (II)

The second approach is to directly derive the gradient of the loss function with respect to $\boldsymbol{W}$. Specifically, suppose the loss function is a function of $\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V}$, denoted $\mathcal{L}(\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V})$; we want to compute $\nabla_{\boldsymbol{W}}\mathcal{L}$ directly, which is a matrix that can be expressed in terms of $\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V},\nabla_{\boldsymbol{U}}\mathcal{L},\nabla_{\boldsymbol{\Sigma}}\mathcal{L},\nabla_{\boldsymbol{V}}\mathcal{L}$ — all of these quantities are just matrices, so there's no need to deal with higher-order tensors.

In the previous section, we already found equivalent functions of $\boldsymbol{U},\boldsymbol{\Sigma},\boldsymbol{V}$ with the same gradient; apart from the parts that are stop_gradient'ed, these equivalent functions are linear in $\boldsymbol{W}$, so the problem essentially becomes one of finding the gradient of a linear composite function. We already discussed the relevant methods in the "Matrix Differentiation" section of Road to Low-Rank Approximation (I): Pseudo-inverse. Specifically, we have

\begin{align} \boldsymbol{X} = \boldsymbol{A}\boldsymbol{B}\boldsymbol{C} &\,\quad\Rightarrow\quad \nabla_{\boldsymbol{B}}f(\boldsymbol{X}) = \boldsymbol{A}^{\top}(\nabla_{\boldsymbol{X}}f(\boldsymbol{X}))\boldsymbol{C}^{\top} \\[8pt] \boldsymbol{X} = \boldsymbol{A}\boldsymbol{B}^{\top}\boldsymbol{C} &\,\quad\Rightarrow\quad \nabla_{\boldsymbol{B}}f(\boldsymbol{X}) = \boldsymbol{C}(\nabla_{\boldsymbol{X}}f(\boldsymbol{X}))^{\top}\boldsymbol{A} \\[8pt] \boldsymbol{X} = \boldsymbol{A}\otimes\boldsymbol{B} &\,\quad\Rightarrow\quad \nabla_{\boldsymbol{B}}f(\boldsymbol{X}) = \boldsymbol{A}\otimes \nabla_{\boldsymbol{X}}f(\boldsymbol{X}) \end{align}

Using these basic formulas together with the chain rule for composite functions, we can write

\begin{equation}\begin{aligned} \nabla_{\boldsymbol{W}}\mathcal{L} \quad = \qquad &\,\boldsymbol{U}(\boldsymbol{F}\otimes(\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{U}}\mathcal{L}) - (\nabla_{\boldsymbol{U}}\mathcal{L})^{\top}\boldsymbol{U}))\boldsymbol{\Sigma}\boldsymbol{V}^{\top} \\[6pt] + &\,\boldsymbol{U}(\boldsymbol{I}\otimes(\nabla_{\boldsymbol{\Sigma}}\mathcal{L}))\boldsymbol{V}^{\top} \\[6pt] + &\,\boldsymbol{U}\boldsymbol{\Sigma}(\boldsymbol{F}\otimes(\boldsymbol{V}^{\top}(\nabla_{\boldsymbol{V}}\mathcal{L}) - (\nabla_{\boldsymbol{V}}\mathcal{L})^{\top}\boldsymbol{V})))\boldsymbol{V}^{\top} \end{aligned}\end{equation}

The whole process just involves repeatedly applying the basic formulas and the chain rule, along with $\boldsymbol{F}^{\top} = -\boldsymbol{F}$; in principle there's no real difficulty, it just requires careful concentration. I recommend that readers work through this derivation themselves — it's a quite practical exercise in matrix differentiation. Finally, let's introduce two pieces of notation

\begin{equation}\newcommand{\sym}[1]{\color{red}{[}#1\color{red}{]_{sym}}} \newcommand{\skew}[1]{\color{red}{[}#1\color{red}{]_{skew}}} \sym{\boldsymbol{X}} = \frac{1}{2}(\boldsymbol{X} + \boldsymbol{X}^{\top}),\qquad \skew{\boldsymbol{X}} = \frac{1}{2}(\boldsymbol{X} - \boldsymbol{X}^{\top})\end{equation}

and we can write the gradient result compactly as

\begin{equation}\nabla_{\boldsymbol{W}}\mathcal{L} = \boldsymbol{U}\Big(2(\boldsymbol{F}\otimes\skew{\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{U}}\mathcal{L})})\boldsymbol{\Sigma} + \boldsymbol{I}\otimes(\nabla_{\boldsymbol{\Sigma}}\mathcal{L}) + 2\boldsymbol{\Sigma}(\boldsymbol{F}\otimes\skew{\boldsymbol{V}^{\top}(\nabla_{\boldsymbol{V}}\mathcal{L})})\Big)\boldsymbol{V}^{\top} \label{eq:w-grad-l}\end{equation}

Gradient (III)

Now let's try a small exercise: finding the gradient of $\boldsymbol{O}=\mathop{\text{msign}}(\boldsymbol{W})=\boldsymbol{U}\boldsymbol{V}^{\top}$, where the concept of $\mathop{\text{msign}}$ was already discussed when we introduced Muon in Appreciating the Muon Optimizer: The Essential Leap from Vectors to Matrices.

According to the definition of $\mathop{\text{msign}}$, we have

\begin{equation}\nabla_{\boldsymbol{U}}\mathcal{L} = (\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V},\qquad \nabla_{\boldsymbol{V}}\mathcal{L} = (\nabla_{\boldsymbol{O}}\mathcal{L})^{\top} \boldsymbol{U}\end{equation}

Substituting into equation $\eqref{eq:w-grad-l}$ gives

\begin{equation}\begin{aligned} \nabla_{\boldsymbol{W}}\mathcal{L} =&\, 2\boldsymbol{U}\Big((\boldsymbol{F}\otimes\skew{\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V}})\boldsymbol{\Sigma} + \boldsymbol{\Sigma}(\boldsymbol{F}\otimes\skew{\boldsymbol{V}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})^{\top}\boldsymbol{U}})\Big)\boldsymbol{V}^{\top} \\[6pt] =&\, 2\boldsymbol{U}\Big((\boldsymbol{F}\otimes\skew{\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V}})\boldsymbol{\Sigma} - \boldsymbol{\Sigma}(\boldsymbol{F}\otimes\skew{\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V}})\Big)\boldsymbol{V}^{\top} \\[6pt] =&\, 2\boldsymbol{U}\big(\boldsymbol{G}\otimes\skew{\boldsymbol{U}^{\top}(\nabla_{\boldsymbol{O}}\mathcal{L})\boldsymbol{V}}\big)\boldsymbol{V}^{\top} \end{aligned}\end{equation}

where $\boldsymbol{G}_{i,j} = 1/(\sigma_i + \sigma_j)$.

Gradient (IV)

Finally, let's consider a commonly used special case: when $\boldsymbol{W}$ is also a positive-definite symmetric matrix, its SVD takes the form $\boldsymbol{V}\boldsymbol{\Sigma}\boldsymbol{V}^{\top}$, i.e., $\boldsymbol{U}=\boldsymbol{V}$. Let's repeat the earlier derivation, first differentiating both sides:

\begin{equation}d\boldsymbol{W} = (d\boldsymbol{V})\boldsymbol{\Sigma}\boldsymbol{V}^{\top} + \boldsymbol{V}(d\boldsymbol{\Sigma})\boldsymbol{V}^{\top} + \boldsymbol{V}\boldsymbol{\Sigma}(d\boldsymbol{V})^{\top}\end{equation}

then left-multiplying by $\boldsymbol{V}^{\top}$ and right-multiplying by $\boldsymbol{V}$:

\begin{equation}\begin{aligned} \boldsymbol{V}^{\top}(d\boldsymbol{W})\boldsymbol{V} =&\, \boldsymbol{V}^{\top}(d\boldsymbol{V})\boldsymbol{\Sigma} + d\boldsymbol{\Sigma} + \boldsymbol{\Sigma}(d\boldsymbol{V})^{\top}\boldsymbol{V} \\[6pt] =&\, \boldsymbol{V}^{\top}(d\boldsymbol{V})\boldsymbol{\Sigma} + d\boldsymbol{\Sigma} - \boldsymbol{\Sigma}\boldsymbol{V}^{\top}(d\boldsymbol{V}) \end{aligned}\end{equation}

from which we see

\begin{gather} d\boldsymbol{\Sigma} = \boldsymbol{I}\otimes(\boldsymbol{V}(d\boldsymbol{W})\boldsymbol{V}^{\top}) \\[8pt] \boldsymbol{V}^{\top}(d\boldsymbol{V})\boldsymbol{\Sigma} - \boldsymbol{\Sigma}\boldsymbol{V}^{\top}(d\boldsymbol{V}) = \bar{\boldsymbol{I}}\otimes(\boldsymbol{V}(d\boldsymbol{W})\boldsymbol{V}^{\top}) \end{gather}

From the second equation, we can further solve

\begin{equation}d\boldsymbol{V} = \boldsymbol{V}(\boldsymbol{K}^{\top}\otimes(\boldsymbol{V}^{\top}(d\boldsymbol{W})\boldsymbol{V}))\end{equation}

where

\begin{equation}\boldsymbol{K}_{i,j} = \left\{\begin{aligned} &\, 1/(\sigma_i - \sigma_j), &\, i\neq j \\ &\, 0, &\, i = j \end{aligned}\right.\end{equation}

Based on this result, we have

\begin{equation}\nabla_{\boldsymbol{W}}\mathcal{L} = \boldsymbol{V}(\boldsymbol{K}^{\top}\otimes(\boldsymbol{V}^{\top}(\nabla_{\boldsymbol{V}}\mathcal{L})) + \boldsymbol{I}\otimes(\nabla_{\boldsymbol{\Sigma}}\mathcal{L}))\boldsymbol{V}^{\top} \end{equation}

Watch out for a trap! The equation above is actually wrong. It's the result of the chain rule, but the chain rule only applies to unconstrained differentiation, whereas here we have the constraint $\boldsymbol{W}=\boldsymbol{W}^{\top}$, so the correct gradient should also include symmetrization:

\begin{equation}\begin{aligned} \nabla_{\boldsymbol{W}}\mathcal{L} =&\, \boldsymbol{V}\Big(\sym{\boldsymbol{K}^{\top}\otimes(\boldsymbol{V}^{\top}(\nabla_{\boldsymbol{V}}\mathcal{L}))} + \boldsymbol{I}\otimes(\nabla_{\boldsymbol{\Sigma}}\mathcal{L})\Big)\boldsymbol{V}^{\top} \\[6pt] =&\, \boldsymbol{V}\Big(\boldsymbol{K}^{\top}\otimes\skew{\boldsymbol{V}^{\top}(\nabla_{\boldsymbol{V}}\mathcal{L})} + \boldsymbol{I}\otimes(\nabla_{\boldsymbol{\Sigma}}\mathcal{L})\Big)\boldsymbol{V}^{\top}\end{aligned}\end{equation}

Another trap is directly substituting $\boldsymbol{U}=\boldsymbol{V}$ into equation $\eqref{eq:w-grad-l}$ to derive the result, which would cause the $\boldsymbol{K}$ term to be doubled. The reason is that equation $\eqref{eq:w-grad-l}$ distinguishes between $\nabla_{\boldsymbol{U}}\mathcal{L}$ and $\nabla_{\boldsymbol{V}}\mathcal{L}$, but under the positive-definite symmetric assumption, $\boldsymbol{U},\boldsymbol{V}$ are identical, so taking the gradient with respect to $\boldsymbol{V}$ effectively sums up what were originally $\nabla_{\boldsymbol{U}}\mathcal{L},\nabla_{\boldsymbol{V}}\mathcal{L}$, leading to double counting. For related literature, see also Matrix Backpropagation for Deep Networks With Structured Layers.

Numerical Issues

Some readers might wonder: I can only guarantee that the singular values of the initialization matrix are pairwise distinct — how can I be sure that the matrix will still satisfy this condition after training? The answer is that the gradient will naturally take care of this for us. From equations $\eqref{eq:dU}$ and $\eqref{eq:dV}$, we can see that the gradient contains $\boldsymbol{F}$, and from $\boldsymbol{F}_{i,j} = \frac{1}{\sigma_j^2 - \sigma_i^2}$ we know that once two singular values get close to each other, the gradient will become extremely large, so the optimizer will automatically push them apart.

However, this property also introduces numerical instability into actual training, mainly manifesting as gradient explosion when $\sigma_i,\sigma_j$ are close. To address this, the paper Backpropagation-Friendly Eigendecomposition proposes replacing exact SVD with "Power Iteration." Later, the paper Robust Differentiable SVD proved that this is theoretically equivalent to a Taylor approximation of $\frac{1}{\sigma_i - \sigma_j}$ in $\boldsymbol{F}_{i,j} = -\frac{1}{（\sigma_i + \sigma_j)(\sigma_i - \sigma_j）}$ (assuming $\sigma_j < \sigma_i$):

\begin{equation}\frac{1}{\sigma_i - \sigma_j} = \frac{1}{\sigma_i}\frac{1}{1-(\sigma_j/\sigma_i)}\approx \frac{1}{\sigma_i}\left(1 + \left(\frac{\sigma_j}{\sigma_i}\right) + \left(\frac{\sigma_j}{\sigma_i}\right)^2 + \cdots + \left(\frac{\sigma_j}{\sigma_i}\right)^N \right)\end{equation}

With this Taylor approximation, at least for the case of $\sigma_j \to \sigma_i$, gradient explosion no longer occurs. Later still, the same authors generalized this to a more general Padé approximation in Why Approximate Matrix Square Root Outperforms Accurate SVD in Global Covariance Pooling?. I'm not too familiar with this line of work, so I won't go into further detail here.

I do have one question, though: if the sole purpose is to avoid numerical blow-up, it doesn't seem necessary to bring in all these tools — wouldn't it be simpler to just add a clipping/truncation to $\sigma_j/\sigma_i$? For instance

\begin{equation}\frac{1}{1-(\sigma_j/\sigma_i)}\approx \frac{1}{1-\min(\sigma_j/\sigma_i, 0.99)}\end{equation}

Wouldn't that simply prevent it from tending to infinity? Or is there something I'm not fully appreciating here? If readers are familiar with the relevant background, please feel free to point it out in the comments — thank you.

General Results

Up to this point, all the matrices that have appeared above have been $n\times n$ matrices, because the derivation was carried out under the assumption stated at the beginning of the article that "$\boldsymbol{W}$ is a full-rank $n\times n$ matrix, with all singular values pairwise distinct." In this section, let's discuss how far this condition can be relaxed.

Simply put, the condition for SVD to be differentiable is that "all nonzero singular values are pairwise distinct" — in other words, the requirement of being square can be dropped, and full rank can also be dropped, but the nonzero singular values must still all be distinct from one another. This is because as soon as there are equal singular values, the SVD is no longer unique, which fundamentally destroys differentiability. Of course, if we only need part of the derivative, we can relax the condition further — for example, if we only want the derivative of the spectral norm, i.e., $\sigma_1$, then we only need $\sigma_1 > \sigma_2$.

So, after relaxing the condition, how does the differential result change? Let's set up the general case $\boldsymbol{W}\in\mathbb{R}^{n\times m}$, of rank $r$, with SVD

\begin{equation}\boldsymbol{W} = \boldsymbol{U}\boldsymbol{\Sigma}\boldsymbol{V}^{\top},\quad\boldsymbol{U}\in\mathbb{R}^{n\times n},\boldsymbol{\Sigma}\in\mathbb{R}^{n\times m},\boldsymbol{V}\in\mathbb{R}^{m\times m}\end{equation}

The reason we allow zero singular values is that here we only care about $d\boldsymbol{U}_{[:, :r]}$ and $d\boldsymbol{V}_{[:, :r]}$, which, under the assumption that the nonzero singular values are pairwise distinct, can be uniquely determined. Starting from equation $\eqref{eq:dU-0}$, the derivation up through equation $\eqref{eq:dU-0}$ is completely general, and equation $\eqref{eq:dU-0}$ also shows why we need to exclude equal singular values: when $\sigma_i = \sigma_j$, $\sigma_j - \sigma_i$ can no longer be inverted. Keeping only the part of equation $\eqref{eq:dU-0}$ related to $d\boldsymbol{U}_{[:, :r]}$, we obtain

\begin{equation}\bar{\boldsymbol{I}}_{[:, :r]}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}\boldsymbol{\Sigma}_{[:, :r]} + \boldsymbol{\Sigma}\boldsymbol{V}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}_{[:, :r]}) = \boldsymbol{E}_{[:, :r]}\otimes (\boldsymbol{U}^{\top}(d\boldsymbol{U}_{[:, :r]}))\end{equation}

Here the definition of $\boldsymbol{E}$ is still $\boldsymbol{E}_{i,j} = \sigma_j^2 - \sigma_i^2$, but $\boldsymbol{E}_{[:, :r]}$ has excluded all the zeros, so it can be inverted without issue, giving the result

\begin{equation}d\boldsymbol{U}_{[:, :r]} = \boldsymbol{U}(\boldsymbol{F}_{[:, :r]}\otimes(\boldsymbol{U}^{\top}(d\boldsymbol{W})\boldsymbol{V}\boldsymbol{\Sigma}_{[:, :r]} + \boldsymbol{\Sigma}\boldsymbol{V}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}_{[:, :r]}))\end{equation}

Going further, let's partition it as follows

\begin{equation}\boldsymbol{U} = \begin{pmatrix}\boldsymbol{U}_{[:,:r]} & \boldsymbol{U}_{[:,r:]}\end{pmatrix},\quad \boldsymbol{F}_{[:, :r]} = \begin{pmatrix} \boldsymbol{F}_{[:r, :r]} \\ \boldsymbol{F}_{[r:, :r]}\end{pmatrix}\end{equation}

and combining this with $\boldsymbol{V}\boldsymbol{\Sigma}_{[:, :r]} = \boldsymbol{V}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]}$, we obtain

\begin{equation}\begin{aligned} d\boldsymbol{U}_{[:, :r]} =&\, \boldsymbol{U}_{[:, :r]}(\boldsymbol{F}_{[:r, :r]}\otimes(\boldsymbol{U}_{[:, :r]}^{\top}(d\boldsymbol{W})\boldsymbol{V}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]} + \boldsymbol{\Sigma}_{[:r, :r]}\boldsymbol{V}_{[:, :r]}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}_{[:, :r]})) \\[6pt] &\,\qquad + \boldsymbol{U}_{[:, r:]}(\boldsymbol{F}_{[r:, :r]}\otimes(\boldsymbol{U}_{[:, r:]}^{\top}(d\boldsymbol{W})\boldsymbol{V}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]})) \end{aligned}\label{eq:dU-r}\end{equation}

By assumption, when $i > r$, $\sigma_i=0$, so every row of $\boldsymbol{F}_{[r:, :r]}$ is $(\sigma_1^{-2},\sigma_2^{-2},\cdots,\sigma_r^{-2})$, and thus $\boldsymbol{F}_{[r:, :r]}\otimes$ is equivalent to right-multiplying by $\boldsymbol{\Sigma}_{[:r, :r]}^{-2}$, so

\begin{equation}\boldsymbol{F}_{[r:, :r]}\otimes(\boldsymbol{U}_{[:, r:]}^{\top}(d\boldsymbol{W})\boldsymbol{V}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]}) = \boldsymbol{U}_{[:, r:]}^{\top}(d\boldsymbol{W})\boldsymbol{V}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]}^{-1}\end{equation}

Finally, using $\boldsymbol{U}\boldsymbol{U}^{\top} = \boldsymbol{I}$ and $\boldsymbol{U} = (\boldsymbol{U}_{[:,:r]},\boldsymbol{U}_{[:,r:]})$, we obtain $\boldsymbol{U}_{[:, r:]}\boldsymbol{U}_{[:, r:]}^{\top} = \boldsymbol{I} - \boldsymbol{U}_{[:, :r]}\boldsymbol{U}_{[:, :r]}^{\top}$, and substituting into equation $\eqref{eq:dU-r}$ gives

\begin{equation}\begin{aligned} d\boldsymbol{U}_{[:, :r]} =&\, \boldsymbol{U}_{[:, :r]}(\boldsymbol{F}_{[:r, :r]}\otimes(\boldsymbol{U}_{[:, :r]}^{\top}(d\boldsymbol{W})\boldsymbol{V}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]} + \boldsymbol{\Sigma}_{[:r, :r]}\boldsymbol{V}_{[:, :r]}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}_{[:, :r]})) \\[6pt] &\,\qquad \color{orange}{+ (\boldsymbol{I} - \boldsymbol{U}_{[:, :r]}\boldsymbol{U}_{[:, :r]}^{\top})(d\boldsymbol{W})\boldsymbol{V}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]}^{-1}} \end{aligned}\end{equation}

This expresses $d\boldsymbol{U}_{[:, r:]}$ as a function of $\boldsymbol{U}_{[:, :r]},\boldsymbol{\Sigma}_{[:r, :r]},\boldsymbol{V}_{[:, :r]}$; under the assumption that the nonzero singular values are pairwise distinct, these three quantities are uniquely determined, so $d\boldsymbol{U}_{[:, :r]}$ is likewise uniquely determined. Compared with equation $\eqref{eq:dU}$, there is one extra term shown in orange. Similarly,

\begin{equation}\begin{aligned} d\boldsymbol{V}_{[:, :r]} =&\, \boldsymbol{V}_{[:, :r]}(\boldsymbol{F}_{[:r, :r]}\otimes(\boldsymbol{V}_{[:, :r]}^{\top}(d\boldsymbol{W})^{\top}\boldsymbol{U}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]} + \boldsymbol{\Sigma}_{[:r, :r]}\boldsymbol{U}_{[:, :r]}^{\top}(d\boldsymbol{W})\boldsymbol{V}_{[:, :r]})) \\[6pt] &\,\qquad \color{orange}{+ (\boldsymbol{I} - \boldsymbol{V}_{[:, :r]}\boldsymbol{V}_{[:, :r]}^{\top})(d\boldsymbol{W})^{\top}\boldsymbol{U}_{[:, :r]}\boldsymbol{\Sigma}_{[:r, :r]}^{-1}} \end{aligned}\end{equation}

Summary

This article has given a fairly detailed derivation of the formula for the derivative of SVD.