<p dir="auto">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:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>

<span class="pl-k">def</span> <span class="pl-en">f</span>(<span class="pl-s1">arr</span>):
    <span class="pl-v">U</span>, <span class="pl-v">S</span>, <span class="pl-v">Vh</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">svd</span>(<span class="pl-s1">arr</span>, <span class="pl-s1">full_matrices</span><span class="pl-c1">=</span><span class="pl-c1">False</span>)
    <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">norm</span>((<span class="pl-v">U</span> <span class="pl-c1">*</span> <span class="pl-v">S</span>) @ <span class="pl-v">Vh</span>)

<span class="pl-v">A</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">diag</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">asarray</span>([<span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float</span>))

<span class="pl-s1">f_g</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">f</span>)

<span class="pl-en">print</span>(<span class="pl-en">f_g</span>(<span class="pl-v">A</span>))</pre></div>
<p dir="auto">Result:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="(DeviceArray(3.8729835, dtype=float32),
 DeviceArray([[nan, nan, nan, nan],
              [nan, nan, nan, nan],
              [nan, nan, nan, nan],
              [nan, nan, nan, nan]], dtype=float32))"><pre class="notranslate">(<span class="pl-v">DeviceArray</span>(<span class="pl-c1">3.8729835</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>),
 <span class="pl-v">DeviceArray</span>([[<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>],
              [<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>],
              [<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>],
              [<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>))</pre></div>
<p dir="auto">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 <code class="notranslate">jax.value_and_grad(jnp.linalg.norm)</code> [2]).</p>
<p dir="auto">After some discussion with <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/frederikwilde/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/frederikwilde">@frederikwilde</a> (the author of the commit <a class="commit-link" data-hovercard-type="commit" data-hovercard-url="https://github.com/google/jax/commit/f443e19141c3d04a6e2947a0dcde8d23039af786/hovercard" href="https://github.com/google/jax/commit/f443e19141c3d04a6e2947a0dcde8d23039af786"><tt>f443e19</tt></a>) we found that the problem is in the calculation of the derivative of the singular matrices <em>U</em> and <em>V<sup>†</sup></em>. The derivation of the derivative calculates in one step the matrix of the inverses <em>1/(S<sub>i</sub> - S<sub>j</sub>)</em> (with <em>S<sub>i</sub></em> the singular values). Obviously, this leads to the described problem.</p>
<p dir="auto">In the source code <a href="https://github.com/google/jax/blob/main/jax/_src/lax/linalg.py#L1291">https://github.com/google/jax/blob/main/jax/_src/lax/linalg.py#L1291</a> 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&gt;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.</p>
<p dir="auto">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 <em>U * S @ Vh</em> 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 <code class="notranslate">jnp.linalg.norm</code> function. [3]</p>
<p dir="auto">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?</p>
<p dir="auto">[1] <code class="notranslate">s_diffs_zeros = jnp.ones((), dtype=s.dtype) * (s_diffs == 0.)</code> (fixing the bug that it has to be <code class="notranslate">s.dtype</code> instead of <code class="notranslate">A.dtype</code>)</p>
<p dir="auto">[2]</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>

<span class="pl-v">A</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">diag</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">asarray</span>([<span class="pl-c1">1</span>, <span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float</span>))

<span class="pl-s1">norm_g</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">linalg</span>.<span class="pl-s1">norm</span>)

<span class="pl-en">print</span>(<span class="pl-en">f_g</span>(<span class="pl-v">A</span>))</pre></div>
<p dir="auto">Result:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="(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))"><pre class="notranslate">(<span class="pl-v">DeviceArray</span>(<span class="pl-c1">3.8729832</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>),
 <span class="pl-v">DeviceArray</span>([[<span class="pl-c1">0.2581989</span>, <span class="pl-c1">0.</span>       , <span class="pl-c1">0.</span>       , <span class="pl-c1">0.</span>       ],
              [<span class="pl-c1">0.</span>       , <span class="pl-c1">0.2581989</span>, <span class="pl-c1">0.</span>       , <span class="pl-c1">0.</span>       ],
              [<span class="pl-c1">0.</span>       , <span class="pl-c1">0.</span>       , <span class="pl-c1">0.5163978</span>, <span class="pl-c1">0.</span>       ],
              [<span class="pl-c1">0.</span>       , <span class="pl-c1">0.</span>       , <span class="pl-c1">0.</span>       , <span class="pl-c1">0.7745967</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>))</pre></div>
<p dir="auto">[3] Result for example code with <a href="https://github.com/google/jax/blob/main/jax/_src/lax/linalg.py#L1291">https://github.com/google/jax/blob/main/jax/_src/lax/linalg.py#L1291</a> replaced by [1]:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="(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))"><pre class="notranslate">(<span class="pl-v">DeviceArray</span>(<span class="pl-c1">3.8729835</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>),
 <span class="pl-v">DeviceArray</span>([[<span class="pl-c1">0.25819886</span>, <span class="pl-c1">0.</span>        , <span class="pl-c1">0.</span>        , <span class="pl-c1">0.</span>        ],
              [<span class="pl-c1">0.</span>        , <span class="pl-c1">0.25819886</span>, <span class="pl-c1">0.</span>        , <span class="pl-c1">0.</span>        ],
              [<span class="pl-c1">0.</span>        , <span class="pl-c1">0.</span>        , <span class="pl-c1">0.5163977</span> , <span class="pl-c1">0.</span>        ],
              [<span class="pl-c1">0.</span>        , <span class="pl-c1">0.</span>        , <span class="pl-c1">0.</span>        , <span class="pl-c1">0.7745966</span> ]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>))</pre></div>
<p dir="auto">If an SVD has degenerate singular values (multiple entries of S are exactly equal), the gradient pass will give <code class="notranslate">nan</code>.</p>
<p dir="auto">This is because the <a href="https://arxiv.org/pdf/1909.02659.pdf" rel="nofollow">AD formula</a> contains something like<br>
<code class="notranslate">1 / (s[i] ** 2 - s[j] ** 2) </code> which becomes <code class="notranslate">nan</code> if <code class="notranslate">s[i] == s[j]</code>. There are also terms with <code class="notranslate">1 / S</code> which become 'nan' if any of 'S' is zero.</p>
<p dir="auto">In particular, this situation arises when multiple singular values are <code class="notranslate">0.0</code>. and it evaluates the entire gradient of a function to <code class="notranslate">nan</code>, even if the function does not depend on the rows of U and columns of V that correspond to the <code class="notranslate">0.0</code> singular values (See example below). Essentially in this case, even if the adjoint gradients of the corresponding slices of U and V are <code class="notranslate">0.</code>, they are multiplied by <code class="notranslate">nan</code> and thus give <code class="notranslate">nan</code>.</p>
<p dir="auto">I don't have a clear vision, what the best case scenario is here.<br>
I just wanted to share that this problem arises for me and how i will dirty-fix it.<br>
Maybe fix "1." would be of interest to implement in jax? thoughts?</p>
<h3 dir="auto">Example</h3>
<p dir="auto">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)</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="mat_degenerate = np.ones([5, 5])
U, S, V = np.linalg.svd(mat_degenerate)
print(S)
&gt;&gt;&gt; [4.9999995e+00 1.9792830e-07 0.0000000e+00 0.0000000e+00 0.0000000e+00]"><pre class="notranslate"><span class="pl-s1">mat_degenerate</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">ones</span>([<span class="pl-c1">5</span>, <span class="pl-c1">5</span>])
<span class="pl-v">U</span>, <span class="pl-v">S</span>, <span class="pl-v">V</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">svd</span>(<span class="pl-s1">mat_degenerate</span>)
<span class="pl-en">print</span>(<span class="pl-v">S</span>)
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> [<span class="pl-c1">4.9999995e+00</span> <span class="pl-c1">1.9792830e-07</span> <span class="pl-c1">0.0000000e+00</span> <span class="pl-c1">0.0000000e+00</span> <span class="pl-c1">0.0000000e+00</span>]</pre></div>
<p dir="auto">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)</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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)"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">truncated_svd</span>(<span class="pl-s1">x</span>, <span class="pl-s1">num</span>):
    <span class="pl-v">U</span>, <span class="pl-v">S</span>, <span class="pl-v">V</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-s1">linalg</span>.<span class="pl-en">svd</span>(<span class="pl-s1">x</span>, <span class="pl-s1">full_matrices</span><span class="pl-c1">=</span><span class="pl-c1">False</span>)
    <span class="pl-k">return</span> <span class="pl-v">U</span>[:, :<span class="pl-s1">num</span>], <span class="pl-v">S</span>[:<span class="pl-s1">num</span>], <span class="pl-v">V</span>[:<span class="pl-s1">num</span>, :]

<span class="pl-k">def</span> <span class="pl-en">foo</span>(<span class="pl-s1">x</span>):
    <span class="pl-v">U</span>, <span class="pl-v">S</span>, <span class="pl-v">V</span> <span class="pl-c1">=</span> <span class="pl-en">truncated_svd</span>(<span class="pl-s1">x</span>, <span class="pl-c1">2</span>)
    <span class="pl-k">return</span> <span class="pl-s1">np</span>.<span class="pl-en">real</span>(<span class="pl-s1">np</span>.<span class="pl-en">trace</span>(<span class="pl-s1">np</span>.<span class="pl-en">dot</span>(<span class="pl-v">V</span>, <span class="pl-v">U</span>))) <span class="pl-c1">+</span> <span class="pl-s1">np</span>.<span class="pl-en">sum</span>(<span class="pl-v">S</span>)</pre></div>
<p dir="auto">And see that the <em>entire</em> gradient becomes <code class="notranslate">nan</code></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="print(value_and_grad(foo)(mat_degenerate)
&gt;&gt;&gt; (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))"><pre class="notranslate"><span class="pl-en">print</span>(<span class="pl-en">value_and_grad</span>(<span class="pl-s1">foo</span>)(<span class="pl-s1">mat_degenerate</span>)
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> (<span class="pl-v">DeviceArray</span>(<span class="pl-c1">6.9999995</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>), <span class="pl-v">DeviceArray</span>([[<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>],
              [<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>],
              [<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>],
              [<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>],
              [<span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>, <span class="pl-s1">nan</span>]], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>))</pre></div>
<h3 dir="auto">Dirty workaround</h3>
<p dir="auto">I have seen people use <code class="notranslate">safe_inverse = lambda x, eps: x / (x**2 + eps)</code> and replace <code class="notranslate">1 / (s[i] ** 2 - s[j] ** 2) </code> by <code class="notranslate">safe_inverse(s[i] ** 2 - s[j] ** 2)</code> as well as <code class="notranslate">1/S</code> by <code class="notranslate">safe_inverse(S)</code> , usually with <code class="notranslate">eps == 1e-12</code>.<br>
So i implemented that. Contact me if anyone is interested</p>