Schwarz-Schur Involution: Lightspeed Differentiable Sparse Linear Solvers
Yu Wang, S. Mazdak Abulnaga, Yaël Balbastre, Bruce Fischl
Abstract
Sparse linear solvers or generalized deconvolution are fundamental to science and engineering, applied in partial differential equations (PDEs), scientific computing, computer vision, and beyond. Indirect solvers possess characteristics that make them undesirable as stable differentiable modules; existing direct solvers, though reliable, are too expensive to be adopted in neural architectures. We substantially accelerate direct sparse solvers by up to three orders of magnitude, violating common assumptions that direct solvers are too slow. We "condense" a sparse Laplacian matrix into a dense tensor, a compact data structure that batchwise stores the Dirichlet-to-Neumann matrices, reducing the sparse solving to recursively merging pairs of dense matrices that are much smaller. The batched small dense systems are sliced and inverted in parallel to take advantage of dense GPU BLAS kernels, highly optimized in the era of deep learning. Our method is efficient, qualified as a strong zero-shot baseline for AI-based PDE solving, and a reliable differentiable module integrable into machine learning pipelines. Our code is available at https://github.com/ wangyu9/Schwarz_Schur_Involution .
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on8
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra et al.NeurIPS 2022 · 5,493 citations
- FlashAttention-2: Faster Attention with Better Parallelism and Work PartitioningTri DaoICLR 2024 · 2,600 citations
- Learning to Simulate Complex Physics with Graph NetworksAlvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying et al.ICML 2020 · 1,439 citations
- Learning Mesh-Based Simulation with Graph NetworksTobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, Peter W. BattagliaICLR 2021 · 1,175 citations
- Multiscale cholesky preconditioning for ill-conditioned problemsJiong Chen, Florian Schäfer, Jin Huang, Mathieu DesbrunSIGGRAPH 2021 · 29 citations
Related papers
- A Deep Conjugate Direction Method for Iteratively Solving Linear SystemsAyano Kaneda, Osman Akar, Jingyu Chen, Victoria Alicia Trevino Kala et al.ICML 2023 · 17 citations
- Learning Interface Conditions in Domain Decomposition SolversAli Taghibakhshi, Nicolas Nytko, Tareq Uz Zaman, Scott P. MacLachlan et al.NeurIPS 2022 · 18 citations
- Neural operators meet conjugate gradients: The FCG-NO method for efficient PDE solvingAlexander Rudikov, Vladimir Fanaskov, Ekaterina A. Muravleva, Yuri M. Laevsky et al.ICML 2024 · 14 citations
- DeepLM: Large-Scale Nonlinear Least Squares on Deep Learning Frameworks Using Stochastic Domain DecompositionJingwei Huang, Shan Huang, Mingwei SunCVPR 2021
- Collapsing Taylor Mode Automatic DifferentiationFelix Dangel, Tim Siebert, Marius Zeinhofer, Andrea WaltherNeurIPS 2025 · 1 citation
