Backward Pass nan for Repeated Eigenvalues

joec: <p>Hi, I am seeing an issue on the backward pass when using torch.linalg.eigh on a hermitian matrix with repeated eigenvalues. I was wondering if there is any way to obtain the eigenvector associated with the minimum eigenvalue without the gradients in the backward pass going to nan. I am performing this calculation as a part of the loss function and there are no learnable parameters after this point.</p>
<p>I have set up the following testcase that illustrates the issue I am having, with the gradients on the torch.linalg.eigh calculation causing issues. I am using pytorch version 1.9</p>
<pre><code class="lang-auto">import torch

matrix = torch.tensor([[
        [
            1.2039, 0, 0.3690, 0
        ],
        [
            0, -1.2039, 0, 0.3690
        ],
        [
            0.3690, 0, -1.2039, 0
        ],
        [
            0, 0.3690, 0, 1.2039
        ],
    ]], requires_grad=True)

e, v = torch.linalg.eigh(matrix)
min_index = torch.argmin(e)
min_eigenvector = v[:, min_index]
min_state = min_eigenvector
min_sum = torch.sum(min_state) # Sum to obtain a scalar for .backward()
min_sum.backward()
assert not matrix.grad.isnan().any()
</code></pre>

KFrank: <p>Hi Joe!</p>
<aside class="quote no-group quote-modified" data-username="joec" data-post="1" data-topic="128049" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/8e8cbc/48.png" class="avatar"> joec:</div>
<blockquote>
<p>using torch.linalg.eigh on a hermitian matrix with repeated eigenvalues.<br>
…<br>
I was wondering if there is any way to obtain the eigenvector associated with the minimum eigenvalue</p>
</blockquote>
</aside>
<p>The short answer:  No, there is not.</p>
<p>The issue is that in the presence of repeated (<em>degenerate</em>)<br>
eigenvalues, “the eigenvector associated with” a degenerate<br>
eigenvalue is not well defined.</p>
<p>An eigendecomposition algorithm (such as <code>torch.linalg.eigh()</code>)<br>
will give you a list of eigenvalues and associated eigenvectors.  But<br>
the two eigenvectors associated with two equal eigenvalues aren’t<br>
well defined (only the two-dimensional subspace spanned by those<br>
two eigenvectors is) – you can mix the two eigenvectors together<br>
to get a different, but equally valid, pair of eigenvectors.</p>
<p>As a result, backwarding through such an eigenvector will give you<br>
ill-defined results that will depend on the particulars of the algorithm<br>
used and that can be <code>nan</code>.  (Backwarding through the <em>eigenvalues,</em><br>
however, is perfectly fine.)</p>
<p>Quoting from the documentation for <a href="https://pytorch.org/docs/1.9.0/generated/torch.linalg.eigh.html" rel="noopener nofollow ugc">torch.linalg.eigh()</a>:</p>
<blockquote>
<p>Gradients computed using the eigenvectors tensor will only be finite when <code>A</code> has unique eigenvalues. Furthermore, if the distance between any two eigvalues is close to zero, the gradient will be numerically unstable, as it depends on the eigenvalues λ_i​ through the computation of 1 / min _i≠j (λ_i − λ_j).</p>
</blockquote>
<p>(And, furthermore, because the eigenvectors are not well defined,<br>
your example loss is also not well defined.)</p>
<p>The following script illustrates the arbitrariness of the choice of<br>
eigenvectors within the subspace defined by a set of degenerate<br>
eigenvalues (and hence of your loss function):</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

matrix = torch.tensor([   # no leading singleton dimension
        [
            1.2039, 0, 0.3690, 0
        ],
        [
            0, -1.2039, 0, 0.3690
        ],
        [
            0.3690, 0, -1.2039, 0
        ],
        [
            0, 0.3690, 0, 1.2039
        ],
    ])

e, v = torch.linalg.eigh (matrix)

print ('e =', e)

vChk = torch.allclose (v @ v.T, torch.eye (4), atol = 1.e-6)
print ('check that v is orthonormal:', vChk)

eChk = torch.allclose (v @ torch.diag (e) @ v.T, matrix, atol = 1.e-6)
print ('check that e, v is an eigendecomposition:', eChk)

sqh = 0.5**0.5

rot = torch.tensor ([
    [ sqh, -sqh,  0.0,  0.0 ],
    [ sqh,  sqh,  0.0,  0.0 ],
    [ 0.0,  0.0,  1.0,  0.0 ],
    [ 0.0,  0.0,  0.0,  1.0 ]
])

vr = v @ rot

vrChk = torch.allclose (vr @ vr.T, torch.eye (4), atol = 1.e-6)
print ('check that vr is orthonormal:', vrChk)

erChk = torch.allclose (vr @ torch.diag (e) @ vr.T, matrix, atol = 1.e-6)
print ('check that e, vr is an eigendecomposition:', erChk)

loss_v = v[:, torch.argmin (e)].sum()
loss_vr = vr[:, torch.argmin (e)].sum()

print ('loss_v  =', loss_v)
print ('loss_vr =', loss_vr)
</code></pre>
<p>This script uses your <code>matrix</code> as its example, but with the leading<br>
singleton dimension removed for simplicity.</p>
<p>Here is its output:</p>
<pre><code class="lang-plaintext">1.9.0
e = tensor([-1.2592, -1.2592,  1.2592,  1.2592])
check that v is orthonormal: True
check that e, v is an eigendecomposition: True
check that vr is orthonormal: True
check that e, vr is an eigendecomposition: True
loss_v  = tensor(0.8408)
loss_vr = tensor(1.1891)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

joec: <p>Thank you for the help and the clear jsutification. Not the answer I was hoping for but very useful nonetheless!</p>