Best way to compute C = A.inv() @ B @ A.inv().T? Linear algebra

dherrera1911: <p>When we have some n-by-n matrix <code>A</code> and we want to compute <code>A.inv() @ b</code> for a vector <code>b</code>, rather than inverting <code>A</code>, we usually use <code>torch.solve(A, b)</code>. This returns the result of <code>A.inv() @ b</code> without actually inverting <code>A</code>, and it is often faster and more numerically stable.</p>
<p>I wonder whether a similar approach exists to computing <code>A.inv() @ B @ A.inv().T</code>, where both <code>A</code> and <code>B</code> are symmetric positive definite.</p>
<p>I tried using <code>torch.linalg.solve</code> twice, but this is often slower than inverting <code>A</code>. I also tried using the <code>torch.cholesky_solve</code>, but this is also slower than just inverting <code>A</code>.</p>
<p>Is there something about this problem that makes inverting <code>A</code> the best option?</p>
<p>Example code:</p>
<pre><code class="lang-auto">import torch

n_dim = 20
n_points = n_dim * 50

X1 = torch.randn(n_points, n_dim)
A = X1.T @ X1

X2 = torch.randn(n_points, n_dim)
B = X2.T @ X2


def fun_invert(A, B):
    A_inv = torch.linalg.inv(A)
    return A_inv @ B @ A_inv


def fun_solve(A, B):
    C = torch.linalg.solve(A, B)
    return torch.linalg.solve(A, C.T).T


def fun_chol(A, B):
    L = torch.linalg.cholesky(A)
    C = torch.cholesky_solve(B, L)
    return torch.cholesky_solve(C.T, L).T

%timeit fun_invert(A, B)
%timeit fun_solve(A, B)
%timeit fun_chol(A, B)
</code></pre>

KFrank: <p>Hi Daniel!</p>
<p>The short story is that <code>cholesky_solve()</code> beats <code>inv()</code> both for larger matrices and when<br>
run on the gpu.  Presumably the theoretical disadvantage of <code>inv()</code> is masked by various<br>
overheads when the dimensions of the matrices are relatively small.</p>
<aside class="quote no-group" data-username="dherrera1911" data-post="1" data-topic="220883" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dherrera1911/48/61905_2.png" class="avatar"> dherrera1911:</div>
<blockquote>
<p>I wonder whether a similar approach exists to computing <code>A.inv() @ B @ A.inv().T</code>, where both <code>A</code> and <code>B</code> are symmetric positive definite.</p>
<p>I tried using <code>torch.linalg.solve</code> twice, but this is often slower than inverting <code>A</code>.</p>
</blockquote>
</aside>
<p>This is not surprising.  The benefits of using <code>solve()</code> instead of <code>inv()</code> are reduced when<br>
you are solving for “multiple right-hand sides,” that is, when your <code>B</code> is a matrix rather than<br>
a vector.  When you invert <code>A</code>, you are solving for the <code>n</code> columns of the inverse of <code>A</code> and<br>
you can formulate this as solving for <code>n</code> right-hand sides.  So calling <code>solve()</code> for <code>B</code>, where<br>
<code>B</code> is a square matrix of dimension <code>n x n</code> is broadly the same as inverting <code>A</code>.  But for your<br>
use case, you have to call <code>solve()</code> twice, whereas you only have to invert <code>A</code> once, reusing<br>
the result.</p>
<p>As an aside, I don’t think <code>A</code> being symmetric positive definite helps for your use case.<br>
It does guarantee that <code>A</code> is invertible, but I’m not aware of any direct (i.e., non-iterative)<br>
solvers that make use of <code>A</code> being SPD.</p>
<aside class="quote no-group">
<blockquote>
<p>I also tried using the <code>torch.cholesky_solve</code>, but this is also slower than just inverting <code>A</code>.</p>
</blockquote>
</aside>
<p>I can reproduce your result of <code>cholesky_solve()</code> being slower (when <code>n</code> is relatively small).<br>
In general, however, <code>cholesky_solve()</code> should be faster.</p>
<p>The harder part of the full cholesky solution is calling <code>cholesky()</code> (the decomposition) and<br>
this part is called once and reused (sort of like reusing the inverse).  You do have to call<br>
<code>cholesky_solve()</code> twice, but this is the easier part of the full cholesky solution.</p>
<aside class="quote no-group">
<blockquote>
<p>Is there something about this problem that makes inverting <code>A</code> the best option?</p>
</blockquote>
</aside>
<p>It looks to me like inverting <code>A</code> wins just due to the vagaries of how the computation for<br>
smaller values of <code>n</code> interacts with the cpu (and gpu) floating-point pipelines.</p>
<p>I find that I can reproduce your result that inverting <code>A</code> is fastest for <code>n = 20</code> (your <code>n_dim</code>)<br>
running on the cpu, but running with significantly larger <code>n</code> (in the example below, <code>n = 200</code>)<br>
and / or on the gpu, the cholesky approach is faster.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">n_dim = 20
...
def fun_invert(A, B):
    A_inv = torch.linalg.inv(A)
    return A_inv @ B @ A_inv

def fun_solve(A, B):
    C = torch.linalg.solve(A, B)
    return torch.linalg.solve(A, C.T).T

def fun_chol(A, B):
    L = torch.linalg.cholesky(A)
    C = torch.cholesky_solve(B, L)
    return torch.cholesky_solve(C.T, L).T
</code></pre>
</blockquote>
</aside>
<p>I’ve timed <code>n = 20</code> vs. <code>n = 200</code> and cpu vs. gpu with a timing script based on your code.</p>
<p>Here is the script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

import time

torch.manual_seed (2025)

device = 'cpu'
if  torch.cuda.is_available():  device = 'cuda'
if  device == 'cuda':
    print (torch.version.cuda)
    print (torch.cuda.get_device_name())

# n_dim = 20
n_dim = 200
n_points = n_dim * 50
n_batch = 10

print ('device:', device, '  n_dim:', n_dim)

n_warm = 5
n_time = 1000

# create batches of SPD A and B
X1 = torch.randn(n_batch, n_points, n_dim, device = device)
A = X1.permute (0, 2, 1) @ X1
X2 = torch.randn(n_batch, n_points, n_dim, device = device)
B = X2.permute (0, 2, 1) @ X2

def fun_batch_invert (A, B):
    for  a, b in zip (A, B):
        a_inv = torch.linalg.inv (a)
        _ = a_inv @ b @ a_inv

def fun_batch_solve (A, B):
    for  a, b in zip (A, B):
        c = torch.linalg.solve (a, b)
        _ = torch.linalg.solve (a, c.T).T

def fun_batch_chol (A, B):
    for  a, b in zip (A, B):
        l = torch.linalg.cholesky (a)
        c = torch.cholesky_solve (b, l)
        _ = torch.cholesky_solve (c.T, l).T

# time invert
for  _ in range (n_warm):
    fun_batch_invert (A, B)
if  device == 'cuda':  torch.cuda.synchronize()
tStart = time.time()
for  _ in range (n_time):
    fun_batch_invert (A, B)
if  device == 'cuda':  torch.cuda.synchronize()
tEnd = time.time()
print ('time_inv: {:.4f}'.format (1000 * (tEnd - tStart) / (n_batch * n_time)))

# time solve
for  _ in range (n_warm):
    fun_batch_solve (A, B)
if  device == 'cuda':  torch.cuda.synchronize()
tStart = time.time()
for  _ in range (n_time):
    fun_batch_solve (A, B)
if  device == 'cuda':  torch.cuda.synchronize()
tEnd = time.time()
print ('time_sol: {:.4f}'.format (1000 * (tEnd - tStart) / (n_batch * n_time)))

# time chol
for  _ in range (n_warm):
    fun_batch_chol (A, B)
if  device == 'cuda':  torch.cuda.synchronize()
tStart = time.time()
for  _ in range (n_time):
    fun_batch_chol (A, B)
if  device == 'cuda':  torch.cuda.synchronize()
tEnd = time.time()
print ('time_cho: {:.4f}'.format (1000 * (tEnd - tStart) / (n_batch * n_time)))
</code></pre>
<p>And here is its output for the four different runs:</p>
<pre><code class="lang-plaintext">2.7.1+cu128
device: cpu   n_dim: 20
time_inv: 0.0244
time_sol: 0.0364
time_cho: 0.0341

2.7.1+cu128
12.8
NVIDIA RTX 3000 Ada Generation Laptop GPU
device: cuda   n_dim: 20
time_inv: 0.1210
time_sol: 0.1847
time_cho: 0.0851

2.7.1+cu128
device: cpu   n_dim: 200
time_inv: 0.5564
time_sol: 0.8801
time_cho: 0.4294

2.7.1+cu128
12.8
NVIDIA RTX 3000 Ada Generation Laptop GPU
device: cuda   n_dim: 200
time_inv: 0.4634
time_sol: 0.8400
time_cho: 0.3604
</code></pre>
<p>A couple of comments:</p>
<p>The cpu timings jump around a lot more than the gpu timings.  I imagine this is because<br>
there are various “random” processes running in the background competing for the cpu.</p>
<p>Empirically, <code>n = 200</code> is still rather small in the context of the gpu.  The benefit of the gpu<br>
becomes dramatically more pronounced for larger <code>n</code>.</p>
<p>Best.</p>
<p>K. Frank</p>