Schwarz–Schur Involution: Lightspeed Differentiable Sparse Linear Solvers
Authors: Yu Wang
Mazdak Abulnaga
Yaël Balbastre
Bruce Fischl
Venue: ICML 2025 poster (2025)
Keywords: differentiable linear solver
sparse solver
generalized deconvolution
partial differential equations
numerical methods
scientific computing
Primary area: optimization
TL;DR: A very fast method for solving sparse linear systems or equivalently generalized deconvolution, qualifying sparse solvers as an efficient reliable differentiable module in neural architectures.

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 batch-wise 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.

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Linear Solvers
Yu Wang 123
S. Mazdak Abulnaga
12
Yaël Balbastre 14
Bruce Fischl
Abstract
Table I. Our substantially faster sparse solver (including the fac- torization and substitution stages) the first direct method to run at interactive rates. It takes our method 10.9 (resp. 220 ms) to solve a Laplaciann system (Dirichlet) on an image of x13
Sparse linear solvers or generalized deconvolution are fundamental science and engineering, ap- plied in partial differential equations (PDEs), sci- entific 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 com- mon assumptions that direct solvers are too slow. We "condense" sparse Laplacian matrix into a dense tensor, a compact data structure that batch- wise stores the Dirichlet-to-Neume reducing the sparse solving to recursively merg- ing pairs of dense matrices that are much smaller. The batched small dense systems 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, quali- fied as a strong zero-shot baseline for Al-based PDE solving, and a reliable differentiable mod- ule integrable into machine learning pipelines. Our code is available at https: / thub. . com/ wangyu9/ Swtp/r/wwaaaaeeeeeeeeeeeaaeeaeea
(resp. 2561 2561). Runtime is reported in milliseconds.
Example
CUDA
SciPy
ours
speedup
2561
36926
253318
220.1
168X 1151X
2049
21354
143100
158.5
135X 903X
1025"
4710
16512
36.45
129X 453X
513"
1036
2051
10,90
95.0X 188X
257
234
355
5.82
40.2X 61.0X
interactive rate algorithm on common-size images. Solving the linear system A -bb is equivalent to the generalized de- convolution of an image b with spatially varying kernel (in rows of) A. Sparse linear solvers underpin modern image processing, computer vision, graphics, as well as numerical methods for PDEs such as finite différeence/llmmet meth- ods (FD/FEM) used in clectrical/mechanica engineering and computational sciences. Users in these areas can also benefit from our approach if efficiency is a primary concern. Our direct solver-even naïvely prototyped in PyTorch-is 60 to 1000 faster than SciPy and 40 to 170x faster than CUDA (cuDSS), which are highly optimized (Table 1).
matrices,
Successes of modern deep learning are largely built on dif- ferentiable namerical linear algebras, particularly matrix as attentions ( (Vaswani et al., multiplications of matrices-nntably convolutions (LeCun 2017) and sparse matrices et al., 1989). High performance Al systems rely on effi- cient implementation of these operations, including dense (Strassen, 1969; Blahut, 2010; Dao et al., 2022; Dao, 2023) and sparse multiplications (Winograd, 1980; Cook, 1966; Lavin & Gray, 2016). In contrast, matrix inversions almost never employed in neural architectures. We believe that inversions are underexplored and can play a prominent role, since (sparse) matrix inversions encompass a large range of operations such as generalized deconvolution, ge- ometry representation, results of physical simulations, de- correlation, and equilibrium states of localized interactions.
Involution.
1. Introduction
In this paper, we propose a conceptually simple approach that reduces solving certain sparse linear system (A, b) to "involuting" tensors (a, B), leading to significant improve- ments. Global linear systems in computer vision and scien- tific computing, such as the Laplacian or Hessian systems image domains, are usually very large yet sparse, as induced by pixel affinity. For a sparse A, direct solvers find the exact solution: x = A for which we present the first
Athinoula Martinos Center, Massachusetts General Hospital, Harvard Medical School 'Massachusetts Institute of Technology, MIT CSAIL, USA 'oony, USA 'University College London, UK. Correspondence to: <(ywang155, biiicss @ www..aarvard.aommmmmmmmm hdd>.
As detailed in $F.2, key obstacle preventing neural archi- tectures from adopting linear solvers is the lack of a sparse solver like ours that is cunsstent--erformann in-the-wild, accurate, robust, and zero-pane pero-parameter. Indirect, a.k.a. iterative, solvers fall short of these goals due to their strong problem dependency, parameter sensitivity, and un-
lightspeed,
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
predictable runtime. The large diversity of PDEs results in laborious workflows requiring expert users in the loop to tweak parameters or preconditioners to employ indirect solvers. Iterative solvers struggle to deconvolute indefi- nite kernels, such as stencils discretizing the Helmholtz PDE (Ernst & Gander, 2011). While existing direct solvers are promising candidates for meeting the desiderata, they are too slow-a limitation that our method overcomes. Cat- alyzed by the surge of interest in physics-infen innsccs- machine learning and AI for PDE/Science (Berens et al., 2023), iden- tification or learning of unknown physical systems open up a unique setting to solve matrix A whose properties are unknown in advance to apply appropriate iterative schemes, necessitating a general-purpose solver that deployable for any A and/or a large varieties of matrices from some training batch. This is similar the generalized deconvo- lution with unknown kernels in computer vision. Applying unsuitable indirect solvers can yield catastrophic failure.
2017) for extensive surveys. Direct solvers first perform the numerical Cianiizaaton-Choleese tion for symmetric A and LU factorization for asymmetric systems-
LDLT factoriza-
by a back substitution to spread out b to yield x. Major backends for direct sparse solvers include SuperLU (Li, 2005), UMFPACK (Davis, 2004), CUDA (Nickolls et al., 2008), Pardiso (Schenk & Gärtner, 2004), Eigen (Guennebaud al., 2014). and Cholmod/- SuiteSparse (Chen et al., 2008). When A comes from spatially constant kernel or point spread function (PSF), the problem becomes an image deconvolution, with efficient frequency domain solvers (Sezan & Tekalp, 1990; Hansen et al., 2006); spectral solvers like FFT/IFFT (Fast Fourier Transform) (Frigo & Johnson, 2005) are limited to solving homogeneous Laplace systems. Instead, we attack PDEs or deconvolution with spatially varying coefficients/kemnlls..
We consider solving indefinite and nonsymmetric square linear systems (Trefethen & Bau, 1997), so methods requir- ing symmetric or positive-definite A do not generally apply, such as CG (Conjugate Gradient) (Hestenes et al., 1952), MINRES (Paige & Saunders, 1975; Liu & Roosta, 2022). Methods apply to our setting include: CGS (Conjugate Gra- dient Squared) (Sonneveld, 1989) and the improved variant biCGstab (Bi-Conjugate Gradient Stabilized) (Van der Vorst, 1992; Sleijpen & Fokkema, 1993), iterative sparse least- squares LSQR (Paige & Saunders, 1 1982b;a) that amounts to CG applied to (ATA) -2 Tb and recent improvement LSMR (Fong & Saunders, 2011), GMRES (generalized min- imal residual) (Saad & Schultz, 1986) and improvements DQGMRES (Saad & Wu, 1996) and FGMRES (Saad, 1993). Other methods include BILQ (Montoison & Orban, 2020; Fletcher, 1976), QMR (Freund & Nachtigal, 1991; 1994), FOM (Saad, 1981), and DIOM (Saad, 1984). The multigrid methods (Bramble, 1993) apply a hierarchical coarse-to-fine strategy, effective as iterative solvers or the preconditioners in CG to yield PCG, especially for elliptic PDEs. We defer discussions on solvers in vision & graphics (Barron & Poole, 2016; Krishnan et al., 2013; Bolz et al., 2003; Jeschke et al., 2009 Horvath & Geiger, 2009; Liu et al., 2016) to $A.
Our improvement is due to "condensing"-transfi GPU's capacity of dense BLAS (Dongarra et al., 1988; 2018) to sparse computation, through design choices with a compact data structure-a tensor storing Dirichlet-to-Neumale ma- trices or sub-systems in batches, and a procedure we coin as "Schwarz-Schur invouution"-rccusivvel recursvvely applying the Schur complement formula to block-wisely contract nodes and keep track of sparsity explicitly. Our approach amounts to a parallel implementation of Gaussian climination under nested dissections (George, 1973) and the induced multi- frontal solvers (Duff& & Reid, 1983), recursively canceling variables. We take advantage of the regularity of image grids to aggressively explore parallelisms, leveraging advances in GPUs sparked by deep learning that shift the best practice towards algorithms better exploiting parallelisms.
Our sparse solvers, with speeds exceeding iterative solvers, retain the accuracy and reliability of direct solvers, efficient enough to be integrated within neural architectures. Owing to the central role of linear solvers, our method enables efficient implementations including but not limited to: generalized deconvolution of spatially varying kernels. identify/reduce gaps between AI & conventional methods
zero-shot baselines for sarming-based PDE solvers.
2. Mathematical Preliminaries
exact Newton solvers made tractable on image domains. mathematical optimization layers embedded in neural nets (Amos& & Kolter, 2017), (physics) solver-in-the-loop. geometric deep learning & algorithms (Litany et al., 2017),
Sparse linear solver = exact generalized deconvolution. In the paper, for an /l x W image with n pixels, := HW, whose boundary has b = 2W +2H pixels, we consider a sparse matrix A = that encodes the affinity of pixels: each row/column of A corresponds to one pixel and Aij is nonzero only if pixels i, E {1,,.., n.. are adjacent in the image, namely their integer coordinates (I), (2,,, e N²2 satisfying that (z, (z,= ( +(,-,, y2, < 2. For a function u(x; y) on the image grid, consider the linear operator or a 3x3 convolution centered (x.1) with al*.0) E R333,
shape/deformation representation, graph/mesh/FFM neu- ral networks (Wang et al., 2019; Pfaff et al., 2020).
eigenbases for spectral neural nets (Bruna et al., 2013), spectral clustering& & segmentation (Shi& & Malik, 2000).
1.1. Related Work
We refer interested readers to standard texts on direct solvers of dense (Davis, 2006) and sparse systems (Duff et al.,
v(z,y) + ΣΣ E EEaddde(--onn (a.)) (6., 8, uu( + da, y + 6,)
2

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Figure 1. Our method immediately accelerates direct solvers in a large variety tasks by orders magnitude, including: FEM for solving PDEs, deconvolution (generalized to spatially varying kernels), eigen solver and spectral method, image segmentation and matting, physical simulation, deformation, geometry processing, shape optimization, Newton's method, and diffeomorphic image registration.
using a 9-point stencil a(x.)) spatially varying kernel:
the anisotropic diffusion coefficient C(x) E R2x2, Vx E Ω the domain 12 subject to either the Dirichlet condition (2) or the Neumann condition (3) at the boundary aΩ:
-- 11
(-11)
a (2.
olo.v (0.-1)
(0.
10.
(1)
(1,
a (1.00
-V - VT(C(x)Vu(xx)) f(x), u(x))an= g(x))... (2)
By flattening the 2-dim array u, into the vector v. E Rax1 (see Figure 25), for Asj a( Ty, Ui yj), the convolution with (2.,)) can be written as the matrix- vector product v t Au. Thus, the linear solve A ''v generalizes deconvolution to the case with spatially vary- ing kernel ( Theoretically, our approach generalizes to larger stencils and 3D, which we leave for future work.
V-(((x)Vu(()) = f(x), n (xxxx lxx)
(3)
Discretizing with first-order piecewise linear FEM yields:
L-GTCG, [Lu]bst>>-f ul1.5g or [Lu]| h,
following the notation and discretization (Wang et al., 2023), in which R,, R=-b g, R discretize the so- lution and right-hand side functions f.g. h; n(x) is the normal direction at the boundary point x; =11.8 +|o+1:n selects rows for boundary or interior pixels, resp.; the ma- trices G, C discretize the gradient operator V and C(x).
Differentiable and repetitive linear solvers. In the learn- ing setting, the system A, b come from some learnable parameters f to be identified with gradient descent. Thus, in addition to solving x A bb., we need to obtain dx/ab and 0x/aA, which have closed-form expressions that can be derived using the adjoint method; evaluating these gradients requires solving the transposed system A see SC.3.
Under Dirichlet boundary condition gxx), the solution to (2) u(x) is unique, which determines the Neumann boundary condition (xx), inducing the Sriichlettt-Neummare (DiN) map: ( h(@n). After discretization, the DtN map becomes the Schur complement of the matrix L. Thus, the ceaffccienttto-oolutii operator: C(x) - u(x), under FD/FEM, can be simply viewed as a linear solver.
It is common to sequentially solve multiple pairs (A, bi), (A, by).... (A, by), i... the same left-hand side A but dif- ferent right-hand sides. Direct solvers, including ours, stand out for allowing separating numerical factorization from back substitution and reuse the former. Numerical factor- ization is the computation involving A only, which is the dominant cost that can be reused for direct solvers. This is possible for indirect solvers, which have to start over for each right-hand side. The typical value of 1: is 2 in differen- tiating linear solvers, h >> 20 when being applied in eigen solvers, or k > 100 in time steppings. Thus, compared to iterative solvers these settings, computationally, it means that iterative solvers get an extra xk slowdown.
What exactly represents- helps to motivate-docs not make any difference to our method. Our method does not assume A symmetric, though this is often the case. Solving Neumann boundary value problems amounts set- ting A L; solving parabolic PDEs amounts to letting A = tL + M, and choosing A = L- K >²M corresponds to solving the Helmholtz equation that arises in wave propaga- tion and electromagnetism, in which M plays role similar to the identity matrix-M is the mass matrix discretized by FEM (Allaire, 2007) with the same sparsity pattern as L.
PDE and FEM preliminaries. Our method applies any invertible A with the aforementioned sparsity. But first let us examine the A that arises from discretizing elliptic PDEs (Gilbarg et al., 1977), situation that motivates our initial development. Surprisingly, the strategy generalizes beyond elliptic PDEs without adaptations to other linear systems in practice, even non positive-sensee semidefinite (PSD). Let us consider solving the Laplace equation defined using
3. Schur Involution for Parallel Elimination
We demonstrate procedure to convert the solution of the linear system x (A, b) to the "involution" of the tensors X + (a. 3), by slicing and inverting many small dense matrices in batches to leverage modern GPU hardware.
3

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
8::0 1-99
(1,,4,25, 26)
(4, 4, 16, 16)
(2, 4, 24)
a (2. 32, 32)
(3) (1,
48 48)
a
:L. 1. 6,6 64)
(1,
g (4, 16 0 (2, 1,,4,21,11
0 (2.2.
( : 1)
g (1,,,,6,111
Figure 2. Schwarz-Schhr involution a 17x17 image. Our method (Algorithm applics batched linear algebras to parallelize Gaussian climination in preseribed aader-pixeh marked in color correspond to nodes to be eliminated in that step. As the initialization, the Schwarz step converts the image into a wire-frame by canceling the interior pixels of each patch. Multiple Schur steps are applied to progressively simplify the wire-frame, until only pixels at the border are left. A Schur step merges every two adjacent subdomains by removing pixels on their shared "edges." We compact subdomains' left-/right-hand sides the tensors a (1) (,,,(() whose shapes evolve as shown. The solutions x with the same shape as , flow reversely: x x(*) + x(() < xxx")...
x( = (a)
3.1. A motivating example: sparse solvers too slow?
3.2.1. 1. SCHWARZ STEP: DECOMPOSE & INITIALIZE DTN
A 4096x4096 image can be divided into 1024x 1024 array of 4 x 4 patches. If adjacent patches overlap by sharing one layer of pixels, then we have a similar division: a 4097x4097 image (or 17x 17 as shown in Figure 2, or x in Figure 3) can be divided into a 1024 1024 (or 4 4, or 1, resp.) array of small patches of size 5 5, since boundary pixels have duplicated representations. On this image, SciPy solver takes more than 20 minutes to solve sparse system A E Rax", = 16785409. However, inverting one million 9x9 matrices in a batch only takes 0.008 seconds using PyTorch a 1rpha.shapp is (1024, 1024,9, 9)
As shown in Figure 3, let P,Q be two subdomains and divide all pixels in the image domain into five (disjoint) groups r, S, t, a, b that: (r, s) is the boundary of P and (s, t) is the boundary of Q-s and Q, and the interior of and Q are a and b, resp. The original sparse system Au v, A _ can be written:
s is the boundary shared between P
Arr
0
A
A
up
VF
Att
A Ass
0
Atb
Ugg
Vt
Aar
Ast
Ass
saa
Aab
Us
Vs
(4)
Aar
0
A as
aa
0
Ua
Va
torch. 1onall. (a (a1phaa) #0. sec.
0
Abt.
A bs
0
Abb
Ubb
Vb
A: sparse matrix
scipy. stapsesparse.... www.sp.....anggga..pppp spsolve (A, b) #1200 sec.
in which there 0 blocks because the interface separates
the pixels into non-adjacent groups, and Ass = ! +
The batched matrix inversion can be viewed as the first step of Gaussian elimination to remove the 9 pixels in the interior 3 x 3 block of each patch, and we already make major progress by removing approximately 9/ z 56.25% variables from the problem, while being 10'x x faster than SeiPy to solve the entire problem! This is strong evidence that the speed of sparse linear solvers has been significantly underestimated and parallelism has been under exploited.
A us (Q) V. = v! + v (Q) can be divided into contributions
(P)
from P and Q resp." Note that matrix A and system (4) are only for illustration purposes and never allocated as an actual matrix in our algorithm (see ŞE). (4) further becomes:
per 0
0 Q
Prs
Up
p'
te
Q
ts
ut
=
q"
(5)
psr
st
p S8
p
+
q"
3.2. Parallel block Gaussian elimination
by exercising Gaussian elimination to cancel Ua: Ub. where
36
37
3
40
P p
p
Are AraA Aar
A
AraA Aus
28
29
p
Aar Ass Aar
A Ans
&
19
20
2
22
R
2
26
3
Quee
Qual
Auu A A Abt
A
A bb Aba
Q
Ast Ast Abt
Aab Abs
Figure 3. Schwarz step removes the interior pixels for each patch.
p p
Ara A-va
q" q
Ath Vbb
A.b A bb Vb.
We introduce a procedure that we term "Schwarz-Schhe involution." " For simplicity of illustration, let us first con- sider a simple case: with only two subdomains, the original problem (4) is first reduced to (5) by the Schwarz step, and then via (9) to (12) by Schur step. Readers unfamiliar with numerical algebra should refer to SF.
CConeretely speaking, s= 14, 32, 40).
r= 10, 1,2,3, 39, 38, 36, 37, 18, 9), [10, 11, 12, 19, 20, 21, 30). t= 6,,,,,,8,1777 26, 35, 44, 43, 42, 41), b- (21lbb1(111111 15, 16, 23, 24, 25, 22, 33, 34).
2The values of ( and ( can be arbitrary, as long their summation the correct Ags. See $E for details.
4

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
using the last equations in $4 eliminate Ua. Up-
(9) the linear system for the joint domain D. Since B, represent the same node, they correspond to the same row/- column; the same applies to ô, ps. Now assemble the new system D by Schur "involuting" the sub-systems P, Q. When eliminating the "wire-frame" nodes ry (i.e. v in re- verse order), introduce the symbols to simplity notation:
Un Ub
Aaa 1
Va
Aar
Aas
A bb
0
Abt
Abs
U1
u1
(6)
p
p
p
Let
P :
par
p
P Paax
P Pas QKE
Oad
P
Pas
acc
Q :
Q
q qi
q"
(8)
P
Pa
P BB +
QuA
P
P
QAR
P Be
SE
X
Oxo
QAR
QAA
A
Ae
The symbol G _ H _ h means matrix G comes from entries of H but after some row/column-vee -wise permutation F PET ppxx, so G == FF = F-th, such that rows/- columns of P. Q list boundary pixels counterclockwise in the ordering of Figure See ŞE for details. The rationale be- hind dividing the domain is that once values the boundary wire-frame r, 8, f. known, the sub-problems solve for ua and Ub become independent: constructions of reduced systems P, Q are independent of each other, concurrently computed in a size-2 batch.
P Sox
P
P5s +
Q
QuA
P
88
Q Pup
P de
Pea
P
OcA
Pet
P
Pay
Po
u
P PBy +
Q
Pp + qk
Ug
Y ==
Q
y
qA
x
u
Poy + Q Q
P8 + q
u
P
Pc
u.
= Pya
P + JQQU
JQA p78
J
P,1
W = P oT + JTQU
: [p JTq.
where = JT is the reverse permutation matrix (see şE), whose action is to reverse the rows (or columns) of a matrix when being multiplied with from left (or right). By plugging in our new definitions, (9) becomes:
3.2.2. SCHUR STEP: MERGE ADIACENT DTN
Schur step merges two subdomains and Q into joint domain D, while converting their left-hand and right-hand sides (P, p) and (Q.q) into new left-hand and right- hand sides (D,d). Divide the nodes in P into contiguous subsets a, B, %, 5, c, such that
X Z
y
(10)
W
u
W
2, G = [4],
= [5, 7]. 5 = [8]. = [9, 10, 11, 12, 13, 14, 15). Di- vide the nodes in Q into contiguous subsets Ks, A4, 1/, such
Schur eliminate u., using:
that: x == [0]. A 2, 4, 5, 6, 7,8,9, 9, 10, 11], 1== [12], 13, 14, 15]. Under the indexing in Figure 4, D's bound- ary consists of the nodes corresponding ax x, Ō, Af- ter merging, B3 k. represent the same node, are ò, uz; Y represents the same set of nodes as y/ but reversely ordered.
u = W-1 (w
Zx)
(back-fill),
(11)
the system (4) to (5) (9) finally becomes Dx d, where:
D D (x YW 'z)
d YW-1 w (12)
In summary, solve the original problem (4) in the spe- cial case of two subdomains, after the Schwarz step that constructs the system matrices P. Q, our final algorithm of the Schur step first solves for the values x D D-1 d, then recovers u, the finally recovers u, Ub -vauuu in the interior of each patch from u, U. u.-values at the patch's boundary.
the solution at the interface- using (11) and
Figure 4. Schur step collapses subdomains P and into D.
The node grouping implies partitioning the matrix P, Q into submatrices, by dividing the rows and columns into subsets.
3.2.3. GENERAL CASES
Summarized in Figure 2. when there are 2 patches of 5 5, we apply the Schwarz step once, and the Schur step & times to recursively glue every two subdomains together. These steps play a role similar to LU factorization (Davis, 2006). See full details in SC, $E and Algorithm 1.
P
P Dry
P.s
P
P
P aao
P ag
P By
P
P ge
aa
Pya
P a
P Pyy
P.s
P
:
P
p
:
p.
Psa
Ps
P
Pas
Poc
Ps
P
P
P&
P
Pcs
Pec
p
Q
Q
Qui
q₂
The Schwarz step eliminate interior nodes in the 3x3 block is basically the same introduced before: for each patch P the interior nodes are eliminated and the system matrix for
Que
Qux
QAA
Qua
QA
Que
Qui
Q
Q
qu
qx
: q.
Q
Qui
qu
the remaining nodes r becomes (A (P) (PP) A-1A-11 (P)) 8A
5

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Pao
Paa
0.A
Pas
Par]
P
Too
0,00
0.AA
One]
Tux Ug
Pa
Paa
Pag
0.x
Pas
Pgc
P
0
Que
Q.
Quy
0.
pp qi
Psa
0xo
0xs
Pag
0xx
05x
Pao
0xs
Psc
0xr
P
+
0sc
Quic
que
Qux
Qxa
Qxy
Qui
0xe
0se
P
0.4
Pas
P
0.g
0.A
0.5
P7a PaB
P75
P74]
[P-r]]
JuQ
JTQ
[J
Py JTq.
Figure 5. The joint linear system that merges two subdomains.
The step transforms the domain to a wire-frame with 2 hol- low subdomains, each of which initially has one patch and is associated with a system matrix (like P. Q.... and a right- hand side (p. q. ...). Again, this stage takes a small fraction of overall runtime but removes the majority of pixels.
theory and computer vision, which means A is positive semidefinite (PSD) and A; = - EEigtt Aij. the Schur complements D, Q are also PSD, and discretizing the Dirichlet-to-Neumaae (DtN) operator, which is well-defined and enjoys numerous nice theoretical properties. See Wang et al. 2018); Levitin al. (2023) and references therein.
A Schur step further simplifies the wire-frame by climinat- ing pixels on some "edges" along the x (or y) direction for j that is odd (or even). The j-th Schur step (j = 1.... k) starts with 2 subdomains that each consists of 2 patches, and ends with 2 subdomains that each has 2 patches. To merge each pair of adjacent subdomains that marked in red and blue, we gather the system matrices for them, also denoted as P, Q, and apply the formula (12).
DtN also plays central role in the study of domain decom- position in numerical methods (Quarteroni & Valli, 1999). Consider the classic example of solving Laplace's equa- tion under Dirichlet boundary condition. After dividing the image domain into subdomains, however, the usual prob- lem arises: Dirichlet data at the interface between subdo- mains are unknown, preventing naive divide-and-conquee But whatever the Dirichlet data at the interface is, the so- lution is linear to it (through the Green's 's function), since we work with a linear PDE: the DtN operator encodes the solution operator that yields the solution for whatever Dirich- let data a simple intuition why DıN is central domain composition (Quarteroni & Valli, 1999). Unlike common practices in larger problems that implicitly realize DtN's (Liu et al., 2016), our algorithm main- tains the DIN/solution operators explicitly as dense matrix. Unlike general-purposs sparse solvers, we do not aim at a solver with minimal floating point operations or memory consumption and reducing fill-in sparsity during the solving, for very large-scale problems with billions of variables. In- stead, we focus on runtime for common image sizes, such as 1024x 1024 with around I million variables. As limitation: our method is more demanding in memory (see $D.4).
Tensorized representations of batched reduced systems are adopted as our method's ' data structure: we never main- tain the global system A as a sparse matrix like all other di- rect solvers. Instead, the j-th Schur step of our algorithm, the linear system is a 4-dim tensor (i) that a ( [k, :,: stores the left-hand side of the (k: I)-th subdomain. For j = (2i + 1), the j-th Schur step takes input ( of size (2 2 16- 24, 16 ), and outputs a (1(+1) of size (24/2-i-11 1 2 24 22, 24 :2); (i)) contains all reduced systems: P. Q are fetched batch- wise from ()) for all red and blue subdomains, e.g., in PyTorch syntax:
P [0:), :, :, ]] # even subdcrains Q = [1:)2 :, :] # odd subdemaing ins
Batched dense linear algebras are employed for automatic parallelisms in PyTorch. For superior performance, it is critical to batch matrix operations including inversion, mul- tiplication, sum, slicing into rows and columns, etc. For
A note on novelty. We remark that key ingredients of our algorithm exist in the literature. We essentially perform Gaussian elimination with a particular pivoting ordering induced by the nested dissection (George, 1973). Similar procedures were applied in early studies of distributing FEM solves among multiple processors, e.g., in civil (Farhat & Wilson, 1987; Farhat et al., 1987; Farhat & Roux, 1991) and ocean engineering (Rakowsky, 1999). However, our setting is different from scientific computing communities, which focus on much larger problems where inter-processor communication rises as a central consideration. Without recent advances in GPUs both in terms of high throughput and memory size, stirred by the needs of deep learning, applying the Schur formula to invert a massive amount of small dense matrices on a single device is either infeasi- ble or uncompetitive compared to alternative direct solvers. Solving a massive amount of dense systems on a single wivice-without needing cross-device communication, is made efficiently only recently, to rise as the best practice.
example, W, X, Y, Z are 4-dim tensors obtained by batch- wise slicing into the last two dimensions of P. Q, e.g., vJJ,, and the following
W P(: 4, ml + code computes the merged system across all subdomains: terch. ainalll 1g. inv (W)
3.3. Discussion
While our method is simple to describe- recursvvlly apply- ing the Schur complement formula in batches, fully under- standing the motivation and appreciating the superiority in performance and numerical behavior, empirically demon- strated in $4, requires familiarity with references in numeri- cal methods and elliptic PDEs. We highlight few aspects.
One question is why the matrices D, P. Q are numerically well-beh and the sum of their sub-blocks W is invert- ible. When the system A comes from discretizing an elliptic PDE, including the "Laplacian matrix" considered in graph
6

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Table 2. Uncertain runtime of iterative solvers. When solving an anisotropic or isotropic Laplacian (Dirichlet) system, our solver is not affected by anisotropy, while Ruge-Steuben AMG (Bell et al., 2022) can be much slower compared solving isotropic systems.
anisotropic Laplacian
isotropic Laplacian
Example
CUDA
SciPy
64
(Ie4))
AMG (le-7)
AMG (le 10)
AMG (le-4)
AMG (le-7)
AMG (le 10)
20/19
21354
143100
158.5
00M
12205
49487
53877
6559
7757
9202
1025"
4710
16512
36.45
175.6
2678
8112
13446
1388
1754
2236
513"
1036
2051
10.90
31.81
832.2
1658
2353
321.9
387.2
515.3
2574
234
355
5.82
7.46
203.0
385.2
556.0
91.0
117.3
347.8
4. Results and Validations
in contrast, iterative solvers, heavily dependent on A, in challenging cases significantly slow down or explode.
We compare our method with direct and iterative solvers and demonstrate significant advantages in runtime. The speedup from our approach can potentially be much more significant than what is reported here. Compared to standard solvers that are already highly optimized, our algorithm is naively prototyped in "TToch-rr-impllmmnntee as low- level operators or optimization the hardware level like CUDA operators will result in further acceleration. Ad- ditionally, as a direct solver, our back-substitution variant enjoys extra speedups in repetitive solves by reusing numer- ical factorization impoossiio for iterative solvers: Table 7.
Table 3. Runtime of a single linear solve on few 513* problems, to the relative tolerance of 10 for iterative solvers. No itera- tive solvers out-of- out-of-the-box can perform well on all tasks. "(s)": solvers henefit from the symmetry assumption when it holds, while other methods do not. "x""": time for iterative solvers should be multiplied by an extra factor of 2 20+ in some applications, since iterative solvers need to run 2 times to differentiate, and 20+ times in eigen solving, while direct solvers do the factorization once and reuse it. "x2": should be multiplied by an extra 2 due to transposed solve to differentiate. "NaN": the solver diverges or fails to yield small error. "*" indicates stopping a larger tolerance due to different stopping criteria; those > 10 5 are:
Comparison with direct solvers. As the performance of
direct solvers, especially our method, is less affected by A, Table I can represent the timing comparison for whatever A. Our method supports both single and double footing-pooit numbers (float-32 and float-64), while existing direct solvers only support float-64, or support float-32 without major speedups (cuDSS). Implemented in float-64, our method often achieves a relative error of 10- 16 - 10-12 even better than SeiPy (Table 9). Our method can use float-32 for further speedup, in which case the error range from 10-6-10-1 typical error tolerance that iterative solvers use, sufficient for many tasks. We evaluate the accuracy of the solution using the standard metric- the (relative) error tolerance:
1.19,
1.
1.4e.4
1.4ee4, 1.41e.4
Lap le-2
Lap le-4
Hel 0.01
Hel 0.1
Hel 1.0 R. Decony
Our-32
12.3
12.2
12.3
12.2
12.2
NaN
0ur-64
35.6
35.4
35.6
35.5
35,5
35,5
SciPy-LL LL
2208
2487
2333
2687
2133
3145
CUDSS-LDU
1011
1029
1040
1037
1020
1038
x2
Our 64 backsub
3.0
3.0
3.0
3.0
3.0
3.0
(s)
65.1
459
563
4142
64412
NaN
(x) MINRES
97.3
1319
(*) 1595
(*) 90541
(9) 11156
NaN
L.MRR
7288
426499
76928
103569
80617
12282
LSQRR
(*a)
2896
(*b)
215758
(*c) 84442
(*d) 20718
(*e) 16158
3290
biCGstab
299
1898
NaN
NaN
NaN
NaN
CGS
58.3
922
NaN
NaN
NaN
DIOM
228
12716
1484
4433
13221
NaN
x
FOM
1157
499930
61344
574742
NaN
NaN
QMR
70.8
570
463
1500
6849
NaN
B
86.5
634
635
2087
6878
NaN
GMRES
954
31017
56478
526975
4675551
NaN
DQGMRES
172
1029
1391
4323
12885
NaN
tol(x) := Ax - bblz/||bl2
(13)
FGMRES
30793
56665
537350
NaN
NaN
AMG
1815
12917
(*)3610256
(*)7840145
(*) 1313919368
NaN
GMRES+AMG
563
502
NaN
NaN
NaN
NaN
which is the stop criterion for most iterative methods. For direct solvers, we compare with SciPy (Virtanen et al., 2020) and CUDA (Nickolls et al., 2008)-the cuDSS LDU method which defaults to the METIS (nested dissection) ordering.
(5) PCG+AMG
408
303
509
(*) 26741
NaN
NaN
In Table 3, we test all methods on solving the image matting
Laplacian A- = L + AM (the L in $B.3, which is highly 10-2 2 10 10-4), 4), the Helmholtz equation A L RMM under wave numbers x² (000, 0.1, 1.0} (Figure 10), and random A-a sparse matrix whose nonzero entries Aij are i.i.d., uniformly drawn from [-1, 1]. We compare with out-of-the-box GPU linear solvers from CuPy (Okuta et al., 2017) and Julia (Bezanson et al., 2017), using the CUDA implementation in the Krylov.jl pack- age (Montoison & Orban, 2023). We are primarily inter- ested in indefinite and nonsymmetric square linear solvers including: LSQR, LSMR, CGS, biCGstab. DIOM. FOM, QMR, BIL.. GMRES, DQGMRES, and FGMRES. While we also report results on CG and MINRES, please note that the comparison may be slightly unfair: they benefit from the symmetry assumption that other methods do not. These oft-of-the-boo iterative solvers can slow down by many
anisotropic,
For iterative solvers, we compare with the algebraic multi- grid solver (Ruge-Stuben) from PYAMG (Bell et al., 2022). For SciPy, we use the "spsolve" function, which defaults to the COLAMD ordering and is backended by SuperLU (Li, 2005). While CUDA can be faster than SeiPy, it is still on the same order of magnitude and cannot achieve interactive rates like ours. Our experiments are collected on an Nvidia GPU A6000 Ada 48GB and an Intel CPU Xeon w9-3475X.
Comparison with iterative solvers. $F.2 extensively dis- cusses why iterative solvers are problematic modules in deep learning pipelines. Although we compare with itera- tive methods and demonstrate advantages, direct methods remain our primary theoretical point of comparison. Di- rect solvers are reliable and robust with consistent runtime;
7

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
orders of magnitude or fail converge for harder prob- lems. PCG+AMG: using Nvidia AMGX (Naumov et al., 2015), we compare with the preconditioned conjugate gra- dient (PCG) with AMG preconditioning; See $F for details of other AMGX solvers. No iterative numerical scheme the parameters that
runtime by only leveraging GPUs, without needing training.
A successful line of research achieves orders-of-magniuee speedup in solving the PDE (2) which is also known as Darcy flow, by learning the coefficient-to-solnnee to-suuttin oper- ator (Li et al., 2020; Wang et al., 2021; Li et al., 2024). Our solver can implement FEM to realize the same opera- tor, which becomes the mapping: c + = A 'b where A G diag(((, cl) We compare with the statc-of-the- art learning solver the task of Darey flows. We test on 1024 Darcy flow examples available with the Python pack- age "neuraloperator" from Li et al. (2024) resampled from 421 to 257². Our method as a direct solver does not need any training and solves the linear system within a tolerance of 10-12 12 (float 64) on all examples. For our method the relative "errors" are 0.869% + 0.0927%, smaller than that of 1.56% in Li et al. (2024). Note that the "error" of our method, i.e., differences between our result and the ground truth, is due solely to the use of different numerical schemes. Our method, if used as the simulation engine, is the ground truth, achieving comparable "inference" speed without any training. Ours at 257² has a running time of 5 - ms, com- parable with many neural auchitectures-Dare few milliseconds for some recent methods (Li et al., 2025).
works well for all problems- generate the problem slows them down significantly. AMG preconditioners improve the CG/PCG and GMRES for el- liptic PDEs but hurt for out-of-scope uses; see Figure 10. In contrast, our solver is wobblemiindependen
Anisotropic Laplacian. We solve Poisson equation with random or constant right-hand sides, with an isotropic or anisotropic kernel. As shown in Table 2, iterative solvers can slow down significantly with anisotropic kernels, while direct solvers are less affected. Table reports the time to solve Laplace equation with Dirichlet condition with either the isotropic kernel 11 or the anisotropic K₂ in Figure 9. The number of iterations increases for anisotropic kernels by a factor of 6x, with the same number of non-zeros (Os in K2 are treated as non-zero zero entries with value 0.0).
Parabolic PDEs and high-precision diffusion kernels.
flows take a
The parabolic PDE du(a.t)
dt = Au(z, t) is solved by our method. The standard practice to approximate u(x, f) is the implicit Euler scheme that solves the system A + fI + M. Heat geodesics is situation in which the high precision of direct solvers within tolerance of 10 10 becomes critical. Crane et al. (2017) approximate geodesic distance from a point To. using the solution under the initial condi- tion u(x,0) = 5(z - za)- The right-hand side is a delta function that is nonzero only at one pixel. The solution u(xx t). known as the heat kernel, a Gaussian-like func- tion centered at Ti, with infinite support. Then geodesics recovered from log ((xxt)) logarrith is of direct inter- ests. x has to be computed within tolerance of 10-19 or even 10 15 to produce non-negative heat kernel values to apply the heat method (Crane et al., 2017). Requiring a tiny tolerance makes iterative methods incompetent, while direct solvers can be simply applied. Figure shows the result of our method. Our method (float-64) solves the system to an error tolerance of 10-16, smaller than that of SciPy. Our method can succeed under extreme numerics even when SciPy fails, by decreasing the stepsize 0-
The results suggest the suitability of our method as an effi- cient solver-in-ther -erriio-h---oop importable by existing architectures, such as solution super-resolution (Ren et al., 2023). As learning-based solvers increasingly import modules from numerical PDEs, it is promising to generalize our involu- tion to procedures with learnable parameters for data-driven discretizations (Wang et al., 2019; Bar-Sinai et al., 2019).
5. Applications
Featured in Figure 1, due to the foundational role of the linear solver, numerous applications immediately benefit from method, removing the need to design problem- specific solvers. For A coming from different applications, our solver consistently demonstrates speedup over other direct solvers to a level similar to the case of solving PDEs in $4. Examples of sparse A encompass:
Mathematical optimization. When A, b are the Hessian and gradient, our solver implements the Newton's method. See $B.1 for detailed discussions and experiments. Despite theoretical superiority, Newton's method has very limited applications in practice due to the expensive Hessian solves, an obstacle that our Similar use cases include shape optimization, Gaussian pro- cesses, and inverse rendering (Nicolet et al., 2021).
4.1. A zero-shot baseline of efficient PDE solvers
Learned PDE solvers emerge in scientific tasks (Lu et al., 2019; 2021; Raissi et al., 2019; Lagaris et al., 1998). In addition to surrogate modeling (e-g., when the governing PDEs are unknown and must be inferred from data), a major advantage of learned solvers over conventional numerical methods is speed: the former can be 3 cndders-of-magniin faster, the main inspiration for our work. Surprisingly, we demonstrate zero-shot approach also capable of improving
speedup removes.
The Laplacian paradigms for graph theory (Chung, 1997). spectral clustering, manifold learning, image/geometty pro- cessing, and vision have been extremely successful; see Wang & Solomon (2019): Chen et al. (2021b) for a sur-
8

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
vey. Early examples include: Laplace equations in optical flow (Hom & Schunck, 1981), semi-supprviised learning with harmonic label propagation (Zhu et al., 2003), Per- ona- ana-Malik model with anisotropic diffusion (Perona & Ma- lik, 1990), and normalized cuts for segmentation (Shi & Malik, 2000). These are examples where PDE or linear al- gebraic solvers play a central role and A -1 or (IA+1) + +I)-+ P 1 instead of the forward operator A, is of primary in- terest. Sparse solvers play an even more prominent role in geometry modeling, where almost every algorithm solves some form of Laplacian systems (Solomon et al., 2014).
Our method is again orders-of-magnitude faster than SciPy: 2572 images, SeiPy has average runtime of 49,, slowing down compared to solving the Laplacian in Ta- ble 2, while it takes 7.5lms for ours (float-64). As expected, PyAMG fails for kernel 02
(2.z) since it is not symmetric.
(Differentiable) physies simulation. When A comes from discretizing a physics-based energy, our method immedi- ately accelerates the physical simulator. $B.2 shows an example. Robot learning engines (Du et al., 2021; Hu et al., 2019) and scientific machine learning (Sanchez- Gonzalez et al., 2020) are increasingly relying on solver- in-the-loop (Amos & Kolter, 2017) using CUDA sparse solvers, which can be replaced with ours.
Image processing: matting, segmentation, and editing. A can be a graph or matting Laplacian L that Li encodes the similarity between adjacent pixels. In this setting, L can be viewed as discretizing an anisotropic Laplace equation. Our solver L-'b can be used for spectral segmentation (Shi & Malik, 2000): see $B.4. Often a constraint term is added to the linear system A= == L+ A in image matting, and Poisson image editing (Pérez et al., 2003). While the con- straint term R.. makes the problem's 's difficulty unevenly distributed across the image and large \ makes the system ill-conditioned, our method in double precision consistently yields an error smaller than SciPy. See details in $B.3.
Handling different boundary conditions. Beyond han- dling the usual Dirichlet or Neumann boundary condition, we even introduce a midpoint reflective boundary condition to handle no-disk topology like spheres. See SC,$D.5.
6. Conclusion and Future Work
It is observed that in many areas-from from medical image anal- ysis to learning for solving Prasittadiiionnl optimization- based methods can be orders-of-magniuude slower com- pared deep learning. We identify that this is partially because the implementation of conventional methods is not specially tailored to a new setting of HPC (high performance computing) on a single workstation, enabled by GPUs. In interactive vision and graphics, considerable effort has been devoted to designing numerical schemes that avoid direct linear solvers. with the assumption that exact linear solvers are too slow. Whether or not the computation time is within tens of milliseconds makes an essential difference for the de- ployability of algorithms in video conferencing. self-driving cars, and virtual reality environment. For example, Bro- Nielsen & Cotin (1996) note "speed is everything" in virtual surgery. Our method eliminates the need for customizing solvers, and revives methods that used to be slow due to (repetitive) linear solvers in a straightforward manner. Our strong baseline for learning based method can serve solver/simulator-nn -ohth-lloo in robot PDEs, or an efficient learning or vision pipelines (Barron & Poole, 2016).
Diffeomorphic image registration. Diffeomorphisms, or smooth bijective maps between image domains (Younes, 2010), are not only a central concept in differential man- ifolds and geometry, but also foundational in image reg- istration (Beg et al., 2005). Recent results in variational suasi-harmonic maps (Wang al., 2023) demonstrate that the map x + (u(x), u(x)) is diffeomorphic if and only if V V. Vu(x)) = 0 (resp. v(x)) for some C(x), with appropriate boundary conditions. Intuitively, this condition states that each pixel must be placed about some weighted average of its neighbors' positions, which is enforced by a sparse linear system. To compute diffeomorphic image registration in Figure 7, we use Wang et al. (2023), which recursively solves anisotropic Laplacian systems. Replacing their linear solver with ours as the backend again leads to speedups. Registering pair of images of size 1024 x 768 requires solving 150 sparse systems se- quentially and takes seconds with our solver, compared to 8 minutes in Wang et al. (2023) if using a SciPy backend.
Due to the foundational role of sparse linear solvers, our method can potentially impact a broad range of applica- tions beyond what presented. Many classical algorithms can be viewed as one or more steps of the Schwarz-Schur involution
Geometry algorithms. $4 shows how to reduce geometric PDEs to anisotropic PDEs: evaluate L, M using curved surface's 's metric, pushed back to the plane. $B.5 handles non-disk topology with appropriate boundary conditions.
a linear Laplacian-like system (Bar- Pérez et al., 2003; Germer et al., 2020; Shi & Malik, 2000). For future work, we plan to design neural architectures based on the "involution." generalizing cotimization-basee Laplacian paradigms, and extend our method work with an arbitrary mesh or graph. Our exact solver can be converted to an approximate solver by down- sampling the DtN and integrated with iterative schemes.
ron & Poole, 2016;
Generalized deconvolution. As shown in Figure 6, we recover an image from its convolution with a spatially vary- ing kernel in the form (1), which reduces to solving linear system (Hansen et al., 2006). To our knowledge, the only algorithm applicable to this setting (equivalent to) using linear solvers (since solving PDE reduces to this problem).
9

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Impact Statement
Algorithm Schwarz-Schur Involution.
This paper presents work whose goal to advance the field of Machine Learning, Optimization, Computer Vision, and Scientific Computing. We call attention the fact that even for fundamental and well-studied task sosk-noliing linear sys- tems, speedups of up to x000 are still achievable, demon- strating the potential for substantial performance gains by developing novel methods like our solver to best leverage hardware capacity. We suspect that perhaps the advances in GPUs have been underutilized terms of improving con- ventional methods, with without relying on deep learning. Possible societal impacts of our work include improved ergy efficiency and reduced carbon footprint: significantly reducing computational time, methods like ours lower power consumption and contribute to more sustainable and envi- ronmentally friendly computing landscape.
(Overall solve: numerical factorization back substitution.)
Require: a x x22x25
Require: g( 122 x=2"
Require: x Ixxxxxx1 only for Dirichlet condition.
(b : 2H + 2W 4 is the number of boundary pixels.)
The Schwarz forward step 200--llorthmm 2:
a 1 , t ,0(),
for in range ()): do The Schur forward step (ooiioonal)-Algoritmmm 4:
(2++1) (2 g(2++ t a (21),
Acknowledgements
The Schur forward step (vertical): a (2++2) g(2it)) + a (21+1),
, g(2tt))
We thank Justin Solomon and Mike Taylor for insightful discussions. We thank the anonymous reviewers, especially for their comments on improving clarity, coff-containmett, and accessibility to the broader ICML community. YB supported by a fellowship from the Royal Society (NIF-R1-232460) Support for this research was provided in part by the BRAIN Initiative Cell Atlas Network (BICAN) the Brain
end for if Dirichlet boundary condition then x (k) is given
X
else
grants U01MH11700 17023 and UM1MH13099 Initiative Brain Connects consortium (U01NS132181, JUMINS132358-011 the National Institute for Biomed- ical Imaging and Bioengineering (1R01EB023281, R21EB018907, R01EB019956, P41EB030006). the Na- tional Institute on Aging (R21 AG082082, 1R01 1R01 AG064027, R01AG016495, IR01AG070988). the National Institute of Mental Health (UMIMH130981, RO1 MH123195, RO1 MH121885, IRFIMH123195),, the National Institute for Neurological Disorders and Stroke, JUU4NS135560 01, ROINS070963, 2R0INS083534, ROINS105820, 158200, R25NS125599), and was made possible by the resources provided by Shared Instrumentation Grants 1S10RR023401, 1S10RR019307, and 1S10RR023043. Much of the computation resources required for this research performed on computational hardware generously provided by the Massachusetts Life Sciences Center In addition, BF is an advisor to DeepHealth, a company whose medical pursuits focus on medical imaging and measurement technologies. BF's interests were reviewed and are managed by Massachusetts General Hospital and Partners HealthCare in accordance with their conflict of interest policies. Support for this research was provided in part by the National Institute of Diabetes and Digestive and Kidney Diseases (1- 1-21-DK-108277
either the improved implementation of Dirichlet other non-disk topology boundary conditions ($C.2)
(R)
or the naïve slow approach: (a) solved by matrix inversion using: a (A) 1xxxxxx66 g Rixxxxx
x + (a (1))
end if for in range(k/22, 0,-0,1111 -1) do The Schur backward step (vertical):
(26-1) (24)
The Schur backward step (horizontal)
Algorithm
x (2i-1) 1)
x (2))
end for The Schwarz backward step Algorithm 3:
(+) x (0)
return (+)) 122"/1x2"* x25x25
10

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
References
Bouaziz, S., Martin, S., Liu, T., Kavan, L., and Pauly, M.
Projective dynamics: Fusing constraint projections for fast simulation. Acm Transactions On Graphics, 33(4): 154, 2014.
Allaire, G. Numerical analysis and optimization: an intro- duction to mathematical modelling and nunerical simu- lation. OUP Oxford, 2007.
Bramble, J. H. Multigrid Methods, volume 294, CRC Press, 1993.
Amos, B. and Kolter,
Optnet: Differentiable opti- mization as a layer in neural networks. In International conference on machine learning, pp. 136-145. PMLR, 2017.
Bro-Niclsed M. and Cotin, S. Real-time volumetric de- formable models for surgery simulation using finite ele- ments and condensation. In Computer graphics forum, volume 15, pp. 57-66. Wiley Online Library, 1996.
Arisaka, S. and Li. Q. Principled acceleration of iterative numerical methods using machine learning. In Interna- tional Conference on Machine Learning, pp. 1041-1059. PMLR. 2023.
Bruna, J., Zaremba, W. Szlam, A., and LeCun, Y. Spec- tral networks and locally connected networks on graphs. arXiv preprint 1rXiv:1312.6203 2013.
Bar-Sinai, Y., Hoyer, S., Hickey, J., and Brenner, M. P. Learning data-driven discretizations for partial differen- tial equations. Proceedings of the National Academy of Sciences, 11661))113344-1544 2019.
Chen, J., Schäfer, F., Huang, J., and Desbrun, M. Multi- scale cholesky preconditioning for ill-conditioned prob- lems. ACM Transactions on Graphics (TOG), 40(4):1-13, 2021a.
Barron, T. and Poole, B. The fast bilateral solver. In European conference on computer vision, pp. 617-632. Springer, 2016.
Chen, Y., Davis, T. A., Hager, W. W., and Rajamanickam, S. Algorithm 887: Cholmod, supernodal sparse cholesky factorization and update/downdate. ACM Transactions on Mathematical Software (TOMS), 35(3):1-14, 2008
Bebendorf, M. Hierarchical matrices. Springer, 2008
Beg, M. F., Miller, M. I., Trouvé, A., and Younes, L. Com- puting large deformation metric mappings via geodesic flows of diffcomorphisms. International journal of com- puter vision, 61: 61:139-157, 2005.
Chen, Y., Chi, Y., Fan, J., Ma, C., al. Spectral methods for data science: A statistical perspective. Foundations and Trends@ in Machine Learning, 14(5):566-806, 2021b.
Chung, F. R. Spectral graph theory, volume 92. American Mathematical Soc., 1997.
Bell, N., Olson, L. N., and Schroder, J. Pyamg: Algebraic multigrid solvers in python. Journal of Open Source Software, 7(72):4142, 2022.
Cook, S. On the minimum computation time for multipli- cation. Doctoral diss., Harvard U., Cambridge, Mass, 1, 1966.
Belongie, S.. Fowlkes, C., Chung, F., and Malik. Spectral partitioning with indefinite kernels using the nyström tension. In Computer Vision Eeion-ECCV 2002: 7th European Conference on Computer Vision Copenhagen, Denmark, May 28-31, 2002 Proceedings, Part III 7, pp. 531-542. Springer, 2002.
Crane, K., Weischedel, C., and Wardetzky, M. The heat method for distance computation. Communications of the ACM, 60(11):90-99 2017.
Dao, T. Flashattention-2: Faster attention with bet- ter parallelism and work partitioning. arXiv preprint arXiv:2307. 0.30006699,,. 2023.
Berens, P., Cranmer, K., Lawrence, N. D., von Luxburg, U., and Montgomery, J. Ai for science: an emerging agenda. arXiv preprint arXiv:2303.04211 2023.
Dao, T., Fu, D., Ermon, S., Rudra, A., and Ré, C. Flashat- tention: Fast and memory exact attention with io-awareness. Advances in neural information processing systems, 35:1 16444466666599 2022.
Bezanson, J., Edelman, A., Karpinski, S., and Shah, V. B. Julia: A fresh approach to numerical computing. SIAM review, 59(1):65-99) 2017.
Blahut, R. E. Fast algorithms for signal processing. Cam- bridge University Press, 2010.
Davis, T. A. Algorithm 832: Umfpack v4. 3-an an unsymmetric-patter multifrontal method. ACM Transac- tions Mathematical Software (TOMS), 30(2):196-19))
Bolz, J., Farmer, I., Grinspun, E.,, and Schröder, P. Sparse matrix solvers on the gpu: conjugate gradients and multi- grid. ACM transactions on graphics (TOG), 22(3):917- 924, 2003.
2004.
Davis, T. A. Direct methods for sparse linear systems. SIAM, 2006.

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Dongarra, J., Duff, I., Gates, M., Haidar, A., Hammarling, S., Higham, N. J., Hogg, J., Lara, P. V., Luszezek, P., Zounon, M., et al. Batched blas (basic linear algebra subprograms) 2018 specification. Technical report, Uni- versity of Tennessee, 2018.
Frigo, M. and Johnson, S. G. The design and implementa- tion of ffww.. Proceedings of the IEEE, 93(2)216-231, 2005.
George, A. Nested dissection of regular finite element mesh. SIAM journal on numerical analysis, 10(2):345- 363, 1973.
Dongarra, J. J., Du Croz, J., Hammarling, S., and Hanson, J. J. An extended set of fortran basic linear algebra sub- programs. ACM Trans. Math. Soft.,, 14(1):1-17, 1988.
Germer, T., Uelwer, T., Conrad, S., and Harmeling, S. Py- matting: A python library for alpha matting. Journal of Open Source Software, 5:2481, 2020.
Du, T., Wu, K., Ma, P., Wah, S., Spielberg, A., Rus, D., and Matusik, W. Diffpd: Differentiable projective dynamics. ACM Transactions on Graphics (TOG), 41(2):1-21, 2021.
Gilbarg, D., Trudinger, N. S., Gilbarg, D., and Trudinger, N. Elliptic partial differential equations of second order, volume 224. Springer, 1977.
Duff, I. S. and Reid, J. K. The multifrontal solution of indefinite sparse symmetric linear. ACM Transactions on Mathematical Software (TOMS), 93):302-3225, 1983.
Grady, L.,, Schiwietz, T., Aharon, S., and Westermann, R. Random walks for interactive alpha-matting, In Proceed- ings of VIIP, volume 2005, pp. 423-429, Citeseer, 2005.
Duff, I. S., Erisman, A. M., and Reid, J. K. Direct methods for sparse matrices. Oxford University Press, 2017.
Grementieri, L. and Galeone, P. Towards neural sparse linear solvers. arXiv preprint arXiv:2203.0660 2022.
Ernst, O. G. and Gander, M. J. Why it is difficult to solve helmholtz problems with classical iterative methods. Nu- merical analysis of multiscale problems, pp. 325-363, 2011.
Guennebaud, G., Jacob, B., et al. Eigen: c++ linear algebra library. URL http://eigen. tuxfamily. org. Accessed, 22, 2014.
Farhat, C. and Roux, F.-X. method of finite element tearing and interconnecting and its parallel solution al- gorithm. International journal for numerical methods in engineering, (2(6)) 1991.
Hansen, C.,, Nagy, J. G., and O'leary, P. Deblurring images: matrices, spectra, and filterin. SIAM, 2006.
He, K., Sun, J., and Tang, X. Fast matting using large kernel matting laplacian matrices. In 2010 IEEE Com- puter Society Conference Computer Vision and Pattern Recognition, pp.
Farhat, C. and Wilson, E. new finite element concurrent computer program architecture. International Journal for Numerical Methods in Engineering, 149977177 1987.
IEEE, 2010.
Hestenes, M. R., Stiefel, E., et al. Methods of conjugate gra- dients for solving linear systems. Journal of research of the National Bureau of Standards, 49(6):409-436 1952.
Farhat, C., Wilson, E., and Powell, G. Solution of finite element systems concurrent processing computers. Engineering with Computers, 2:157-165, 1987.
Horn, B. K. and Schunck, B. G. Determining optical flow. Artificial intelligence, 17(1-3):185 1981.
Fletcher, R. Conjugate gradient methods for indefinite system. Springer-Veriee Berlin. Heidelberg. New York, pp. 73, 1976.
Horvath, C. and Geiger, W. Directable, high- resolutoolt simulation of fire on the gpu. ACM Transactions on Graphics (TOG), 28(3):1-8, 2009.
Fong, D. C.-L. and Saunders, M. Lsmr: An iterative algo- rithm for sparse suustes problems. SIAM Journal on Scientific Computing, 33(5):2950-2971, 2011.
Hovland, P. and Hückelheim, Differentiating through linear solvers. arXiv preprint arXiv:2404.17039 2024.
Freund, R. W. and Nachtigal, N. M. Qmr: a quasi-minimal residual method for non-hermitian linear systems. Nu- merische mathematik, 60(1): 60(1):315-333 1991.
Hu, Y., Anderson, L.,, Li, T.-M., Sun, Q., Carr, N., Ragan- Kelley, J., and Durand, F. Difftaichi: Differentiable programming for physical simulation. arXiv preprint NKKv:/910.0093333 2019.
Freund, R. W. and Nachtigal, N. M. An implementation of the qmr method based on coupled two-term recurrences. SIAM Journal on Scientific Computing, 15(2): 15(2):313-33) 1994.
Jeschke, S., Cline, D., and Wonka, P. gpu laplacian solver for diffusion curves and poisson image editing. ACM Transactions on Graphics, 28(5):1 2009.
12

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Krishnan, D., Fattal, R., and Szeliski, R. Efficient precondi- tioning of laplacian matrices for computer graphics. ACM transactions on Graphics (1OG), 32(4):1-15, 2013.
Litany, O., Remez, T., Rodola, E., Bronstein, A., and Bron- stein, M. Deep functional maps: Structured prediction for dense shape correspondence. In Proceedings of the IEEE international conference on computer vision, pp. 5659-5667, 2017.
Lagaris, I. E., Likas, A., and Fotiadis, D. I. Artificial neu- ral networks for solving ordinary and partial differential equations. IEEE transactions on neural networks, 9(5): 987-1000, 1998.
Liu, H., Mitchell, N., Aanjaneya, M., and Sifakis, E. A scalable schur-complement fluids solver for heteroge- neous compute platforms. ACM Transactions on Graph- ics (TOG), 1- 2016.
Lavin, A. and Gray, S. Fast algorithms for convolutional neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4013- 4021, 2016.
Liu, Y. and Roosta, Minres: from negative curvature detection to monotonicity properties. SIAM Journal on Optimization, 32(4):2636-266) 2022.
LeCun, Y., Boser, B., Denker, J. S., Henderson, D., Howard, R. E., Hubbard, W., and Jackel, L. D. Backpropaga- tion applied to handwritten zip code recognition. Neural computation, 1(4):541-555 1989.
Lu, L.,, Jin, P., and Karniadakis, G. E. Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of opera- tors. arXiv preprint arXiv: /900/990//.3// 2019.
Levin, A., Lischinski, D., and Weiss, Y. A closed-form solution to natural image matting. IEEE transactions on pattern analysis and machine intelligence, 30(2):228- 242, 2007.
Lu, L.,, Jin, P., Pang, G., Zhang, Z.,, and Karniadakis, G. Learning nonlinear operators via deeponet based on the universal approximation theorem of operators. Nature machine intelligence, 3(3):218- (0:))218229, 2021.
Levitin, M., Mangoubi, D., and Polterovich, Topics in spectral geomelry, volume 237. American Mathematical Society, 2023.
Montoison, A. and Orban, D. Bilq: An iterative method for nonsymmetric linear systems with quasi-minimum error property. SIAM Journal Matrix Analysis and Applications, 41(3): 145)::11 2020.
Li, H., Miao, Y., Khodaei, Z. S., and Aliabadi, M. An architectural analysis of deeponet and a general exten- sion of the physics-informer deeponet model on solving nonlinear parametric partial differential equations. Neu- rocomputing, 611:128675, 2025.
Montoison, A. and Orban, D. Krylov. jl: A julia basket of hand-picked krylov methods. Journal of Open Source Software, 8(89):5187) 2023.
Mullen, P., Tong, Y., Alliez, P., and Desbrun, M. Spectral conformal parameterization. In Computer Graphics Fo- rum, volume 27, pp. 1487-1494. Wiley Online Library, 2008
Li, M., Lian, X.-C., Kwok, J. T., and Lu, B.-L. Time and space efficient spectral clustering via column sampling. In CVPR 2011, pp. 2297-2304. IEEE, 2011.
Li, X. S. An overview of superlu: Algorithms, implementa- tion, and user interface. ACM Transactions on Mathemat- ical Software (TOMS), 31(3):302-325 2005.
Naumov, M., Arsaev, M., Castonguay, P., Cohen, J., De- mouth, J., Eaton, J., Layton, S., Markovskiy, N., Reguly, I., Sakharnykh, N., al. Amgx: library for gpu ac- celerated algebraic multigrid and preconditioned iterative methods. SIAM Journal on Scientific Computing, 37()): S602-$626, 2015.
Li, Y., Chen, P. Y., Du, T., and Matusik, W. Learning preconditioners for conjugate gradient pde solvers. In In- ternational Conference on Machine Learning, pp. 19425- 19439. PMLR, 2023.
Négiar, G., Mahoney, M. W., and Krishnapriyan, Learning differentiable solvers for systems with hard con- straints. In ICLR, 2023.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint arXiv:2010.08893 2020.
Nickolls, J., Buck, I., Garland, M., and Skadron, K. Scalable parallel programming with cuda: Is cuda the parallel pro- gramming model that application developers have been waiting for? Quene, 6(2):40-53, 2008
Li, Z.,, Zheng, H., Kovachki, N., Jin, D., Chen, H., Liu, B., Azizzadenesheli, K., and Anandkumar, A. Physics- informed neural operator for learning partial differential equations. ACM/JMS Journal of Data Science, 1(3):127,, 2024.
Nicolet, B., Jacobson, A., and Jakob, W. Large steps in inverse rendering of geometry. ACM Transactions on Graphies (TOG), 40(6):1 2021.
13

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Nocedal, J. and Wright, S. J. Numerical optimization. Springer, 1999.
Ren, P., Rao, C., Liu, Y., Ma, Z., Wang, Q. Wang. J.- X., and Sun, H. Physr: Physics-informed deep super- resolution for spatiotemporal data. Journal of Computa- tional Physics, 492: 12438, 2023.
Okuta, R.,. Unno, Y., Nishino, D., Hido, S., and Loomis, C. Cupy: A numpy-compatible library for nvidia gpu caleu- lations. In Proceedings of workshop machine learning systems (LeurningSys) in the thirly-first annual confer- ence on neural information processing systems (NIPS), volume 6, 2017.
Reuter, M., Wolter, F.-E., and Peinecke, N. Laplace- beltrami spectra as 'shape-dna of surfaces and solids. Computer-Aided Design, 38(4):342-366, 2006.
Saad, Y. Krylov subspace methods for solving large unsym- metric linear systems. Mathematics of compulation, 37 (155):105-126, 1981.
Ovsjanikov, M., Ben-Chen, M., Solomon, J., Butscher, A., and Guibas, L. Functional maps: flexible representation of maps between shapes. ACM Transactions on Graphics (ToG), 31(4):1-11, 2012.
Saad, Y. Practical use of some krylov subspace methods for solving indefinite and nonsymmetric linear systems. SIAM journal on scientific and statistical computing, (1):203-228, 1984.
Paige, C. C. and Saunders, M. A. Solution of sparse in- definite systems of linear equations. SIAM journal on numerical analysis, 12(4)6617-629, 1975.
Saad, A flexible inner-outer preconditioned gmres al- gorithm. SIAM Journal Scientific Computing, 14(2): 461-469, 1993.
Paige, C. C. and Saunders, M. A. Algorithm 583: Lsqr: Sparse linear equations and least squares problems. ACM Transactions on Mathematical Software (TOMS), 8(2): 195-209, 1982a.
Saad, Y. and Schultz, M. H. Gmres: A generalized mini- mal residual algorithm for solving nonsymmetric linear systems. SIAM Journal on scientific and statistical com- puting, 7(3):856-888 86356686,,.00 1986.
Paige, C. C. and Saunders, M. A. Lsqr: An algorithm for sparse linear equations and sparse least squares. ACM Transactions on Mathematical Software (TOMS), 8(1):
43 43-71,11 1982b.
Saad, Y. and Wu, K. Dagmres: direct quasi-mininat residual algorithm based on incomplete orthogonalization. Numerical linear algebra with applications, 3(4):329 343, 1996.
Pérez, P., Gangnet, M., and Blake, A. Poisson image editing. ACM Transactions on Graphics, 22(3):313-318, 2003.
Perona, and Malik, Scale-space and edge detection using anisotropic diffusion. JEEE Transactions on pattern analysis and machine intelligence, 12(7):629-639) 1990.
Sanchez-Gonzaler, A., Godwin, J., Pfaff, T., Ying, R., Leskovec, J., and Battaglia, P. Learning to simulate complex physies with graph networks. In International conference on machine learning, pp. 8459-8468, PMLR, 2020.
Pfaff, T., Fortunato, M., Sanchez-Gonzan Gonzalez, A., and Battaglia, P. Learning mesh-based simulation with graph networks. In International conference on learning repre- sentations, 2020.
Schenk, O. and Gärtner, K. Solving unsymmetric sparse sys- tems of linear equations with pardiso. Future Generation Computer Systems, 20(3):475-477, 2004.
Praun, E. and Hoppe, H. Spherical parametrization and remeshing. ACM transactions graphics (TOG), 22(3): 340-349, 2003.
Sezan, M. I. and Tekalp, A. M. Survey of recent develop- ments in digital image restoration. Optical Engineering, 29(5):393 4043, 1990.
Quarteroni, A. and Valli, A. Domain decomposition methody for partial differential equations. Oxford University Press, 1999.
Shi, and Malik, J. Normalized cuts and image segmenta- tion. IEEE Transactions on pattern analysis and machine intelligence, 2288):888-905, 2000.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physics- informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Sleijpen, G. L. and Fokkema, D. R. Bicgstab (1) for linear equations involving unsymmetric matrices with complex spectrum. Electronic Transactions on Numerical Analysis, ((11):2000, 1993.
Rakowsky, N.
The schur complement method as fast par- allel solver for elliptic partial differential equations in occanography. Numerical linear algebra with applica- tions, 66))497-- 1999.
Solomon, J. Numerical algorithms: methods for computer vision, machine learning, and graphics. CRC press, 2015.
14

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Solomon, J., Crane, K., and Vouga, E. Laplace-belrrami: The swiss army knife of geometry processing. In Sympo- sium Geometry Processing Graduate School (Cardiff, UK, 2014), volume 2, 2014.
Wang, Y., Kim, V., Bronstein, M., and Solomon, J. Learning geometric operators meshes. In Representation Learn- ing on Graphs and Manifolds 2019 (ICLR workshop), 2019.
Sonneveld, P. Cgs, a fast lanczos-type solver for nonsym- metric linear systems. SIAM journal scientific and statistical computing, 10(1):36-52. 1989.
Wang. Y., Guo, M., and Solomon, J. Variational quasi- harmonic maps for computing diffeomorphisms. ACM Transactions Graphics (TOG), 42(4))1-26, 2023.
Strassen, V. Gaussian climination is not optimal. Nu- merische mathematik, 13(4):354-356 1969.
Winograd, S. Arithmetic complexity of computations, vol- ume 33, Siam, 1980.
Trefethen, L. and Bau, II, Numerical linear algebra, 1997.
Yi, L., Su, H., Guo, X., and Guibas, L. J. Syncspecenn: Synchronized spectral enn for 3d shape segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2282-2290, 2017.
Van der Vorst, H. A. Bi-egstab: fast and smoothly con- verging variant of bi-cg for the solution of nonsymmetric linear systems. SIAM Journal on scientific and Statisticul Computing, 13(2):631-644 1992.
Younes, L. Shapes and diffeomorphisms, volume 171.
Springer, 2010.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention all you need. Advances in neural information processing systems, 30, 2017.
Zhu, X., Ghahramani, Z., and Lafferty, J. D. Semi- supervised learning using gaussian fields and harmonic functions. In Proceedings of the 20th International con- ference on Machine learning (ICML-033, pp. 912-919, 2003.
Virtanen, P., Gommers, R., Oliphant, T. E., Haberland, M., Reddy, T., Cournapeau, D., Burovski, E., Peterson, P., Weckesser, W., Bright, et al. Scipy 1.0: fundamental algorithms for scientific computing in python. Nature methods, 1613):211-272, 2020.
Vladymyrov, M. and Carreira-Perpinant M. The variational nystrom method for large-scale spectral problems. In International Conference on Machine Learning, pp. 211 2 220. PMLR, 2016.
Vladymyrov, M., Von Oswald, J., Miller, N. A., and Sandler, M. Efficient linear system solver with transformers. In Al for Math Workshop @ ICML, 2024.
Wang, S., Wang, H., and Perdikaris, P. Learning the solution operator of parametric partial differential equations with physics-infrrmed deeponets. Science advances, 7(40): eabi8605, 2021.
Wang, An efficient spectral solver for image segmentation and normalized cuts. 2025.
Wang, Y. and Solomon, Intrinsic and extrinsic operators for shape analysis. Handbook of numerical analysis, volume 20, pp. 41-115. Elsevier, 2019.
Wang, and Solomon, Fast quasi-harmond weights for geometric data interpolation. ACM Transactions on Graphics (TOG), 40(4):1-15, 2021.
Wang, Y., Ben-Chen, M., Polterovich, I., and Solomon, Steklov spectral geometry for extrinsic shape analysis. ACM Transactions on Graphics (TOG), 38(1):122 2018.
15

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Table of Contents
C.2.2
Solution 2
25
C.2.3 Solution 3
25
Introduction
C.3
Differentiable linear solvers and derivatives
25
1.1 Related Work
2
C.4
Transposed solve with reused Schur com- plement
26
2
Mathematical Preliminaries
2
C.5
Two settings in linear solvers
27
3
Schur Involution for Parallel Elimination
3
C.6 Involution: exact inverse convolution (spa- tially varying kernel)
27
3.1
A motivating example: sparse solvers too
slow?
C.7
Inverse problems, optimal control of PDEs, and consttoondd optimization
27
3.2
Parallel block Gaussian climination
C.8 Generalize to larger kernels
28
3.2.1
Schwarz step: decompose & initial-
ize DiN
4
D Discussions and Implementation Details
28
3,2.2
Schur step: merge adjacent DtN
5
D.1 Details the algorithm
28
3.2.3
General cases
5
D.2 Distributed representations of sparse systems
28
3.3
Discussion...
6
D.3 Details on experiment setups
28
D.4 Memory complexity....
28
4
Results and Validations
7
D.5
Midpoint reflective boundary condition
29
4.1
A zero-shot baseline of efficient PDE solvers
8
E Detailed Method Description with Necessary Background Information
5
Applications
8
30
E.1
Case study: two subdomains
30
Conclusion and Future Work
E.2
Graph algorithm perspectives
35
Extended Discussions on Related Work
17
E.3
Parallel Schwarz climination step
36
E,4
Details on Schur step
37
B Visualization, Applications, and Experiments
17
E.5
Final algorithm
39
B.1 Newton's method and interactive graphics
18
E.6
Illustration 3x7 image
42
B.2 Physical simulation and shape optimization
19
B.3 Image matting and segmentation
F
Additional Details
45
B.4 Fast backends for eigen solvers in spectral
F.1 Details on FEM and PDE discretization: ad- dibility of parallel linear elements
methods
20
45
B.5
PDEs domains with non-disk topology
21
F22 Issues with incorporating iterative solvers in deep learning
45
B.6
Timing details
22
E3 Differentiable sparse solvers: widely de- sired, yet absent
B.7 Numerical stability
23
50
B.8
Complexity analysis...
Method: Extended Discussions
24
C.1 Dirichlet-to-Neumaan factorization
C.2 Solvers for Neumann boundary condition
C.2.1 Solution
25
16

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
A. Extended Discussions on Related Work
Bolz et al. (2003) pioneer the application of GPUs accel- erate iterative solvers for computer graphics applications. Barron & Poole (2016) apply preconditioned conjugate gra- dients for fast bilateral filtering; their solvers operate in lower dimensional bilateral space, unlike our method which aims for pixel-space solvers. Krishnan et al. (2013) use a Schur complement formula to construct a Laplacian precon- ditioning scheme. For very small-scale problems, it well- known that dense direct solvers can be faster (Bro-Nielsen & Cotin, 1996). Problem-specific iterative solvers have been designed for image processing (Jeschke et al., 2009), physi- cal simulation (Horvath & Geiger, 2009; Liu et al., 2016), and geometry processing (Krishnan et al., 2013). Direct solvers such as CHOLMOD have been used in differen- tiable rendering and inverse geometry design (Nicolet et al., 2021).
We focus on direct solvers-exact solutions that are com- puted up to numerical errors accumulated due to machine precisions, rather than iterative ones that only yield approxi- mate solutions under given error bound. The performance of direct solvers is more stable, whereas iterative solvers can greatly slow down for difficult A and may fail to converge if used without care. In contrast to direct solvers, itera- tive methods such as multigrid solvers only yield solutions within a prescribed error tolerance. Even so, they are still much slower than ours in scenarios that they are specifi- cally designed for, and many orders of magnitude slower challenging examples. Direct solvers are praised for their robustness and reliability, while being much slower. We im- plement a direct solver to achieve both accuracy/reiaabiiity and speed, combining the best of both worlds.
Li al. (2023) apply preconditioner learned from data, and Arisaka & Li (2023) develop learning-based acceleration of iterative methods under a meta-learning framework.
Despite a well-studied problem, general-purposs sparse di- rect solvers, introduced decades ago and having remained stable since, are not specifically designed for today's 's vi- sion and learning applications. We focus on problems that fit on a single GPU, that nonetheless encompass common image sizes, avoiding considerations in cross-device com- munication that arise in larger scale problems in scientific computing. In addition, advances in GPUs sparked by deep learning shift the best practice towards algorithms exploit- ing parallelisms and that modern GPU BLAS kernels are extremely good at solving large number of small problems. We observe that the high throughput offered by modern GPUs has been under-exploited in direct sparse solvers at the common scale of problems in vision.
Domain decomposition. Our approach conceptually fol- lows the divide-and-conuuer strategy, dating back to the cel- ebrated Schwarz alternating method in domain decomposi- tion (Quarteroni & Valli, 1999). The Dirichlettto-Neumaan (DtN) operator is a standard object used to "glue" solutions between subdomains (Quarteroni & Valli, 1999). Typically, it is implicitly maintained through its matrix-vector prod- uct, realized by solving sparse systems on subdomains (Liu et al., 2016). In contrast, we make a different design choice to explicitly maintain the discrete DtNs as dense matrices.
Sparse solvers: missing in differentiable programming. Differentiable sparse linear or eigen solvers have been as popular requested feature in deep learning packages: please refer to $F.3 for a list of examples of community requests and discussions on sparse linear or eigen solvers. In a related but different effort, recent works (Grementieri & Galeone, 2022; Négiar al., 2023; Vladymyrov et al., 2024) learn to solve linear systems, for which our method can serve as strong baseline or an efficient training engine. Our analysis and discussion of FEM in $2 suggest that the learning-basee based PDE solvers and learning-base based linear solvers are in fact attacking the same underlying problem- yet they are cur- rently studied as separate fields with little cross-reeerencingg
Sparse linear solvers in vision & graphics. computer vision and graphics, previous developments have focused on iterative solvers, and we are not aware of prior effort to design direct solvers for Laplacian-like systems-a missing setting we address. A major component of contributions in optimization-basee methods is to develop problem-specific iterative solvers. Our approach can bypass the need for iterative schemes in many tasks.
Hierarchical approaches or multigrid methods are popular especially when A arises from elliptic PDEs (Allaire, 2007). However, in situations where the solution is non-smooth, such as the Helmholtz equation, common preconditioned it-
erative solvers can work very poorly (Ernst & Gander, 201 I).
B. Visualization, Applications, and Experiments
Incomplete LU factorizations are another family of precon- ditioning schemes. For positive semi-definite systems, in- complete Cholesky factorization can be applied (Chen et al., 2021a). Depending the properties of A, popular choices of iterative methods include Jacobi methods, Gauss-Seidel, Krylov subspace including (preconditioned) conjugate gradi- ents (PCG) for positive-definite A, and generalized minimal residual (GMRES) for un- un-symmetric A (Solomon, 2015).
Note that our goal in this section is not to compete on ev- ery task with of but to pick classic methods already requiring a solution to some Laplacian-like systems, collecting sparse matrices A covering a broad array of cases to test our methods with. Details on the timing are provided in SB.6.
17

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
kemel on smooth surface.
on surface.
Figure 8. Our solver accelerates by orders of magnitude the com- putation geodesic distances. The heat kernel computed by our method is shown (in log scale). This is task where iterative meth- ods struggle: often the kemell solution must within tolerance of 10 10 for accurate geodesic distances (Crane al., 201 7).
K
0.0
-1
0.0
Input by: image. X1 AT bb.
Recovered image:
Input bg:
Recovered image:
convolved image.
X2 A bbg.
Ka
-0.25
-0.5
0.25
Figure 6. Our solver is applicd to image deconvolution with A. We test with two cases of spatially varying kernels: (x.x) (in(88))" K1 + K2. K1, K2 E gaxa, ag is randomly drawn (z,, y)- where K1, K2 defined in Figure
-0.5
2
-0.5
Kemel (*1))) e x2
from b
from random b
Figure 9. Our solutions to isotropic and anisotropic Poisson, fair comparison, Os in K₁ are treated non-zero entries with value 0.0 all solvers.
(a) Moving image. (b) Target image. (c) Moved (ours).
Figure 7. For computing diffeomorphisms and image registration,
x² 0.01/64
xe 0.1/64
k² 1.0/64
previous optimization-base methods take few minutes high-res images. Switching the backend to our solver immediately reduces the runtime few seconds.
fowrs) = 35.6
fours) 35.6
jours) 35.5
GCG) 563
(CG) = 4121
(CG) = 6442
(PCG) 220
(PCG) = 26741
(PCG) inf
Figure 10. The solutions the Helmholtz equation under different wave number K, computed using our method. This is the example used in Table 3. Unlike the performance of both CG and PCG (with AMG preconditioning in $F) significantly degrades under higher frequency (larger or even fail to converge (inf). Without problem-specific precon- ditioner, PCG be less efficient than CG.
B.1. Newton's method and interactive graphies
method whose runtime consistent,
As an immediate benefit, fast solver potentially unlocks the capacity of Newton's 's method (Nocedal & Wright, 1999) relying Hessian solves for applications in interactive physics, computer vision, and graphics.
When = H is the Hessian and b = g is the gradient of some energy of function on the image domain, the solution
18

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
fixed boundary- minimal surface problem. The face is parameterized goal is to find the one minimizing the total area, which is nonlinear objective. Minimal surface area can be used as prior in shape completion. Newton's method sequentially solves H H.'. Thanks to the second-order convergence, using Newton's method with our solver as the backend, converges within 20 iterations, while gradient descent using L-BFGS takes more than 1000 iterations to get to the same accuracy. Replacing SciPy in Newton's 's method with our solver immediately yields speedup of 400x.
height field u(2; y). and the
Newton solver made efficient. Despite superiority in the convergence rate, we find that Newton's 's method has very limited applications in computer vision. Perhaps it is par- tially because of the expensive Hessian solves, a barrier that method removes.
This phenomenon is even more prominent in interactive computer graphics. In the seminal work Projective Dynam- ics (PD) (Bouaziz et al., 2014), as the central assumption when it comes to fast simulators, it is made explicit that Newton's 's method requires orders-of-mand offeenitu-- fewer itera- tions while solving the Hessian each iteration is expensive, making it less competitive.
Cnefficients
Solution
Difference with the "ground truth"
+ G diag(c.e)G
=
Figure 11. Our solvers realize the FEM coefficiemt-to-ooutton map.
H yields the descent direction in Newton's method. It is very common that H and g come from an energy in which pixels interact locally with their neighbors, rather than other pixels in the image. In this setting, the Hessian H has the same sparsity the Laplacian matrix apply our method; in fact, for some physics-based energy, H equals exactly to an anisotropic Laplacian, H +- L. In this case, the sparse linear solver is the sole bottleneck in applying Newton's 's method: constructing entries in H, g involves only per-element/pixel computation, which can happen in parallel.
The assumption that linear solve is expensive necessitates the development of alternative methods. A common pattern of the alternative fast methods is to reuse the factorization of the Laplacian matrix and run tens of thousands of such cheaper iterations, while Newton's method can converge in of iterations. This assumption also governs the best practice in differentiable simulation, e.g., in robot learn- ing (Du et al., 2021).
Our method challenges this long-standing assumption, po- tentially unlocks the capacity of Newton's 's method, and eases the design of efficient algorithms.
B.2. Physical simulation and shape optimization
Figure 12. The minimal surface problem searches for the surface whose minimized under prescribed boundary, having been a prior in, c.g., shape completion. The minimal surface computed by Newton's 's method, which is accelerated by x00 with our Hessian solver compared to the SciPy backend.
restpose
Figure 13. Some shape deformation produced by our method.
Physical simulation already can benefit from our method: Laplace equation governs physical phenomena in electro- magnetism, gravity, heat diffusion, fluid dynamics, and shape deformation. Incorporating the coefficients C(x)
400x x faster Hessian solvers for minimal surfaces. As an application, in Figure 12, we apply our solver to Newton's 's method to compute the surface that minimizes the area with
19

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
allows pushing forward the PDEs in some irregular geo- metric domain onto a canonical domain, and accounts for the distortion induced by the geometric mapping-liik how
ful direct solvers potentially follow a numerically unstable order when performing Gaussian elimination.
A purpose of this validation is to emphasize that viewed as Gaussian elimination, our method cancels variables in an order that is dependent on the entries in A, avoiding many numerical issues. A common misperception of our method is that it performs Gaussian elimination that is blind to data A. This is not the case. Our method does prescribes some block structures, which are independent of A. However, within each block, the canceling order of nodes are still permuted, leading to robust numerical behaviors. This is a critical detail hidden in how to compute matrix inversions we use mhorch.linalg.iii is backended by cuBLAS.. Thus, when inverting batch of dense matrices, it performs careful numerical schemes on each matrix individually to avoid usual pitfalls in numerical operations.
we compute distances on surfaces in $4. In addition, a straightforward way to utilize our solver is to write down some quadratic objective whose minimizer yields a linear system solve, The quadratic objective may come from the linearization of some nonlinear energy, such as the Hes- sian. Our solver currently only supports one variable per pixel; we plan to generalize it to multiple variables per pixel, e.g., two variables can represent the J. y coordinates of deformation. This can be done by generalizing the "DtN matrix" to be twice as large and viewing our method sim- ply as a block Gaussian elimination. Instead, we present simple trick to directly leverage our current solver in the setting of two variables per pixel. We put both the z-1/- components in complex number u, Uy + ity. For in- stance, we can squeeze the 2n x 2n block matrix arising in a conformal deformation energy (Mullen et al., 2008) into a single n X n complex-value valued Hermitian matrix that A = conj(AT) Without modifying the code, our solver applies to complex-valuee linear systems. Figure 13 shows the deformation obtained in this way.
function which
B.3. Image matting and segmentation
In image matting (Grady et al., 2005; He et al., 2010; Levin et al., 2007), the matrix A encodes both an affinity matrix L, Laplacian with values decided by pixels' color, and soft constraint term. It typically minimizes for an alpha function x that is smooth as measured by L, and constrained to be 0 or 1 (stored in k) in the foreground or background, respectively. R is a binary matrix selecting those constrained pixels.
Lxx + A||RTx x - kl| conntant
(14)
In this case, A := L + XRRT. = ARk. Most of the change in x is restricted to the band between two layers of the usev-provided masks (often called the trimap). We validate using the linear system obtained by PyMatting on the examples shown in Figure 14. We use the image matting framework PyMatting (Germer al., 2020) and its imple- mentation of the random walk Laplacian (Grady al., 2005) with a small 3 x 3 stencil as L. The error of our method in double precision is on the scale of 10 10-16, and is slightly smaller than SciPy in every image.
Input - L
Trimap - R. Alpha 4- our x
Our output
Figure 14. The input image is used to generate an affinity matrix A, the trimap for A, b. and the output alpha is A -b.
We choose to test it with our method for a particular reason: the difficulty of solving concentrates at the band, not evenly distributed across the image. The problem can be converted to easy task for geometry-awand PDEs if solving only within the band region with Dirichlet condition at known pixels. But (14) represents a common way to incorporate the constraint softly: add a penalty term weighted by a very large coefficient X. It makes some entries Aiy at constrained pixels very large compared to unconstrained pixels. Uneare-
B.4. Fast backends for eigen solvers in spectral methods
For many practical applications, often it is the low-frequency components of the Laplacian that are relevant, rather than the high-frequency ones: x + A magnifies the low- frequency components, while x + Ax does the opposite. This observation makes our method relevant for many com-
20

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
puter vision tasks.
eigenvalue problem for MATLAB, which the code of Shi & Malik (2000) relies, while it takes as short 55 milliseconds if using our linear solver, yielding accelera- tion of 455 This is also many orders of magnitude faster compared to the Nyström methods that only approximate the spectrum (Vladymyrov & Perreira-erppnan, 2016; Li et al., 2011; Belongie et al., 2002). This eigenvalue problem requires step of numerical factorization (11 ms for our solver), and tens of the back substitution steps.
In many vision applications, matrix A comes from pixel values and can be considered as surrogate representation of an image. but it is not clear what constitutes relevant b. Then, the eigenvalue problem Ax = Ax becomes the right computational model to extract information from A.
We demonstrate how common eigen solvers can directly benefit from the fact that linear solvers can be made much more efficient than what people thought previously. Design- ing standalone eigen solver is beyond the scope of this paper. The concurrent work (Wang, 2025) uses our efficient linear solver as the inverse iteration oracle to implement a fast eigen solver, reducing the design of fast eigen solvers to that of linear solvers. Our linear solver can serve as the atomic linear operator solvers. For example, our solver can immediately speed up the spectral image segmentation method, normalized cuts (Shi & Malik, 2000). In this case, we choose the matrix A as (I - L), where L is the matting Laplacian with the 3 x 3 kernel in $B.3 that becomes L after having its rows and columns normalized, following Shi & Malik (2000).
Recall that the back substitution can be much faster than the numerical factorization since it skips calculating any new left-hand sides. Thus, our linear solver can immediately accelerate by orders of magnitude the spectral segmenta- tion in the full pixel space. This is different from previ- ous methods like Barron & Poole (2016), where a major source of speedup comes from reducing the computation to lower-dimensional space. Future work might combine our approach with the dimension-reduucton strategy for further improvement, and leverage our solver as parameter-free differentiable layer similarly to Barron & Poole (2016).
A A11 for iterative eigen
B.S. PDEs on domains with a non-disk topology
Our solver applics shape analysis and geometry pro- cessing for curved surfaces as well. The heat kernel computed on two surfaces with the spherical topology are shown, accelerated by orders of magnitude using method. Using a corner reflective an octahedral parameterization (Praun & Hoppe, 2003), our method can solve geometric PDE that is defined surface with spherical topology ($B.5).
Figure 16.
boundary condition
Figure 15. Our linear solver, when being used by concurrent work to realize the inverse iteration oracle in iterative eigen solvers, is immediately transferred into an orders-of-me maggnniude faster cigen solver-455x faster than MATLAB's sparse eigen solver for 513 513 images, with applications to spectral image segmen- tation (Shi & Malik, 2000). The top 5 eigenfunctions of the astro- naut, cameraman, coffee, pepper, and lemur images, computed hy 20 calls to A -'x, are shown.
We can solve PDEs on a surface with a non-disk topology, by leveraging the surface parameterization techniques, such as the octahedral parameterization (Praun & Hoppe, 2003) to cut and map spherical topology surface onto an image domain. Our method can easily incorporate midpoint reflective boundary condition induced by the octahedral parameterization (Praun & Hoppe, 2003), by modifying the last step to fill in the Dirichlet boundary condition in the Neumann solver (C.2). It allows us to process surfaces with spherical parameterization. Figure 16 shows the heat kernel computed on curved surfaces that have the spherical topology using our method.
For examples shown in Figure 15, iterative eigen solvers require only 13 ~ 20 calls to A '',,, to compute the top cigenvectors within the tolerance of 10 3 ~ 10-2, 2, For an image of size 5132, it takes 25 seconds to solve the top-6
21

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Adding the ability to handle curved surfaces also makes our solver applicable to shape analysis and geometric deep learning, such as spectral methods for shape corre- spondence (Ovsjanikov et al., 2012; Litany et al., 2017; Yi et al., 2017), Laplacian eigenvalues for shape classi- fication (Reuter et al., 2006), and many other applica- tions (Solomon et al., 2014).
when solving elliptic PDEs, though our method does not optimize for aceuracy like the SuperI.U method does. We had suspected that the more careful pivoting in SuperLU might help to achieve a better precision, which nonetheless is not observed for solving Laplacian systems.
Table 11 reports the runtime when solving the matting sys- tem. Interestingly, SuperL from SciPy significantly slows down in the matting example, due to the unbalanced con- centration in spatial diffieulty, as we have discussed in $B.3. However, our method (float 64) still has an error smaller than SciPy.
2
3
4
5
6
Table 4. The experiments same to the ones Table 3 but for 1025 x 1025 images.
Lap le-2
Lap le.4
Hell Te:2
Hel 1.0
R-Decony
Ours64
175
175
175
175
175
Ours-32
36.4
36.3
36.4
36.4
SciPe-LUi
16001
15883
15899
16008
23175
CUDSS-iDD
4566
4592
5005
4554
4687
X2
Ours Gu-s-64-back-eu
3.1
3.1
3.1
3.7
10
20
30
40
50
CG
92.1
2250
21473
66 MINRES
113
1292
5037
33864
NuN
Figure 17. Our linear solver with the midpoint reflective bound- ary condition can be used as the reverse itcration for computing eigen functions of a surface with the the octahedral paramcter- ization (Praun & Hoppe, 2003). Some Laplace-Beltrane cigen functions of the Stanford bunny are shown. Eigen functions of the Laplace-Beltrami Beltrami operator are foundational for manifold data analysis and geometric deep learning.
LSMK
8244
1042012
749451
400585
1871 (")
LSQR
4520
NaN
NoN
118307
827
823
NAN
NON
CGS
95
538
NAN
NIN
NIN
DIOM
2263
17349
39340
347352
NuN
X
POM
5934
863361
667058
NGN
QMR
4268
2444
NEN
BILQ
150
4228
2094
NGN
N&N
GMRES
1028
279081
285348
NEN
NIN
DOGMKEES
13018
49156
RGMRES
1051
40096
NoN
NoN
NIN
B.6. Timing details
Table 5. The error as measured by relative tolerance corresponding to Table 4. Note that random deconvolution (R (R-Deconv) may not be a very well-defined task to evaluate accuracy, since the matrix A can be close to singular due to randomness.
We provide details on the timing comparison. Again, un- like the baseline solvers that are already highly optimized with official support from Nvidia, our method is simply prototyped in PyTorch with great potential for further accel- eration. We have used torch.compile to reduce the overhead of PyTorch, which makes Py PoTch code 1.2 to 1.8x x faster. In addition to Table 2, Table 8 and 9 report the average runtime (in milliseconds) and the error tolerance to solve an isotropic or anisotropic system with the Dirichlet condi- tion, respectively. The runtime of direct solvers including CUDA and our method are not affected by the anisotropic coefficients that make the problem much harder for iterative solvers as demonstrated in Table 2. SciPy does slow down for the anisotropic coefficients, but is still comparable to the isotropic case. In all experiments, we do not treat any matrix as symmetric even though they are.
Earor
Lap le-2
Lap 1o-4
Ie-2
R-Decom NaN
Qurs-32
1.44-005
1:21F-05
5.358-03
3.57F-04
Ours-64
1899-14
2.268-14
1.378-11
8.838-15
1.89-409
S:Py-LU
2.755-14
2.825-14
2.165-13
4.325-13
1.78E-15
CUDSS-LDU
2,038-14
2005-14
7,378-14
9.868-13
541F-13
Table 6. The error as measured by relative tolerance corresponding to Table 3. The back-substitution variant our solver shares the same error with (the full solve variant of) our method.
Laroe
Lap 1e-2
Lap 1o44
Hel 1c:2
Hel 1.0
R-Decane
Ours-32
1,52E-05
1.71E-05
1.94E-01
2.305-03
NuN
Qurs-64
2.715-14
3.03--14
5345-113
4.305-12
3:30E-06
SaPy-E1
4.069.14
4.19F.14
7.228.13
1,03F.12
2.06-115
CUDSS-LDU
2
3.04--14
1.56-113
2.868-12
6.24B-12
Table and 10 compare the average runtime when solv- ing the Dirichlet or Neumann boundary condition. Since reuse the Dirichlet solver to the Neumann problem, it can take slightly more time for the Neumann problem, but still comparable to the Dirichlet case. Future work may implement $C.2 for a better Neumann solver.
Table 7. The back substittion time for our method to solve Neu- mann problem. Surprisingly, the speed for 1025² image is similar to that of a 2572 image, suggesting that there are probably large room for low-level optimization for the scale of 2372.
Resolution
Time (milliseconds)
Ours (float-64)
257
2.56
Ours (float-64)
513
3.01
Surprisingly, the error tolerance of our method in float 64 consistently smaller than that of SuperL. from SciPy
Ours (float-64)
1025
3.13
22

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Table 8. Average runtime (in milliscconds) solve isotropic Laplacian with Dirichlet condition.
time
error tolerance
Example
CUDA
SciPy
ours
ours-64
SciPy
ours-64
ours
2049a
21297
138219
158.5
0OM
1.92e-14
1.07e-14
6.86e-6
1025"
4663
15526
37.4
172.9
1.27e-14
7.56e-15 15
5,28c-6
513"
1053
2223
10.9
31.9
8,52e-15
5,43e-15
3,87e-6
2572
235
334
5.11
7.55
5.49e-15
3.79e-15
2.53e-6
Table 9. Average runtime (in milliseconds) to solve an anisotropic system with Dirichlet condition.
time
error tolerance
Example
CUDA
SciPy
ours
ours-64
SciPy
ours-64
ours
2049
21354
143100
159.3
0OM
3.43e-14 14
2.03e-14 14
1.07e-5
1025
4710
16512
37.5
171.3
2.54e-14 14
1.57c- 14
7.97e-6
513"
1036
2051
10.9
31.77
1.36e-14 14
8.85e- 15
4.44e-6
257
234
355
5.82
7.43
2.24e-16
1.57e-16
2.45e-6
Table 10. Average runtime (in milliseconds) to solve an anisotropic system with Neumann condition.
time
error tolerance
Example
CUDA
SciPy
ours
ours-64
SciPy
ours-64
ours
2049
21396
186231
171.6
0OM
2.41e-15
1.47c-15
6.73e-7
1025
4779
17511
42.2
187.4
2.40c-15
1.47e-15
6.74e-7
5132
1038
2174
12.2
35.4
2.37e-15
1.47e-15
6.6le-7
2574
235
389
6.47
8.82
2.28e-15
1.47e-15
6.56e-7
Table 11. Average runtime (in milliseconds) to solve the matting system with Neumann condition.
time
error tolerance
Example
CUDA
SciPy
ours-64
SciPy
ours-64
2049
22601
244125
OOM
2.89c-16
1 1,80c-16
1025*
4721
16969
187.6
2.73e-16
1,77e-16
5132
1046
6223
35.8
2.50e-16
1. 1.65e-16
257
232
1045
8,79
2.24e-16
1.57e-16
B.7. Numerical stability
well-defined objects in some sense. It will be interesting for future work to study the numerical behavior of our method thcoretically. For Helmholtz equations that are not elliptic, our errors in float-64 64 are still comparable to the baseline methods. For the deconvolution with random kernel, our method can have larger error; it can fail if using float-32, but note that this is not well-defined task, since under ran- dom spatial kernels with large probability the sub-system in some patch can be quite close to singular to make the global system also close to singular, so it is expected that SuperLU with global pivoting will be the most numerical safe option. Still, in numerically challenging situations, method with float-64 can offer accuracy that is sufficient for deep learning systems.
Our method prescribes structure of pivoting and elimina- tion ordering for embarrassingly parallel elimination, which means that the row- and column-wise pivoting is limited to switching rows and columns corresponding to nodes within the same sub-domain. Theoretically, such restriction on the pivoting scope could compromise the numerical stabil- ity for some problems. However, we consistently observe that when solving Laplacian systems or elliptic PDEs, our method even has a smaller error consistently than that of SuperLU or cuDSS (METIS ordering), which employ more conservative pivoting strategies. These should be related to the fact that the sub-systems that our method explicitly main- tains are discrete Brrcclee-to-Neumane operators which are
23

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
The row-wise and column-wise pivoting in the LU factor- ization is meant to improve the numerical stability of the right solve, not for the left solve c cA equivalently, the transposed solve A cc, for the back propagation step. Our Birichlet-to-Neumeee rectly to both the left and right solves, without the need of a separate factorization for AT.
solution from many small dense matrices are very efficiently inverted.
In summary, our "Schwarz-Schur involution" recursively merges subdomains and eliminates nodes/pixels at the in- terface between subdomains, and updates the maintained system matrix by applying a Schur complement formula.
factorization can be applied di-
Following the nested dissection hierarchy (George, 1973). recursively divide the domain into two subdomains and reduce the degrees of freedom: nodes at the interface be- tween subdomains have duplicated copies in each subdo- main. This yields a quad-tree in 2D: as shown in Figure 18, recursively applying this step reduces the number of subdo- mains following the sequence: (512,512) - (256, 256) + (128, 128) + (64, 64)... + (1, 1). This procedure visually corresponds Figure 18 but in the reverse order.
B.8. Complexity analysis
To emphasize, unlike the standard numerical analysis, in our setting, we are less concerned about the asymptotic complexity of the numerical algorithm, but rather the actual runtime common image sizes. These are very small- scale problems by the standard of scientific computing, having small constant in the complexity can be equally or even more important than the asymptotic rate. The constant and thus the best algorithm are hardware dependent, so the actual runtime as we have reported in the paper should the only criterion, in addition to which we still provide an asymptotic complexity analysis.
C.1. Dirichlet-to-Neumaan factorization
We call the Schwarz and Schur steps as the Dirichlet-to- Neumann factorization, a numerical factorization procedure similar to the LU factorization. The role of the dense matri- ces we saved in axl)) and the overall hierarchy is analogous to the L and U factors in standard LU factorization (Davis, 2006). In the Schur step, small patches are recursively merged, during which the new left-hand side-the DtN matrices for the merged subdomains are constructed by in- verting sub-block of the DtN matrices P, Q at subdomains using the equations in $3.2.2; every matrix there will be saved for later usage-in in the back substitution stage to ap- ply to the right-hand side.
The time complexity of our algorithm is dominated by the parallel matrix inversion in the Schur steps:
log v
2
12)3 3n
3/22 +
3/2
21.5
21.5
1=1
(15)
where we have assumed the cubic cost of inverting dense matrices and that the parallelized inversions take the same time as inverting one matrix of the same size.
The output in the last Schur step is the tensor (*) of size (1, b, 6). where again is the number of pixels on the border of the image domain. Denoting (*) [0. 0, :,:] as the
C. Method: Extended Discussions
dense matrix D Roxb the original linear system Ax b at the last level of the hicrarchy becomes Du d, where D E hxd and d E R
Back substitution is the standard step to fill in solutions for all rows of x. For Dirichlet boundary condition, we directly fill in rows that correspond to boundary pixels, x| x|I++++ g. Then values at the interface can be recursively backed filled by solving (11), so the values in x at the wire-frame of the image domains can be obtained. Finally, we recover the values in x that correspond to the interior 3 x 3 block of each patch.
Figure 18. After the Schwarz step that eliminates interior nodes/pixels for each the subdomains, the domain becomes a "wire-frame" that is progressively simplified in the Schur steps. The Schur step recursively collapses adjacent hollow subdomains by applying Schur formula merge two DtN matrices.
Condense sparsity and restrict the number of variables. For 103 x x image. its discrete Laplacian A is sparse matrix of 10 10 which is much too large deal with as dense matrix. If we can instead somehow work with one (dense) 10 x 103 matrix and a few smaller ones, then we can leverage the efficient dense linear algebraic kernels BLAS (Dongarra et al., 1988; 2018). Indeed, the largest dense matrix we have to invert is at the scale of 10 x 103. Our algorithm is a procedure to assemble the final
C.2. Solvers for Neumann boundary condition
For Neumann boundary condition, there are multiple ways to directly apply our Dirichlet solvers: 1) solve some DIN system for the final domain to fill in the missing Dirichlet data and call the Dirichlet solver ($C.2.1); 2) the actual option we use in the paper for better performance: modify the first Schwarz step to also eliminate pixels the boundary
24

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
of the domain (SC.2.3).
We use the Dirichlet boundary value problem to explain our method as it arises naturally: our method is a hierarchical construction of DIN matrices that are directly applicable under the known Dirichlet condition. The linear condition (11) that we enforced at the interface says that the Dirichlet data should be chosen so that the resulting Neumann data on either side of the interface must match each other- flux into the interface should equal to flux outward the other side of the interface.
Figure 19. The wire-frame structure resulted in the Neumann clim- nnttoon, obtained by modifying the Dirichlet elimination in Fig- ure 18.
C.2.3. SOLUTION
While Solution should the theoretically optimal way in this paper, it requires implementing new edge collapse procedure and handling boundary patches differently. In
Most tasks in computer vision have the natural (Neumann) boundary condition. The Neumann boundary value prob- lem can solved in a similar fashion. In principle, one can recursively work with the ciinnaanntt--iiiiiii ma- trix instead of the Dirichlet-to-Neumann matrix, which are (pseudo) inverses of each other. Then we fill in the miss- ing Neumann boundary condition in a hierarchical fashion. Instead, we present simple solution that directly leverages the Dirichlet solver. The goal is to minimize effort and reuse Dirichlet solvers.
Figure 20. The final hierarchy for Neumann solver.
order to reuse the Dirichlet solver like Solution does, propose further simplification: the boundary of the im- age domain we only eliminate wire-frames without two end points. The nodes at the boundary of the image domain are still kept, though any rows and columns corresponding to them in the left-hand side are zeros, and any rows corre- sponding them the right-hand side are also zeros. Then, we can reuse the same equations in $3.2.2. This solution has the advantage that at the last level, the system solve will be smaller than D: while D has the same size, we know that most rows and cols in D are zeros and can be thrown away. So we only need to solve for, e.g., D 1:4.d,,:4.66. d1:45 when patch size is x enlarged to 5 x so (+) selects one from every 4 pixels along the border. The new system matrix D 1:4:0,1o 1:46 E TR0/1x611 records the interplay among the scatter points on the border of the image domain.
C.2.1. SOLUTION 1
The original linear system Ax b the last level of the hierarchy becomes Du d, where D E R and d E R From the Gaussian elimination viewpoint, the DIN matrix D is nothing more than the linear system matrix with the interior nodes in the image eliminated. Thus, the Dirichlet boundary condition that is missing in order to apply the Dirichlet solver for the Neumann problem can be simply found by solving D --.d.
Indeed, this immediately gives Neumann solver. How- ever, there is one drawback with this strategy of reusing the Dirichlet solver in a Neumann problem. The system matrix D 16x as large as the largest matrix that one has to solve the Dirichlet problem. (D 4x as large in rows and columns so a total of 16). This makes the Neumann solver for large-scale problem noticeably slower than the Dirichlet solver, even though in principle they should be equally difficult.
C.3. Differentiable linear solvers and derivatives
Differentiable linear solvers are useful in many scientific applications (Hovland & Hückelheim, 2024). Following notations of Wang & Solomon (2021), let us derive the closed-form formulas for the partial derivatives x = A'bb: we need to obtain 8x/ab and 8x/0A.
C.2.2. SOLUTION 2
As shown in Figure 19, to avoid the need to invert a larger matrix, we design pre-processing step to further eliminate nodes at the beginning.
It convenient to introduce the notation to flatten the matrix A into a vector:
a = FLATTEN(()) E T (
In general, it better to eliminate nodes as early possible, since that can happen in parallel, rather than deferring it to the final stage to increase the size of the large dense matrix. In some sense, the Neumann boundary value problem is easier than the Dirichlet counterpart, since the former allows node elimination at an earlier stage.
(16)
which stores the the nonzero entries of A in a vector a, where (===) is the number of nonzero entries, such that the matrix product
A D[)
(17)
25

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Figure 21. The node elimination order conesponds to Figure 19. An ideal elimination order for the Neumann boundary conditions.
The node elimination order corresponds to Figure 20. The elimination order actually use for the Neumann boundary conditions for the case implementation.
Figure 22.
recovers the matrix A in the coordinate format, where J2 [=""")/", Here (-AG(-) converts a vector into a diagonal matrix and notice that
algorithm already gathers all ingredients to simply apply the left multiplication A -1 in licu A Tc. This can lead to an extra 2x speedup of our method compared other direct solvers, when differentiating linear solvers.
DIAG( DIAG(V)u=
(18)
C.4. Transposed solve with reused Schur complement
for two vectors u, v, where O v is the element-wise mul- tiplication of two vectors. Namely, their i-th rows, (i, :) and J (1, :), are one-hot sparse vectors that put the entry a; in the correct row and column in the matrix.
Schur complement reduces solving the x 2 block matrix
x
x
(23)
Then, any infinitesimal deviation (šA, 8b, 8x) from the stationary (A, b, x) such that Ax b must satisfy that:
Z
W
z
w
to the solution of a smaller matrix Dx = d, with the formula to update the left-hand side (LHS):
(fA)x + Adx = Hb DIDAG(6 + + J2 = öb
D = X- (X-YW
(LHS update),
(24)
J[DIAG(J2x)áa + Adx åb
the formula to update the right-hand side (RHS):
So we have:
--y-YY--M (RHS update),
(25)
8x/0a = A - [[
(19)
(20) and the formula for back substitution (after receiving calculated from the coarser level):
8x/ab abb = =A
Thus, provided with F, the gradient of the energy F(x) w.r.t. x, the chain rule gives the gradients V E and E.
z = W W Zx]
(back-fill).
(26)
V.E = (J2x) @ (JI(AT)
(21)
Now let us consider the transposed solve:
((T) VxE
(22)
X
Y
[
y
(27)
Z
W
z
w
Remark: to evaluate the gradients, we need to solve the system AT instead of A in general, unless A is symmetric. Our method an exactly "symmetric" algorithm for A and
or
AT: For our method, it is not necessary to run the Dirichlet- Neumann factorization one more time for AT, since our
X
Zr
y
(28)
YT
WT
W
26

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Instead of directly calculating the solution using =D-id
generalized deconvolution with spatially varying kernel, but they are usually not called deconvolution task. To distinguish from the common setting of deconvolution, we term our process Schwarz-Schur involution (Sinv2D), the exact inversion of the convolution with a spatial kernel. We call the set of dense matrices that are applied across all subdomains as the Schwarz-Schur involution kernel at that level, which are stored in the tensor (
D : - TYT) = DJ à: W Tw,
(29)
we would like to reuse W, D the updated left-hand side for the untransposed system. Though the formula D DT = (D-1)T allows avoiding the second matrix inversion with a cheaper matrix transpose, we adopt an even more efficient approach that avoids transposing the dense system matrix W. D.
Theoretically, our method can enjoy further improvements when the kernel a(*,)) is a spatially constant matrix (the timing reported in the paper does not take advantage of this: we always treat (**.)) as spatially varying even it is constant for fair comparisons). In this case, at each subdivision level, the S-involution kernels are the same across the domain. Thus, there is no need to maintain a dense matrix for each subdomain individually; instead, we can maintain one S-involution kernel that applies to every subdomain.
x= [(y WwW--ZZZD ZZDD11
(RHS update: left)
(30)
z (ww - W-
(back-fill: left). (31)
These left solves are straightforward to apply since we main- tain W-1 _D- as dense matrices.
What is even better: for a prescribed kemel, the S-involution kernel has a closed-form formula that can be derived and calculated ahead of time, bypassing the need of numerical factorization. When A is the isotropic Laplacian, the result- ing DIN matrix approximates continuous DtN operator which has known low-rank approximation, the structure that may be leveraged in future work (Bebendorf, 2008).
In summary, in the final algorithm to solve the transposed
system A -b,, we simply replace the updating formulas (25).(26) with (3)).(3), whenever a Schur complement elimination step is applied.
C.5. Two settings in linear solvers
C.7. Inverse problems, optimal control of PDEs, and PDE-constrained optimization
We clarify that there are two different settings when solving (multiple instances of) x A bb::
Our method accelerates not only the forward linear solver, but also the "backward" process of differentiating into the linear solver, which is useful in, c.g. inverse problems, opti- mal control of PDEs, and PDE-constrained optimization.
A is unchanged and fixed. This is when the common trick can be applied, e.g. in interactive computer graph- ics (Bouaziz et al., 2014; Du et al., 2021): pre-factorize = BTB once and at runtime reuse the fixed factoriza- tion for speedup, only performing back substitution. A is changed over iterations, solving systems with co- efficients unknown in advance. This is more common situation for nonlinear systems. For example, in Newton's method, A is the Hessian that changes over iterations.
For x A- bb, back-propagating the gradient w.r.t. x to that w.r.t b is straightforward since dx A A-1ob, so will focus on how to differentiate into the coefficients of A. There are two ways:
In fact, our method implemented in PyTorch, is already end-to-end differentiable w.r.t. entries of A. However, it is generally advised against differentiating through the low- level implementation details: see Hovland & Hückelheim (2024) for empirical studies and references therein for dis- cussions. We also observe that indeed this naive approach is inefficient to implement the backward step to obtain gradients.
In this paper, we focus on the second setup, where the system must be solved from scratch and the numerical fac- torization stage-which is the main bottleneck-is therefore included. Nonetheless, both settings can benefit from our method significantly.
C.6. Involution: exact inverse convolution (spatially varying kernel)
Alternatielly, as described in SC.3, without relying on the automatic differentiation feature of PyTorch, we can explicitly calculate the gradient by solving the adjoint system A -Tc, where e depends on x, similarly to the dif- ferentiable bilateral solver (Barron & Poole, 2016). When A is symmetric, this involves solving a linear system with the same left-hand side, see, e.g., Wang et al. (2023).
Many classical algorithms can be viewed as single step of Schwarz-Schur involution-solfe like system. Image filtering (Barron & Poole, 2016), edit- ing (Pérez et al., 2003), matting (Germer et al., 2020), and segmentation (Shi& & Malik, 2000) are classical examples. The linear solving in these methods can be viewed as the
a linear Laplacian-
27

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Thus, our solver immediately speeds up the optimization of some distributed parameters A, thich-writted in the equivalent form a(*.))_can represent convolutional ker- nel, local material property, or the Jacobian of a deformation field. In the literature of inverse PDEs, this is often referred to as the problem of distributed parameter identification.
D.2. Distributed representations of sparse systems
In our method, we never have to construct global sparse matrix as other direct sparse solvers do, leading to extra sav- ing in time. In the initialization step the Schur complements are stored in distributed fashion.
In $5, we apply the differentiable solver to search for the A that minimizes an objective E(-) This an inverse problem
D.3. Details on experiment setups
Compared to Python, the Julia ecosystem has much better support for sparse linear algebraic CUDA APIs that are criti- cal in scientific computing. For this reason, when comparing our method with CUDA, we use the up-to-date Julia binding that provides APIs backended by the most recently released cuDSS with cuSolver, cuSparse and CUDA. cuDSS is the ofate-of-the-ar official CUDA library that supports sparse linear solvers including sparse LU (LDU) factorization, be- coming our direct point of comparison.
min E(x)
(32)
s.t. x = A b
In $5, the linear system comes from FEM: A = GTC to yield PDE-constrained optimization (Wang & Solomon, 2021; Wang et al., 2023), and the goal is to optimize C that encodes a deformation field that minimizes the image matching loss.
D.4. Memory complexity
C.8. Generalize to larger kernels
We emphasize that the reported memory for our solver is the maximum GPU usage for its workspace storing all in- termediate matrices and can be released after the solution is obtained: if the solver is included as a layer in the neural networks, the memory can be released after the forward pass (though storing the workspace can skip the numerical factorization stage the back-propagation step to improve performance). For example, if the solver layer is included 10 times in the neural architecture, the peak memory usage for the model is same to including the solver only 1 time. Our solver currently treats symmetric A as asymmetric ones: future work can leverage the symmetry to save half of the memory usage.
Our approaches can be generalized to a larger kernel size, e.g., to deconvolute 5 x 5 or 7 x convolutional kemels,, with a more complicated implementation. In the current 3 x 3 case, one layer of boundary pixels can separate two subdomains P. Q, and it will require two layers of boundary pixels in the case of 5 x 5 to separate two subdomains P.Q. (or 3 layers of boundary pixels in the case of 7 x 7). Similarly, the boundary pixels (at the border of the whole image) in the last Schur step will have two layers of pixels. The parallel elimination procedure will be similar, but will have to account for the fact that the "rame-frame" has two layers of pixels.
In fact, in some sense, our current method already supports kernels larger than 3 x 3. The actual constraint that we have is that every pixel can only contribute to pixels in the same x 5 patch (and recall that pixels at the patch boundary belong to multiple patches). Thus, the convolution window for pixel can cover the entire 1/2/4 patches it belongs to: for example, pixels at the patch boundary can use a local window of 9 x 9 or 9 X 5, and it is 5 x 5 for an interior pixel, though the window may not be centered at it. Also, recall that the x 5 patch size is a hyperparameter that we are free to change arbitrarily in the current method.
Table 12. Memory usage (MB) comparison. Note we do not op- timize our memory usage and there is likely large room for im- provement. Note that if used in differentiable solver setting, the memory of CUDA might have to be multiplied by an extra factor of due to transposed solve, while there is no need for our method. (*: while it fits in memory to run the code, using "torch.compile" for an optimized runtime exceeds the limit.)
GPU memory usage (MB)
Example
CUDA
ours-32
ours-64
40972
33330
OOM
OOM
2049a
8104
16094
32188 (*)
1025*
2234
3880
7244
D. Discussions and Implementation Details
513²
880
890
1736
D.1. Details on the algorithm
257
548
242
420
A hyperparameter in our method is the size of patches, which we choose as 4 x 4 (enlarged to 5 x 5 for dupli- cating boundary pixels at interfaces), except for the image of size 2561² where we choose the size to be 5 x 5 (enlarged to 6).
Our method achieves speed at the cost of memory usage. Table 12 compares the GPU memory usage with the CUDA solver. Note that our current implementation is generous in memory usage and we believe there is likely plenty of
28

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
room for improvement. We save every intermediate dense matrices allow quick experimentation of ideas, and do not use any PyTorch in-place operations, which we find may cause counterintuitive behaviors for sliced matrices, making debugging challenging.
such that
Ua
Ua ((a+b) Ub
(1a=+b) Ub
(b-+cc)
(EV(U(())))) Un U(c-dd) Ud REV (()
In addition, the memory usage can be significantly reduced or even eliminated, if the kemee a (m-1)) is constant spatially: future implementation may take advantage of this. In this case, all dense matrices the same across the domain, and even have known closed form that can be derived ahead of time.
Ua
(aa+b) Ub ((o+d)
(c-+d)
Ud (ed-se)
V
Ut
U
Ua U(e+ff) uf REV (((=)))
u(e++f) Uf u(g+b))
U(c+f))
uf
D.5. Midpoint reflective boundary condition
U(f+g)
Up
Ug
Ua U(g->h) Up (EV(
(gg+hh)
Un
[(h+aa)
in which REV(-) indicates flipping a vector, i.e. the vector with the same elements appearing the reverse order. In this case, we can reduce the number of variables using Fv. F is sparse binary matrix, describing how entries in be copied from entries in v, such that 1 F1.
h d
0
System (33) then can be solved via = (FTDF) FFTd FFd and the solution can be recovered from u + Fv.
Figure 23. Name pixels the boundary of the image domain in counter-clockwise ordering.
To handle spherical topology, we introduce midpoint re- flective boundary condition, induced by the octahedral parameterization (Praun & Hoppe, 2003) that maps from a surface with spherical topology to the image domain. Shown Figure 23, the boundary of the image domain is split into the chain a + d + + + The problem becomes
min 1u duu
2
s.f. U, = UI = u. Ug
((a-b)) = ((cbb)
(33)
((c+d) = ((e-dd)
U(e+f) U(g++())
(g [((a>s))
The sub-vector U(a=5)) does not include two endpoints b.
The constraints come from how the octahedral parameteriza- tion (Praun & Hoppe, 2003) specifically sets the boundary condition to enforces the spherical topology. To solve the problem, we stack the unique values in into a vector v,
29

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
E. Detailed Method Description with Necessary Background Information
In this section, we provide a detailed description of our method along with the necessary background, supplementing the brief explanation in the main text (kept short due to space constraints). Nonetheless, understanding the internal details is typically unnecessary for most users, as our method is designed be used black-box module. Readers familiar with numerical methods can skip this section.
Plug-and-play for end users Just as convolution and matrix multiplication accessed via CUDA (euDNN/cuBLAS)) that envision solver can be similarly exposed through low-level APIs (e.g., via cuDSS-like interface) with further opportunities of optimization, allowing users to easily integrate it without needing to implement anything themselves.
Notation: matrix slicing For matrix A E R we use the matrix slicing notation A,s to denote the submatrix of A that takes the rows specified by and columns specified by For example, for r [2. 1, 3], [6, 7] Ars refers to:
2
2
A.s
A
A1,7
(34)
A3,6
A3.7_
the 3 x 2 submatrix that is obtained by selecting the first three rows, in the order of 2, 1,3, and 6-th, 7-th columns from matrix A. The lower scripted A,s should not be confused with upper scripted A which is just a variable that contains r, as part of its name and does not necessarily have anything to do with slicing.
Reverse permutation matrix
Define J so that JTX (or XJ) permutes rows (or columns) X reverse order.
J =
JJ=
(35)
1
Schur complement
Schur complement reduces the problem of solving the 2 block matrix,
x
(36)
Z
W
to that of solving a smaller matrix Dx = d, assuming an invertible W, where
D = X YW-Iz)
= YW-lw
(37)
by canceling using the second line of the equation that rewrites z W !w Zx] as a function of x.
While equations in the paper might look bit dense, at the high level the overall procedure is quite simple: recursively applying the Schur complement the system many times reduce the problem a smaller system. There are a few extra considerations: 1) implicitly or explicitly re-order rows and columns of the matrix as necessary, so that we know the part we eliminate is located at the sub-block W (or anywhere we prescribed); 2) apply many Schur complements in parallel.
E.1. Case study: two subdomains
We review some standard concepts before explaining our method for accessibility larger audience. The reader may refer $E.6, smaller image for clarity of illustration.
36
39
40
27
2
30
9
H S
G
Figure 24, The connectivity graph resulted from the use of the 3 x 3 kernel size, for a 5 x image under the lexicographic sweeping order.
30

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Mau
Figure 25. For the x 9 image, using the pixel indexing in Figure 24, visualization of the sparsity pattern of A E B 45x44 for the original sparse linear system E, A.j = vi, Vi. As we can sec, the original Laplacian-like a banded matrix with three banded diagonals.
For the 5 x 9 image shown in Figure 24, each pixel corresponds to a node in the graph and is assigned with an index i {0, 1,...,445- 43-1). Two nodes i and where j E (0, 1, ...., 45-1 1) }. are connected by edge if the x 3 kernel centered at one node will cover the other node. Aij 0 if node and node j are not connected by an edge in the graph.
For the original sparse linear system, A IR 45x45
45-1
Σ
A = v., (0, 1,..,, 45 45-1).
(38)
1=0
Figure 25 visualizes the sparsity pattern of A 45x445 In this visualization of A, only where there (i,j) corresponds an entry that A.; # () namely and j are connected by edge in Figure 24, and all other entries that omitted zeros. The original A is Laplacian-like matrix with a 9-point stencil. A is a banded matrix with three banded diagonals. Previous direct solvers have to explicitly maintain the sparse matrix A shown in Figure 25, while our solver does not. Instead, we work with a more compact representation that takes advantage of the regular grid structure: we start with the dense block submatrices shown in Figure 26.
31

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Vt
Vb
Figure 26. After reordering rows and columns, the original system becomes the 5 x 5 block system:
A
0
Art
Ars
Ara
0
A
u,
Vrr
Atr
Att
Ats
0
Ata
Atb
ut
Aar
Ast
Asa
Asa
Aab
u,
Vs
Aar
Aat
Aas
Aan
Aab
ua
Va
0
Abr
Abt
Abs
0
Abn
Abb
Lub
Vbb
By pivoting the rows and columns of A, we arrive at the linear system in Figure 26. It easy to verify that the pivoting does not change the problem: it only re-orders the equations and variables. From Figure 26, now it is visually clear that after the row- and column- pivoting, some submatrices are zero. We treat the nonzero block matrices dense matrices though they can be sparse.
32

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Now, let start with:
A
A
up
VF
A
tt
A ts
A tb
Vt
Agr
Ast
Aga sa
A sa
Agb
Us
Vs
(39)
Aar
0
A as
A na
0
Va
A
Abs
Abb
Ub
Vb
where
[4, 13, 31, 40), r [0, 2,,,3, 37, 27, 18,
a= [10, 11, 12, 19, 20, 28, 30),
(40)
t= 5, 6, 6,7,8, 17, 35, 44, 43, 41], [14, 15, 16, 23,24, 24, 25, 33, 34).
Patch division and patch-wise finite element methods
Recall that
Ass (P)) + A (Q) v! +v!
(41)
(42)
can be divided into contributions from P and Q, resp. With this partition, the computations for patch P and Q are made independent from each other, and thus can be done in parallel. The values of ( and ( can be arbitrary. as long as
their summation is the correct Ass. In fact, their values are never used standalone, only their sum A (P) A (Q) used (they become parts of the matrices P, Q which are summed later).
Especially for applications like PDEs (discretized with first-order piecewise linear FEM), each patch P and Q will contribution (P) the value of Aas independently: their portions of contributions A and A only depend on information within the patch P and Q, respectively. This makes the computation within P and Q exactly independent from each other, easing parallel implementation.
System partitioning
Arr
0
A
A ra
0
up
0
Att
Ats
0
Atb
Ut
Agr
Ast
Ass
(P)
+ Ass
(Q)
A
Asb
llg
Vs
(P)
+
Vs
Q
(43)
Aar
Aas
A
0
Ua
V
0
Abt
Abs
0
Abb
Vb
Agr
0
A
A
0
up
0
0
0
0
0
Up
VF
0
0
0
Ut
0
A
A
A
0
V
A ATr
0
A
88
(P)
Agn
Ug
+
O
Ast
A
(Q)
Agh
Ul
(P)
+
Q
(44)
A ar
0
Aas
Aga
0
Un
O
0
0
0
Un
Va
0
0
0
0
Ub
O
Abt
A bs
0
Apb
Ub
0
Instead of working with the sparse matrix A and right-hand side b, method works with the dense blocks that compactly put in two tensors a. 333 as will introduce
come from the nonzero parts of A:
A
A
u
V,
?
Agr
A
Asa
ug
Ve
(45)
A
A
as
A Aaa
11g
Va
Att
A
A
Ut
Ve
Ast
A (Q)
Agb
V:
(Q)
(46)
Abt
Abs
A Abb
Vb.
33

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
The symbol "L"" indicates the above systems cannot be solved directly, since our assumption is that the values Ass and A can be arbitrary, and only their summation is a known value Agx. In fact, the solution depends every patch any local computation cannot yield the correct solution. But we can still climinate Ua up by writing their values function of Ur Uy: Up.
un
Va Vb
Aar
0
As
Ur
nm
Ut
(47)
A
Ab
0
Abt
Abs
Us
and plugging it into the joint system:
A
A
Up
A
0
Vr
u,
0
Att
Ats
+
0
tb
(48)
A sT
Ast
st
Ass
A
Agb
ub
which leads to:
Arr
0
A
Ara
0
A
0
Va
Aar
0
Aas
Vr
0
A
A ts
0
A
(49)
Agr
Agt
Ags
Aga
Aab
0
A
-1
bbJ
Vb
0
Abe
Abs
V
which further simplifies to:
Arr
Am aa
A -1
Aar
0
A
Ayn
A
Up
82
0
Att
Ath A hb
Abt
Ats
Atb
bb
Abs
Ut
A
A
aa
Anr
Ast
Agh A bb
Abt
Ass A
-1
88
bb
ug
(50)
Av1a Vn
Vr
Vt
A Vb
Vs
A Va Ash A Vb
solving which yields the solution to the original problem.
Arr Ara
nm
A
ar
A Ara
ur
0
0
0
0
+
0
Att Atb -1 Abt
bb
Ats
Arb bb
Abs
Aar
AsaA-1 Aar
(P)
Aas
0
Ast
Asb Abt
A
Asb Abs
(51).
V A Va
Vt
Aib Vb
A-Iva
(Q)
Asb A bb V..
The derivations demonstrate that the computation can be split into contributions from each subdomain.
34

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
E.2. Graph algorithm perspectives
Gaussian elimination, including our method, can be simply understood as a graph algorithm that removes nodes from graph, while updating the edge weights accordingly. The graph algorithm perspective make the overall algorithm significantly casier to understand.
Weighted graph. The matrix A plays the same role as the adjacency matrix in graph theory. As in Figure 21, initially each pixel in the image is node in the graph, and two nodes are connected by an edge with weights Aij. if the pixels and j are adjacent (as defined in $2).
Node elimination Sequential Gaussian elimination removes nodes by one from the graph. When removing node from the current graph, only need to apply simple modification to matrix A: for all pair of nodes i, j that are both adjacent to k, update Asj by subtracting the term A A -1
Auj from it.
kk
Aij + Aij A
Akjs
Vi,j
Ask f- 0,
Vi
(weights updating)
(52)
A f O,
Vi
A&k f-
Since Ak.: A. the row and column that correspond to node & become zero after the step, we can actually delete them from the matrix; note we will need to relabel and nodes after the deletion. That is, when deleting a node k, the indirect influence node j on via k, is attributed through direct influence of node j on
While the sequential Gaussian elimination is sufficient solve the linear system and is mathematical equivalent to our method, is inefficient. As an improvement, BLAS kernel.
remove multiple nodes in single elimination, leverage dense CUDA
Block (multiple nodes) elimination The updating formula generalizes to the case when i, j, k each is not a single node, but each consists of a set of nodes. Then, we have block Gaussian elimination: first divide the domain into three sets of nodes as r, 8, and f, such that any node in not connected to any node in as separated by The x 3 block:
Apr
pr
Apss
0
Art
A sT
Ags
Ast
Atr
Ata
Att
becomes the 2 block:
A
Asr A
Ars
Ast
Ats
Au
Namely, the update rule is to subtract the adjustment term A A Ays-
A f- A.r Ag Ars:
Ays + 0,
(weights updating)
(53)
Agr + 0,
A +-
Parallel block elimination For efficiency, we apply many Schur complements in parallel. In $3, we simply do two block Gaussian climination at the same time. Then our algorithm is basically parallel block Gaussian climination in which many groups of nodes (marked in yellow Figure 3) are removed by concurrently subtracting many adjustment terms. The major effort in the code is "index- -trockinra.iin carefully track what the indices remaining nodes become after some nodes are removed.
35

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
E.3. Parallel Schwarz elimination step
Previous discussions of the Schwarz step in $3.2.1 are only for illustrative purposes and do not match our implementation of the Schwarz step. Instead, in this section provide full details for the Schwarz step that eliminates the interior of each patch.
The boundary-fistt ordering. As shown in Figure 27, we adopt boundary- ordering that uses 7) {0 15) to index the boundary nodes, and w [16 ~ 24} for interior nodes.
7 0- 1, 1515
(54)
w (1 17, 18, 18....2
(55)
:.:.
Figure 27. The boundary-fing first ordering. Schwarz step removes the interior pixels for each patch in parallel.
In contrast, Figure 28 shows the lexicographic sweeping order.
16
77
19
IS
I6
I1
12
10
Figure 28, The lexicographic sweeping ordering.
Define the permutation vector ya that is useful to map array stored in the lexicographic sweeping order to the boundary-first ordering, and its inverse permutation vector v:
#: 00, 4,9 14, 19, 24, 23, 22, 21, 20, 15, 10, 6,8, 11, 12, 13, 16, 17, 18]
(56)
V = 0,, 4, 15, 16, 17, 18, 5, 14, 19, 20, 21, 6, 13, 22, 23, 24,7, 12, 11, 10
(57)
Patch-wise subsystems. This step directly constructs the matrices P, Q E R 16x16 The input to the Schwarz step is the sparse system matrices (P)) (Q) E 25 and right-hand sides (P) h (Q) E 22x1 Although H (Q) are sparse, we treat them as dense matrices in our algorithms. (H(P), h(P)) and (Q), (Q)) play the role of sub-systems previous derivations.
Air
A
Ara
Vr
(P) P
Agr
A S8
A sa
(P) :
V:
(P)
(58)
A
A as
Ana
Va
Att tt
A ts
A&b
VE
(
P
Agt
A (
Agb
9
P
(Q)
(59)
Abt
bs
Abb
Vh
In fact, our algorithm always adopts H as the direct representation of A, and never explicitly constructs A as a matrix. Using H is much more natural choice work with parallel linear solvers.
So, our method directly constructs matrix H. the following way: the entry H represents some interaction or affinity between node and node in the subdomain P, under the boundary-fore ordering shown in Figure 27. For example,
36

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
common that the patch-wise sub-system is given in the lexicographic sweeping order shown in Figure 28 as matrix and vector E, of the same size. in this case, should first map it to the boundary-firs ordering using:
H pl. h c|
(60)
The symbol " amounts using some permutation vector in this way.
We will also be able to map it back to the lexicographic sweeping order if necessary using the inverse permutation vector:
E H(v,v), h(v
(61)
We store the patch-wise sub-systems in the tensors (+)) E 22xx255x25, 3 E 1221125x1
6 [.00. = H
3 10 0,,0 h
a [1,0, H
3 0,, = (
Then, we call Algorithm for parallel implementation of the Schwarz step:
Algorithm 2 The parallel Schwarz step. Forward pass.
Require: (+()) E TRd xd2 x25x25 ((() E RR xdxx
(d1,d2) is (2//2, 2*22), or (2, 1) in the 2-patch illustrative example.
X 0 n]
Y = = a n, w
y =
(62)
Z a [1
W ==a
w :[:
(63)
D: (x 'z)
d= YW
(64)
a: D, ==d
(65)
return (0) E RRdxxddx xd2x1616 g0(0) E Rd1xdxx16x1
Algorithm 2 simultaneously calculates the new sub-systems for all subdomains-P and Q in the case of two subdomains. and the sub-systems can be fetched from the tensor a.
P az (0, 0,:,:]
Q ax [1, 0,:
Algorithm 2 still applies when there is an di d2 array of patches, instead of an 2 array of patches that we used the illustration example. Algorithm conducts batch-wise computation that simultancously calculates the new sub-systems for all patches.
E.4. Details Schur step
Now provide details the Schur step that merges every two adjacent sub-systems, to supplement $3.2.2.
The Schur step assembles the new system D by Schur "involuting" the sub-systems P. Q. The step takes input the left-hand and right-hand sides (P,p) and (Q, q). and outputs new left-hand and right-hand sides (D,d).
We store the sub-systems in four dimensional tensor (i) that a ( [c, d, :. stores the system matrix at (c, d) in the array of subdomains.
37

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Algorithm 3 The parallel Schwarz step. Backward pass.
Require: x ( tensor of size (d1, d2, 16, 16).
x Un
(66)
W = W - (w- - Zx)
(back-substitution),
(67)
( = 22,,/////,,2//,,,,, 25.522
(68)
(s n.: == un
(69)
( [, = uw
(70)
return (*) tensor of size (d1, d2, 25, 25).
()) contains all reduced systems. For j = (2i + 1). the j-th Schur step takes as input (i) of size (2//2-i, 2k/2-1 16 - 2"), and outputs o (i+1) size (2+/2-i-1 102---1,,2/22111 24 2', 24 22))
22/2
For j (2i 1), i 0, 1, ..., is 1, the j-th Schur step merges subdomains the horizontal direction.
For -2i 2) 0, ..., k 1, the j-th Schur step merges subdomains in the vertical direction.
0.:
Figure 29, Figure 4: Schur step collapses subdomains P and Q into D.
Horizontal merge
Divide the nodes into contiguous subsets o., (9 7, 6, c, such that
00, 1,2,3), [4], [5, 6,, [8], [9, 10, 11, 12, 13, 14, 15).
(71)
Divide the nodes in Q into contiguous subsets X, 14, such that:
K = [0], x = [1,2, 3, 4, 6, 9, 10, 1-], == [13, 14, 15).
(72)
Under the indexing in Figure 4, D's boundary consists of the nodes corresponding to a, B, A, 8, B, represent the same node after merging, so are 8, μ..
It makes sense define their indices in the merged domain.
à = [0. 1, [4], Å [5. 6, 8, 10, 11, 12, 13, 14, [16]. 2 = [17, 18, 19, 20, 21, 22, 23].
(73)
Note that the values of these indices depend on the shape of the patches: we only give their values the first Schur step. The node grouping implies partitioning the matrix P, Q into submatrices, by dividing the rows and columns into subsets.
38

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
P Paa
Pa
0A
Pas
Pa
Pac
P
P Pa.
P Ppg
0
P
P9
Be
P g
Oxa
0x
Oxx
OA8
Oxc
Oxy
P 6oc
P 5/
OS
P85
P
P
P cC
P c/
0,A
P
P
P FT
P yyc
P
A
P y6
P
Pre
P
(74)
my
au
A
0 0a
0
0 Dey
Up
Pa
QKI
Q
Q QNIa
Q
Ug
Da + qu qA ps P&+qu qu
Oxa
QXi
QAA
QA
Oxe
Q
ux
0aax
ju
Q/14
0.c
Ug
0 ecx
0g3
0aA
0
Oec
eey
u
Pe Py + JTq
[0. yo
JT
Q
JTQ
JT
Q pyp
yE
Wyy
is the linear system for the joint domain D. Since B, represent the same node, they correspond to the same row/column; the same applies to ò, j.. Note that y represents the same set of nodes as z but reverse order. Let J = JT to be the reverse permutation matrix (see ŞE), whose action to reverse the rows (or columns) of matrix when being multiplied with from left (or right).
E.5. Final algorithm
Algorithm specifies the overall algorithm with both the numerical factorization and back substitution stages.
The numerical factorization variant can be implemented by only keeping the computations involving a, and returns and saves the values of (9), (1) ... a for the use of the back substitution step.
The back substitution variant can be realized by only keeping the lines involving 3, and using the ("), (z)), .... a cached in the numerical factorization step.
39

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Algorithm 4 The parallel Schur step, j-th step. j (2i+1). Forward pass. Horizontal.
Require: a of size 02//2-1,2//2-11 101166-2 2', 16. 2'),
Require: (00 of size 1///--,,2/2-- 10-16-14 , 1).
P, Q are fetched batches from aن for all red and blue subdomains.
P =
0 0 : 244 = 10 2,444
(75)
Q a 00 [1 2
q = 3 [1 i,
(76)
Poo
P
Pa
Pay
Pos
ac
P
Pa
P
P By
P
De
P
P
Pyaa
P Y/B
P Py
py8
P yc
P
Py
p
(77)
Poa
P 8B
P B
P85
Pae
Ps
Pes
P c/3
P ecy
Ped
Pec
Pe
Q KR
Q&A
QKI
kP
QAx
QAa
Qu
QA
qx
Que
QuA
Quje
Q
qu
q.
(78)
Qui
QuA
Quj
QuI
qu
P aa
P
a
Oa
Pas
P Pac
P Pax
33
KK
Q
P
Pas
Q QuIa
P Se
X:
Oxa
Ax
QAA
QAI
O
(79)
Pôa
P
P
83
Qux
P +
Q
Poe
Pea
P
cB
Oax
Pe8
P
cc
P
Pay
Pa
P
P
+
Q
Ps + 9N
Y
Q
y
(80)
P
Poy
+
Q
J
P8 + qu
P
Ps
= Pya
P P
JTQ P.s + J'Qu
P,d
(81)
W = P my + JTQ
w P + JTq
(82)
D : (x YW
d: YW w
(83)
a: D, B:
(84)
return (ا+ن) size 124/2-i-1,gh/2-11 24- 24.22,4.
gl++1) of size (2k/2-i-1 2//2-1 (22/2---1122/2----24--0 21,1
40

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
Algorithm 5 The parallel j-th Schur step, j (2i + 1). Backward pass. Horizontal.
Require: x ((++1) of size (2#/2-1, 22/222,,2 (2 24.4.22) 24,22,))
ua B ā.
U
3 B (A:,
u
= ((5) or
UA = a (t [,;;AA1) A
(85)
Ug
u 3 1,,,,8.1
ua
3 [A;,,,]
ua
3
X ==
(86)
ug
u:
x
ZER0S(2 24//2-1, 116-2% 21,
(87)
W-1 (w Zx)
(back-substitution))
(88)
where we use upper script to emphasize that it does not correspond to slice into a vector.
0, =
x 0 k,,,,,,,, k. u
u&
ug
x [1 2,:;; A, =uر
2,444
u
"0::::111=""" or
(89)
:1:1;;;;144 =u
u
20002;;0;;-u
ua
2.;
return ( of size (2t/2-6, 24/2-1 16. 1).

---

Schwarz-Schur Involution:
Lightspeed Differentiable Sparse Lincar Solvers
E.6. Illustration 3x7 image
An example with detailed visualization Let consider smaller 7 image for the convenience of illustration, as shown in Figure 24. Each pixel corresponds to node in the graph and is assigned with an index E {0, 1,... 21 1}. Two nodes and j. where j E {0. 1, .... 21 1), connected by edge if the 3 x 3 kernel centered at node will cover the other node. Aū = 0 if node and node j are not connected by an edge in the graph. Each of the color includes the following subsets of nodes:
[0, 1, 2, 7, 14, 15, 16), s [3, 10, 17J, [4, 5, 6, 13, 18, 19, 20). a [8, 9J. b [11, 12).
The original sparse linear system, A E 22
21-1
Σ
u v. {0, (0.1..... ..., 21 1}.
(90)
j=0
Figure The connectivity graph resulted from use the 3 kernel size, for T'maage under lexicographic sweeping order.
42

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
U1
U1
112
U3
[1
U5
UG
U7
Ug
Ug
U10
U11
U12
U13
U14
U15
W16
U17
W18
U19
U20
(0.0)
(0.1)
(0.7)
(0.8]
Vo
(1.0)
(1.1)
(3.2)
(1.7)
(..8
(1.9)
V1
(2.2)
(2,9)
V2
(3.2) (3.3) (3,
(3.9)
(3, (3)10)
(3. [3)
V3
(4.5)
(4.10)
(4. (1) (4.12)
V4
(5, 4) (6,
(5.6)
(6. 181 12)
[5.
V5
(5,5)
(6.6)
(6, (3.
V6
(7.0)
(7.7)
(5)
V7
0)
(8.1)
(8.2)
(8. 7)
(8.8) (8.9)
(8. 14) (8.15) (8. 16)
V8
(9.2)
(9.10)
(9.15) (9. 16) (7)
V9
(1 (10.3) (10.4)
(10,9) (10, 10) 11]
(10, 16) (10, (10, 181
V10
(11,3)
(11,4) (11,5)
(11, 10)111 11](11,
(1), (1) 19)
V11
(12.4) (12.5 5) (12.6)
(12 11) 12)
(12, (82 192
V12
113.5 5) (10,6)
(13, (2) 18;
(13, (13, 20)
V13
(14.7)
(14,,
V14
(15.1 7) (15.8) (15,9)
(15. (15, (1)
V15
(16.8) (16,9) (16, 10)
(16, (1) (6) 17)
V16
(17.9) 10)
(17. 16), (17,
V17
(18, 10)18, 11)
C1N, (18, (18, 19;
V18
11)(10, (19,
(19, (19 (19, 20)
(20), 12) (20 18:
(20, 19) (20.20)
V20
Figure 31, Same as Figure for 3 7 image.
43

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
u
u!
U1
U2
144
U15
U16
U1
U5
U6
113
U18
U1
U3
1.0
U17
Ug
U12
(0.0)
(0.8)
Vo
(1,0)
(1.1)
(.2) (1.7)
(1,8)
(1.9)
V1
12
(2.2)
(2.3) (2,
(2.8)
(2,9)
V2
(7.9)
(7.1)
(7.7) (1.14) (7.15)
(7.88
V7
(14,7) (40,44 18)
(14,8)
V14
(16.7) (4) 13)(16,1))
(16.8)
V15
(16. 13) 10)
(16, 10) 17) (16.8) (16.0)
V16
(4,4) (4.5)
(4.3) (4,
(4. (1,11)
V4
(5.4)
(3.5) (5.0) (5.18)
(5.
V5
(6.5) (6,6)
(6, 13)
(6, 12)
V6
(13.5)
(13,18)
(13,12)
V13
(38. (18,18 19)
(18, 10) (38. 17)
(38, 11) (18.12)
V18
13) 19, (19, (1)
(19, 12)
V19
(30, 13)
(20, 19) (20, 200
(20,12) 12)
V20
(3.2)
(3.4)
(3.3) (3.
(3.9) (1.11)
V3
(10, (10,4)
(10,3) 10)(10
(10,9)
(17.16) 16)
(37. 18)
(17, (17. 17)
(17.9) (1)
V17
(8.0)
(8.2)
(8.7) (8.14) (8)15) (8, 16)
(8.8)
(8.0)
V8
19.1) (9.2)
(9, 15) (9, (6)
(9.3) (9) (9)
(9:8) (9.9)
Vg
(11,4) (31,5)
(11, (81(11,19))
(1133 (11, 10)
(11,1)
V11
Vb
(12.4) (12.5) (12.0) 13 (12 (12.
(12.
V12
Figure 32. Same as Figure 26 but for a 3 7 image.
44

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
F. Additional Details
Please note that our use of the term "involution" is not related to the mathematical concept that refers to a self-inverse function as an "involution," though deploying our solver in an eigen solver resembles similar concepts. We use "involution" to refer to the exact, generalized, and inverted convolution-generaaiiee in the sense of spatially varying kemels. Our use of "involution" is close to its biomedical meaning- the shrinking or reduction in the size of an organ, but applied to an image.
E.1 Details on FEM and PDE discretization: addibility of parallel linear elements
Although where the matrix A or the tensor a comes from does not matter to apply our sparse solver (though the numerical error can depend on A), as long as A is invertible with the sparsity pattern have described, supplement with implementation details of assembling matrix A for reproducibility. Despite that theoretically one could use, e.g., bilinear bases as well, in this paper, we use the first-order piecewise- linear finite element method (Wang et al., 2023) to discretize the partial differential equations. There are multiple advantages: 1) The discretized system preserved certain algebraic properties of the continuous PDEs that are particularly valuable in geometry tasks (Wang et al., 2023). 2) Most importantly, the addibility of parallel linear elements. It is easy to verify that if each of the patch (i,j) puts in (i, the Laplacian matrix L (f-5) that is discretized with the Neumann boundary condition, then assembling together these patch- wise sub-matrices yields the Laplacian matrix L for the whole image domains, under the same Neumann boundary condition. Here "assembling patch-wise sub-matrices" means summing together the sub-matrices to a large sparse matrix L by indexing into the rows and columns of L that the patch nodes correspond to. In other words, we know it holds that:
L = Σ (1.) Raxx
(91)
(b-f)
where (+J) E TR 25x25 comes from the Laplacian matrix discretizing an elliptic PDE with the Neumann boundary condition, and L E Rax= comes from discretizing the same PDE & boundary condition but over the entire image domain. Here J(i) E 225x is the binary matrix, playing the role of putting the patch sub-system L (4.5 in the correct rows and columns of the large sparse system.
The fact can be verified by noticing that using first-order piecewise-linear FEM (Wang et al., 2023) to discretize the elliptic PDE (3) which stores the anisotropic tensor C(x) as a 2 matrix per triangle, each triangle will contribute to the entries of A independently using information only within that triangle. The fact that patch sub-systems can be "seamlessly" added together thanks to the use of the Neumann boundary condition- the natural boundary condition. Using certain high-order finite element methods might necessitate information exchange between adjacent patches, which is still doable but requires extra care in the implementation.
double covered FEM mesh layout for images. During the tessellation of the regular grid as triangular meshes, to avoid bias in choosing the orientation of the long edge of the triangles, we simply use double covered triangular mesh-ead square has four cross-edge triangles whose edges connect pixels.
Figure 33, Our double-covered piccewise- -incar triangle mesh layout.
F.2. Issues with incorporating iterative solvers in deep learning
Our method is Sparse solver that is Consistent-performmacee Hyperspeed, in-the Wild, Accurate, Robust, and Zero- parameter. In this section, we explain why indirect, a.k.a. iterative, solvers fall short of these goals. We identify key obstacle preventing neural architectures from adopting linear solvers is the lack of method like ours.
Deploying iterative solvers appropriately comes with a tedious workflow, requiring users in the loop with PhD-level expertise in numerical analysis, PDEs, GPU optimization, or image processing. Iterative solvers require woublem-specine solvers with many parameters, preconditioners which also have parameters. Although solvers such as LSMR and LSQR seem to
45

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
"config_versian" 2, "solver": {
"scope": "main", pprint_grid_ssatss
"store_res seccre_rs__history"" history"; 1, "sovver; "FGMRES", "primt_selve_stats" 1,
"obtain_timings": 1,
10
"precendittioer":: "interpolator": "D2", "print_grid_ssatst" 1, Seggressive_leveis"" 1, "sclver": "AMG",
11
12
15
"smocther";
16
"relaxation_factor""" 1, "scope": "jacobi", "solver": "JACOBT_I_1 b
17
18
19
20
"presweeps": 1, "selector": "HMIS",
21
"coarsest "coarsest_sweeps":
"coarse solver": : "NOSOLVER", "max_iters"; 1, "max_rov_sum": 0.9, "strength_thrreholdd:: 0.25,
24
25
26
"min_coarse_roxx 2, "scope": "amg_solver", "max "max_levels" 50, "cycle": "V", "postsweeps";
29
30
33
m 100, monitor_residual 1, gmres_n_restart": 10, "cnvergence"; "RELATIVE INI", "tolerance" 1e-06, "norm": "L2"
34
35
36
38
39
40
Figure 34, The FGMRES+AMG solver (with HMIS selector) in Table 3.
be applicable to all matrices, especially when it comes to random deconvolution, we observe that they can be inefficient compared to our method in $4.
The large diversity of possible PDEs makes it difficult for user master all the methods well and implement the entire arsenal of numerical PDEs. Especially, we lack automatic pipeline analyze the type of A to deploy the appropriate iterative solvers, and to automatically set the parameters. preconditioner, and the parameters of the preconditioner for good performance by analyzing entries of A. In the setup of PDE learning, the exact form of PDE is unknown. To algorithmically search for the PDE (that is, the matrix represents it), one will have to implement iterative solvers for all possible PDEs, which is an infeasible task for unstudied PDEs. Iterative solvers can perform poorly for highly singularly, stiff, anisotropic coefficients, oscillatory solutions such as the Helmholtz equation under relatively large frequency (Ernst & Gander, 201 1).
Even restricting to solving some well-studied PDEs, there are efficient iterative numerical schemes, the optimal values of these parameters also depend on the entries and properties of A. The many parameters in iterative solvers must be appropriately set to yield a correct result: sigma shift values, preconditioning schemes, and parameters such as error tolerance, maximum iterations. For example, Figures 34, 35 and 36 are the FGMRES+AMG and GMRES+AMG, and PCG+AMG solvers we test in Table 3, under the Nvidia AMGX framework (Naumov et al., 2015).
46

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
"config_versian"" 2,
"solver". {
"preconditioner":
stats":
"print_ vis data": 0,
"sclver"; "AMG",
"smocther";
"scope": "jacobi", "solver": "BLOCK JACCBI" pnonitor_residuat 0, "print_solve_stats""
10
11
b
14
"print__sove_statts": 0, "presweeps": 1,
15
16
"interpolator"" "D2", "max_iters": 1, "menitor_ressddua"" "store_res" res his_oory 0,
18
19
20
"scope": "amg", "max levels": 50, "cycle"; "y" "postsweeps";
23
24
},
25
"sclver": "PCG", ierint_sclve_stateee 1,
"obtain_timings"": 1,
28
"max_iters"; 100,
29
"monitor_residuall;; 1,
30
*convergence": INI'', "scope": "main", "tolerance" : 1e-06,
32
33
"norm": "12"
34
Figure 35, The PCG+AMG solver in Table 3.
47

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
"config_version"" 2,
#a 1,
eexception_handle ing" 1,
"solver"; {
"scope"; "main",
grid stats";
sstorerres : 1,
"solver: "GMRES", "print_solve_stats":: 1,
10
"obtain__timngss: 1,
12
"precenditiioeer": "interpolatorr: "D2", "print__rid_statss::: 1, "sclver": "AMG" "smoother": "JACOBI_LI", "presweeps": 2, "selector": "PMIS",
13
14
15
17
18
secsarsest_sweepr""s 2,
"coarse_solver": "NOSOLVER", "max_iters": 1, "interg_max_eeeentt::" 4,
22
23
"nin__oarre_rooss" 2, "scope": "amg_ssoveer" "max_levels": 50, "cycle": "V", "postsweeps";
26
27
28
}
"max maxx_iters"; 100,
30
mmnittor_residuul1: 1,
11
ggmres_n_restart": 10,
32
"cnvergencc"; "RELATIVE INI",
33
"tolerance" : 1e-06, "norm": "I2"
35
36
Figure 36. The GMRES+AMG solver in Table
48

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
That said, even to solve single instance of PDE with iterative solvers, the user may have manually repeat the trial-and- enror steps until convergence: tweak the parameters and then the iterative solver with them; if not converged the result is not desirable, the user has go back to change the parameters. The tedious user-in-the-loop workflows prevent training with solver in the loop massive amount of data in scientific machine learning in way similar to the practice of computer vision and natural language processing.
Unpredictable runtime: When the input problem changes, the runtime can change dramatically, which a critical drawback for interactive applications such as self-driving cars or other time-sensitive scenarios. Conservative parameter settings significantly slow down iterative solvers, and aggressive parameter settings risk insufficient iterations or unconverged results.
Catastrophic failures can incur when applying unsuitable iterative solvers: it can obtain totally wrong results to ruin trained models, calling for user intervention: manually fixes for certain training examples and restoration from trained checkpoints.
49

---

Schwarz-Schur Involution: Lightspeed Differentiable Sparse Lincar Solvers
E33 Differentiable sparse solvers: widely desired, yet absent
Here links examples of feature requests and discussions on differentiable linear solvers in the community, Despite of the frequent requests, linear solvers have only very limited support in mainstream deep learning packages, which perhaps makes sense because existing sparse linear solvers differentiable modules and our method. Naïvely porting existing solvers to, e.g., PyTorch, makes them the sole bottleneck in the training pipelines. For example, JAX has a sparse solver, which is CPU-only and wraps up a naive callback of Concurrent to
quite limited in applicability, and inefficient compared other
SciPy's "spsolve" function. PyToreh still dees not have a sparse linear solver until Oet. (version the development of our work, since version 2.5, Py Torch recently starts to supporting a sparse solver backended by cuDSS (which requires compilation from the source; as of August 2025, the pip binary release of PyTorch does not ship with the cuDSS solver). However, the cuDSS solver is much slower than our method, demonstrated throughout the paper.
tttpes//githup. wtw.sttww.ccoo//oooc////stee//ee http://pooo...ppooot///uesee///s
Sctps: http:://gtttbb... con/ hytps/////pptoors//ss///////a//a
tttpe://// thob. con/ http:///www....ooo./////////////
attps:
attpe: /ttpe://ao ///: http::///tttooooooom www.://///ttteeece http:///ww.....oooooooo////////p
attps: /ttps://githoe www.c/oo////eeeeeee http:///w......aaaaaaa/aa//a/aaa
tp/:///11 //git whb.. eensor www.s//elllo///eoooooo/o////////
attpe: rr
Lensor- layout- gradcchheck/
443309.
Lpa ht/g///thiiaa won/ ond atcorrs wow.////onnnnn//m/
50


Reviews:
[{"review_id": "FUypr2Y3kt", "reviewer": "Reviewer_eGtJ", "summary": "The paper proposes a method for accelerating sparse linear and PDE solvers by transforming sparse Laplacian matrices into dense tensors. This procedure uses dense GPU BLAS kernel to batch and run such system in parallel. This method is differentiable, which can be potentially useful for machine learning pipeline integration.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# General responses\nWe thank the reviewers for their careful examination of our paper. We note that all reviewers appreciate the broad impact that our method will have (eGtJ: “can be applied to extensive areas”, D4Wa: “greatly accelerate scientific computation in many domains”, QcrL: “both novel and clever”, “significant potential for applications in various domains”).\n\nWe also acknowledge that our paper is rather unusual, in that it revisits an area that is often taken for granted (linear solvers), and in its current form assumes textbook knowledge about classical methods, in particular Gaussian elimination. It’s now clear that by bypassing some steps in our derivations (for reasons of readability and space), we have made the method more difficult to follow for some readers. However, our derivatives are correct (as confirmed by QcrL and numerical experiments), and we are happy to add intermediate derivations in Appendix. These extra steps, along with a more extensive graphical presentation of the problem and its specific sparsity, are available in https://anonymous.4open.science/r/icml-rebuttal-268F and will be added to Appendix.\n\nThe reviewers' comments highlight a specificity of our paper, in that it identifies a problem that is *critical yet long overlooked and considered having little-room-to-improve*. It is generally accepted that existing direct solvers are close-to-optimal, and haven’t been improved in decades. However, we do show that the approach we propose accelerates the inversion of large systems with Laplace-like sparsity by several orders of magnitude. There is nothing magical though: we simply make the observation–which has already been made for *forward* operators and is critical to the success of deep networks–that extremely parallelized problems can be implemented extremely efficiently on GPUs, such that they outperform “optimal” (flop-wise) sequential algorithms. Our *main contribution* is in transforming a problem that is not parallel by construction into a form that enables the use of dense parallel operations. Rather than focusing on the mathematics (which build on simple ideas) of our method, we have chosen to focus on its potential applications and on quantifying its performance gains on an array of real-world applications.\n# Broad impact\nDue to the foundational role of linear solvers, our method can broadly impact:\n* We provide the *first method to invert convolution* efficiently and exactly (in the generalized case of spatially-varying kernels) at interactively rate, for vision & image processing.\n* Immediate improvements across many areas: PDE solvers, learning-for-PDEs, image segmentation, physical simulation, shape optimization, and so on.\n* Since solvers are used in all subareas of science and engineering, many computing, solver-in-the-loop and learning systems can benefit from ours by substituting their existing solvers.\n* Our finding---sparse solvers can be 1000x faster than current algorithms---qualifies them as a module in the neural networks; especially those involving geometry/physics, since conventional methods in e.g. physics/geometry processing heavily rely on sparse solvers.\n* Sparse solvers have been a long requested feature in e.g. PyTorch; the lack of efficient implementation hinders its adoption in neural architectures, especially for scientific ML.\n* Optimization: by making Newton’s method much faster on images.\n* Spectral methods are accelerated: vision, graphics, geometric deep learning.\n* Partially explain and reduce the performance gap between conventional methods and deep learning for many tasks.\n# Other PDEs\nWe appreciate the insightful suggestions on more PDEs like the Navier–Stokes. We note that *Laplace FEM solvers are effectively sparse linear solvers*, making them ideal for evaluating our method. Many nonlinear PDEs—minimal surfaces, deformation, heat flow, optimal control of PDEs (diffeomorphisms), and eigenvalue problems—*repeatedly call Laplace solvers as subroutines*. Shown in Appendix, our method *directly accelerates these different types of PDEs*.\n\nSome classic Navier–Stokes solvers involve solving Laplace equations with *constant* kernels—these can often be handled efficiently by FFT-based methods, and do not require general-purpose sparse solvers like ours. In contrast, our method targets **spatially-varying kernels** (e.g., inhomogeneous materials so FFTs no longer apply) or **generalized deconvolution** with spatially-varying kernels. While our method *can* handle more general fluid models with varying diffusivity or anisotropy, we are not aware of widely-used PDE learning benchmarks involving such settings.\n\nWe discuss in Appendix (line 962) on how future work can employ **a simplified variant** of our method to accelerate constant-kernel PDEs like regular Navier–Stokes, and we will try to implement the discussed variants and apply it to Navier–Stokes experiments if time permits. This will further extend the applicability of our method.", "claims_and_evidence": "The major claims of this paper include the following: First, the significant speedup compared to the existing solutions. This is supported by Table 1, where it shows the average runtime to solve Laplacian systems under the proposed method, CUDA, Scipy, and AMG. Second, the paper claims that the proposed method is applicable to more problems than PDE, which is demonstrated in examples such as image segmentation.", "methods_and_evaluation_criteria": "The evaluation contains multiple PDE problems, including anisotropic and isotropic Laplacian systems, along with Darcy flow experiments compared to neural operator baseline. Another evaluation aspect is the comparison of average runtime, performed on SciPy, CUDA, and AMG. The author uses the relative error tolerance as the main evaluation metric, which is a standard measurement solution accuracy.", "theoretical_claims": "I have no comments on the theoretical claims due to a lack of expertise in this specific area.", "experimental_designs_or_analyses": "The experimental designs in this paper cover multiple domains, including PDE solving and real-world applications in graphics and vision. In Section 5.2, “A zero-shot baseline: learn-to-solve PDEs,” the author evaluates the proposed method against a state-of-the-art neural operator on the Darcy Flow dataset. While this comparison demonstrates the method’s applicability to PDE solving, it primarily focuses on a single PDE equation. While the paper shows that the method is widely applicable, it is suggested to include experiments on additional PDE types, such as Navier-Stokes equations.", "relation_to_broader_scientific_literature": "The paper builds on literature in sparse linear solvers, domain decomposition methods, such as the Dirichlet-to-Neumann methods operator, and the backends for direct sparse solvers.", "essential_references_not_discussed": "I have no comments on the essential references.", "other_strengths_and_weaknesses": "Strength: \nThe proposed method significantly speeds up the PDE solver and the method can be applied to extensive areas more than PDE solving.", "other_comments_or_suggestions": "I do not have specific comments on this aspect as it is outside my primary research area.", "questions_for_authors": "I have no further questions to the authors.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pOZvs8Fqeb", "reviewer": "Reviewer_D4Wa", "summary": "The authors propose an efficient method to solve sparse linear systems. Current algorithms for solving such systems are slow, which hinders their applications in real-time scenarios such as interactive graphics. The authors propose a direct solver, which uses a divide-and-conquer strategy to efficiently solve sparse linear systems. The proposed method is differentiable, thus can be integrated into modern machine learning framework.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the thoughtful feedback on improving our paper. Please also refer to Re:eGtJ for extra descriptions of our methods.\n# Method motivations and descriptions\nWe are committed to improving the exposition and we are confident we can make the paper significantly more accessible to a broader ICML audience.\n\nWe have added supporting figures in the following link:\nhttps://anonymous.4open.science/r/icml-rebuttal-268F/rebuttal-figures.pdf\nThis includes a step-by-step illustration of our method.\n\nWe realize our presentation assumed too much background in Gaussian elimination which we are happy to expand upon. To address this, we have added a brief summary and will expand it into a detailed explanation in the revised paper.\n## Parallel Gaussian elimination\nSchur complement reduces solving the 2x2 block matrix, \n[ X  Y; Z  W ]\nto solving a smaller matrix [X - YW^{-1} Z]\n\nWhile equations in the paper might look a bit dense, at the high level the overall objective is quite simple and intuitive: recursively applying the Schur complement many times to reduce the problem to a smaller system.\nTwo extra considerations have to be added: 1) always first re-order rows and columns of matrices, so that we know the part we want to eliminate is located at W (or anywhere we prescribed); 2) apply many Schur complements in parallel.\n## Graph algorithm perspectives\nOur method, as Gaussian elimination, can be simply understood as a *graph algorithm* that removes nodes from a graph. The actual procedure looks involved because the data is put in a tensor to leverage BLAS for parallel computing, and we avoid explicitly constructing the graph.\n\nAs in Figure 21, initially each pixel in the image is a node in the graph, and two nodes are connected by an edge with weights $A_{ij}$, if the pixels i and j are adjacent (as defined in Sec. 2). In other words, the matrix A plays the same role as the adjacency matrix in graph theory.\nGaussian elimination removes nodes one by one from the graph: when removing a node $k$ from the current graph, the only modification we need to make to matrix A is: for all pair of  nodes i,j that are both adjacent to k, update $A_{ij}$ by subtracting the term $A_{ik}A_{kk}^{-1}A_{kj}$ from it; 2) delete the row and column correspond to node $k$.\nI.e., the indirect influence of node j on i via k, is attributed through a direct influence of node j on i, when deleting $k$.\n## Block elimination\nThe updating formula generalizes to the case when $i,j,k$ each is not a single node, but each consists of a set of nodes. For example, for  r=[1,2,3], s=[6,7], $A_{rs}$ refers to the 3x2 submatrix that selects the first three rows and 6-th, 7-th columns from matrix A.\n\nThen, we have block Gaussian elimination: first divide the domain into three sets of nodes as r, s, and t, such that any node in r is not connected to any node in t as separated by s.\nThe 3x3 block: \n[Arr, Ars, Art==0;\n Asr, Ass, Ast; \n Atr==0,Ats, Att;]\nbecomes the 2x2 block: \n[Ass - Asr Arr^{-1} Ars, Ast; \n Ats, Att;]\nNamely, the update rule is to subtract the adjustment term: Asr Arr^{-1} Ars.\nIn Sec 4, we simply do two block Gaussian elimination at the same time.\n\nThen our algorithm is basically a parallel block Gaussian elimination in which many groups of nodes (marked in yellow in Figure 3) are removed by concurrently subtracting many adjustment terms.\n\nThe major effort in the code is ``index-tracking’’: carefully track what the indices of remaining nodes become after some nodes are removed.\n\nIllustrated in Figure 24 as a step-by-step breakdown diagram of algebraic manipulation: we apply a few steps of transformations: equations in Gray box are the transformed equations, and definitions in White box (including P, Q and their submatrices) are intermediate symbols we introduce to simplify the notation. The steps simply apply Schur complements many times and in parallel.\n# Re: Questions\nWe add visualization, sample code, and derivation diagrams to help to understand the method in the anonymous link above.\n> how do we arrive at equation 4\n\nPlease refer to Figure 21-23, for a visualization of the matrix partitioning process used to derive Equation (4), demonstrated on a 3×7 image.\n\n> what are the submatrices A_rr, A_rs, P_rr, P_rs represent\n\nPlease refer to Re:QcrL-Clarifications for definitions of submatrices.\n\n# Plug-and-play for end users\nOur method is designed to be *black-box usable* for the vast majority of users—understanding the internal details is often not required (though we will clearly document them). Just as convolution and matrix multiplication are accessed via CUDA (cuDNN/cuBLAS), we envision our solver being similarly exposed through low-level APIs (e.g., via a cuDSS-like interface), allowing users to easily integrate it without needing to implement anything themselves. Our code will be released publicly. A demo and code that include all implementation details are already available at the anonymous link.", "claims_and_evidence": "The authors claim their method is faster than current baselines, Table 1 shows that indeed the runtimes are orders of magnitude faster.", "methods_and_evaluation_criteria": "The major problem right now with the manuscript is the description of the method. It is currently hard to follow, sections 4.2.1 and 4.2.2. Although the authors tried to put visuals to make the texts easier to follow, certain symbols are introduced without enough details. For example, how do we arrive at equation 4, what are the submatrices A_rr, A_rs representing exactly. Same question for equation 5, what does P_rr, P_rs represent. It may be better if the authors start with even smaller P,Q in Figure 4, say 3x7 and show exactly how the submatrices in equation 4 are composed of. If there are too many details, some can be put into a supplementary. Some of the results can also be moved to supplementary, it is essential that the method is clear enough to understand so that it can be reproduced by interested researchers.", "theoretical_claims": "None", "experimental_designs_or_analyses": "The experimental analysis looks good.", "relation_to_broader_scientific_literature": "Solving sparse linear systems is a general problem, with wide applications to scientific domains, some of which are mentioned in the manuscript. The authors innovation on efficiently solving such sparse systems can greatly accelerate scientific computation in many domains.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "None", "other_comments_or_suggestions": "None", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rrvaW8pKHA", "reviewer": "Reviewer_QcrL", "summary": "This paper investigates efficient methods for solving sparse linear equations that commonly arise in applications related to partial differential equations (PDEs) and convolutional neural networks. The key insight of the study is the exploitation of hidden structures within convolutional kernels, allowing the division of an image into small, independent patches. Within each patch, computations can be performed independently by leveraging the concept of Gaussian elimination.\n\nThis technique follows a divide-and-conquer approach, where the problem is progressively broken down into smaller subproblems. The solution process involves solving these subproblems independently and then performing back-substitution from the final step to the initial step. By structuring the problem in this manner, the authors effectively harness the parallel computing capabilities of GPUs to handle multiple small but dense square matrices efficiently.\n\n## update after rebuttal\n\nI want to thank the authors for updating the codebase to include a comparison with scipy's direct solve method. I can verify scipy is significantly slower than the proposed method.\n\nHowever, I also implemented indirect method using cupy+GPU, here are the preliminary results:\n\n```\n# GPU sparse solve alternative\nA_cupy_sparse = cp.sparse.csr_matrix(A_scipy) # Convert SciPy sparse to CuPy sparse\nlb_cupy = cp.asarray(lb.flatten()) # Keep RHS as dense CuPy array\nx_cupy_sparse = None\nprint(\"Timing CuPy sparse solve (with synchronization):\")\nfor _ in range(5):\n  time_start_sparse = time.time()\n  # x_cupy_sparse = cupyx.scipy.sparse.linalg.spsolve(A_cupy_sparse, lb_cupy) # method 1, 89 second, norm difference with gsol is 0.955\n  # x_cupy_sparse = cupyx.scipy.sparse.linalg.splu(A_cupy_sparse).solve(lb_cupy) # method 2, 17 seconds, norm difference with gsol is 0.005\n  # x_cupy_sparse = cupyx.scipy.sparse.linalg.lsqr(A_cupy_sparse, lb_cupy) # method 3, 84 second, norm difference with gsol is, nan\n  x_cupy_sparse = cupyx.scipy.sparse.linalg.cg(A_cupy_sparse, lb_cupy) # method 4, 0.15 second, norm difference with gsol is 0.085\n  # x_cupy_sparse = cupyx.scipy.sparse.linalg.gmres(A_cupy_sparse, lb_cupy) # method 5, 0.54 second, norm difference with gsol is 1.40\n  # x_cupy_sparse = cupyx.scipy.sparse.linalg.cgs(A_cupy_sparse, lb_cupy) # method 6, 0.14 second, norm difference with gsol is 0.606\n  # x_cupy_sparse = cupyx.scipy.sparse.linalg.minres(A_cupy_sparse, lb_cupy) # method 7, 0.066 seconds, norm difference with gsol is 141.08\n  # x_cupy_sparse = cupyx.scipy.sparse.linalg.lsmr(A_cupy_sparse, lb_cupy) # method 8, 14 seconds, norm difference with gsol is 165.35\n  cp.cuda.Stream.null.synchronize() # Wait for GPU to finish\n  print(f\"cupy sparse solve time: {time.time() - time_start_sparse}\")\n```\n\nFor comparison, the proposed method (on GPU) runs in 0.032 seconds. The scipy direct method runs in 11.12 seconds. The norm difference between gsol and x_scipy is 0.021.\n\nHence, we can see the proposed method is still faster than these cupy indirect solve implementation, but the speedup is now much smaller than anticipated. The accuracy for the indirect method is kind of okay in my stand, since the norm is calculated as np.linalg.norm(x_vec - y_vec), and the vector has length 263169. So, the speedup over indirect method is more around 5-10 I would say.\n\nAlso, for this current example image size, transfer A_scipy to the GPU doesn't seem to take that long. Therefore, having the claim of 1000x speedup in the title is \\textit{a little bit} of an exaggeration. I would probably add a quantifier in the title saying that this impressive speedup is only for direct solve.\n\nPersonally, I still support the acceptance of this paper due to this nice observation/idea of being able to do lots of local parallelization for direct solve, although a little less excited than first reading this work (I thought this was a HUGE breakthrough). My current score is more like 3.5 than 4, but I am going to keep it as 4. For revision, I would strongly suggest the authors should compare thoroughly with indirect cupy methods (as I show above) and also modify the title and abstract to reflect that there is a place for the indirect cupy method.\n\nIf the authors can show the speedup of the proposed method over the indirect cupy method is much more significant (2 or 3 orders of magnitude) on large images (2561 x 2561), that would be wonderful. However, since I do not have that information, I cannot champion for something I haven't seen or experimented by myself.\n\nIf the authors feel strongly about the advantage of their method on large image instances, they are welcome to update the codebase in the anonymous code repo and send a message to the AC with clear evidence so that AC may check it out.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the insightful comments and efforts to verify our method & derivation!\n\nWe appreciate the detailed feedback on improving clarity, and will be sure to elaborate on full details to reduce readers’ efforts to reconstruct the steps.\n\nPlease also refer to Re:eGtJ for extra descriptions, and refer to the Re:D4Wa for background information and the visualizations in https://anonymous.4open.science/r/icml-rebuttal-268F/rebuttal-figures.pdf (Figures 21-24)\n# Clarifications (+D4Wa):\n> Clarify the meaning of the color-coded subscripts for matrix A in Section 4.2.1: \n\nplease refer to Figures 22,23 in the anonymous link: each color indicates a group of pixels: r = [0, 1, 2, 7, 14, 15, 16], s = [3, 10, 17], t = [4, 5, 6, 13, 18, 19, 20], a = [8, 9], b = [11, 12], in the case of 3x7 image (Figure 21).\n\nA_rr, A_rs represent the 7x7 and 7x3 submatrices (since the group r, s have 7, 3 nodes, resp.):\nA_rs is the slicing into matrix A that uses r as indices to select rows and s to select columns: A_rs = A[r, s]. A_rr = A [r,r].\n\nSee Figure 24 for the definitions of P_rr, P_rs:  \nP_rr = Arr − Ara Aaa^{-1} Aar,\nP_rs = Ars − Ara Aaa^{-1} Aas.\n\nP, Q  are 10×10 matrices (since after eliminating nodes in a, b, each of the\nsubdomains have 10 nodes remaining, as shown in Figure 21.)\n\nP_rr is a 7x7 matrix, as a submatrix of P, selecting the rows and columns that correspond to the red nodes (there are 7 of them). \nP_rs is a 7x3 matrix, similar to P_rr except that its columns correspond to the green nodes (3 of them).\n\nIn Figure 24, we provide a step-by-step breakdown of Gaussian elimination from Equation (4.2) to (5), including the derived value of $u_a$ same as suggested.\n# Re: Questions\nQ1: Yes, it is exactly right that because the kernel size is 3x3, a 3x3 window centered at the a brown dot cannot incorporate information from the light and dark blue dots. We add a demonstration in Figure 21.\n\nQ2: The values of $A_{SS}^P$ and $A_{SS}^Q$ can be arbitrary, including $1/2 A_{SS}$, as long as their sum is $A_{SS}$. In fact the values of $A_{SS}^P$ and $A_{SS}^Q$ are never used standalone, only their sum $A_{SS}^P + A_{SS}^Q$ is used (they become parts of the matrices P,Q which are summed later).\n\nThe terms of $A_{SS}^P$ and $A_{SS}^Q$ represent the partition of the value $A_{SS} = A_{SS}^P + A_{SS}^Q$, i.e. into contributions from the patch P and Q separately. This partition is because then the computation for patch P and Q can be done in parallel. Especially for applications like 1st-order FEM discretized PDEs, each patch P and Q will contribution to the $A_{SS}$, and their contributions $A_{SS}^P$ and $A_{SS}^Q$ only depend on information within the patch P and Q, respectively, This makes the computation within P and Q exactly independent from each other so it can be done in parallel.\n\nQ3: The overall runtime of L-BFGS is 20x slower than our method. Note that L-BFGS is a very different type of method---Newton's method using sparse Hessian solver with SciPy backend remains our primary point of comparison, which is 400x slower than our method.\n\nQ4: Only the baseline method [Shi & Malik 2000] uses MATLAB---the “eigs” function and it takes theirs/MATLAB 25s to solve the top-6 Eigs.  In contrast, it takes less than 0.3s for our method (since the major computation is sequentially calling 20 times of A^{-1} b, that each takes 0.011s using our method).\n\nQ5: For fair comparison, both ours and all baseline direct solvers (Scipy, CUDA) do not use symmetry. Theoretically, methods that use symmetry can roughly reduce half of the computation. Note the baselines AMG and (effectively) neural operators do exploit symmetry: In other words, the actual improvement of our method is more significant than what is presented in the paper. In theory, our method could use symmetry to get another 2x speedup. Some relevant discussions are in Appendix A.6, C.4.\n\nQ6: Our basic routine A^{-1} b is a forward problem: solves for x given both A,b. When A is not known, we need to search for the value of A, using the gradient of E w.r.t. entries of A, which further depends on the Jacobian of x over entries of A. As explained in B.5, our solver is already differentiable in PyTorch so no extra work needed to make it differentiable to obtain the Jacobian. Otherwise, one has to explicitly implementing the adjoint method that solves for A^{-T}, as done in [Wang et al. 2023] and also here:\nhttps://arxiv.org/pdf/2404.17039v2\n\nQ7: We provide sample code in the link: https://anonymous.4open.science/r/icml-rebuttal-268F  and promise the final code will be released. In the code repo, the major effort to implement our solver is “index-tracking”: i.e., in Figure 3, when some nodes are eliminated, what will be the new indices of these nodes (so that subsequent matrix slicings can find the correct submatrices). These are tedious details that can be found in the code. Actually, the 2-lines code in line 179 is a minimum validation to trace down why our solver can be much faster.", "claims_and_evidence": "The claims made in the paper are supported by clear and compelling evidence, particularly in Section 4.2, where the authors use color-coded visualizations and mathematical derivations to illustrate key computational concepts. While the explanations are well-structured, those unfamiliar with convolutional neural networks may find some parts challenging to follow. However, after manually performing the Gaussian elimination steps from Equations (4.2) to (5) and verifying the authors' calculations, the divide-and-conquer approach became evident.\n\nThe computational results presented are particularly impressive. Compared to CUDA-based implementations, the proposed method achieves a speedup of two orders of magnitude. Additionally, when compared to SciPy implementations, the authors report a computational speedup exceeding three orders of magnitude. These results strongly support the effectiveness of the proposed approach.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria (running time comparison) are valid for the problem and application.", "theoretical_claims": "I have manually checked the Gaussian Elimination step from Equation 4 to Equation 5 and can verify that they are correct.", "experimental_designs_or_analyses": "I have checked the experimental designs and analyses. It's mainly solving the same sparse linear systems and comparing the running times, which is very straightforward.", "relation_to_broader_scientific_literature": "Yes, this work is at the intersection of scientific computing and differential equations. Different branches of literature reviewers are discussed extensively in this work.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The proposed method is both novel and clever, demonstrating significant potential for applications in various domains. However, one notable weakness lies in Section 4, particularly Subsection 4.2, where the methodology is discussed. The explanation lacks sufficient detail, making it difficult for readers unfamiliar with convolutional neural networks to grasp the key ideas immediately. In particular, understanding the notation and subscripts used in the matrices required considerable effort. The color-coded elements in Subsection 4.2.1 are not well explained, which adds to the confusion. Additionally, the process of Gaussian elimination from Equation (4.2) to (5) is not elaborated on sufficiently, requiring extra effort from the reader to reconstruct the steps.\n\nAnother point of concern is Figure 4, where the rationale behind dividing patches in the presented manner is unclear. A suggestion for improvement would be to clarify in the caption that the convolutional kernel is 3x3, which influences the selection of points included in the computation. Specifically, when the kernel focuses on a brown dot, it only includes nearby red and green dots while never incorporating the light blue dots. A clearer explanation would help readers understand why certain dots incorporate information while others do not. Adding such clarifications would significantly improve the paper’s clarity and accessibility.", "other_comments_or_suggestions": "Thank you very much for writing this paper. I have enjoyed reading, understanding, and learning from your proposed idea. Below are some writing suggestions:\n\nClarify the meaning of the color-coded subscripts for matrix A in Section 4.2.1.\n\nProvide a step-by-step breakdown of Gaussian elimination from Equation (4.2) to (5). At least mention that $u_a = A_{aa}^{-1}(v_a - A_{ar} u_r - A_{as} u_s)$.\n\nImprove the caption of Figure 4 to explicitly mention the kernel size (3x3) and its impact on patch selection.\n\nConsider adding a brief explanation or visual aid to illustrate why the green dots incorporate information from all other colored dots while the brown and light blue dots do not incorporate information from all colored dots.", "questions_for_authors": "1. Regarding Figure 4, is our understanding correct that the kernel size is , and that this is why the brown dots only incorporate information from red, brown, and green dots but not from the light and dark blue dots?\n\n2. In Line 218, what do $A_{SS}^P$ and $A_{SS}^Q$ represent? Do they both equal 1/2 of $A_{SS}$?\n\n3. In Line 678, you discuss the running time of your method but do not mention the running time of the gradient descent method using L-BFGS. How long does L-BFGS take in comparison?\n\n4. In the appendix, Line 749 (right column), you mention that it takes 25 seconds to solve the top-six eigenvalue problem in MATLAB. Which method is this referring to? Is it your proposed method or the baseline? If it is your method, how long does the baseline take to solve the same problem?\n\n5. In Line 813 (right column), you state that matrices are not treated as symmetric, even when they are. Do the baseline methods exploit this prior knowledge of symmetry?\n\n6. In Appendix Section B5, you discuss optimizing matrix A for Problem 11. Could you clarify how this optimization is related to solving the sparse linear system?\n\n7. There is no code submission for this paper. While the overall methodology is clear on a high-level, verification through implementation would be helpful. Would it be possible to provide a sample implementation during rebuttal, specifically for the deconvolution operation shown in Figure 7? An anonymous GitHub repository or similar would greatly enhance reproducibility and allow for further validation of your approach.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors provide a practical method for solving sparse linear system, following a divide and conquer approach that allows parallelization and efficient use of GPU.
Reviewers agreed that the method seems to offer, in multiple settings, interesting speedups. Nevertheless reviewer `QcrL`, with additional experiments, suggested that the 1000x claim was too bold, and this needs to be toned down for the camera ready version. 
As acknowledged during the rebuttal, the authors must also include in appendix additional derivations to ease understanding of the paper for the ICML audience.
Decision: Accept (poster)