Linear solver for sparse matrices

Sehane: <p>Hello,<br>
I’d like to solve a linear system Ax=b where A is not square, but I know that there is exactly one solution.</p>
<p>The matrix A is represented as a sparse matrix that cannot be densified because it is too large. I also want the autograd to work on A.</p>
<p>The problem is that the only solutions I found so far are either computing a dense representation of A (which doesn’t work since A is too big) or using scipy (which is not compatible with autograd).</p>
<p>Is there a solution availiable on pytorch ?<br>
Thank you !</p>

l5rk: <p>I haven’t got experience working with sparse matrices in particular, but a quick peruse of the docs shows that matrix multiplication operations are supported. With that in mind, if your matrix is very tall or very wide, but the other dimension is reasonable, you could try to explicitly form the <a href="https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse#Linearly_independent_columns" rel="noopener nofollow ugc">Moore Penrose Inverse</a>. This involves taking (A^T A)^(-1). If you have a very tall matrix with a reasonable number of columns, A^T A may be a reasonable shape, and you could convert that intermediate step into a dense matrix to find either the explicit inverse (not usually recommended due to numerical stability issues) or any other linear solver supporting dense matrices.</p>

Sehane: <p>Both dimensions of the matrix are too large. The thing is that I’m shure that solving a linear system on a very big sparse matrix is differenciable and not computationally very expensive (using the algorithm of Gauss for example). That’s why I’d expect pytorch to have such a solver, but I didn’t find any.</p>

KFrank: <p>Hi Séhane!</p>
<aside class="quote no-group" data-username="Sehane" data-post="1" data-topic="200289" 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/s/71c47a/48.png" class="avatar"> Sehane:</div>
<blockquote>
<p>I’d like to solve a linear system Ax=b where A is not square, but I know that there is exactly one solution.</p>
</blockquote>
</aside>
<p>Because <code>A</code> is not square, your system of equations must be either<br>
under-determined or over-determined.  If it is under-determined, it<br>
has, by definition, more than one solution.</p>
<p>If it is over-determined, you will, in general, have <em>no</em> solutions, so <code>A</code><br>
must have some special structure that permits it to have a solution.</p>
<p>Could you explain that structure, as well as why it leads to <em>exactly<br>
one</em> solution?</p>
<aside class="quote no-group">
<blockquote>
<p>The matrix A is represented as a sparse matrix<br>
…<br>
using scipy (which is not compatible with autograd).</p>
</blockquote>
</aside>
<p>Let’s say that you do compute <code>x</code> somehow.  Do you need <code>x</code> to be<br>
differentiable with respect to <code>A</code> or to <code>b</code> or to both?  (If you only need<br>
differentiability with respect to <code>b</code>, you could quite easily write a<br>
<a href="https://pytorch.org/docs/2.2/autograd.html#function" rel="noopener nofollow ugc">custom autograd function</a> that uses scipy’s sparse solver for both<br>
the forward and backward passes.)</p>
<p>Best.</p>
<p>K. Frank</p>

Sehane: <p>Hello,<br>
So I want to differenciate regarding some parameters in A.<br>
The setting is the following.<br>
I have a matrix P describing a markov decision process ( which means that each line of P is a probability distribution, so it has positive values that sums up to one).<br>
I want to find the stationnary state, which is defined as the limit of (P^n) when n tends to infinity. It is caracterized as the only matrix w such that :</p>
<ul>
<li>wP=w</li>
<li>w  is a probability distribution, so it has positive values that sums up to one</li>
</ul>
<p>To solve it I decided to define a matrix X as P-Id and a line of ones and Y is a row of zeros with a one at the bottom.</p>
<p>To make this work you additionally need to transpose the matrices correcty. I implemented it and it worked with dense matrices. The thing is that P is sparse, so X is spparse too, and I want to scale the problem. It is only possible using sparse matrices.</p>
<p>Thank you for your answer !</p>

KFrank: <p>Hi Séhane!</p>
<aside class="quote no-group quote-modified" data-username="Sehane" data-post="5" data-topic="200289" 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/s/71c47a/48.png" class="avatar"> Sehane:</div>
<blockquote>
<p>I have a matrix P describing a markov decision process<br>
…<br>
I want to find the stationnary state, which is defined as the limit of (P^n)</p>
</blockquote>
</aside>
<p>The short answer is use the power-iteration method to compute your<br>
stationary state, <code>w</code>.  <code>torch.sparse.mm (sparse, dense)</code> supports<br>
backpropagation, so you will be able to compute your gradients without<br>
materializing a dense matrix with the size of <code>P</code>.</p>
<p>Beware that there is some nuance here, as discussed below.</p>
<aside class="quote no-group">
<blockquote>
<ul>
<li>wP=w</li>
</ul>
</blockquote>
</aside>
<p>Note that this is the eigenvector equation for <code>w</code> (with eigenvalue one), so<br>
calling the proposed scheme “power iteration” is appropriate.</p>
<aside class="quote no-group">
<blockquote>
<p>To solve it I decided to define a matrix X as P-Id and a line of ones and Y is a row of zeros with a one at the bottom.</p>
</blockquote>
</aside>
<p>This is also legitimate, except, as you’ve noted, <code>torch.linalg.lstsq()</code><br>
(and similar pytorch tools) don’t (yet?) support sparse matrices.</p>
<p>Further detail:</p>
<p>We have two schemes for computing your stationary vector:  power iteration<br>
and <code>lstsq()</code>.  We can understand both of these as functions that map a<br>
matrix <code>P</code> to its stationary vector, <code>w</code>.</p>
<p>As seen in the script given below, these give the same result, but not the<br>
same gradient.  This is because these two schemes aren’t actually the<br>
same function – they only give the same result when applied to a<br>
legitimate markov-process transition matrix.</p>
<p>So you have, conceptually, a bunch of different functions that all give the<br>
same result when applied to a legitimate transition matrix, but that produce<br>
different gradients.</p>
<p>(As a simple example of this, consider <code>f (x) = 1</code>, that is, the constant<br>
function, and <code>g (x) = x**2</code>.  These two different functions have the same<br>
value when applied to <code>x = 1</code>, but have different derivatives.  In this example<br>
the “condition” that <code>x = 1</code> is analogous to the condition that <code>P</code> be a valid<br>
transition matrix in your real problem.)</p>
<p>But which gradient is the “correct” one?  In some sense that’s up to you<br>
and the details of your use case.</p>
<p>One approach to addressing this problem is to map the input matrices<br>
to these functions, which may or may not be transition matrices, to valid<br>
transition matrices and then proceed to the stationary-state computation.<br>
(I call this “normalization” in the script.)  If we ensure that this normalization<br>
map doesn’t change a matrix that is already a transition matrix, we will still<br>
get the same result as from the unnormalized functions.</p>
<p>With this scheme, the normalized versions of power iteration and <code>lstsq()</code><br>
are the same function as one another and produce the same gradient.</p>
<p>(Note that pytorch’s support for sparse tensors is a moving target.  It’s gotten<br>
a lot better recently, but still has a lot of missing functionality and probably<br>
still some bugs.  I’ve run my example script with the pytorch’s latest stable<br>
version, 2.2.2, and it wouldn’t surprise me if earlier versions don’t work the<br>
same way.  You will see some oddities in the example script that are required<br>
due to missing sparse functionality.)</p>
<p>Consider the following script that implements these schemes, including<br>
(normalized) power iteration for sparse <code>P</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2024)

n = 100
sparseMult = 3
powerCnt = 100

# create random, sparse, transition matrix
#   this matrix need not be ergodic nor have unique
#   stationary state, but will "probably" be okay

P = torch.rand (n, n)                  # random probability matrix (before masking and normalizing)
rowCounts = torch.poisson ((sparseMult - 1) * torch.ones (100)).long() + 1
m = torch.zeros (n, n)                 # to become mask matrix 
for  i, k in enumerate (rowCounts):
    m[i][torch.multinomial (torch.ones (n), k)] = 1

P *= m                                 # mask out most probabilities
P /= P.sum (dim = 1, keepdim = True)   # normalize each row to sum to one

# check sparsity
print ('sparsity (non-zero fraction):     ', (P != 0.0).sum() / P.numel())

# check valid transition matrix
print ('\nverify valid transition matrix:')
print ('any less than zero:               ', (P &lt; 0.0).any())
print ('any greater than one:             ', (P &gt; 1.0).any())
print ('max row-sum discrepancy:          ', (P.sum (dim = 1) - 1.0).abs().max())

PSparse = P.to_sparse()                # also create sparse version of P

# create random initial state vector -- u is a global variable
u = torch.rand (n)
u /= u.sum()

def statePower (T):                    # use "power iteration" to find stationary state
    w = u.clone()
    for  i in range (powerCnt):  w @= T
    return w

def stateLstsq (T):                    # use lstsq to solve linear equation for stationary state
    w = u.clone()
    # form augmented matrix and "result" vector
    A = T.T - torch.eye (n)
    A = torch.cat ((A, torch.ones (1, n)), dim = 0)
    b = torch.zeros (n + 1)
    b[-1] = 1.0
    w = torch.linalg.lstsq (A, b)[0]
    return w

def tNorm (T):                         # "normalize" T to ensure valid transition matrix
    eps = 1.e-12
    t = T.abs()
    t /= (t.sum (dim = 1, keepdim = True) + eps)
    return t

def statePowerNorm (T):                # normalize T before "power iteration"
    return statePower (tNorm (T))

def stateLstsqNorm (T):                # normalize T before "power iteration"
    return stateLstsq (tNorm (T))

def statePowerSparse (T):              # version of statePower for sparse T
    w = u.clone().unsqueeze (-1)
    TT = T.T
    for  i in range (powerCnt):  w = torch.sparse.mm (TT, w)
    return w.squeeze()

def tNormSparse (T):                   # "normalize" sparse T to ensure valid transition matrix
    eps = 1.e-12
    t = T.abs()
    rowSum = torch.sparse.sum (t, dim = 1)
    rowSumInv = 1.0 / (rowSum.to_dense() + eps)
    # t *= rowSumInv                   # can't backpropagate through broadcast mul(), so do it "by hand"
    expVals = rowSumInv[t.indices()[0]]
    rsiExpand = torch.sparse_coo_tensor (t.indices(), expVals, t.shape).coalesce()
    t = t * rsiExpand
    return t

def statePowerNormSparse (T):          # version of statePowerNorm for sparse T
    return statePowerSparse (tNormSparse (T))

wPwr = statePower (P)
wLst = stateLstsq (P)
wPwrN = statePowerNorm (P)
wLstN = stateLstsqNorm (P)

print ('\ncheck agreement between state-vector algorithms:')

# computed state vectors agree
print ('max diff between wPwr and wLst:   ', (wPwr - wLst).abs().max())

# "norm" versions of state vectors also agree
print ('max diff between wPwrN and wPwr:  ', (wPwrN - wPwr).abs().max())
print ('max diff between wPwrN and wLstN: ', (wPwrN - wLstN).abs().max())

# compute gradients
P.requires_grad = True

# compute grad from power iteration
wPwr = statePower (P)
(wPwr**2).sum().backward()             # call backward on some scalar function of wPwr
gradPower = P.grad.clone()

# compute grad from lstsq solution
P.grad = None
wLst = stateLstsq (P)
(wLst**2).sum().backward()             # call backward on the same scalar function of wLst
gradLstsq = P.grad.clone()

print ('\ncheck disagreement / agreement for gradients of algorithms:')

# but gradients from the two versions disagree
print ('max diff between grads:           ', (gradPower - gradLstsq).abs().max())

# compute gradients using "norm" versions

# compute powerNorm grad
P.grad = None
wPwrN = statePowerNorm (P)
(wPwrN**2).sum().backward()            # backward on wPwrN
gradPowerNorm = P.grad.clone()

# compute lstsqNorm grad
P.grad = None
wLstN = stateLstsqNorm (P)
(wLstN**2).sum().backward()            # backward on wLstN
gradLstsqNorm = P.grad.clone()

# however the gradients from the "norm" versions agree
print ('max diff between "norm" grads:    ', (gradPowerNorm - gradLstsqNorm).abs().max())

# compute state vector and gradient using sparse power "norm" version

PSparse.requires_grad = True

wSprN = statePowerNormSparse (PSparse)
(wSprN**2).sum().backward()
gradSparseNorm = PSparse.grad.clone().coalesce()

print ('\ncheck state vector and gradient for sparse algorithm:')

# sparse-version of state vector agrees
print ('max diff between wSprN and wLstN: ', (wSprN - wLstN).abs().max())

# sparse gradient agrees with corresponding elements of dense gradient (using "norm" versions)
print ('max diff for sparse "norm" grad:  ', (gradSparseNorm.values() - gradLstsqNorm[PSparse.indices().unbind()]).abs().max())
</code></pre>
<p>Here is the script’s output:</p>
<pre><code class="lang-plaintext">2.2.2
sparsity (non-zero fraction):      tensor(0.0275)

verify valid transition matrix:
any less than zero:                tensor(False)
any greater than one:              tensor(False)
max row-sum discrepancy:           tensor(1.1921e-07)

check agreement between state-vector algorithms:
max diff between wPwr and wLst:    tensor(7.4506e-08)
max diff between wPwrN and wPwr:   tensor(1.8626e-08)
max diff between wPwrN and wLstN:  tensor(7.5554e-08)

check disagreement / agreement for gradients of algorithms:
max diff between grads:            tensor(0.2146)
max diff between "norm" grads:     tensor(5.9372e-08)

check state vector and gradient for sparse algorithm:
max diff between wSprN and wLstN:  tensor(8.5808e-08, grad_fn=&lt;MaxBackward1&gt;)
max diff for sparse "norm" grad:   tensor(8.4285e-08)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

alexj: <p>Sorry, you mentioned that sparse.mm supports autograd. What types of sparse tensors are supported?</p>

eacuellarq: <pre><code class="lang-auto">import torch
from scipy.sparse import csr_matrix
from scipy.sparse.linalg import splu

class SuperLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, b):
        # Ensure A is a CPU COO sparse tensor and b is a CPU dense tensor
        assert A.layout == torch.sparse_coo, "A must be a COO sparse tensor"
        assert A.device == torch.device('cpu'), "A must be on CPU"
        assert b.device == torch.device('cpu'), "b must be on CPU"
        
        # Extract indices, values, and shape from A
        A = A.coalesce()
        indices = A.indices()
        rows_np = indices[0].numpy()
        cols_np = indices[1].numpy()
        values_np = A.values().numpy()
        shape = A.shape
        
        # Convert to SciPy CSC matrix for SuperLU
        csr_A = csr_matrix((values_np, (rows_np, cols_np)), shape=shape)
        csc_A = csr_A.tocsc()
        
        # Solve Ax = b using SciPy's SuperLU
        b_np = b.numpy()
        invA = splu(csc_A)
        x_best_np = invA.solve(b_np)
        x_best = torch.tensor(x_best_np, dtype=b.dtype)
        
        # Save necessary variables for backward pass
        ctx.invA = invA
        ctx.rows_np = rows_np
        ctx.cols_np = cols_np
        ctx.x_best = x_best
        ctx.A_shape = shape
        
        return x_best

    @staticmethod
    def backward(ctx, grad_output):
        # Retrieve saved variables
        invA = ctx.invA
        rows_np = ctx.rows_np
        cols_np = ctx.cols_np
        x_best = ctx.x_best
        shape = ctx.A_shape
        
        # Solve A^H grad_b = grad_output (adjoint solve)
        grad_output_np = grad_output.numpy()
        grad_b_np = invA.solve(grad_output_np, trans='H')
        grad_b = torch.tensor(grad_b_np, dtype=grad_output.dtype)
        
        # Compute gradient for A: - grad_b[rows] * x_best[cols].conj()
        rows_tensor = torch.from_numpy(rows_np).long()
        cols_tensor = torch.from_numpy(cols_np).long()
        
        grad_b_rows = grad_b[rows_tensor]
        x_best_cols_conj = x_best[cols_tensor].conj()
        
        grad_A_values = - grad_b_rows * x_best_cols_conj
        
        # Construct sparse gradient tensor for A
        indices = torch.stack([rows_tensor, cols_tensor], dim=0)
        grad_A = torch.sparse_coo_tensor(indices, grad_A_values, shape)

        return grad_A, grad_b
        
</code></pre>