Issue 1:
Calculating a gradient of the SVD of a matrix where at least two of the singular values are equal lead to a gradient with only NaNs. Let me illustrate the problem by a sample code and result:

import jax
import jax.numpy as jnp

def f(arr):
U, S, Vh = jnp.linalg.svd(arr, full_matrices=False)
return jnp.linalg.norm((U * S) @ Vh)

A = jnp.diag(jnp.asarray([1, 1, 2, 3], dtype=float))

f_g = jax.value_and_grad(f)

print(f_g(A))

Result:

(DeviceArray(3.8729835, dtype=float32),
DeviceArray([[nan, nan, nan, nan],
[nan, nan, nan, nan],
[nan, nan, nan, nan],
[nan, nan, nan, nan]], dtype=float32))

This sample is nothing but a very complicated expression of the norm of a matrix whose gradient is defined (easy to see if replacing the gradient call by jax.value_and_grad(jnp.linalg.norm) [2]).

After some discussion with @frederikwilde (the author of the commit f443e19) we found that the problem is in the calculation of the derivative of the singular matrices U and V†. The derivation of the derivative calculates in one step the matrix of the inverses 1/(Si - Sj) (with Si the singular values). Obviously, this leads to the described problem.

In the source code https://github.com/google/jax/blob/main/jax/_src/lax/linalg.py#L1291 there is already the commented out code part [1] which would implement a pseudo-inverse in this case. As far as I understand the reason why this code is not used at the moment is the general problem that there is a gauge-freedom of the singular subspace (which has n>1 if there are two equal singular values) is not solved and therefore the derivative of only one of the two unitaries alone would make no sense in this case.

As far as I understand the problem, it should be ok to do the pseudo-inverse of the differences, if users are not interested in the derivative of a single unitary but of the complete product U * S @ Vh or variants of this one. To illustrate that enabling the code path [1] leads to correct results in some cases, one can replace the code line and rerun the example above. The resulting gradient is the same one as if one used the VJP-rule of the jnp.linalg.norm function. [3]

What is your opinion on this problem? Can the gauge-freedom somehow be fixed so it is no problem? Is the use case of using only the derivative of one unitary very unlikely so one could stick to the pseudo-inverse?

[1] s_diffs_zeros = jnp.ones((), dtype=s.dtype) * (s_diffs == 0.) (fixing the bug that it has to be s.dtype instead of A.dtype)

[2]

import jax
import jax.numpy as jnp

A = jnp.diag(jnp.asarray([1, 1, 2, 3], dtype=float))

norm_g = jax.value_and_grad(jnp.linalg.norm)

print(f_g(A))

Result:

(DeviceArray(3.8729832, dtype=float32),
DeviceArray([[0.2581989, 0. , 0. , 0. ],
[0. , 0.2581989, 0. , 0. ],
[0. , 0. , 0.5163978, 0. ],
[0. , 0. , 0. , 0.7745967]], dtype=float32))

[3] Result for example code with https://github.com/google/jax/blob/main/jax/_src/lax/linalg.py#L1291 replaced by [1]:

(DeviceArray(3.8729835, dtype=float32),
DeviceArray([[0.25819886, 0. , 0. , 0. ],
[0. , 0.25819886, 0. , 0. ],
[0. , 0. , 0.5163977 , 0. ],
[0. , 0. , 0. , 0.7745966 ]], dtype=float32))

Issue 2:
If an SVD has degenerate singular values (multiple entries of S are exactly equal), the gradient pass will give nan.

This is because the AD formula contains something like

1 / (s[i] ** 2 - s[j] ** 2) which becomes nan if s[i] == s[j]. There are also terms with 1 / S which become 'nan' if any of 'S' is zero.

In particular, this situation arises when multiple singular values are 0.0. and it evaluates the entire gradient of a function to nan, even if the function does not depend on the rows of U and columns of V that correspond to the 0.0 singular values (See example below). Essentially in this case, even if the adjoint gradients of the corresponding slices of U and V are 0., they are multiplied by nan and thus give nan.

I don't have a clear vision, what the best case scenario is here.

I just wanted to share that this problem arises for me and how i will dirty-fix it.

Maybe fix "1." would be of interest to implement in jax? thoughts?

Example

I contruct a matrix that has only a few non-zero singular values (actually, only the first in non-zero, but due two numerical error, the second is non-zero too)

mat_degenerate = np.ones([5, 5])
U, S, V = np.linalg.svd(mat_degenerate)
print(S)
>>> [4.9999995e+00 1.9792830e-07 0.0000000e+00 0.0000000e+00 0.0000000e+00]

I design a function that only depends on the first two singular values (which are clearly neither degenerate nor zero), and the corresponding rows (columns) of U (V)

def truncated_svd(x, num):
U, S, V = np.linalg.svd(x, full_matrices=False)
return U[:, :num], S[:num], V[:num, :]

def foo(x):
U, S, V = truncated_svd(x, 2)
return np.real(np.trace(np.dot(V, U))) + np.sum(S)

And see that the entire gradient becomes nan

print(value_and_grad(foo)(mat_degenerate)
>>> (DeviceArray(6.9999995, dtype=float32), DeviceArray([[nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan],
[nan, nan, nan, nan, nan]], dtype=float32))

Dirty workaround

I have seen people use safe_inverse = lambda x, eps: x / (x**2 + eps) and replace 1 / (s[i] ** 2 - s[j] ** 2) by safe_inverse(s[i] ** 2 - s[j] ** 2) as well as 1/S by safe_inverse(S) , usually with eps == 1e-12.

So i implemented that. Contact me if anyone is interested