🤖 AI

torch-sla: Differentiable Sparse Linear Algebra with Adjoint Solvers and Sparse Tensor Parallelism for PyTorch

This paper introduces torch-sla, an open-source PyTorch library that unifies differentiable sparse linear algebra by providing a single autograd-aware API for diverse solvers across multiple CPU and GPU backends, while supporting batched operations and scalable distributed multi-GPU execution through an O(1)-graph adjoint framework.

Original authors: Mingyuan Chi, Shizheng Wen

Published 2026-05-07

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Mingyuan Chi, Shizheng Wen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive, complex puzzle. In the world of computer science and machine learning, this puzzle is often a "sparse linear system"—a giant grid of numbers where most of the cells are empty (zero), but the few filled-in cells hold the key to everything.

For a long time, the popular programming tool PyTorch (used by many AI researchers) had a problem: it was great at handling dense puzzles (where every cell is filled), but it was terrible at solving these giant, mostly-empty puzzles while also learning from the mistakes. If you tried to solve one, the computer would either crash, run incredibly slowly, or forget how to learn from the result.

Enter torch-sla, a new open-source library that acts like a universal translator and a super-efficient mechanic for these puzzles. Here is how it works, using simple analogies:

1. The "Smart Dispatcher" (Unified Backend)

Imagine you are a delivery driver with a package. You need to get it to a destination, but the road conditions change: sometimes it's a small city street, sometimes a massive highway, and sometimes you're in a different country entirely.

Before, you had to manually choose a different vehicle for every single trip (a bicycle for small jobs, a truck for big ones, a boat for water). If you picked the wrong one, you wasted time or got stuck.

torch-sla is like a smart dispatcher. You just say, "I need to solve this puzzle," and it automatically checks:

Where are you? (Is your computer using a CPU or a powerful graphics card/GPU?)
How big is the puzzle? (Is it a tiny 100-piece puzzle or a 100-million-piece one?)
What shape is it? (Is it perfectly symmetrical?)

Based on these clues, it instantly picks the best "vehicle" (solver) from a garage of five different options (including tools from SciPy, NVIDIA, and PyTorch itself). You don't have to know which one to pick; the library does it for you.

2. The "Magic Shortcut" (Adjoint Differentiation)

This is the library's most magical trick.

The Old Way (Naive Backpropagation):
Imagine you are walking through a dark maze, taking 1,000 steps to find the exit. To learn from your journey, you decide to write down every single step you took, the exact position of your feet, and the air pressure at each moment. When you get to the end, you have a notebook with 1,000 pages of notes. If you want to go back and see where you made a mistake, you have to flip through all 1,000 pages. If the maze gets bigger, your notebook gets huge, and your computer runs out of memory (RAM) just trying to hold the notes.

The torch-sla Way (Adjoint Solver):
Instead of writing down every step, torch-sla uses a "magic shortcut" based on a mathematical principle called the Implicit Function Theorem.
Think of it like this: You walk the maze once. When you reach the exit, instead of retracing your steps, you instantly calculate the "reverse map" in your head. You only need to remember the start, the finish, and the rules of the maze.

Result: No matter if the maze took 10 steps or 10,000 steps, your "notebook" stays the same small size.
Benefit: The paper shows this saves 195 times more memory than the old way. It allows you to solve massive puzzles on a single computer that would have previously crashed the system.

3. The "Team Huddle" (Distributed Solvers)

What if the puzzle is so big that one computer can't hold it? You need a team of computers (like a group of friends working on a giant mural).

Usually, when friends work on separate parts of a mural, they have to shout back and forth to make sure the edges match up. In computer terms, this is called "halo exchange."

The Problem: When you try to learn from the final picture, the old tools didn't know how to "shout back" the corrections correctly to the right people.
The torch-sla Solution: It teaches the computers how to pass the "correction notes" in reverse. If Computer A sends a piece of data to Computer B, torch-sla ensures that when the learning happens, Computer B sends the correction back to Computer A perfectly. This allows the library to scale up to 400 million unknowns across three powerful GPUs.

4. What Can It Actually Do?

The paper demonstrates that this library can handle:

Direct Solvers: Cracking the code instantly for medium-sized puzzles.
Iterative Solvers: Taking step-by-step guesses for massive puzzles.
Nonlinear & Eigenvalue Solvers: Handling puzzles where the rules change as you solve them, or finding specific "vibrations" (eigenvalues) in the data.
Batching: Solving thousands of these puzzles at the same time (like solving 1,000 different mazes simultaneously).

The Real-World Test

To prove it works, the authors ran a test called "Inverse Coefficient Learning."
Imagine you see the final shape of a shadow (the solution) and you want to figure out what object (the input) cast it.

They used torch-sla to "learn" the shape of a hidden object by looking at the shadow it cast.
The library successfully figured out the hidden object with 99.77% accuracy in under a minute.
Crucially, the researcher didn't have to write any complex math code to make the learning happen; they just called a simple command: A.solve(f).

Summary

torch-sla fills a huge gap in the PyTorch ecosystem. It takes the difficult, memory-hungry task of solving giant, sparse math puzzles and makes it:

Automatic: It picks the best tool for the job.
Efficient: It uses a "magic shortcut" to save massive amounts of memory.
Scalable: It lets teams of computers work together seamlessly.
Learnable: It allows AI models to learn directly from these complex math problems without breaking.

It's essentially the missing link that lets PyTorch users tackle the same kinds of scientific and engineering problems that were previously only possible in other, more specialized programming environments.

Technical Summary: torch-sla

Problem Statement

Differentiable sparse linear algebra is a foundational requirement for scientific machine learning, including graph neural networks (GNNs), neural operators on unstructured meshes, and differentiable simulation pipelines. However, the PyTorch ecosystem lacks a unified library to support these workloads. Native PyTorch support is fragmented: torch.sparse provides low-level kernels and a non-differentiable, CPU-only spsolve (a SuperLU wrapper), while torch.linalg is restricted to dense matrices.

Existing workarounds are insufficient for large-scale or complex training scenarios:

Breaking the Graph: Calling external solvers (e.g., cuDSS, PETSc) breaks the autograd graph, forcing users to drop gradients or write custom torch.autograd.Function wrappers for every problem.
Inefficient Differentiation: Naively backpropagating through $k$ iterations of an iterative solver creates an $O(k)$ -node computational graph. For a problem with $10^6$ degrees of freedom (DOF) and 1,000 iterations, storing intermediate vectors requires roughly 80 GB of memory, rendering large-scale training infeasible.
Scalability Limits: Problems with $10^7$ – $10^9$ unknowns require distributed domain decomposition. Existing PyTorch distributed primitives do not natively support the transposed halo exchange required for sparse-operator adjoints in the backward pass.
Backend Fragmentation: Direct solvers (e.g., cuDSS) are fast for small problems ( $<10^5$ DOF) but suffer from $O(n^{1.5})$ fill-in and memory exhaustion for larger systems. Iterative solvers scale better but incur kernel-launch overhead on small problems. No unified interface exists to dispatch across these heterogeneous backends based on device and problem size.

Methodology

The authors present torch-sla, an open-source library that fills this gap via three core components:

1. Unified Backend Abstraction

torch-sla exposes a single autograd-aware API (.solve, .matvec, .eigsh, .det) that dispatches across five interchangeable backends based on device and problem size:

CPU: SciPy (SuperLU, UMFPACK, CG, BiCGStab) and Eigen (CG, BiCGStab).
GPU: NVIDIA cuDSS (LU, Cholesky, LDLT), CuPy (iterative and direct), and a PyTorch-native iterative solver.
Auto-Dispatch: The system automatically selects the backend: direct solvers for small-to-medium problems (e.g., $<2 \times 10^6$ DOF on GPU) and iterative solvers for larger problems to avoid memory overflow. Symmetry and positive-definiteness are detected to optimize factorization (e.g., upgrading LU to Cholesky).

The library supports four sparse tensor types organized by locality and batching: SparseTensor (single matrix/shared pattern), SparseTensorList (distinct patterns per batch), and their distributed variants DSparseTensor and DSparseTensorList.

2. Adjoint Differentiation Framework

To avoid the $O(k)$ memory and graph complexity of naive backpropagation, torch-sla employs an implicit function theorem approach.

Linear Systems: For $Ax=b$, the backward pass solves a single adjoint system $A^\top \lambda = \partial L / \partial x$ . The gradient with respect to the matrix is computed as $-\lambda_i x_j$ only on the sparsity pattern.
Nonlinear Systems: For $F(u, \theta)=0$ , the backward pass requires one adjoint linear solve ( $J^\top \lambda = \partial L / \partial u$ ) and a vector-Jacobian product, independent of the number of nonlinear iterations.
Eigenvalue Problems: Gradients for eigenvalues utilize the Hellmann–Feynman theorem ( $\partial \lambda / \partial A_{ij} = v_i v_j$ ), requiring only an outer product on the sparsity pattern.
Complexity: This framework reduces the autograd graph to $O(1)$ nodes and memory to $O(n + \text{nnz})$ , independent of the number of solver iterations ( $k$ ).

3. Distributed Layer with Halo Exchange

For multi-GPU execution, torch-sla implements domain decomposition where each process owns a row block and halo metadata.

Forward Pass: Processes exchange halo values before local Sparse Matrix-Vector (SpMV) operations.
Backward Pass: The library implements a transposed halo exchange. While the forward pass scatters owned values to neighbor halos, the backward pass sums incoming halo gradients back into the owning process's boundary entries.
Implementation: Both forward and backward communications are wrapped as torch.autograd.Function primitives using torch.distributed, ensuring end-to-end differentiability across distributed solves.

Key Contributions

Unified API: A single library providing autograd-aware interfaces for direct, iterative, nonlinear, and eigenvalue solvers across five backends (SciPy, Eigen, cuDSS, CuPy, PyTorch-native).
Scalable Adjoint Framework: An $O(1)$ -graph differentiation method that eliminates the memory bottleneck of iterative solver backpropagation, enabling training on problems with millions of unknowns.
Distributed Autograd: The first implementation of distributed sparse solvers in PyTorch that natively supports end-to-end differentiation via transposed halo exchange.
Extensibility: A backend abstraction that allows integrating external libraries (e.g., PETSc, Trilinos, hypre) or learned preconditioners by implementing a common interface.

Experimental Results

The authors evaluated torch-sla on 2D Poisson benchmarks using NVIDIA H200 GPUs:

Scalability: The library scales to 169M DOF on a single GPU and 400M DOF across three GPUs.
Performance: For a 1M DOF problem, the torch-sla CG backend (190 ms) outperforms cuDSS (7.3 s) by 38 $\times$ . Above 2M DOF, direct solvers run out of memory, while iterative solvers maintain near-linear scaling.
Memory Efficiency: In a comparison of naive backpropagation vs. the adjoint method for 1,000 CG iterations, the naive approach required 64 GB (OOM at 2,000 iterations), while the adjoint method maintained a constant **328 MB** usage, achieving a 195 $\times$ memory reduction.
Gradient Verification: Gradients were verified against finite differences for nonlinear and eigenvalue problems, showing relative errors below $10^{-5}$ .
End-to-End Application: An inverse coefficient learning task (recovering spatially varying conductivity in a Poisson equation) successfully converged to a 0.23% relative L2 error in under one minute using a standard PyTorch Adam loop.

Significance

The paper claims that torch-sla closes a long-standing gap in the PyTorch ecosystem, providing capabilities analogous to JAX's jax.lax.custom_linear_solve and jax.scipy.sparse.linalg but within a unified PyTorch library. It enables researchers to:

Train models involving sparse linear solves without dropping gradients or porting code to JAX.
Scale scientific machine learning workloads to hundreds of millions of degrees of freedom on distributed GPU clusters.
Treat sparse solvers as differentiable layers, facilitating the research of learned preconditioners and end-to-end optimization of physical simulations.

The authors note limitations, including the current reliance on Jacobi preconditioning for the PyTorch-native backend (which limits convergence speed for very large systems) and a focus on 2D Poisson benchmarks, with 3D PDEs and indefinite systems identified as future work.