Matrix Inversion in PyTorch

Overview

This document provides a comprehensive guide to matrix inversion using PyTorch, focusing on the torch.linalg.inv function. It covers API usage, practical examples (including GPU acceleration), performance considerations, and numerical stability tips.

Function Signature

torch.linalg.inv(A, *, out=None) → Tensor

Computes the inverse of a square matrix (or batch of matrices) A. The inverse satisfies A @ A.inv() = A.inv() @ A = I, where I is the identity matrix.

Parameters

• A (Tensor) – Input tensor of shape (*, n, n) where * is zero or more batch dimensions. Must be square and invertible.
• out (Tensor, optional) – Output tensor. Ignored if None.

Returns

• Tensor – Inverse of A. Same shape and dtype as A.

Raises

• RuntimeError – If A is not square or is singular (non-invertible) within numerical tolerance.

Basic Usage

CPU Example (100×100)

import torch

# Create a random 100x100 matrix (always full rank with probability 1)
n = 100
A = torch.randn(n, n, dtype=torch.float64)

# Compute inverse
A_inv = torch.linalg.inv(A)

# Verify
I = torch.eye(n, dtype=torch.float64)
error = torch.norm(A @ A_inv - I, p='fro')
print(f"Frobenius error: {error:.2e}")

GPU-Accelerated Example (100×100)

import torch
import time

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float64
n = 100

A = torch.randn(n, n, dtype=dtype, device=device)

start = time.time()
A_inv = torch.linalg.inv(A)
torch.cuda.synchronize()  # Wait for GPU to finish
end = time.time()

print(f"Inversion time on {device}: {end - start:.6f} s")

# Verification (on GPU)
I = torch.eye(n, dtype=dtype, device=device)
error = torch.norm(A @ A_inv - I, p='fro')
print(f"Reconstruction error: {error:.2e}")

Expected Performance (100×100)

Hardware	dtype=float64	dtype=float32
NVIDIA V100 (GPU)	~0.3 ms	~0.2 ms
NVIDIA RTX 3090	~0.2 ms	~0.1 ms
Intel Xeon CPU	~1.5 ms	~1.0 ms

Note: For small matrices, GPU kernel launch overhead may diminish speedups.

Batch Matrix Inversion

The same function handles batches seamlessly:

batch_size = 64
n = 100
A_batch = torch.randn(batch_size, n, n, device='cuda')
A_batch_inv = torch.linalg.inv(A_batch)  # shape: [64, 100, 100]

The inversion is performed independently for each matrix in the batch.

Numerical Considerations

1. Condition Number

Ill-conditioned matrices (large condition number) lead to large numerical errors. Estimate condition number:

cond = torch.linalg.cond(A)
if cond > 1e6:
    print("Matrix is ill-conditioned, inverse may be inaccurate")

2. Singular Matrices

For exactly singular matrices, torch.linalg.inv throws a RuntimeError. For near‑singular cases, consider:

• Adding a small diagonal perturbation: A_reg = A + 1e-6 * torch.eye(n, device=device)
• Using pseudo‑inverse: torch.linalg.pinv(A, rtol=1e-5)

3. Data Type Recommendations

• float64 – Default for high accuracy, especially when condition number > 1e4.
• float32 – Faster and memory efficient, but may lose precision for ill‑conditioned matrices.
• complex64/complex128 – Supported for complex-valued matrices.

4. Gradient Computation

If A.requires_grad = True, the inverse operation supports automatic differentiation:

A = torch.randn(n, n, requires_grad=True, device='cuda')
A_inv = torch.linalg.inv(A)
loss = A_inv.norm()
loss.backward()  # Gradients w.r.t. A are computed correctly

The gradient formula used is:
$$\frac{\partial X^{-1}}{\partial X}=-X^{-\top }\otimes X^{-1}$$
where $\otimes$ denotes the Kronecker product.

Alternative Inversion Methods

PyTorch does not expose low‑level LU or Cholesky solvers for direct inversion, but you can implement custom inversion if needed (e.g., for educational purposes). However, torch.linalg.inv is highly optimized (calls LAPACK on CPU and cuSOLVER on GPU) and should be the default choice.

Error Handling

def safe_inv(A, eps=1e-6):
    """Attempt inversion; fallback to pseudo-inverse if singular."""
    try:
        return torch.linalg.inv(A)
    except RuntimeError:
        print("Matrix singular, using pseudo-inverse instead")
        return torch.linalg.pinv(A)

Memory Footprint

For a single (n, n) matrix, the inversion requires temporary storage roughly proportional to n². For large n (>5000), memory may become a bottleneck.

Example memory usage (float64):

• n=1000 → matrix: 8 MB, workspace: ~30–50 MB
• n=5000 → matrix: 200 MB, workspace: ~1 GB

References

• PyTorch Documentation: torch.linalg.inv
• LAPACK routine ?GETRF/?GETRI (CPU)
• cuSOLVER cusolverDnXgetrf/cusolverDnXgetri (GPU)

Conclusion

torch.linalg.inv provides a robust, GPU‑accelerated matrix inversion routine. For most practical sizes (up to a few thousand), it delivers high performance and accuracy. Always verify invertibility and consider conditioning when numerical stability is critical. For large‑scale problems, prefer solving linear systems directly with torch.linalg.solve rather than computing the full inverse.

Source

https://gitee.com/waterruby/ANNA.git

(torch_env) x@x-X99:~/pro/ANNA/experiments$ python inv_100.py
使用设备: cuda
数据类型: torch.float64
求逆耗时: 0.147686 秒
重建误差 (Frobenius 范数): 3.567055e-13
反向误差 (Frobenius 范数): 5.542944e-13

逆矩阵左上角 5x5 子块:
tensor([[-0.1154, -0.0390, -0.0728,  0.4767, -0.5041],
        [-0.0144,  0.1074,  0.1905, -0.3552,  0.3985],
        [-0.1725, -0.1471,  0.2727, -0.1940,  0.5182],
        [ 0.0482,  0.0323, -0.0815, -0.1578,  0.0574],
        [-0.0433, -0.0402,  0.0622,  0.0675, -0.0541]], device='cuda:0',
       dtype=torch.float64)