Pytorch diagonalization of Tri diagonal complex symmetric matrix

username: Fabio_Anselmi
cooked: <p>Hi all!</p>
<p>I am looking for a backprop-friendly algorithm in pytorch to find the eigenvectors and eigenvalues of a tridiagonal symmetric matrix (only the main and upper and lower diagonal different from zero).<br>
The variables on which I am running the grad are the matrix elements and the loss is a backprop-friendly function of the eigenvalues and eigenvectors.</p>
<p>Any idea?<br>
torch.eig and variations do not work since the grad flow does not pass through.</p>
<p>Thank you!</p>
<p>Fabio</p>

username: AlphaBetaGamma96
cooked: <p>Hi <a class="mention" href="/u/fabio_anselmi">@Fabio_Anselmi</a>,</p>
<p>Do you have a minimal reproducible example?</p>
<p>Also, the eigenvalues won’t have a derivative by definition so you can’t backprop through the eigenvalues, you should be able to backprop through the eigenvectors (I think).</p>

username: Fabio_Anselmi
cooked: <p>Yes, here is my tentative to solve the problem with QR algortithm… but the results are not correct, I tried to test it with not backprop friendly algos and it gives different results  <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<p>import numpy as np<br>
import torch<br>
import torch.nn as nn<br>
import torch.nn.functional as F</p>
<p>import matplotlib.pyplot as plt<br>
import os<br>
os.environ[‘KMP_DUPLICATE_LIB_OK’] = ‘True’</p>
<p>SEED = 467<br>
torch.manual_seed(SEED)</p>
<p><span class="hashtag">#Constants</span> of the system<br>
n=4<br>
Gamma = torch.rand(n)<br>
Omega = torch.randn(n)# could be negative<br>
g = torch.rand(n-1)</p>
<p>def tridiag_eig(g,Gamma,Omega):<br>
n = Gamma.shape[0]<br>
diag = -torch.relu(Gamma)/2 +1j<em>Omega <span class="hashtag">#impose</span> positiveness constraints<br>
upper_diag = -1j</em>torch.relu(g)<br>
tridiag_mat = torch.diag(diag) + torch.diag(upper_diag,1) + torch.diag(upper_diag,-1)</p>
<pre><code># Initialize variables
eig_vals = tridiag_mat.new_empty(n)
left_eig_vecs = torch.eye(n, dtype=tridiag_mat.dtype, device=tridiag_mat.device)
right_eig_vecs = torch.eye(n, dtype=tridiag_mat.dtype, device=tridiag_mat.device)

# QR Algorithm
for i in range(100):
    # QR decomposition
    Q, R = torch.linalg.qr(tridiag_mat) 
    
    # Compute the new matrix
    tridiag_mat = torch.mm(R, Q)

    # Update eigenvectors
    left_eig_vecs = torch.mm(left_eig_vecs, Q)
    right_eig_vecs = torch.mm(right_eig_vecs, torch.inverse(Q))

    # Check convergence
    if torch.max(torch.abs(torch.triu(tridiag_mat, diagonal=1))) &lt; 1e-6:
        break

# Extract eigenvalues
eig_vals = torch.diag(tridiag_mat)

return eig_vals, left_eig_vecs, right_eig_vecs
</code></pre>
<p>eig_vals, left_eig_vecs, right_eig_vecs = tridiag_eig(g,Gamma,Omega)</p>

username: KFrank
cooked: <p>Hi Fabio!</p>
<aside class="quote no-group" data-username="Fabio_Anselmi" data-post="1" data-topic="180769" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fabio_anselmi/48/60413_2.png" class="avatar"> Fabio_Anselmi:</div>
<blockquote>
<p>torch.eig and variations do not work since the grad flow does not pass through.</p>
</blockquote>
</aside>
<p><code>torch.linalg.eig()</code> works for me (as does <code>torch.linalg.eigh()</code>):</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; n = 5
&gt;&gt;&gt;
&gt;&gt;&gt; d_main = torch.randn (n)
&gt;&gt;&gt; d_lower = torch.randn (n - 1)
&gt;&gt;&gt; m_lower = torch.diag (d_main) + torch.diag (d_lower, -1)
&gt;&gt;&gt; m_tri = m_lower + torch.diag (d_lower, 1)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.linalg.eig (m_tri)                   # general algorithm
torch.return_types.linalg_eig(
eigenvalues=tensor([ 1.4005+0.j,  0.5639+0.j, -1.2217+0.j, -1.1171+0.j, -0.7208+0.j]),
eigenvectors=tensor([[ 0.0027+0.j, -0.0882+0.j, -0.9948+0.j,  0.0512+0.j,  0.0054+0.j],
        [ 0.0450+0.j, -0.9944+0.j,  0.0899+0.j,  0.0295+0.j,  0.0168+0.j],
        [ 0.5250+0.j, -0.0086+0.j, -0.0403+0.j, -0.7924+0.j, -0.3077+0.j],
        [ 0.8454+0.j,  0.0575+0.j,  0.0246+0.j,  0.5224+0.j,  0.0922+0.j],
        [ 0.0875+0.j,  0.0098+0.j, -0.0114+0.j, -0.3092+0.j,  0.9468+0.j]]))
&gt;&gt;&gt;
&gt;&gt;&gt; torch.linalg.eigh (m_lower)                # takes advantage of symmetry of matrix
torch.return_types.linalg_eigh(
eigenvalues=tensor([-1.2217, -1.1171, -0.7208,  0.5639,  1.4005]),
eigenvectors=tensor([[ 0.9948,  0.0512,  0.0054,  0.0882, -0.0027],
        [-0.0899,  0.0295,  0.0168,  0.9944, -0.0450],
        [ 0.0403, -0.7924, -0.3077,  0.0086, -0.5250],
        [-0.0246,  0.5224,  0.0922, -0.0575, -0.8454],
        [ 0.0114, -0.3092,  0.9468, -0.0098, -0.0875]]))
&gt;&gt;&gt;
&gt;&gt;&gt; m_tri.requires_grad = True
&gt;&gt;&gt; loss = torch.linalg.eig (m_tri)[0][2]      # loss is a single eigenvalue
&gt;&gt;&gt; loss.backward()
&gt;&gt;&gt; m_tri.grad
tensor([[ 9.8957e-01, -8.9380e-02,  4.0040e-02, -2.4454e-02,  1.1315e-02],
        [-8.9380e-02,  8.0730e-03, -3.6165e-03,  2.2087e-03, -1.0220e-03],
        [ 4.0041e-02, -3.6166e-03,  1.6201e-03, -9.8947e-04,  4.5784e-04],
        [-2.4454e-02,  2.2088e-03, -9.8947e-04,  6.0430e-04, -2.7962e-04],
        [ 1.1315e-02, -1.0220e-03,  4.5784e-04, -2.7962e-04,  1.2938e-04]])
&gt;&gt;&gt;
&gt;&gt;&gt; m_tri.grad = None
&gt;&gt;&gt; loss = torch.linalg.eig (m_tri)[1][1, 2]   # loss is a single element of a single eigenvector
&gt;&gt;&gt; loss.backward()
&gt;&gt;&gt; m_tri.grad
tensor([[ 0.0634, -0.0057,  0.0026, -0.0016,  0.0007],
        [ 0.5604, -0.0506,  0.0227, -0.0138,  0.0064],
        [-0.2186,  0.0197, -0.0088,  0.0054, -0.0025],
        [ 0.1320, -0.0119,  0.0053, -0.0033,  0.0015],
        [-0.0590,  0.0053, -0.0024,  0.0015, -0.0007]])
</code></pre>
<p>[Edit:  I overlooked that you were asking about the complex case.]</p>
<p>The situation is similar for the complex case, but even in the (hermitian)<br>
tri-diagonal case, you can now have complex eigenvectors, so the<br>
eigenvector phase ambiguity makes the gradient of (phase-dependent)<br>
functions of the eigenvectors not well defined.  (Gradients of eigenvalues<br>
are still fine.)</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; n = 5
&gt;&gt;&gt;
&gt;&gt;&gt; d_main = torch.randn (n, dtype = torch.complex64)
&gt;&gt;&gt; d_lower = torch.randn (n - 1, dtype = torch.complex64)
&gt;&gt;&gt; m_lower = torch.diag (d_main) + torch.diag (d_lower, -1)
&gt;&gt;&gt; m_tri = m_lower + torch.diag (d_lower, 1).conj()
&gt;&gt;&gt;
&gt;&gt;&gt; torch.linalg.eig (m_tri)                   # general algorithm
torch.return_types.linalg_eig(
eigenvalues=tensor([-2.4665+0.3781j, -1.5699+0.4308j,  1.8091+0.4371j,  0.7244+0.4485j,
         0.0592+0.0190j]),
eigenvectors=tensor([[-4.1827e-01+0.1134j,  6.5241e-01+0.0000j, -1.4480e-01+0.1424j,
         -2.9093e-01+0.4490j,  9.6288e-02+0.1866j],
        [-4.0111e-01-0.3376j,  2.1858e-01+0.2747j,  3.9774e-01+0.0802j,
          6.3391e-01+0.0000j,  1.5024e-02-0.1544j],
        [ 6.0069e-01+0.0000j,  4.2588e-01+0.0718j,  4.8217e-01-0.3556j,
         -5.6944e-02+0.0620j,  6.8369e-04+0.2855j],
        [-3.9204e-01-0.1196j, -4.4602e-01-0.1976j,  6.2306e-01+0.0000j,
         -4.1192e-01+0.1344j, -7.1238e-03+0.1854j],
        [ 2.1511e-02+0.0899j,  4.8561e-04+0.1631j,  1.8057e-01+0.1198j,
         -3.3740e-01+0.0566j,  9.0328e-01+0.0000j]]))
&gt;&gt;&gt;
&gt;&gt;&gt; torch.linalg.eigh (m_lower)                # takes advantage of hermeticity of matrix
torch.return_types.linalg_eigh(
eigenvalues=tensor([-2.4809, -1.5880,  0.0173,  0.7716,  1.8363]),
eigenvectors=tensor([[-0.4259+0.0000j, -0.6590-0.0000j, -0.2628+0.0000j,  0.5261+0.0000j,
         -0.1965+0.0000j],
        [-0.2990-0.4251j, -0.2088-0.2968j,  0.0988+0.1405j, -0.3690-0.5247j,
          0.2281+0.3243j],
        [ 0.5983+0.0911j, -0.4138-0.0630j, -0.3100-0.0472j,  0.0325+0.0049j,
          0.5924+0.0902j],
        [-0.3422-0.2370j,  0.3956+0.2740j, -0.1670-0.1156j,  0.3280+0.2272j,
          0.5165+0.3577j],
        [ 0.0084+0.0943j, -0.0147-0.1647j,  0.0777+0.8694j,  0.0347+0.3877j,
          0.0200+0.2234j]]))
&gt;&gt;&gt;
&gt;&gt;&gt; m_tri.requires_grad = True
&gt;&gt;&gt; loss = torch.linalg.eig (m_tri)[0][2]      # loss is a single eigenvalue
&gt;&gt;&gt; loss.backward()                            # well defined for eigenvalues
&gt;&gt;&gt; m_tri.grad
tensor([[ 0.0416+0.0051j, -0.0550+0.0633j, -0.1237+0.0026j, -0.1020+0.0785j,
         -0.0145+0.0424j],
        [-0.0409-0.0731j,  0.1669+0.0141j,  0.1500+0.1966j,  0.2469+0.0719j,
          0.0854-0.0266j],
        [-0.1189-0.0344j,  0.1900-0.1583j,  0.3616+0.0509j,  0.3340-0.1806j,
          0.0621-0.1166j],
        [-0.1094-0.0679j,  0.2347-0.1050j,  0.3477+0.1526j,  0.3852-0.0868j,
          0.0949-0.0992j],
        [ 0.0067-0.0443j,  0.0657+0.0606j, -0.0012+0.1321j,  0.0805+0.1113j,
          0.0447+0.0168j]])
&gt;&gt;&gt;
&gt;&gt;&gt; m_tri.grad = None
&gt;&gt;&gt; loss = torch.linalg.eig (m_tri)[1][1, 2]   # loss is a single element of a single eigenvector
&gt;&gt;&gt; loss.backward()                            # fails because of phase ambiguity in eigenvectors
Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  File "&lt;path_to_pytorch_install&gt;\torch\_tensor.py", line 487, in backward
    torch.autograd.backward(
  File "&lt;path_to_pytorch_install&gt;\torch\autograd\__init__.py", line 200, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: linalg_eig_backward: The eigenvectors in the complex case are specified up to multiplication by e^{i phi}. The specified loss function depends on this quantity, so it is ill-defined.
&gt;&gt;&gt; m_tri.grad
&gt;&gt;&gt; 
</code></pre>
<p>Note that <code>torch.eig()</code> was deprecated and removed several versions ago.</p>
<p>From recollection, <code>torch.eig()</code> was not as smart about recognizing that<br>
the eigenvectors were real so that the backward pass would be well defined.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Fabio_Anselmi
cooked: <p>Hi Frank, thank you for your answer!  Can I fix somehow the phase to solve the ambiguity and reintroduce the grad?</p>

username: KFrank
cooked: <p>Hi Fabio!</p>
<aside class="quote no-group" data-username="Fabio_Anselmi" data-post="5" data-topic="180769" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fabio_anselmi/48/60413_2.png" class="avatar"> Fabio_Anselmi:</div>
<blockquote>
<p>Can I fix somehow the phase to solve the ambiguity and reintroduce the grad?</p>
</blockquote>
</aside>
<p>Without knowing the details of your use case, I can’t really speculate what<br>
an appropriate fix might be.</p>
<p>In general, the right approach is to call <code>grad()</code> only on quantities that are<br>
independent of the (ambiguous) phase of the eigenvectors.  Quoting from<br>
the <a href="https://pytorch.org/docs/2.0/generated/torch.linalg.eigh.html" rel="noopener nofollow ugc">documentation for <code>eigh()</code></a>:</p>
<blockquote>
<p>For this reason, the loss function shall not depend on the phase of the eigenvectors, as this quantity is not well-defined.</p>
</blockquote>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.0.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; n = 5
&gt;&gt;&gt;
&gt;&gt;&gt; d_main = torch.randn (n, dtype = torch.complex64)
&gt;&gt;&gt; d_lower = torch.randn (n - 1, dtype = torch.complex64)
&gt;&gt;&gt; m_lower = torch.diag (d_main) + torch.diag (d_lower, -1)
&gt;&gt;&gt; m_lower.requires_grad = True
&gt;&gt;&gt; eigs = torch.linalg.eigh (m_lower)
&gt;&gt;&gt; eigs
torch.return_types.linalg_eigh(
eigenvalues=tensor([-2.4809, -1.5880,  0.0173,  0.7716,  1.8363],
       grad_fn=&lt;LinalgEighBackward0&gt;),
eigenvectors=tensor([[-0.4259+0.0000j, -0.6590-0.0000j, -0.2628+0.0000j,  0.5261+0.0000j,
         -0.1965+0.0000j],
        [-0.2990-0.4251j, -0.2088-0.2968j,  0.0988+0.1405j, -0.3690-0.5247j,
          0.2281+0.3243j],
        [ 0.5983+0.0911j, -0.4138-0.0630j, -0.3100-0.0472j,  0.0325+0.0049j,
          0.5924+0.0902j],
        [-0.3422-0.2370j,  0.3956+0.2740j, -0.1670-0.1156j,  0.3280+0.2272j,
          0.5165+0.3577j],
        [ 0.0084+0.0943j, -0.0147-0.1647j,  0.0777+0.8694j,  0.0347+0.3877j,
          0.0200+0.2234j]], grad_fn=&lt;LinalgEighBackward0&gt;))
&gt;&gt;&gt; loss = eigs[1][1, 2].abs()                 # loss is a phase-independent function of the eigenvectors
&gt;&gt;&gt; loss.backward()                            # backward() is well defined
&gt;&gt;&gt; m_lower.grad
tensor([[-0.1912+5.1455e-10j,  0.0693-9.8580e-02j, -0.0954+1.4522e-02j,
         -0.0802+5.5560e-02j,  0.0248-2.7710e-01j],
        [ 0.0693+9.8580e-02j, -0.0758+0.0000e+00j,  0.0400+4.1699e-02j,
          0.0566+2.0807e-02j, -0.1446+8.3531e-02j],
        [-0.0954-1.4522e-02j,  0.0400-4.1699e-02j,  0.0420+4.6566e-10j,
         -0.0132+6.4294e-03j, -0.0251+1.0246e-01j],
        [-0.0802-5.5560e-02j,  0.0566-2.0807e-02j, -0.0132-6.4294e-03j,
         -0.0366+3.7253e-09j,  0.0309-3.7008e-02j],
        [ 0.0248+2.7710e-01j, -0.1446-8.3531e-02j, -0.0251-1.0246e-01j,
          0.0309+3.7008e-02j,  0.2617+7.4506e-09j]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: Fabio_Anselmi
cooked: <p>Thank you, I see.<br>
I suspect I do not have a way around the problem.</p>
<p>My loss, as below, is a function of the right (U) left (V) eigenvectors and the eigenvalues and a fixed TCorr function, that I approximate :</p>
<p>class LossCorr(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(LossCorr, self).<strong>init</strong>()<br>
self.Gamma = nn.Parameter(torch.rand((n), dtype=torch.float32), requires_grad=True)<br>
self.Omega = nn.Parameter(torch.rand((n), dtype=torch.float32), requires_grad=True)<span class="hashtag">#could</span> be negative<br>
self.g = nn.Parameter(torch.rand((n-1), dtype=torch.float32), requires_grad=True)<br>
self.C = nn.Parameter(torch.rand((n,n), dtype=torch.complex64), requires_grad=True)# C can be complex</p>
<pre><code>def forward(self):
    W=torch.zeros((n,n,n))
    E,V,U = Diagonalization_algo(self.g,self.Gamma,self.Omega) #Calculate the eigenvectors and eigenvalues
    
    for i in range(n):
        UV = torch.mm(U[:,i].unsqueeze(dim=0).T,V[:,i].unsqueeze(dim=0))
        W[i,:,:] = (self.C.T@UV)@self.C.conj()
    
    Cor=torch.zeros((n,n),dtype=torch.complex64)
    
    loss=0
    for t in range(Tmax):
        for i in range(n):
             Cor += W[i,:,:]*torch.exp(E[i]*t)
          
        loss += (1/Tmax)*torch.norm(Cor-TCorr(t),2)        
     
    return loss,Gamma,Omega,g
</code></pre>

username: KFrank
cooked: <p>Hi Fabio!</p>
<aside class="quote no-group quote-modified" data-username="Fabio_Anselmi" data-post="7" data-topic="180769" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/fabio_anselmi/48/60413_2.png" class="avatar"> Fabio_Anselmi:</div>
<blockquote>
<p>I suspect I do not have a way around the problem.<br>
…<br>
self.C = nn.Parameter(torch.rand((n,n), dtype=torch.complex64), requires_grad=True)# C can be complex</p>
</blockquote>
</aside>
<p>I’m not sure that I see your problem.  Based on what you have posted, it<br>
looks like <code>C</code> is the only complex entity and it does not appear to participate<br>
in <code>Diagonalization_algo()</code>.  (As an aside, why not use something provided<br>
by pytorch such as <code>torch.linalg.eig()</code> instead of writing your own?)</p>
<p>The code you posted before suggests that you are diagonalizing a real,<br>
symmetric, tri-diagonal matrix.  You should therefore have real eigenvectors<br>
(as well as eigenvalues), so you shouldn’t have any phase ambiguity, and<br>
backpropagation should work fine.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Fabio_Anselmi
cooked: <p>Hi Frank!</p>
<p>Gamma Omega and g are real but the matrix defined from them is complex:</p>
<p>def tridiag_eig(g,Gamma,Omega):<br>
n = Gamma.shape[0]<br>
diag = -torch.relu(Gamma)/2 +<strong>1j</strong>*Omega <span class="hashtag">#impose</span> positiveness constraints<br>
upper_diag = -<strong>1j</strong> *torch.relu(g)<br>
tridiag_mat = torch.diag(diag) + torch.diag(upper_diag,1) + torch.diag(upper_diag,-1)</p>