CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra
Authors: Andres Potapczynski
Marc Anton Finzi
Geoff Pleiss
Andrew Gordon Wilson
Venue: NeurIPS 2023 poster (2023)
Keywords: Machine Learning
Numerical Linear Algebra
partial differential equations
Gaussian processes
equivariance
graph learning
spectral analysis

Abstract: Many areas of machine learning and science involve large linear algebra problems, such as eigendecompositions, solving linear systems, computing matrix exponentials, and trace estimation. The matrices involved often have Kronecker, convolutional, block diagonal, sum, or product structure. In this paper, we propose a simple but general framework for large-scale linear algebra problems in machine learning, named CoLA (Compositional Linear Algebra). By combining a linear operator abstraction with compositional dispatch rules, CoLA automatically constructs memory and runtime efficient numerical algorithms. Moreover, CoLA provides memory efficient automatic differentiation, low precision computation, and GPU acceleration in both JAX and PyTorch, while also accommodating new objects, operations, and rules in downstream packages via multiple dispatch. CoLA can accelerate many algebraic operations, while making it easy to prototype matrix structures and algorithms, providing an appealing drop-in tool for virtually any computational effort that requires linear algebra. We showcase its efficacy across a broad range of applications, including partial differential equations, Gaussian processes, equivariant model construction, and unsupervised learning.

CoLA: Exploiting Compositional Structure for
Automatic and Efficient Numerical Linear Algebra
Andres Potapczynski
Mare 2inzi
Geoff Pleiss
Andrew Gordon Wilson
New York University, "Cameegi Mellon University, University of British Columbia,
Vector Institute
Abstract
Many areas of machine learning and science involve large linear algebra prob- lems, such eigendecompositions, solving linear systems, computing matrix exponentials, and trace estimation. The matrices involved often have Kronecker, convolutional, block diagonal, sum, product structure. In this paper, we pro- pose a simple but general framework for large-scale linear algebra problems machine learning, named CoLA (Compositional Linear Algebra). By combining linear operator abstraction with compositional dispatch rules, CoLA automatically constructs memory and runtime efficient numerical algorithms. Moreover, CoLA provides memory efficient automatic differentiation, low precision computation, and GPU acceleration both JAX and PyTorch, while also accommodating new objects, operations, and rules in downstream packages via multiple dispatch. CoLA can accelerate many algebraic operations, while making easy to prototype matrix structures and algorithms, providing an appealing drop-in tool for virtually any computational effort that requires linear algebra. We showcase efficacy across broad range of applications, including partial differential equations, Gaussian processes, equivariant model construction, and unsupervised learning.
Introduction
The framework automatic differentiation has revolutionized machine learning. Although the rules that govern derivatives have long been known, automatically computing derivatives was nontrivial process that required (1) efficient implementations of base-case primitive derivatives, (2) software abstractions (autograd and computation graphs) to compose these primitives into complex computations, and (3) a mechanism for users to modify or extend compositional rules to new functions. Once libraries such PyTorch, Chainer, Tensorflow, JAX, and others [l, 8, 30, 31, 38, 47J figured out the correct abstractions, the impact was enormous. Efforts that previously went into deriving and implementing gradients could be repurposed into developing new models.
In this paper, we automate another notorious bottleneck for ML methods: performing large-scald scale linear algebra (e.g. matrix solves, eigenvalue problems, nullspace computations). These ubiquitous operations are the heart of principal component analysis, Gaussian processes, normalizing flows, equivariant neural networks, and many other applications [2. 12, 13, 17, 18, 27, 28, 34, 37, 39]. Modeling assumptions frequently manifest themselves as algebraic structure- -ucchsschh as diagonal dominance, sparsity, or low-rank factorization. Given structure (e.g., the sum of low-rank plus diagonal matrices) and a linear algebraic operation (e.g., linear solves), there is often a computational routine (e.g. the linear-time Woodbury inversion formula) with lower computational complexity than seneral-purposs routine (e.g., the cubic-time Cholesky decomposition). However, exploiting
Equua contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Simple Operators
Composition Operators
Base
Pr
Σ
II
Case
A
(A)
(A)
Tr(A)
exp(A)
(A)(A)
Table Many structures have explicit composition rules to exploit. Here show the existence of a dispatch rule that can be used to accelerate linear algebraic operation for some matrix structure what is possible with the dense and iterative base cases. Many combinations (shown with are automatically accelerated as a consequence of other rules, since for example Eigs and Diag used in other routines. In absence of rule, the operation will fall back the iterative and dense base case for each operation (shown in Columns are basic linear operator types such as Diagonal, T: Triangular, P: Permutation, C: Convolution, S:Sparse, Pr: Projection and composition operators such as sum, product, Kronecker product, block diagonal and concatenation. All compositional rules can be mixed and matched and are implemented through multiple dispatch.
structure for faster computation is often an intensive implementation process. Rather than having an object A in code that represents Now-rank-plle plus-diagona matrix and simply calling solve(A, b), a practitioner must instead store the low-rank factor F as a matrix, the diagonal as a vector, and implement the Woodbury formula from scratch. Implementing structure-awee learning models is often seen as a major research undertaking. For example, a nontrivial portion of the Gaussian process literature devoted deriving specialty inference algorithms for structured kernel matrices (e.g. 7, 11, 19, 25, 29, 46, 52, 53, 24).
routines in machine
As with automatic differentiation, structure- aware linear algebra is ripe for automation. We introduce a general numerical framework that dramatically simplifies implementations efforts while achieving a high degree of computational efficiency. In code, we represent structure matrices as LinearOperaton objects which adhere to the same API as standard dense matrices. For example, a user can call A -ib or rig(A) on any LinearOperator A, and under-the-hood our framework derives computationally efficient algorithm built from set of compositional dispatch rules (see Table 1). If little is known purpose base case (c.g. Gaussian elimination or
about A, the derived algorithm reverts to a GMRES for linear solves). Conversely, if A is known to be the Kronecker product of a lower triangular matrix and positive definite Toeplitz matrix, for example, the derived algorithm uses specialty algorithms for Kronecker, triangular, and positive definite matrices. Through this compositional pattern matching, our framework can match or outperform pperial-purpose implementations across numerous applications despite relying only small number of base LinearOperator types.
Furthermore, our framework offers additional novel functionality that is necessary for ML applicatonns (see Table 2). In particular, we automatically compute gradients, diagonals, transposes and adjoints of linear operators, and we modify classic itcrative algorithms to ensure numerical stability in low precision. We also support specialty algorithms, such SVRG [23] and novel variation of Hutchinson's diagonal estimator [22], which exploit implicit structure common matrices in machine learning applications (namely, the ability express matrices as large-scale sums amenable to stochastic approximations). Moreover, our framework is easily extensible in both directions: a user implement
linear operator (i.e. column in Table 1), (i.e. one row in Table 1). Finally, our routines benefit from GPU and TPU acceleration and apply to symmetric and non-symmetric operators for both real and complex numbers.
linear algebraic operation
We term our framework CoLAA (Compositional Linear Algebra), which we package in a library that supports both PyTorch and JAX. We showcase the extraordinary versatility of CoLA with

---

broad range of applications in Section 3.2 and Section 4, including: PCA, spectral clustering, multi-task Gaussian processes. equivariant models, neural PDEs, random Fourier features, and PDEs like minimal surface or the Schrödinger equation. Not only does CoLA provide competitive performance to specialized packages but it provides significant speedups especially in applications with compositional structure (Kronecker, block diagonal, product, etc). Our package is available at attps://gt tew.tttub.ooo www.:wwwwwaawaewaaeeeea cola.
2
Background and Related Work
Structured matrices Structure appears throughout machine learning applications, either occurring naturally through properties of the data, or artificially as a constraint to simplify complexity. A nonexhausitve list of examples includes: (1) low-rank matrices, which admit efficient solves and determinants [54]] (2) sparse matrices, which admit fast methods for linear solves and eigenvalue matrices, which admit efficient spectral decompositions; (4) Toeplitz or circulant matrices, which admit fast matrix-vector products. See Section 3 and Section for applications that these structures. Beyond these explicit types, we also consider implicit structures, such as matrices with clustered eigenvalues or matrices with simple unbiased estimates. Though these implicit structures do not always fall into straightforward categorizations, is possible to design algorithms that exploit their inherent properties (see Section 3.3).
problems [14, 44]: (3)
Iterative matrix-free algorithms Unlike direct methods, which typically require dense instan- tiations of matrices, matrix-free algorithms only access matrices through routines that perform matrix-vector multiples (MVMs) [e.g. 44]. The most common matrix-free algorithms- -uchh as conjugate gradients, GMRES, Lanczos and Arnoldi iteration- under the category of Krylov subspace methods, which iteratively apply MVMs to refine a solution until a desired error tolerance achieved. Though the rate of convergence depends on the conditioning or spectrum of the matrix, the number of iterations required is often much less than the size of the matrix. These algorithms often provide significant computational speedups for structured matrices that admit sub-quadratic MVMs (e.g. sparse, circulant, Toeplitz, etc.) when using accelerated hardware (GPUs or TPUs) designed for efficient parallel MVMs [c.g. 10, 20, 51].
Multiple dispatch Popularized by Julia [6]. multiple dispatch is functional programming paradigm for defining type-specific behaviors. Under this paradigm, a given function (c.g. solve) have multiple definitions, each of which are specific to particular set of input types. A base-case definition solve [LinearOperator]] would use a generic matrix-vector solve algorithm (e.g. Gaus- sian elimination GMRES), while type-specific definition (e.g. solve [Sum], for sums of matrices) would use special purpose algorithm that makes use of the subclass' structure (e.g. SVRG, see Section 3.3). When a user calls solve( A, b) runtime, the dispatcher determines which definition of solve to use based on the types of A and b. Crucially, dispatch rules can be written for composi- tional patterns types. For example, a solve [Sum [LowRank, Diagonal]] function will apply the Woodbury formula to a Sum operator that composes LowRank and Diagonal matrices. (In contrast, under an inheritance paradigm, one would need detine a specific SumOfLowRankAnde agonal sub-class that uses the Woodbury formula, rather than relying on the composition of general purpose types.)
Existing frameworks for exploiting structure Achieving fast computations with structured matrices often a manual effort. Consider for example the problems of second order/natural gra- dient optimization, which require matrix solves with (potentially large) Hessian/Fisher matrices. Researchers have proposed tackling these solves with matrix-free methods [33], diagonal approx- imations [e.g. 4), low-rank approximations [e.g. 42], or Kronecker-factorizable approximations [34]. Despite their commonality-relying on structure for fast solves-all methods currently require different implementations, reducing interoperability and adding overhead to experimenting with new structured approximations. As an alternative, there are existing libraries like SeiPy Sparse [50], Spot [49], PyLops [41], or GPyTorch [20], which offer a unified interface for using matrix-fres free algorithms with any type of structured matrices. user provides an efficient MVM function for a given matrix and then chooses the appropriate iterative method (e.g. conjugate gradients or GMRES) to perform the desired operation (e.g. linear solve). With these libraries, a user can adapt to different structures simply by changing the MVM routine. However, this increased interoperability comes at the cost of efficiency, as the iterative routines are not optimal for every type of structure. (For example, Kronecker products admit efficient inverses that are asymptotically faster than conjugate
3

---

gradients; see Figure 1.) Moreover, these libraries often lack modern features (e.g. GPU acceleration or automatic differentiation) or are specific certain types of matrices (see Table 2).
Package
GPU Support
Autograd
Non-symmetric
Complex Numbers
Randomized Algorithms
Composition Rules
Matrices
Scipy Sparse
x
PyLops
*
*
v
X
GPyTorch
/
/
x
X
X
CoLA
V
V
Table 2: Comparison of scalable linear algebra libraries. PyLops only supports propagating gradients through vectors but not through the linear operator's parameters. Moreover, PyLops has limited GPU support through CUPY, but lacks support for Py PyTorch, JAX or TensorFlow which are necessary for modern machine learning applications.
3 CoLA: Compositional Linear Algebra
We now discuss all the components that make CoLA. In Section 3.1 we first describe the core MVM based LinearOperator abstraction, and in Section 3.2 we discuss our core compositional framework for identifying and automatically exploiting structure for fast computations. In Section 3.3, we highlight how CoI.A exploits structure frequently encountered in MI. applications beyond well- known analytic formulae (e.g. the Woodbury identity). Finally, in Section 3.4 we present CoL.A's machine learning-spi speciifc features, like automatic differentiation, support for low-precision, and hardware acceleration.
3.1 Deriving Linear Algebraic Operations Through Fast MVMs
Borrowing from existing frameworks like Scipy Sparse, the central object of our framework the LinearUperator: linear function finite dimensional vector space, defined by how it acts on vectors via matrix-vector multiply MVMA v + Av. While this function has a matrix representation for given basis, we do not need store or compute this matrix to perform MVM. Avoiding the dense representation of the operator saves memory and often compute.
Some basic examples of LinearOperatorr by a 2-dimensional array and use the standard MVM routine [Av], = =E== A.j Sparse matrices, which can be represented by key/value arrays of the nonzero entries with the standard CSR-sparse MVM routine; Diagonal matrices, which represented by I-dimensional array of the diagonal entries and where the MVM is given by (a)v = d,u;; Convolutin operators, which are represented by convolutional filter array and where the MVM is given by Conv(a)v the Jacobian represented implicitly through an autograd Jacobian Vector Product represented by function and input and where the MVM is given by Jacobian(f..r x)v JVP(f. v). In CoLA, each of these examples are sub-classes of the LinearOperator superclass.
unstructured Dense matrices, which represented
JVP operators-
Through the LinearOperator" 's MVM, it is possible to derive other linear algebraic operations. As a simple example, we obtain the dense representation of the LinearOperator by calling MVM(e1). .... MVM(ex). each unit vector e,, We describe several key operations supported by our framework, some well-established, and others novel CoLA.
Solves, eigenvalue problems, determinants, and functions of matrices As a base case for larger matrices, CoLA uses Krylov subspace methods (Section 2, Appendix C) for many matrix operations. Specifically, we use GMRES [43] for matrix solves and Arnoldi [3] for finding eigenvalues, determinants, and functions of matrices. Both these algorithms can be applied to any non-symmetric and/or complex linear operator. When LinearOperaten (e.g. self-adjoint, positive semi-definite) we use more efficient Krylov algorithms like MINRES, conjugate gradients, and Lanczos (see Section 3.2). As stated in Section 2, these algorithms are
are annotated with additional structure

---

11,4 A.
Σ A.
BlockDiag (A, B)
Kron(A, B)
MVM (r)
Σ, Σ, log M T log 11,
E (1 #1+
TA + TM
TAND + NaTu
Solve
AA + 4B
SANB + Nase
Eigs (E)
EA + Es
Ea + En
Table 3: CoLA selects the best rates for each operation or structure combination. Asymptotic runtimes resulting from dispatch rules on compositional linear operators in our framework. Listed operations matrix vector multiplies, linear solves, and eigendecomposition. Here denotes error tolerance. For a given operator of size x N, we denote T as its MVM cost, s its linear solve cost, E' its cigendecomposition cost and its condition number. A lower script indicates to which matrix the operation belongs to.
matrix free (and thus memory efficient), amenable GPU acceleration, and asymptotically faster than dense methods. See Section C.2 for a full list of Krylov methods used by CoLA.
Transposes and complex conjugations In alternative frameworks like Scipy Sparse must mamually define a transposed MVM v + A for linear operator objects. In contrast, CoLA uses novel autograd trick derive the transpose from the core MVM routine. We note that A is the vector-Jacobian product (VJP) of the vector v and the Jacobian A 0Aw DAw/ww. Thus, the function transpose( A) returns Linear0perator object that uses VJP this idea Hermitian conjugates, using the fact that A'v Tv (Al)..
v) its MVM. We extend
Other operations In Section 3.3 we outline how to stochastically compute diagonals and traces of operators with MVMs, and in Section 3.4 we discuss novel approach for computing memory-efficient derivatives iterative methods through MVMs.
Implementation CoLA implements all operations (solve, eig, logdet, transpose, conjugate, etc.) following a functional programming paradigm rather than as methods of the LinearOperator object. This is not minor implementation detail: demonstrate in the next section, it is crucial for the efficiency and compositional power our framework.
3.2 Beyond Fast MVMs: Exploiting Explicit Structure Using Composition Rules
While the GMRES algorithm can compute solves more efficiently than corresponding dense methods such the Cholesky decomposition, especially with GPU parallelization and preconditioning, it is not the most efficient algorithm for many LinearOperators, For example, if A = Diag(a), then we know that A A Diag(a-1) without needing to solve a linear system. Similarly, solves with be inverted efficiently through back substitution, and solves with eirculant
triangular matrices matrices can be computed efficiently in the Fourier domain (a) F (piag(Fa). (where F the Fourier transform linear operator). We offer more examples in Table (left).
As described Section 2, example, we implement the solve [Diagonal], solve [Tr TTriangulari ar], and solve [Circulant] dispatch rules using the efficient routines described above. a specific LinearOperator subclass does not have specific solve dispatch rule then we default to the base-case solve rule using GMRES. This behaviour also applies to other operations, such as logdet, eig, diagonal, etc.
we use multiple dispatch to implement these special case methods. For
The dispatch framework makes it easy to implement one-off rules for the basic LinearOperaton sub-classes described in Section L. However, its true power lies in the use of compositional rules, which we describe below.
Compositional Linear Operators In addition to the base LinearOperator sub-classes (c.g. Sparse, Diagonal, Convolution). our framework provides mechanisms to compose multiple LinearOperators together. Some frequently used compositional structures are Sum ( A;). Product AII; Kronecker (A 3 B), KroneckerSum (A + B). BlockDiag [A, O; 0, B and Concatenation [A, B). Each of these compositional LinearOperators are defined by (1) the base LinearOperator objects to be composed, and (2) corresponding MVM routine, which is typically written in terms of the MVMs of the composed LinearOperators. For example, MVMsun V Σ, (VM,(v), where MVM, are the MVM routines for the component LinearOperators.
5

---

Dispatch rules for compositional operators are especially powerful. For example, consider Kronecker products where we have the rule (A @ B) = A @ B Though simple. this rule yields highly efficient routines for numerous structures. For example, suppose we want to solve (A where A is dense, B is diagonal, and C is triangular. From the rules, the solve would be split over the product, using GMRES for A, diagonal inversion for B, and forward substitution for C. This breakdown is much more efficient than the base case (GMRES with MVMKroa).
@ B C)x
When exploited to their full potential, these composition rules provide both asymptotic speedups (shown in Table 3) as well as runtime improvements on real problems across practical sizes (shown in Figure 1). Splitting up the problem with composition rules yields speedups in surprising ways even in the fully iterative case. To illustrate, consider one large CG solve with the matrix power B = AA; in general, the runtime is upper-bounded by O(nT k" log }). where T is the time for a MVM with A, is the condition number of A, and is the desired error tolerance. However, splitting the product via a composition rule into a sequence of solves has a much smaller upper-bound of O(T log "). We observe this speedup in the solving the Bi-Poisson PDE shown in Figure 1(b).
Additional flexibly and efficiency via parametric typing A crucial advantage of multiple dis- patch is the ability to write simple special rules for compositions of specific operators. While general purpose solve[Sum] method (SVRG; see next section) yields efficiency over the GMRES base case, it is not the most efficient algorithm when the Sum operator is combining a LowRank and a Diagonal operator. In this case, the Woodbury formula would be far more efficient. To account for this, CoLA allows for dispatch rules on parametric types; that is, the user defines solve [Sum [LowRank, Diagonal]] dispatch rule that is used if the Sum operator is specifically combining LowRank and Diagonal linear operator. Coding these rules without multiple dispatch would require specialty defining sub-classes like NowRankPlusDiagaae object, increasing complexity and hampering extendibility.
over the LinearOperator
Decoration/annotation operators Finally, we include several decorator types that annotate existing LinearOperators with additional structure. For example, we define SelfAdjoint (Her- mymeeric.,metiic).... Unitary (orthonormal), and PSD (positive semi-definite) operators, each of which wraps an existing LinearOperator object. None of these decorators define a specialty NVM; however, these decorators can be used to define dispatch rules for increased efficiency, For example solve [PSD] can use conjugate gradients rather than GMRES, and solve [PSD [Tridiagonal]] can use the linear time tridiagonal Cholesky decomposition [see e.g.. 21, Sec. 4.3.6].
Taken together Our framework defines 16 base linear operators, 5 compositional linear operators, decoration linear operators, and roughly 70 specialty dispatch rules for solve, eig, and other operations. (See Table 1 for short summary and Appendix A for complete list of rules.) We note that these numbers are relatively small compared with existing solutions yet-as we demonstrate in Section 4- these operators and dispatch rules are sufficient to match exceed performance of specialty implementations in numerous applications. Finally, we note that CoLA is extensible by users in both directions. A user can write their own custom dispatch rules, either to (I) define new LinearOperator and special dispatch rules for it, or (2) to define a new algebraic operation for all LinearOperators, and crucially this requires no changes to the original implementation.
3.3 Exploiting Implicit Structure in Machine Learning Applications
So far we have discussed explicil matrix structures and composition rules for which there are simple analytic formulas easily found well-known references [e.g. 21, 44, 48). However, current large systems-especially those found in machine learning- often have implicit structure and special properties that yield additional efficiencies. In particular, many ML problems give rise to lincar operators composed of large summations which are amenable to stochastic algorithms. Below outline two impactful general purpose algorithms used in CoLA to exploit this implicit structure.
Accelerating iterative algorithms on large sums with SVRG Stochastic gradient descent (SGD) widely used for optimizating problems with very large or infinite sums to avoid having to traverse the full dataset per iteration. Like Monte Carlo estimation, SGD is very quick to converge to a few decimal places but very slow to converge to higher accuracies. When an exact solution is required on a problem with a finite sum, the stochastic variance reduced gradient (SVRG) algorithm [23] is much more compelling, converging on strongly convex problems (and many others) at an exponential rate, with runtime O((1 + k/M) log !) where is the condition number and is the desired accuracy.

---

10"
Runtime (sec)
10-1
Dense
iterative
10
Runtime (sec)
Iterative
10--
COLA
Runtime
10
CoLA
10
10-1
10-1
Mult-grid
Mult--gri COLA
10-1
10"²
10a
25
50
75
100
25
50
75
100
50
100
150
Size (10')
Size (10+)
Size (10²)
(a) Multi-task GPs (8)
(b) Bi-Poisson (II)
(c) EMLP
Figure 1: Empirically, our composition rules yield the best runtimes across applications consisting of linear operators with different structures (more application details Section 4). We plot mean runtime (over repetitions) for different methods (dense, iterative and ours (CoLA)) against the size of the linear operator. (a) Computing solves on multi-takk GP problem [7] for linear operator having Kronecker structure K @ Kx, where K is kernel matrix containing the correlation between the tasks and X is RBF kernel on the data. For this experiment we used a synthetic Gaussian dataset which has dimension D = 33, N = 1K and we used 7 = 11 tasks. (b) Computing solves on the 2-dimensional Bi-Poisson PDE problem for the composition A composed with itself on grid of sizes up to N 10003. We use CG with a multi-grid aSA preconditioner [9] to solve the linear system required in this application. (c) Finding the nullspace of equivariant MLP of linear operator having block diagonal structure. Here, NullF refers to the iterative nullspace finder algorithm detailed in [16]. We ran a 5-node symmetric operator S(5) as done in [16] with MLP sizes up to 15K. See Appendix D for further details.
the Laplacian operator
Convergence Criteria
102
CoLA (SVRG)
10a
CoLA (SVRG)
10
CoLA (SVRG)
109
iterative
Residual
10-1
Iterative
Residual
10
10-1
10-2
Mynd ww www
10-1
10-4
10-
50
100
0
50
100
150
200
150
MVMs (10*)
MVMs (10*)
(a) PCA (Eig)
(b) GPs RFF (Solve)
(c) Neural-IVP (Solve)
Figure 2: CoLA exploits the sum structure of linear operators through stochastic routines. (a) Eigenvalue convergence criteria against number of MVMs for computing the first principal component on Buzz (N = 430K, D = 77) using VR-PCA [45]. (b) Solve relative residual against number of MVMs for random Fourier features (RFFs) approximation [40] to a RBF kernel with
-=11 1K features on Elevators (N = 12.5K, D = 18). (c) Solve relative residual against number of MVMs when applying Neural-IVP [17] to the 2-dimensional wave equation equation done in [17]. See Appendix D for further details.
When the condition number and the number of elements in the sum is large, SVRG becomes desirable alternative even classical deterministic iterative algorithms such as CG or Lanezos whose runtimes are bounded by O(Vk log !). Figure 2 shows the impact of using SVRG to exploit the structure of different linear operators that are composed of large sums.
Stochastic diagonal and trace estimation with reduced variance Another case where we exploit implicit structure is when estimating the trace the diagonal of linear operator. While collecting the diagonal for dense matrix is trivial task, it is a costly algorithm for an arbitrary LinearOperator
defined only through its MVM-it requires computing Diag(A) (A) > e; Ae; where O is the Hadamard (elementwise) product. If we need merely an approximation or unbiased estimate of the diagonal (or the sum of the diagonal), we can instead perform stochastic diagonal estimation [22] Diag(A) A) = 1 Σ 2j @ Az, where the 2j E R are any randomly sampled probe vectors with covariance T. We extend this randomized estimator to use randomization both in the probes, and random draws from a sum when A = Σ 1= A;:
Diag EM A.) Eig O Azaij
7

---

1.0
0.75
Runtime
0.5
0.25
Tre(20K)
MHD(4.8K)
Dcs18(11.9K)
Solves
Eig
Logdet
Torch
CoLA
CPU
GPU
Figure For sufficiently large problems, switching from dense iterative algorithms provides consistent runtime reductions, especially on GPU, where matrix multiplies can be effectively parallelized. We plot the ratio between the runtime of a linear algebra operation using CoLAA or PyTorch on different hardware (CPU and GPU) divided by the runtime of using PyTorch CPU. For the linear solves, we use the matrix market sparse operator Trefethen; for the eigenvalue estimation, we use the matrix market sparse operator mhd4800b and, finally, for the log determinant computation, we use the matrix market sparse operator besstk 18. We provide additional details in Section D.4.
In Section B.1 we derive the variance of this estimator and we show that it converges faster than the base Hutchinson estimator when applicd Sum structures, We validate empirically this analysis in Figure 5.
3.4 Automatic Differentiation and Machine Learning Readiness
Memory efficient auto-differenniation In ML applications, we want to backpropagate through operations like A -1 Eigss()) (A). Tr(A), exp(A), log (A). To achieve this, in CoLA we define novel concept of the gradient of dinear0perator which detail in Appendix B. For routines like GMRES, SVRG, and Arnoldi, we utilize custom backward pass that does not require backproagating through the iterations of these algorithms. This custom backward pass results in substantial memory savings (the computation graph does not have to store the intermediate iterations of these algorithms), which we demonstrate in Appendix B (Figure 6).
Low precision linear algebra By default, all routines CoLA support the standard float32 and float64 precisions. Moreover, many CoLA routines also support float16 and bfloat16 half precision using algorithmic modifications for increased stability. In particular, we use variants of the GMRES, Arnoldi, and Lanczos iterations that are less susceptible to instabilities that arise through orthogonalization [44, Ch. 6] and we use the half precision variant of conjugate gradients introduced by Maddox et al. [32]. See Appendix C for further details.
Multi framework support and GPU/TPU acceleration CoLA is compatible with both Py Torchh and JAX. This compatibility not only makes our framework plug-and-play with existing implemented models, but also adds GPU/TPU support, differentiating from existing solutions (see Table 2). CoLA's iterative algorithms are the class of linear algebra algorithms that benefit most from hardware accelerators as the main bottleneck of these algorithms are the MVMs executed at each iteration, which can easily be parallelized hardware such as GPUs. Figure 3 empirically shows the additional impact of hardware accelerators across different datasets and linear algebra operations.
Applications
We now apply CoLA to an extensive list of applications showing the impact, value and broad applicability of our numerical linear algebra framework, as illustrated in Figure 4. This list of applications encompasses PCA, linear regression, Gaussian processes, spectral clustering, and partial differential equations like the Schrödinger equation or minimal surface problems. In contrast to Section (Figure & Figure 2), the applications presented here have basic structure (sparse, vector-product, etc) but not compositional structure (Kronecker, product, block diagonal, etc). We choose these applications due to their popularity and heterogeneity (the linear operators have different properties: self-adjoint, positive definite, symmetric and non-symmetric), and to show that CoLA

---

10
100
CoLA
sk
CoLA
Runtime (sec)
5k
Runtime (sec)
10a
CoLA (GPU)
75
GP (ele)
CoLA (CPU)
Epochs
CoLA (kin)
10a
50
GP (kin)
10
25
10
k=5
k=10 Components
k=20
10a
10
10"
10a
10
20
40
50
Dasetset Size
Runtime (sec)
(a) PCA
(b) Linear Regression
(c) GPs
(B)
SciPy
Runtime (sec)
10a
sk (L)
CoLA
Runtime (sec)
10a
CoLA (L)
10a
CoLA.
Runtime
10a
SciPy
10
CoLA
10-1
SciPy (AX)
10+
10a
10a
10
10
104
10a
10*
Edges
Grid Size
Grid Size
(d) Spectral Clustering
(e) Schrödinger Equation
(f) Minimal Surface
Figure 4: CoLA is easily applied to numerous applications with competitive performance.
Here sk: skleamn GP: GPyTorch and the tuple (N, D) denotes dataset size and dimensionality. (a): Runtime for PCA decomposition on Buzz (437.4 4K, 77). (b): Linear regression runtime on Song (386.5K, 90), where epochs) on exact GP inference on Elevators (14K, 18) and Kin (20K, 8) on GPU. (d): Spectral clustering runtime sk(L) denotes sklearn's 's implicitly restarted Lanczos implementation and sk(A) denotes sklearn's 's LOBPCG with an algebraic multi-graph preconditioner (PyAMG) 26]. CoLA(L) denotes our Lanczos implementation and CoLA(B) our LOBPCG implementation. (e): Runtimes for finding the smallest eigenfunctions expanding grids of Schrödinger equation with expanding finite difference grid. (f): Runtimes for solving the minimal surface equation via root finding expanding grids. Here SciPy utilizes the ARPACK package, highly-optimized Fortran implementation of the Arnoldi iteration, while SciPy JAX (the SciPy version integrated with JAX) and CoLA utilize python Amnldi implementations. Appendix expands on the experimental details.
CoLA both GPU and CPU. (c): Training efficiency (measure in
citations graph (cit-HepPh) consisting on 34.5K nodes and 842K edges.
performs in any application. We compare against several well-known libraries, sometimes providing runtime improvements but other times performing equally. This is remarkable our numerical framework does not specialize in any of those applications (like GPyTorch) nor does it rely on Fortran implementations of high-level algorithms (like sklearn SciPy). Below describe each of the applications found in Figure 4.
Principal Component Analysis PCA is classical ML. technique that finds the directions in the data that capture the most variance. PCA can be performed by computing the right singular vectors of X E NN When the number of data points N is very large, stochastic methods like SVRG in VR-PCA [45] can accelerate finding the eigenvectors over SVD or Lanczos, as shown in Figure 2(a).
Spectral Clustering Spectral clustering [36] finds clusters of individual nodes in a graph by analyzing the graph Laplacian L D W where D denotes diagonal matrix containing the degree of the nodes and W the weights on the edges between nodes. This problem requires finding the smallest eigenvectors of L. We rum this experiment on the high energy physics arXiv paper citation graph (cit-HepPh).
Gaussian processes GPs are flexible nonparametric probabilistic models where inductive biases are expressed through a covariance (kernel) function. At its core, training a GP involves computing and taking gradients of the log determinant of a kernel log K| and of a quadratic term y (where y is the vector of observations).
Schrödinger Equation
this problem we characterize the spectrum of an atom or molecule by finding the eigenspectrum of PDE operator in a Schrodinger equation Hy Etf. After discretizing v to grid, we compute the smallest eigenvalues and eigenvectors of the operator H which for this experiment is non-symmetric perform compactfying transform.

---

Minimal Surface Here we solve set of nonlinear PDEs with the objective of finding the surface that locally minimizes its area under given boundary constraints. When applied to the graph of a function, the PDE can be expressed as f(z) (1 + =²)= Zy + (1 + +22))===
0 and solved by root finding on discrete grid. Applying Newton- Rapison,,, we iteratively solve the non-symmetric linear system Z 2 J
f()) where is the Jacobian of the PDE operator.
Bi-Poisson Equation The Bi-Poisson equation A uu = p is a linear boundary value PDE relevant in continuum mechanics, where A is the Laplacian. When discretized using grid, the result is large symmetric system to be solved. We show speedups from the product structure in Figure 1(b).
Neural PDEs Neural networks show promise for solving high dimensional PDFs. One approach for initial value problems requires advancing an ODE on the neural network parameters 0, where 0 = M(0)==F(() where M is an operator defined from Jacobian of the neural network which decomposes as the sum over data points M "/ Σ; M, and where I' is determined by the governing dynamics of the PDE [15, 17]. By leveraging the sum structure with SVRG, we provide further speedups over Finzi et al. [17] shown in Figure 2(c).
Equivariant Neural Network Construction As shown in [16]. constructing the equivariant layers of a neural network for given data type and symmetry group is equivalent to finding the nullspace of a large linear equivariance constraint Cv 0, where the constraint matrix C is highly structured, being block diagonal matrix of concatenated Kronecker products and Kronecker sums of sparse matrices, In Figure ((c) we show the empirical benefits of exploiting this structure.
Discussion
We have presented the CoLA framework for structure-aware linear algebraic operations in machine learning applications and beyond. Building on top of dense and iterative algorithms, we leverage explicit composition rules via multiple dispatch to achieve algorithmic speedups across wide variety of practical applications. Algorithms like SVRG and a novel variation Hutchinson's 's diagonal estimator exploit implicit structure common to large-scale machine learning problems. Finally, CoLA supports many features necessary for machine learning research and development, including memory efficient automatic differentiation, multi-framewor support of both JAX and PyTorch, hardware acceleration, and lower precision.
While structure exploiting methods are used across different application domains, domain knowledge often does not cross between communities. We hope that our framework brings these disparate communities and ideas together, enabling rapid development and reducing the burden of deploying fast methods for linear algebra at scale. Much like how automatic differentiation simplified and accelerated the training of machine learning models-with custom autograd functions as the exception rather than the rule- CoLA has the potential to streamline scalable linear algebra.
Acknowledgements
This work is supported by NSF Award 1922658, NSF CAREER IIS-2145492, BigHat Biosciences, Capital One, and an Amazon Research Award.
References
[1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp. Geoffrey Irving, Michael Isard. Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large- scale machine learning on heterogeneous systems, 2015. URL https: //ww://wwww/cc tensorflo.. org/ Software available from tensorflow.org.
[2] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable Second Order Optimization for Deep Learning. Preprint arXiv 2002.09018:2 2020.
10

---

[3] Walter Edwin Arnoli. The principle of minimized iterations in the solution of the matrix eigenvalue problem. Quarterly of applied mathematics, 9(1):17-29, 1951.
[4] S Becker and Yann Lecun. Improving the convergence of back- proopppppagttoon learning with second-order methods. In Proceedings of the 1988 Connectionisl Models Summer School, San Mateo, pages 29-37. Morgan Kaufmann, 1989.
[5] Nathan Bell, Luke N. Olson, Jacob Schroder, and Ben Southworth. PyAMG: Algebraic Multigrid Solvers in Python. Journal of Open Source Software, 2023.
[6] Jeff Bezanson, Alan Edelman, Stefan Karpinski, and Viral B. Shah. Julia: A Fresh Approach to Numerical Computing. arXiv preprint arXiv: 1607, 2014.
[7] Edwin V. Bonilla, Kian Ming A. Chai, and Christopher K. I. Williams. Multi-task Gaussian Process Prediction. Advances in Neural Information Processing Systems (NeurIPS), 2007.
[8] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations Python+NumPP programs. SoftwareX, 2018. URL
http: // /ttp:://iihhiin http//ttttuuut.e/ee cwm/ google/ jax
[9] M. Brezina, R. Falgout, S. MacLachlan, T. Manteuffel, S. McCormick, and J. Ruge, Adaptive Smoothed Aggregation (aSA) Multigrid. SIAM Review, 2005.
[10] Benjamin Charlier, Jean Feydy, Joan Alexis Glaunes, François- David Collin, and Ghislain Durif. Kernel operations on the GPU, with autodiff, without memory overflows. Journal of Machine Learning Research, 2211):35573444
2021.
[11] John P Cunningham, Krishna V Shenoy, and Maneesh Sahani. Fast gaussian process methods for point process intensity estimation. In International Conference on Machine Learning (ICML), pages 192-199, 2008.
[12] Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in Neural Information Processing Systems (NeurIPS), 2013.
[13] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations. International Conference on Machine Learning (ICML), 2019.
[14] Timothy A Davis. Direct methods for sparse linear systems. SIAM, 2006.
[15] Yifan Du and Tamer A Zaki. Evolutional Deep Neural Network. Physical Review E, 104(4): 045303, 2021.
[16] Mare Finzi, Max Welling, and Andrew Gordon Wilson, A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups. International Conference on Machine Learning (ICML), 2021.
[17] Mare Finzi, Andres Potapezynski, Matthew Choptuik, and Andrew Gordon Wilson. A Stable and Scalable Method for Solving Initial Value PDEs with Neural Networks. International Conference on Learning Representations (ICLR), 2023.
[18) Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry Hungry Hippos: Towards Language Modeling with State Space Models. Preprint arXiv 2212.14052v3, 2023.
[19] Jacob Gardner, Geoff Pleiss, Ruihan Wu, Kilian Weinberger, and Andrew Wilson. Product kernel interpolation for scalable gaussian processes. In International Conference Artificial Intelligence and Statistics, pages 1407-1416. PMLR, 2018.
[20] Jacob R. Gardner, Geoff Pleiss, David Bindel, Kilian Q. Weinberger, and Andrew Gordon Wil- son. GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration. Advances in Neural Information Processing Systems (NeurIPS), 2018.
11

---

[21] Gene H Golub and Charles F Van Loan. Matrix Computations. The Johns Hopkins University Press, 2018. Fourth Edition.
[22] Michael Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in 1076, 1989.
and Computation, 18(3):005.
[23] Rie Johnshon and Tong Zhang. Accelerating Stochastic Gradient Descent using Predictive Variance Reduction. Advances in Neural Information Processing Systems (NeurIPS), 2013.
[24] Sanyam Kapoor, Marc Finzi, Ke Alexander Wang, and Andrew Gordon Gordon Wilson. SKI- ing Simplices: Kernel Interpolation the Permutohedral Lattice for Scalable Gaussian Processes. International Conference on Machine Learning (ICML), 2021.
[25] Matthias Katzfuss and Joseph Guinness. A general framework for vecchia approximations of gaussian processes. Statistical science, 36(1):124 141, 2021.
[26] Andrew Knyazev. Toward The Optimal Preconditioned Eigensolver: Locally Optimal Block Preconditioned Conjugate Gradient Method. SIAM Journal on Scientific Computing, 2000.
[27] Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural Operator: Learning Maps Between Function Spaces. Preprint arXiv 0888008481v,, 2021.
[28] Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the Intrinsic Dimension of Objective Landscapes. International Conference on Learning Representations
(ICLR), 2018
[29] Jackson Loper, David Blei, John Cunningham, and Liam Paninski. general linear-time inference method for gaussian processes on one dimension. The Journal of Machine Learning Research, 22(1):
2021.
[30] Dougal Maclaurin. Modeling, inference and optimization with composable differentiable procedures. PhD thesis, School of Engineering and Applied Sciences, Harvard University, 2016.
[31] Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient- hyperparameter opti- mization through reversible learning. In International conference on machine learning, pages 2113-2122. PMLR, 2015.
[32] Wesley J. Maddox, Andres Potapczynski, and Andrew Gordon Wilson. Low-Precision Arith- metic for Fast Gaussian Processes. Conference on Uncertainty in Artificial Intelligence (UAI), 2022.
[33] James Martens. Deep learning via hessian-free optimization. In Proceedings of the 27th International Conference on International Conference on Machine Learning, pages 735-742, 2010.
[34] James Martens and Roger Grosse. Optimizing neural networks with fronecker-factoree approx- imate curvature. In International conference machine learning, pages 2408-2417. PMLR, 2015.
[35] Ber-Gunnar Martinsson and Joel Tropp. Randomized Numerical Linear Algebra: Foundations & Algorithms. arXiv 2002.01387v3, 2020.
[36] Andrew Y. Ng. Michael I. Jordan, and Yair Weiss. On Spectral Clustering: Analysis and an algorithm. Advances in Neural Information Processing Systems (NeurIPS), 2001.
[37] Eric Nguyen, Karan Goel, Albert Gu, Gordon W. Downs, Preey Shah, Tri Dao, Stephen A. Baccus, and Christopher Ré. S4ND: Modeling Images and Videos Multidimensional Signals Using State Spaces. Preprint arXiv 2210.06583v2. 2022.
[38] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
12

---

[39] Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: Visual reasoming with a general conditioning layer. In Proceedings of the Thirty-Secone AAAI Conference Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (1AAI-18), and the 8th AAAI Symposium on Educational Advances Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 3942-3951. AAAI Press, 2018.
[40] Ali Rahimi and Ben Recht. Random Features for Large-Scale Kernel Machines. Advances in Neural Information Processing Systems, 2007.
[41] Matteo Ravasi and Ivan Vasconcelos. PyLops-A linear-operator Python library for scalable algebra and optimization. SoftwareX, 11: 11:10661, 2020. ISSN 2352-7110. doi: https://oii.rr// 10. 1016/j.softx ht16//sooftx2211.100033 2019.103611. URL https: //www. sciencedirect. www.cooo//ccaoo /tmm//iinnn//rrtccee/
pi1//$332710019800666
[42] Nicolas Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algorithm. Advances in neural information processing systems, 20, 2007.
[43] Youcef Saad and Martin H Schultz. Gmres: A generalized minimal residual algorithm for solving nonsymmetrie linear systems. SIAM Journal on scientific and statistical computing, 7 (3):856-866 869, 1986.
[44] Yousef Saad. Iterative methods for sparse linear systems. SIAM, 2003.
[45] Ohad Shamir. A Stochastic PCA and SVD Algorithm with an Exponential Convergence Rate. arXiv preprint arXiv: /t..:.4992222 20942848v 2015.
[46] Edward Snelson and Zoubin Ghahramani. Sparse gaussian processes using pseudo-inputs. Advances in neural information processing systems, 18, 2005.
[47] Seiya Tokui, Kenta Oono, Shohei Hido, and Justin Clayton. Chainer: next-generation open source framework for deep learning. In NeurlPS Workshop on Machine Learning Systems (Learningyss), volume 5, pages 1, 2015.
[48] Lloyd N. Trefethen and David Bau. Numerical Linear Algebra. SIAM, 1997.
[49] Ewout van den Berg and Michael Friedlander. Spot - Linear-Operator Toolbox. SoftwareX, cs. cs. ubc. ca/labs/scl/sppa spot/.
2013. URL http: //www//www//.o
[50] Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, CJ Carey, Ilhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-2722 2020. doi: 10. 10011//41192-0199900000
[51] Ke Wang, Geoff Pleiss, Jacob Gardner, Stephen Tyree, Kilian Q Weinberger, and Andrew Gor- don Wilson. Exact gaussian processes on a million data points. Advances neural information processing systems, 32, 2019.
[52] Andrew Wilson and Hannes Nickisch. Kernel interpolation for scalable structured gaussian processes (kiss-gp). In International conference machine learning, pages 1775-1784. PMLR, 2015.
[53] Andrew G Wilson, Elad Gilboa, Arye Nehorai, and John P Cunningham. Fast kernel learning for multidimensional pattern extrapolation. Advances in neural information processing systems, 27, 2014.
[54] Max A Woodbury. Inverting modified matrices. Department of Statistics, Princeton University, 1950.
[55] Peng Xu, Bryan He, Christopher De Sa, Ioannis Mitliagkas, and Chris Re. Accelerated stochastic power iteration. International Conference on Artificial Intelligence and Statistics, pages 58-67. PMLR, 2018.
13

---

Appendix Outline
This Appendix organized as follows:
In Appendix A we describe various dispatch rules including the base rules, the composition rules and rules derived from other rules.
In Appendix B we provide an extended discussion of several noteworthy features of CoLA, such as doubly stochastic estimators and memory- autograd implementation.
In Appendix C we include pseudo-code various of the iterative methods incorporated in CoLA and discuss modifications to improve lower precision performance.
In Appendix D we expand the details of the experiments in the main text.
A
Dispatch Rules
We now present the linear algebra identities that we use to exploit structure in CoLA.
A.1
Core Functions
A.1.1
Inverses
We incorporate several identities for the compositional operators: product, Kronecker product, block diagonal and sum. For product we have (AB) = (B AB''A-1 and for Kronecker product we have
(A
B
A
x B
In terms of block compositions we have the following identities:
1
A
A
A
B
A
A
'BD*
and
0
D
0
D-1
D
0
D
0
D
D-1
A
B
-
A -1-
B
A
0
I
0
C
D
0
I
0
D - CA
CA-BB
BB
B
CCA
I
Finally, for sum we have the Woodbury identity and its variants. Namely, for Woodbury we have
(A + UBV) A - A U
U + VA- VA-1
the Kailath variant where
(A -CA-
A A A CA CA
and the rank one update via the Inerman-Morrisee formula
(A + bc = A-1 cTAbAb A A-
Besides the compositional operators, have some rules for special operators. For example, for A Diag (a) we have - Diag (a-1) Also, if Q is unitary then Q Q* or if Q is orthonormal then Q = QT.
A.1.2 Figendecomposition
We now assume that the matrices in this section are diagonalizable. That is, Eigs (A) = AA, VA where A VAAA V In terms of the compositional operators, there is not a general rule for product or sum. However, for the Kronecker product have gs A & B) A & AB, V A & B and for the Kronecker sum we have (A + B) = AA + AB. VA & VB. Finally, for block diagonal we have
A
0
VA
0 VD
Eigs
D
0
AD
14

---

A.1.3 Diagonal
As base case, if we need to compute Diag (A) for general matrix A we may compute each diagonal element by cf Ac; Additionally, is large enough we switch to randomized estimation Diag(A) (Z AZ)1/N with NV( 1)dx N where N is the number of samples used ap- proximate the diagonal. In terms of compositional operators, we have that for sum Diag (A + B) Diag (A) + Diag (B). For Kronecker product we have Diag(A A @ B) vec(Diag (A A)Diaa(B)T) and for Kronecker sum Diag(A A + B) (Dca (A) 1T + IDiag (B)T). Finally, for block composition we have
B D D
Diag
[iiag(A),iia(())
C
A.1.4 Transpose / Adjoint
As explained in Section 1, as a base case we have an automatic procedure to compute the transpose or adjoint any operator
autodiff. However, we also incorporate the following rules. For sum we have (A + B)* A + B and (A + B) AT + BT, For product we have (AB) B* and (AB) = BTAT, Kronecker product we have A
s B) A' B and (A @ B) T A &
For the Kronecker sum we have (A 4 B)* A block composition we have
B and (A $ B) AT + BT. In terms of
(c
A
B)
B D
A
C*
)
A
A CT
and
D
B'
D
B
DJ
A if is self-adjoint and Finally for the annotated operators we have the following rules. AT = A if A is symmetric.
A"
A.1.5
Pseudo-inverse
As base case, if we need compute A+. we may use SVD (A) U, and therefore set A+ VV+V*, where Σ inverts the nonzero diagonal scalars. If the size of A is too large, then we may use randomized SVD. Yet, is uncommon to simply want A+, usually we want to solve least-squares problem and therefore
solvers that are as expensive SVD. For the compositional operators have the following identities. For product (AB)*
as
(A AB) (ABB*)* and for Kronecker product we have (A + B)* A diagonal we have
x B
For block
(A
A D
)
= A 0 O
D
Finally, we have some identities that are mathematically trivial but that are necessary when recursively exploiting structure as that would save computation. For example, if Q is unitary we know that Q Q and similarly when is orthonormal. If A self-adjoint, then A+ A and also if symmetric and PSD.
A.2 Derived Functions
Interestingly, the previous core functions allow to derive multiple rules from the previous To illustrate, we have that Tr (A) Σ, Diag (A),. Additionally, if is PSD we have that
(A) = VAJ (AA) V and A is both symmetric and PSD then (A) = V (Aa) Vİ- where in both
used Eigs (A) AA, V A. Some example functions for PSD matrices
Sqrt (A) = V A 1/2 v-1 or Log (A) = V log A Which also this rules allow us to define
1/2
LogDet (A) Tr (Log ())).
A.3 Other matrix identities
We emphasize that there are myriad more matrix identities that we not intentionally include such Tr B) = Tr(A) + (B Tr(AB) Tr(BA) when A and B are squared. These
15

---

additional cases are not part of our dispatch rules either they are automatically computed from other rules (as in the first example) or they do not yield any computational savings (as in the second example).
B Features in CoLA
B.1 Doubly stochastic diagonal and trace estimation
Singly Stochastic Trace Estimator
Consider the traditional stochastic trace estimator:
7
Tr(Base](A) 1 ñ 2 Azj
(I)
j=1
with each 2j ~ (N,, where A a D D matrix. When A is itself a sum A = 1 m A,, we can expand the trace as [r[Base] (A) = 1 across elements of the sum.
2,A, with probe variables shared 2) >
Consider the quadratic form Q = z Az. which for Gaussian random variables has cumulant generating function of (t) = log Ef('") - log det A. From the generating function we can derive the mean and variance of this estimator: [[Q] K3 (0) Tr(A) and Var[Q]
(0) 2Tr(A?) )) Since Tr (A) is a sum of independent random draws of Q, we see:
2 n Tr (A )
E Tr (ase|a)) Tr(A) and Var[Tr[Base|A)]
(2)
Doubly Stochastic Trace Estimator For the doubly stochastic estimator, we choose probe vari- ables which are sampled independently for each element of the sum:
n
m
Tr Sum] (A)
z A.ij.
(3)
TEVTA
Separating out the elements of the sum, we can write the estimator as Tr Summ (A) = , Σラ R where Rg independent random samples of the value R generating function is merely Kn(t) Ei Kq. (t/m) where Q; = z Azz. Taking derivatives we find that,
1 Σm z. A,Z, The cumulant m
E|R] = (2((0) = 1 Tr(A,) = Tr(A),
(4)
1=1
m
m
Var R k
2Tr(A
2 TT m Σ A
(5)
m²
m
Assuming bounded moments on A,, then both to fixed values as m + co. Given that Tr [r(Sum] (A) variance of the doubly stochastic estimator:
=
1 E, A, and S(A)
m E, A will converge
m
1 Σ Ry. we can now write the mean and
Tr[Sum]] (A)] TT(A) and Var [T[[Su]] (A)]
(r((A))
(6)
mn
As the error of the estimator can be bounded by the square root of the variance, showing that while the error for Tr Base is O(1/vn) (even when applied to sum structures), whereas the error for Tr [Sun] ((1/ nmm significant asymptotic variance reduction.
The related stochastic diagonal estimator
T
m
Diag(Sum|
(A)
Zi C A,zij-
ram
achieves the same O(1/ / /m)) convergence rate, though we omit this derivation for brevity as follows the same steps.
In Figure 5 empirically how our doubly stochastic diagonal estimator outperforms the standard Hutchinson estimator.
16

---

Diagonal Estimate Relative Error
Hutchinson
100
Doubly Stochastic
10
10-2
10-2
100
102
Passes through dataset
Figure 5: Improved convergence of doubly stochastic diagonal estimator. Convergence of our doubly stochastic diagonal estimator in evaluating the diagonal of the UCI Buzz empirical covariance matrix (batch size 100). Shown the relative eтгоr of the estimate vs the number of passes through the n data points of the dataset. Our diagonal estimator has lower variance and converges faster than the standard Hutchinson estimator.
a(A -1b)
a(A -b))
102
Unrolled autodiff
Runtime (sec)
Memory (MB)
102
CoLA vjp rules
10
101
Unrolled autodiff
100
100
CoLA vjp rules
102
103
10a
102
103 CG iters
10a
CG iters
aloglAl
alogiAl
102
Runtime (sec)
Memory (MB)
101
Unrolled autodiff
10
Unrolled autodiff
CoLA vjp rules
CoLA vjp rules
100
10"
102
10
102
10
Lanczos iters
Lanczos iters
(a) Backwards pass runtime
(b) Backwards pass memory
Figure 6: Our autograd rules allow for fast and memory efficient backpropagation. For two different linear algebra operations A 'b and log Agl. we show the runtime and peak memory utilization required to compute the derivatives as we increase the size of the problem. In all plots, we compare CoLA'' autograd rules against the autograd default backpropagating through each iteration of the solver (unrolled autodiff). Notably, using the custom autograd rules allows us to save substantial memory and runtime when performing the backwards pass.
17

---

B.2 Autograd rules for iterative algorithms
For machine learning applications, we want seamlessly interweave linear algebra operations with automatic differentiation. The most basic strategy is to simply let the autograd engine trace through the operations and backpropagate accordingly. However, when using iterative methods like conjugate gradients or Lanczos, this naive approach is extremely memory inefficient and, for problems with many iterations, the cost can be prohibitive (as seen in Figure 6). However, the linear algebra operations corresponding to inverse, eigendecomposition and trace estimation have simple closed form derivatives which we can implement to avoid the prohibitive memory consumption and reduce runtime.
Simply put, for operation like = CGSolve, CGSolve(A,b) = A-1b we must define Vector Jacobian Product: VJP(S; (A, b), v) = (vT 8f
8A v fff.) However, for matrix-free linear operators, we cannot afford to store the dense matrix A, and thus neither can we store the gradients with respect to each of its elements! Instead we must (recursively) consider how the linear operator was constructed in terms of its differentiable arguments. In other words, we must flatten the tree structure of possibly nested differentiable arguments into a vector: 0 = flatten[A] For example for A Kron (Diag(01) (01). Conv(62)). flatten[A] 02. From this perspective, we consider A as a container or tree of its arguments 0, and define uT Øf ;= unflatten juT 0f which coincides 8A with the usual definition for dense matrices. Applying to inverses, we can now write simple VJP:
30
V
8A = inflatten (0 + unflatten(f)A "b,
af
0, 'b,
(8)
for v (T 1 (OgAa) -b, and we will adopt this notation below for brevity. Doing so gives a memory cost which is constant in the number of solver iterations, and proportional to the memory used in the forward pass. Below we list the autograd rules for some of the iterative routines that we implement in CoLA with their VJP definitions.
1. y Solve(A, b):
W By
- (A -1 W) ((800))
80
2. A, V Eigs(A) (A):
:
W 01
00
= wiDiag V (8gAs)
3. X, V Eigs A):
WT dv: = WT( A) AgAov.
4. y = log |A|:
By 80
Tr A (gAa)
5. y = Diag A)
w 80
= wiDiag (0gAa)
In Figure 6 show the practical benefits of our autograd rules. We take gradients of different linear solves A that were derived using conjugate gradients (CG), where each solve required an increasing number CG iterations.
C Algorithmic Details
In this section we expand upon three different points introduced in the main paper. For the first point we argue why SVRG leads to gradients with reduced variants. For the second points we display all the iterative methods that upon CoLA's strategy for dealing with the different numerical precisions that we support.
base algorithms in CoLA. Finally, for the third point expand
C.1 SVRG
In simplest form, SVRG [23] performs gradient descent with the varianced reduced gradient
w -- w - (1..(w g, (wo) + g(wo))
where 9s represents the stochastic gradient evaluated at only single element or minibatch of the sum, and g g(wo) is the full batch gradient evaluated the anchor point wo which is recomputed at the end of each epoch with an updated anchor.
With different loss functions, we can use this update rule to solve symmetric or non-symmetric linear systems, to compute the top eigenvectors or even find the nullspace of a matrix. Despite the fact that the corresponding objectives are not strongly convex in the last two cases, it has been shown that
18

---

Symmetric Solve Aw b
Top-k Figenvectors AW= WA
Nullspace AW- 0
9:(w))
A,W-1 b
A,W + WW [55]
A,W [16]
Table 4: SVRG gradients for solving different linear algebra problems.
gradient descent and thus SVRG will converge at this exponential rate [55, 16]. Below we list the gradients that enable us to solve different linear algebra problems: In each of the three cases listed above, we can recognize that if the average of all the gradients g(w) ) is O, then the corresponding linear algebra solution has been recovered.
While may seem that we need take three complete passes through {A.} per SVRG epoch (due the three terms in Equation 9), we can reduce this cost to two complete passes exploiting the fact that the gradients are linear in the matrix object, replacing A, A, A.Wo with A;(W Wo) where appropriate. In all of the Sum structure experiments where we leverage SVRG, the x-axis measures the total number of passes through {A; two for each epoch for SVRG.
C.2 Iterative methods
In Table we list the different iterative methods (base cases) that for different linear algebraic operations as well as for different types of linear operators. As seen in Table 5, there are many alternatives to our base cases, however we opted for algorithms that are known to be performant, that are well-studied and that are popular amongst practitioners. A comprehensive explanation of our bases cases and their alternatives can be found in Golub and Loan [21] and Saad [44].
Linear Algebra Op
Base Case
Alternatives
Ax b (non-symmetric)
GMRES
BiCGSTAB, LGMRES, QMR
Ax = (self-adjoint)
MINRES
GMRES
Ax b (PSD)
CG
GMRES
Eigs(A) (non-symmetricc
Arnoldi
IRAM, Bi-Lanczos
Eigs(A) (self-adjoint)
Lanezos
LOBPCG
A+ +
CG
LSQR, LSMR
A = UEV*
Lanczos, rSVD
Jacobi-Davidson
f(A) (self-adjoint)
SLQ
Amnolii
Table CoLA's base case iterative algorithm and some alternatives. now expand on
acronyms. GMRES: Generalized Minimum RESidual, BiCGSTAB: BiConjugate Gradient STABi- lized, QMR: Quasi-Minimal Residual, MINRES: MINimum RESidual, CG: Conjugate Gradients, IRAM: Implicitly Restarted Arnoldi Method, LOBPCG: Locally Optimal Block Preconditioned
Conjugate Gradients, Bi-Lanczos: Bidiagonal Lanczos, CGS: Conjugate Gradient Squared, LSQR: Least squares QR, LSMR: Least squares Minimal Residual iteration, LGMRES; Least squares Gener- alized Minimum RESidual, rSVD: randomized Singular Value Decomposition, and SLQ: Stochastic Lanczos Quadrature.
C.3 Lower precision linear algebra
The accumulation of round-off error is usually the breaking point of several numerical linear algebra (NLA) routines. such, common to use precisions like float64 or higher, especially when running these routines on CPU. In contrast, in machine learning, lower precisions like float32 or float16 are ubiquitously used because more parameters and data can be fitted into the GPU memory (whose memory is usually much lower than CPUs) and because the MVMs can be done faster (the CUDA kernels are optimized for operations on these precisions). Additionally, the round-off error incurred on MVMs is not detrimental when training machine learning models (as are already
19

---

running noisy optimization algorithms) as when solving linear algebra problems (where round-of error can lead to poor solutions). Thus, is an active area of research in NLA to derive routines which utilize lower precisions than float64 or that mix precisions in order to achieve hetter runtimes without a complete degradation of the quality of the solution.
In CoLA we take a two prong approach to deal with lower precisions in our NLA routines. First, we incorporate additional variants of well-known algorithms that propagate less round-off error the expense of requiring more computation, as seen in Figure 7. Second, we integrate novel variants of algorithms that are designed to be used on lower precisions such as the CG modification found in Maddox et al. [32]. We now discuss the first approach.
As discussed in Section C.2, there two algorithms that key for eigendecomppsitions. The first is Aroodi (applicable to any operator), and the second is Lanczos (for symmetric operators) - where actually Lanczos can be viewed as a simplified version of Arnoldi. Central to these algorithms is the use of an orthogonalization step which is well-known to be a source of numerical instability. One approach to aggressively ameliorate the propagation of round-off error during orthogonalization is to use Householder projectors, which is the strategy that we use in CoLA. Given a unitary vector u, Householder projector (or Householder reflector) defined the following operator R 2uu". When applied to a vector x the result Rx is basically a reflection of x over the u space. To easily visualize this, suppose that x E RR² and u 01. Hence,
I1
F1
11
Rx
2
I2
which is exactly the reflection of the vector across the axis generated by 62. Most notably, R. is unitary R. = I which can be easily verified from the definition. Being unitary is crucial as under the usual round-off error model, applying R to another matrix A does not worsen the already accumulated error E. Mathematically, || (A E) RAJ [RE] |E|| where the last equality results from basic properties of unitary matrices. We are going to use Arnoldi as an example of how Householder projectors are used during orthogonalization. In Figure 7 we have an example of two different variants of Amoldi present in CoLA. The implementations are notably different and also it is easy to see how Algorithm is more expensive than Algorithm 1. First. note that for Algorithm 2 we have two for loops (line 6 and line 8) whereas for Algorithm we only have one (line 4-6). Worse, the two for loops Algorithm require more flops than the only for loop in Algorithm Note that we do not always favor the more expensive but robust implementation of an algorithm as in some cases, like when running GMRES, the round-off error not as impactful to the quality of the solution, and shorter runtimes are actually more desirable.
D Experimental Details
In this section we expand upon the details of all the experiments ran in the paper. Such details include the datasets that were used, the byperparameters of different algorithms and the specific choices of algorithms used both for CoLA but also for the alternatives. We run each of the experiments times and compute the mean dropping the first observation (as usually the first run contains some compiling time much not too large). We do not display the standard deviation as those numbers are imperceptible for each experiment. In terms of hardware, the CPU experiments were run on an IRtel(R) Core(TM) i5-9600K CPU @ 3.70GHz and the GPU experiments were run on a NVIDIA GeForce RTX 2080 Ti.
D.1 Datasets
Below we enumerate the datasets that we used in the various applications. Most of the datasets are sourced from the University of California at Irvine's (UCI) Machine Learning Respository that can be found here: https: //trr//aaeeeeeeeee ics. uci hosting these UCI benchmarks can be found here: https: / github. com/treff http//effrevaass//e/ datasets (we have no affiliation).
php. Also, a community repo
1. Elevators. This dataset is a modified version of the Ailerons dataset, where the goal is to to predict the control action on the ailerons of the aircraft. This UCI dataset consists of N = 14K observations and has D = 18 dimensions.
20

---

Algorithm Arnoldi iteration
Algorithm Householder Arnoldi iteration
1: Inputs: A, Op Mu/ ||oll where possibly Vo ~ N (0, I), maximum number of itera- tions and tolerance (0, 1).
1: Inputs: A. F 0 where possibly
N (0, I), and maximum number of iterations
T.
2: for = 0 7
do
2: for j = 0 to '' do u, GET SEO.HOOSSOOLDDE_EECCC..
3:
Mj+1 < Aq;
4:
for i 0 to do
4:
R I I- 1-2u,u;
5:
h4. - q (Aqz)
h, R,
6:
Vj+1 j+1 hi.9.
6:
q RRu R.,.-. R2
7:
end for
7:
if < ''' then
8:
h1+1, |+1
1ト+1 R, Ro(Aq,)
9:
if h1+1h < e then
end if
10:
stop
10: end for
11: else
11: return H. Q = WEWEEHOLDDR_EC(W...... k))
12:
9i+1 V 1
12: function GET HLHHHOUSEHLDDRRRR.
13:
end
13: W4 = 0 for < and u, = wE for i >> k.
14:
end for
14: uk w/ ||]]
15:
return H. Q (qo) ... 1...1/9--11977
15: return u
16: end function
Figure 7: Different versions of the same algorithm, but the Householder variant being more numeri- cally robust.
2. Kin40K. The full name of this UCI dataset is Statlog (Shuttle) Data Set. This dataset contains information about NASA shuttle flights and we used subset that consists of 40K observations and has D 8 dimensions.
3. Buzz. The full name of this UCI dataset is Buzz in social media. This dataset consists of examples of buzz events from Twitter and Tom's Hardware. We used subset consisting of
N = 430K observations and has D = 77 dimensions.
4. Song. The full name of this UCI dataset is YearPredictionMSS. This dataset consists of N -386.5K observations and it has D 90 audio features such as 12 timbre average features and 78 timbre covariance features.
5. cir-HepPh. This dataset is based arXiv's HEP-PH (high energy physics phenomenology) citation graph and can be found here: https: / /snap.aapp.a.. /ttasaaaaaaaaa..aaaeaaa.aaaaaaaa html. The dataset covers all the citations from January 1993 April 2003 of V 34, 549 papers, ultimately containing E] 421, 578 directed edges. The notion of relationship that we used in our spectral clustering experiment creates a connection between two papers when at least one cites another (undirected symmetric graph). Therefore the dataset that we used has the same number of nodes but instead E = 841, 798 undirected edges.
stanf orrd edu/ /ata/ htte/cit--tt cit-He/h..
D.2 Compositional experiments
This section pertains to the experiments of Section 3.2 displayed in Figure 1. We now elaborate on each of Figure I's panels.
(a) The multi-task GP problem exploits the structure of the following Kronecker operator Kr @ Kx where K is a kernel matrix containing the correlation between the tasks and Kx is a RBF kernel on the data. For this experiment, we used a synthetic Gaussian dataset where the train data x; ID) which has dimension D 33, N 1K and used T 11 tasks (where the tasks basically set the size of Kr). We used conjugate gradients (CG) the iterative method, where set the hyperparameters to tolerance of 10-6 and to a maximum number of iterations to 1K. We used the exact same hyperparameters for CoLAA
(b) For the bi-poisson problem set up the maximum grid to be N = 1000 Since this PDE problem involves solving symmetric linear system, we used CG as the iterative method with tolerance of 10 and maximum number of iterations of 10K. The previous
21

---

parameters also apply for CoLA. We note that PDE problems are usually solved to higher tolerances as the numerical enror compounds as we advance the PDE.
(c) For the EMLP experiment we consider solving the equivariance constraints to find the equivariant linear layers of graph neural network with nodes. To solve this problem, we need to find the nullspace of large structured constraint matrix. We use the uniformly channel heuristic from [16] which distributes the N channels across tensors of different orders. We consider our approach which exploits the block diagonal structure, separating the nullspaces into blocks, as opposed to the direct iterative approach exploiting only the fast MVMs of the constraint matrix. We use tolerance of 10
D.3 Sum structure experiments
This section pertains to the experiments of Section 3.3 contained in Figure 2. We now elaborate on each of Figure 2's panels.
(a) In this experiment we computed the first principal component of the Buzz dataset. For the iterative method we used power iteration with a maximum number of iterations of 300 and a stop tolerance of 10-7 CoLA used SVRG also with the same stop tolerance and maximum number of iterations. Additionally, we set SVRG's batch size to 10K and the learning rate to 0.0008. We note that single power iteration roughly contains 43/2 21.5 times more MVMs than a single iteration of SVRG. In this particular case, the length of the sum is given by the number of observations and therefore SVRG uses 430/10 10 43 times less elements per iteration, where 10 comes from the 10K batch size. Finally, the 2 is explained by noting that SVRG incurs in full sum update every epoch.
(b) In this experiment we trained a GP by estimating the covariance RBF kernel with J = 1K random Fourier features (RFFs). The hyperparameters for the RBF kernel are the following: length scale (& = 0.1), output scale (a = 1) and likelihood noise (o² 0.1). Moreover, we used CG the iterative solver with tolerance of 10- and 100 the maximum number of iterations (the convergence took much less iterations than the max). For SVRG we used the same tolerance but set the maximum number of iterations to 10K. batch size of 100 and learning rate of 0.004. We note that a single CG iteration roughly contains 10/2 = 5 times more MVMs than a single iteration of SVRG. In this particular case, the length of the sum is given by the number of RFFs and therefore SVRG uses 1000/100~ = 10 times less elements per iteration, where 100 comes from the batch size.
(c) In this experiment we implemented the Neural-IVP method from Finzi et al. [17]. We consider the time evolution of wave equation in two spatial dimensions. At each integrator step, a linear system M(0) F(0) must be solved to find é, for a d 12K 12K dimensional matrix. While Finzi et al. [17] use conjugate gradients to solve the linear system, we demonstrate the advantages of using SVRG, M(O))
Σ (f) is sum over the evaluation at m = 50K distinct sample locations within the domain. In this experiment we use a batch size of 500 for SVRG, and employ rank 250 randomized Nyström preconditioning for both SVRG and the iterative CG baseline.
D.4 Hardware speed-up comparisons
This section pertains to the experiments of Figure 3. For all these experiments we computed the runtime reduction as a fraction between the time that it takes CoLA to run some linear algebra operation and Py PoTchh using the same hardware. As an example, assume that PyTorch takes 200 seconds to compute a solve using a CPU and 100 seconds to compute the same solve but now using GPU, Moreover, assume that CoLA's iterative algorithm takes 100 seconds to compute the same solve on a CPU and 40 seconds on a GPU. Thus, the runtime reduction would be 100/200 = 0.5% for the CPU column whereas 40/ 40/100 100 0.4 for the GPU column.
Solves. this experiment we calculated the % runtime reduction when running torch. linalg. Siinalg.oove solve on the Trefethen N = 20K matrix market sparse operator. In this experiment, CG was run with a tolerance of 10-11 and a maximum number of iterations equal to the operator size.
22

---

2. Eigenvalue estimation. In this experiment we calculated the % runtime reduction when running torch. linalg. eigh on the mhd4800b . 4.8K matrix market sparse operator. In this experiment, Lanczos was run with tolerance of 10-9 and maximum number of iterations equal to 100.
3. Log determinant computation. In this experiment we calculated the % runtime reduction when running torch. linalg.logdet on the besstk18 N = 11.9K matrix market sparse operator. In this experiment, the stochastic Lanezos quadrature was run using 30 Lanczos probe estimates and 25 samples.
D.5 Applications
This section pertains to the experiments of Section 4 displayed in Figure 4. We now elaborate on each of Figure 4's panels.
(a) In this experiment we compute 5, 10 and 20 PCA components for the Buzz dataset. We compared against sklearn which uses the Lanczos algorithm through the fast Fortran-based ARPACK numerical library. In this case, CoLA uses randomized SVD [35] with a rank 3000 approximation.
(b) In this experiment we fit Ridge regression on the Song dataset with a regularization coefficient set to 0.1. We compared against sklearn using their fastest least-square solver lsqr with a tolerance of 10-4, In this case, CoLA uses CG with the same tolerance and with a maximum number of iterations set to 1K. Additionally, we CoLA using CPU and GPU whereas we used only CPU for sklearn as it has no GPU support. We observe how in the arguably most popular ML method, CoLA is able to beat leading package such as sklearn.
(c) In this experiment we fit GP with RBF kernel on two datasets: Elevators and Kin40K. We only used up to 20K observations from Kin40K as that the maximum number of observations that would fit the GPU memory without needing to partition the MVMs. We compare against GPyTorch which uses CG and stochastic Lanezos quadrature (SLQ) to compute and optimize the negative log- marrinnaa likelihood (loss function). Both experiments were
GPU for 100 iterations using Adam as an optimizer with learning rate of 0.1 with the default values of 31 = 0.9 and 32 = 0.999. Additionally, for both GPyTorch and CoLA, the CG tolerance 250 and 20 probes were used for SLQQ Note that both CoLAA and GPy Torch have similar throughputs, for example GPyTorch runs 100 iterations on Elevators on 43 seconds whereas CoLA runs 100 iterations on 49 seconds. When training GP, we solve block of 11 linear systems (1 based on y and 10 based on random probes) where one key difference is that the CG solver for GPyTorch has a stopping criteria based on the convergence of the mean solves whereas CoLA has stopping criteria based on the convergence of all the solves.
to 10-1 with a maximum number of CG iterations of
(d) In this experiment we run spectral clustering on the cit- eit-HepPh dataset using an embedding size of 8 and also clusters for k-means (with only run of k-means after estimating the embeddings). We compare against sklearn using two different solvers, one based on Lanezos iterations using ARPACK and another using an Algebraic Multi-Grit solver AMG. In this case, CoI.A also uses Lanezos iterations with default tolerance of 10-6, We see how sklearn's 's AMG solver runs faster than CoLA's but this mostly the algorithmic constants as they have similar asympiotical behavior (similar slopes).
(e) In this experiment we solve the Schrödinger equation to find the energy levels of the hydrogen atom on a 3-dimensional finite difference grid with up to N = 5K points. In order to handle the infinite spatial extent, we compactify the domain by applying the arctan function. Under this change of coordinates, the Laplacian has a different form, and hence the matrix forming the discretized Hamiltonian is no longer symmetric. We compare against SciPy's 's Arnoldi implementation with 20 iterations where CoLA also uses Arnoldi with the same number of iterations. Surprisingly. CoLA's JAX jitted code has competitive runtime when compare to SciPy's runtime using ARPACK.
(f) In this experiment we solve minimal surface problem on a grid of maximum size of N 100 points. To solve this problem we have run Netwon- Netwon-Rhapson where each
23

---

inner step involves linear solve of an non-symmetric operator. We compare against SciPy's GMRES implementation as well as JAX's integrated version of SciPy. The main difference between the two is that SciPy calls the fast and highly-optimized ARPACK library whereas SciPy (JAX) has its only Python implementation of GMRES which only uses JAX's primitives (equally it is done in CoLA). The tolerance for this experiment was 5e-3. We see how CoLA's GMRES implementation competitive with SciPy (JAX) but still does not beat ARPACK mostly due to the faster runtime of using lower level GMRES implementation.
24


Reviews:
[{"review_id": "jI8FlcpLe3", "reviewer": "Reviewer_qUas", "summary": "The paper presents a library to automate the efficient execution of numerical linear algebra kernels commonly occurring in ML applications. The library recursively exploits compositional structure beyond versus standard packages which treat numerical matrix kernels as back boxes. The proposed framework provides memory efficient automatic, differentiation, low precision computation, GPU acceleration in both JAX and PyTorch.", "questions": "-) Is there any new numerical method involved in the library? My understanding is that every single numerical method used for linear systems, eigenvalue problems, etc, is known, and the main novelty of the paper is to wrap them together in an efficient manner for composite tasks. \n-) Judging COLA as a library requires some reasonable level of insight in the library itself. I understand that the authors must stay anonymous, but it is rather difficult to fully understand the benefits of COLA without looking at the code itself.", "limitations": "N/A", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "-) The quality of the manuscript is good. All concepts are communicated clearly and the paper is very well-written.\n-) The authors have put a lot of effort to provide automation for several important numerical kernels and structures. A long list of important applications are listed, and COLA can be of major significancy in several ML areas, essentially speeding up innovation.\n-) The numerical results indicate that COLA can be faster than baseline alternatives on a wide range of numerical tasks.", "weaknesses": "-) One aspect I found confusing is the lack of information regarding the numerical algorithms featured in some of the results. For example, in Figure 1, it is not clear why the library is faster for the Kronecker problem. Likewise, the same is true for the Bi-Poisson problem. Can the underlying algorithms be found online?\n-) In similar spirit, it is not clear whether COLA is compared against state-of-the-art numerical approaches. For example, Multigrid is the de-factor choice for elliptic Poisson problems, is this what the authors list as 'iterative'? If not, what is the point of listing a comparison against a non-optimal iterative (or direct) algorithm?\n-) In Figure 3, scipy is competitive with COLA, if not better in some tasks. This furthers complicates the message of the paper. What is the main reason for publication? The fact that COLA includes a wide-range of solvers for composite tasks or that it can be faster or more memory efficient in general?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are appreciative of your thoughtful feedback. In our response, we provide important clarifications, and new results inspired by your comments. Although contributions of this type can be hard to evaluate, we believe CoLA should be judged by its strong potential for scientific impact, and hope you can consider raising your score in light of our response.\n\n**Why CoLA composition rules produce algorithmic speedups and Figure 1**. The main objective of Figure 1 is to show how exploiting the compositional structure provides an algorithmic improvement. For example, for matrices with Kronecker product structure, it is common practice to use an iterative algorithm like CG to perform a linear solve (see e.g. [6]). However, we show it is more efficient to split the problem into two using the Kronecker structure: in particular, decompose $(K_T \\otimes K_X)^{-1} \\mathrm{vec}(Y) = \\mathrm{vec}(K_T^{-1} Y K_X^{-1})$ and the complexity is reduced from $O\\big(\\sqrt{\\kappa_T \\kappa_X}(m^2n + mn^2)\\log1/\\epsilon\\big)$ to $O\\big(\\sqrt{\\kappa_T}m^2n\\log1/\\epsilon+\\sqrt{\\kappa_X}mn^2\\log1/\\epsilon\\big)$. The computational burden is likewise reduced for when using dense LU or Cholesky based solvers.\n\nFor the BiPoisson problem, ($\\Delta^2 x = \\rho$), it is more efficient to use the product structure and perform the two linear solves separately. As you mention, a specialized multigrid method can very efficiently solve the discretized elliptic differential operator here. Just as in the dense and iterative approaches that we discussed, the multigrid method also benefits from splitting up the problem with CoLA’s composition rules. We have run the multigrid method both with the CoLA decomposition (solving the PDE by inverting $\\Delta$ twice) and without (solving by inverting $\\Delta^2$) and show in the attached PDF (Figure A (Right)) that applying that doing so yields significant runtime improvements. In other words, CoLA’s approach of recursively breaking up structure provides runtime benefits independent of the algorithm being used to solve the problem (CG or multi-grid). \n\n**Why CoLA should be published at NeurIPS**. The research behind CoLA serves to substantially reduce the bottleneck of deriving efficient algorithms that exploit the structure commonly found in machine learning and scientific computing. Regarding the scientific impact of frameworks such as CoLA, we would like to draw an analogy to PyTorch and reverse mode automatic differentiation (NeurIPS 2019). From a narrow point of view, autograd is merely an application of the chain rule, and yet its impact on machine learning research has been almost immeasurable. Autograd obviates the need for deriving backpropagation rules for each model separately, and bespoke autograd rules (such as our own for iterative solvers in CoLA) can be slotted into an existing language of differentiable functions only when necessary, without requiring the whole structure to be constructed anew. Likewise with CoLA, we have developed an approach such that new rules can be slotted into an existing linear algebra ecosystem, and that ecosystem need not be rewritten for each use case. \n\nThe potential impact of this framework is extremely large. In a sense, machine learning is largely linear algebra, and common modeling assumptions give rise to structure that can be exploited for significant computational savings. CoLA will help make researchers more broadly aware of the structure they can exploit, and significantly reduce the bottleneck to implementing methods that exploit structure, as well as prototyping various structures for their problems. The speed gains when using CoLA depend on the degree of compositional structure. For some problems like the Schrodinger equation (Figure 3 (e)), CoLA achieves parity with SciPy; in more complex problems like equivariant neural networks (Figure 1(c)) CoLA achieves remarkable speedups over existing solutions by exploiting compositional structure (see Figure 1). Of course, CoLA will not be a magic bullet for every problem, and we believe it is actually to the paper’s credit that it provides an honest and comprehensive presentation, including results where CoLA is essentially on par with alternatives.\n \n**On the purpose of Figure 3**. We would like to clarify the purpose of Figure 3. While Figures 1 and 2 demonstrate the efficiency gains on problems with compositional structure, Figure 3 demonstrates the breadth of CoLA’s applicability in real-world applications, including on problems with no compositional structure. On these problems without compositional structure, we do not expect CoLA to outperform existing specialized methods. However, as seen in Figure 3, CoLA remains competitive with these specialized methods, demonstrating that even in the “worst case” scenario (no compositional structure) there is no downside to using CoLA. We will clarify this point about Figure 3 in the main text.\n\n**Regarding new numerical methods**, we believe you may have missed the novel doubly stochastic trace / diagonal estimator that we introduce Section 3.3. This algorithm, though not the centerpiece of our paper, is an important methodological contribution as it yields lower variance than the standard Hutchinson estimator (see Appendix B.1). Moreover, we provide a novel procedure to automatically compute gradients through iterative methods, as well as an automatic procedure to compute diagonals and transposes / adjoints of linear operators (Sections 3.1 & 3.4) through their matrix-vector product routine.\n\n**CoLA’s code**. We have been careful to retain anonymity, but we are beyond excited to publicly announce the library, as we feel the community would value it greatly. However, we argue that the key ingredients behind CoLA (a pleasingly simple programmatic mechanism for exploiting compositional structure, as well as the algorithms covered in Section 3) are sufficiently general purpose concepts that can be evaluated independent of implementation."}, {"review_id": "AY7CurB2jm", "reviewer": "Reviewer_pi1f", "summary": "This works presents CoLA, a framework for extending the linear algebra interface of modern numerical packages to take advantage of the structural properties of linear operators present in machine learning and other applications. By adding adaptive multi-type based dispatching CoLA adaptively exploits both dense and iterative methods to decrease the computational costs of performing a given operation. CoLA is also capable of exploiting the compositional structure that is present by combining different linear operator properties, creating a large collection of possible applications. Additionally, CoLA provides all these features in an extensible framework capable of backpropagation to ensure integration in modern machine learning and deep learning applications. Evaluations support the authors assertion that providing structural information during processing facilitates better performance on many benchmark applications.", "questions": "- It probably doesn't help that one of the applications I know the most about, spectral clustering, is the application where CoLA seems to provide mixed results. Would closing the performance gap between CoLA and the sk (PyAMG) backend be a simple process of extending the multi-dispatch interface?\n- I found the added ability to efficiently backprop through iterative solve methods interesting, were evaluation results for this feature presented in any of the experiments?\n", "limitations": "The limitations of the work are clear from the presentation and no issues require further acknowledgment, to the best of my knowledge.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The work clearly addresses an ongoing issue in many numerical linear algebra applications that require specialized structural operator structures to be individually implemented and exploited on a per-application basis. This process is not on tedious and time-consuming but also ripe with opportunities for error during the computation of the backward pass updates required for integration in a modern machine learning application.\n- Overall the writing and exposition of the problem, proposed solution and evaluations are clear and well articulated in the text.\n- The proposal dovetails naturally with the well-known LinearOperator interfaces that exist in Scipy and provide an easier route to define further extensions for developers to provide application-specific knowledge for further customization.\n- An interface for composing different linear operator properties for CoLA to exploit seems to be a novel and interesting extension over other existing implementations.\n- Performance results sufficiently demonstrate that having access to more structural information naturally supports more opportunities for application performance improvement while lowering the cost of the developer to exploit that structure through an intuitive and simple interface.\n", "weaknesses": "- Though interesting the work presented may be a better fit in a venue that focuses on numerical methods and software. The intended audience is quite broad.\n- The core contributions and methods are relatively straightforward and are well-known to the numerical linear algebra community. This work seeks to make the process of exploiting the structural properties of linear operators easier for developers to use and integrate into a modern machine-learning application but it's not clear whether this contribution would have an appreciable impact on the developer community. Though this is my personal opinion it seems that problems encountered by developers are simple enough to be solved manually in most cases.\n- The evaluation section focuses on applications that exhibit basic structure but it's hard to see the value of the additional features mentioned: backprop and lower precision.\n- There seems to be a strong reliance on the appendix to fill in missing explanations due to size constraints.\n- Although the interface provides more flexibility regarding basic and/or compositional structure this still leaves a host of additional options to select the appropriate iterative approach to solve a sparse matrix. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. In our response, we provide substantial clarifications, as well as experimental results catalyzed by your comments. We appreciate your feedback, and hope you can consider increasing your score in your final evaluation. We strongly believe this effort will have a significant impact on the machine learning community, comparable to highly impactful libraries such as GPyTorch [6] and BoTorch [2], both of which appeared at NeurIPS in previous years. We would be happy to engage if there are further questions.\n\n**Fit to NeurIPS**. Our library, though by no means limited to machine learning problems, was specifically designed with machine learning applications in mind. Consider the following design decisions:\n\n1. _Machine learning specific features_. Unlike other frameworks, CoLA offers backpropagation (a nontrivial contribution of our paper), GPU acceleration, and low-precision operations—three necessary features for modern machine learning applications. While these features are broadly applicable, their impact has unquestionably been dominated by ML in recent years.\n2. _Algorithms suited for the “implicit structure” of machine learning problems_. As outlined in Section 3.3, many of the algorithms used by CoLA (e.g. randomized diagonal estimation, randomized preconditioning, SVRG, etc.) are especially well-motivated for ML objective functions, which often feature large summations over data that are amenable to randomized algorithms. \n3. _Ability to rapidly prototype with different structures_. The flexibility that results from the use of dispatch rules is of particular value to ML researchers who are more inclined to prototype different structures (diagonal, convolution, low-rank) without a strong a priori sense of what might provide a good approximation.\n4. _Evaluation on machine learning applications_. The applications in the paper are dominated by machine learning relevant topics (GPs, equivariant neural nets, neural PDEs, spectral clustering, PCA…). \n\nWe thus argue that the machine learning community has the most to gain by using CoLA. We would also note that, as stated above NeurIPS has been a venue for similar software frameworks like GPyTorch [6] and BoTorch [2]. CoLA is similar in nature to these other frameworks, but has arguably an even greater potential for impact in ML as its applicability spans more applications.\n\n**Impact on the developer community**. We respectfully disagree that “problems encountered by developers are simple enough to be solved manually in most cases.” In machine learning applications (e.g. second order optimization), it is often common to prototype with different types of matrix structure (e.g. block diagonal versus low-rank versus Kronecker approximations of the Hessian matrix). While the rules that govern these different structures are not necessarily complicated, switching between these different structures is a tedious and error-prone process. Indeed, this is a pain point that we encountered in many of our own projects, which inspired our development of this library. CoLA will automate away this process, enabling more rapid prototyping.\n\nThe usefulness of such automation should not be underestimated. Automatic differentiation frameworks have significantly impacted ML research and development, even though computing gradients is conceptually straightforward. We argue that CoLA will have a similar effect by targeting a different bottleneck in the ML pipeline.\n\n**Value of backprop and lower precision.** We note that these two features are critically important to our library. Optimization through backpropagation is now the dominant paradigm in machine learning, and low-precision arithmetic is becoming increasingly prominent to improve speed and memory consumptions. We also emphasize that our implementation of these features are nontrivial contributions. A naive implementation of backpropagation (i.e. directly backpropagating through numerical methods) would incur significant memory (see Figure 4 in the supplementary), and the standard implementations of numerical methods are well known to be unstable for low-precision arithmetic.\n\nFor a specific experiment that demonstrates the value of backpropagation, we would draw your attention to the Gaussian process application in Figure 3. The parameters of the kernel are chosen by backpropagating the negative log marginal likelihood function, which requires computing a solve and log determinant of the kernel matrix, as well as the gradients of these operations. To demonstrate the value of low precision, we added a linear regression experiment in float16 to show the runtime efficiencies that can be gained. Please see Figure B (Right) of the attached rebuttal pdf.\n\n**Selecting the appropriate interactive approach**. We believe that there may be some confusion about how the appropriate numerical method is selected. Most users will likely use the *high-level CoLA interface* (i.e. calling `cola.solve`, `cola.eigs`, `cola.trace`, etc.), in which CoLA automatically determines an appropriate default numerical method based on the underlying structure and size of the linear operator. “Power users”, who may want to specify the underlying algorithm, can use the *low-level CoLA interface* (i.e. directly calling `cola.cg`, `cola.gmres`, etc.) or pass a keyword argument to the high level interface solve(A, b, method=”cg”). We will clarify this point in the paper.\n\n**Performance of CoLA on Spectral clustering**.  As you note, there is a performance gap between CoLA and scikit-learn on the spectral clustering application. This gap is the result of CoLA not using the same algorithm as scikit-learn. We have now incorporated LOBPCG into CoLA and you can find the results on the attached PDF, Figure B (Left). Now that we have included LOBPCG (sk(B) vs CoLA(B)), CoLA again achieves better runtimes.\n\nWe appreciate your thoughtful questions and we are happy to engage further!"}, {"review_id": "kkr10szfHF", "reviewer": "Reviewer_Wjxi", "summary": "In this paper, the authors introduce CoLA, a library which streamlines the use of various linear algebra routines within frameworks of relevance for machine learning applications. The library revolves around the LinearOperator object, and provides numerous implementations of various numerical algorithms involving operations with such object, from inversion (system solution), to eigenvalue computation, to operators manipulation.\n\nMost notable features of the library are:\n- Structure-awareness and automatic dispatch: the CoLA framework is able to leverage information regarding the relevant structure of the LinearOperator considered (such as positive-definiteness, symmetry, sparsity, composition of operators), so to automatically identify and utilise the most apt algorithm specialisation for the target operation\n- Integration with existing frameworks: CoLA can interface with both JAX and PyTorch codebase, supports both GPU and TPU acceleration, and has some support for low-precision arithmetic and complex numbers\n- Autograd for iterative routines: CoLA implements some relevant iterative routines for linear algebra, and defines some ad-hoc rules for efficiently performing automatic differentiation through them\n\nThe main reported results pertain the performance comparison of the CoLA-implemented routines versus alternative implementations. Overall, CoLA performance is shown to be comparable to that of the baselines; moreover, when there is gain to be had from leveraging the structure of the underlying linear operator, CoLA is shown to be able to do this effectively.", "questions": "Overall I’m quite satisfied with the paper. One minor doubt / curiosity I still have is:\n- When comparing CoLA with other existing baselines in Fig3, it seems like your implementation underperforms in both (a) PCA and (d) Spectral Clustering. You elaborate a bit for (d) in Appendix D4, but could you expand on this? In particular, is the difference in performance simply due to the lack of an optimised implementation on your side, or are there some structural causes? And what are the main implementation differences?", "limitations": "Limitations of CoLA are not explicitly commented upon, but at the same time the framework proposed seems very flexible and efficient (as showcased in the experiments). The main limitations are hence connected to what methods are readily implemented in the framework, rather than being structural.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The authors propose a very useful framework. Most noticeably, it automates away the need for manually tuning the method choice depending on the properties of the operator, when numerical linear algebra algorithms are involved. Moreover, it effectively leverages some clever design choices (such as multiple dispatch). Overall, it can become a valid tool for numerous applications in the field of ML (and other fields as well)\n- The authors propose an interesting solution to efficiently performing autograd on iterative procedures, which are ubiquitous in linear algebra applications\n- The paper is reasonably well-written. Even though the breadth of applications considered in their work is indeed rather large, the authors still manage to present the key advantages of their framework in a clean manner, without being dispersive", "weaknesses": "- The main weakness I see, is the lack of actual “novelty” in the work being proposed - at least in the classical sense of the word. Apart from the autograd rule in Appendix B2, in fact, the various methods proposed and implementation choices are not new. This notwithstanding, the main goal of the project consists in collecting available linear algebra routines into a unified, ready-to-use, efficient library which can be easily encapsulated within existing ML frameworks, and as such it is still valuable to the research community", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful and supportive review. We were glad to see your appreciation of the broader ways in which CoLA can help contribute to the scientific community. \n\n**Performance gap on PCA and Spectral clustering**. On PCA, we believe the runtime gap is primarily the result of CoLA’s Python overhead. On the spectral clustering example we use Lanczsos to compute the smallest eigenvectors whereas scikit-learn typically uses LOBPCG. When comparing against scikit-learn using Lanczos (sk(L) vs CoLA(L)), we perform slightly better due to minor differences in the Lanczos implementation (scikit learn uses implicit restarts and we do not). We have now incorporated LOBPCG into CoLA and we compare the results (sk(B) vs CoLA(B)) in Figure B (Left), with the CoLA implementation coming out slightly ahead again. \n\nPlease let us know if we can assist with any other questions."}, {"review_id": "hA4jkEw5WT", "reviewer": "Reviewer_ZAWv", "summary": "This work proposes a new library for solving linear algebra problems involving structured matrices. The library incorporates highly efficient linear algebra kernels tailored to handle specific types of structured matrices. Moreover, it uses compositional rules to tackle problems involving matrices with composition structures. By employing these rules, the library eliminates the need for manual implementation of many efficient algorithms applicable to such matrices.", "questions": "n/a", "limitations": "n/a", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "1. Designing efficient numerical linear algebra libraries is important to the machine learning community.\n\n2. This work is well-written and the presentation is clear.", "weaknesses": "Overall I believe the novelty of the work is very limited. The numerical linear algebra algorithms discussed in the paper do not appear to be novel, and the proposed library essentially consolidates these existing algorithms. While this library may be the first to utilize compositional rules for automated algorithmic development, the rules themselves are straightforward and uncomplicated, lacking significant novelty. In summary, this work represents a combination of various engineering efforts, but its scientific contributions are mostly incremental.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We appreciate the positive remarks regarding the clarity of our work. Below, we clarify the novel methodological contributions of our submission. At the same time, we also note that many impactful papers at NeurIPS have been centered around frameworks where algorithmic innovation is not the focus (see references in the general remarks). \n\n**Novelty in Composition Rules**: We respectfully but strongly argue that the use of compositional dispatch rules to recursively subdivide linear algebraic operations is a substantial innovation. The rules may seem “uncomplicated,” but in this context, simplicity adds to the practical utility of the approach, and we ask that our research is judged by its impact rather than its complication. (We also note that the NeurIPS reviewing guidelines explicitly endorse new combinations of known techniques.) The potential impact of this approach can be seen in Figures 1 and 2, which thoroughly demonstrate scenarios where this compositional approach is not only advantageous but also atypical. \n\n\n**Additional Contributions Overlooked**: We’d like to highlight several novel contributions of our paper that may have been overlooked. We have introduced a novel algorithm for stochastic diagonal and trace estimation of sum objects that we prove (in appendix B.1) and empirically verify that it achieves a speedup over the standard Hutchinson estimators (in the attached PDF, Figure A - Left). Moreover, we provide memory efficient autograd rules (section 3.4) to be used in conjunction with the iterative algorithms.\n\n**Significance of Software Libraries in Scientific Contribution**: We'd like to emphasize that many impactful papers presented at NeurIPS and other leading conferences have been centered on the development and presentation of software libraries. This community has recognized that the practical implementations, when done effectively, propel research forward, as evidenced by works such as [1, 2, 3, 4, 5] (see general remarks). These libraries, even without introducing radically new algorithms, provide tangible benefits to the research community.\n\nIn conclusion, while individual rules may appear obvious, the framework for integrating these structures in a general way, and the cumulative impact of the many features we introduce help move the research community forward. Given these points and the precedence of NeurIPS accepting software-centric papers as crucial contributions, we kindly request you to reconsider your assessment of our work."}, {"review_id": "b3vI28htAs", "reviewer": "Reviewer_V5p6", "summary": " Similarly to automatic differentiation which enabled to shift the focus from computing derivatives to deriving new algorithms, the authors propose to automate large scale linear algebra, with \"structure-aware\" linear algebra. Various algorithms leverage problem structure to expedite the evaluation of operations on linear operators. This toolbox aims to bridge the gap, allowing individuals without extensive knowledge of algorithms and tricks to still benefit from them. The successful accomplishment of this paper's objective could have a fair impact on the machine learning community by automating a significant portion of model optimization.", "questions": "**Porting to Julia.**\n\nIt is a nice use case of multiple dispatch, and I like the functional paradigm used. Are there any plans to port it in Julia? There already exists in the `LinearAlgebra` module \"structure aware\" types (`Symmetric`, `Hermitian`), and the multiple dispatch-functional paradigm are built in, so I believed it would certainly meet its public there. \n\n**Using SVRG.**\n\n`l. 245` and `248`, the complexity of SVRG seems bigger than Lanczos (at least in $\\kappa$; what is $M$?). A reference for the convergence rate of CG and Lanczos would be welcomed. Why is an accelerated version of SVRG not considered, to speed up convergence and further improve the runtime performances? The cyanure toolbox holds some [benchmarks](http://thoth.inrialpes.fr/people/mairal/cyanure/benchmarks.html) along with some [references](http://thoth.inrialpes.fr/people/mairal/cyanure/references.html). Is it because of the memory limitations? \n\n**Arnoldi Iterations.**\n\nIn Fig. 3, it is shown that `Arpack` performs better than the Python implementation of the Arnoldi operations. Given the modularity of your approach, wouldn't it be worth to add a dispatch rule which specifically uses `Arpack`? \n\n**Typos.**\n\n* `l. 82` : matrix vector **multiples** ? \n\n* Fig. 1 could benefit from having one legend for all 3 subplots. Contrary to what is asserted in the paper's checklist, there are no error bars. Given the moderate runtime ($10^2$s), using more than 3 repetitions (e.g 10) with error bar is necessary. \n\n* `245`: missing parenthesis ", "limitations": "The article has no Limitations section but seems compared extensively to other libraries in sec. 3.4. ", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "This article is well written. The problem is well posed, and the proposed solution is convincing. \n\nImportantly, the library's well-thought-out and modular structure is crucial for its growth and positive impact on the community.", "weaknesses": "The number of dispatch rule is relatively small. The authors mention 70 dispatch rules on `l. 224`, but in appendix A only a dozen such rules are shown (perhaps we reach 70 dispatch rules by counting all the possible combination?). If I had to write a new model, I'm wondering if (1) I would use this framework, with the (arguably small) associated overhead of using a new library or (2) simply scroll through appendix A to see what are the rules relevant for my implementation. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your thoughtful and supportive review! Below we hope to bring clarity to your questions.\n\n> If I had to write a new model, [could I] simply scroll through appendix A to see what are the rules relevant for my implementation?\n\nYes, implementing the necessary rules and algorithms for a specific problem is always possible. However, while one could simply use Appendix A as a lookup table to select appropriate algorithms, there are many scenarios where the best rules may not be immediately obvious, especially as the compositional structure becomes increasingly nested and/or complex. Additionally, we note that—beyond the dispatch rules—our library includes memory-efficient gradients of iterative operations that are difficult to implement correctly, particularly while retaining the matrix free LinearOperator abstraction. We expect that most researchers writing their own implementations of common algorithms such as GMRES or stochastic Lanczos quadrature will not find it high on their priority list to implement memory efficient backpropagation rules and thus leave the compute and memory savings unexploited. We also would like to think that having a large existing linear algebra ecosystem (that new rules can be slotted into if necessary) will help free up researcher time for other pursuits.\n\n**SVRG:** We have listed the complexity of CG and Lanczos $O(\\sqrt{\\kappa}\\log1/\\epsilon)$ in terms of the number of matrix-vector-multiplies, equivalent to full passes through the sum when applied to a sum linear operator $A=\\sum_{i=1}^M A_i$, where M here is the number of elements in the sum [7]. The $O((M+\\kappa)\\log1/\\epsilon)$ iteration complexity of SVRG (as it is usually expressed) becomes $O((1+\\kappa/M)\\log1/\\epsilon)$ measured by these full passes. For large values of M, SVRG can still be faster than CG/lanczos even without the acceleration.\nWe found that well chosen momentum values for accelerated SVRG can result in speedups; however, compared to the learning rate, we found it difficult to automate the selection of this hyperparameter from the data in a robust and inexpensive way.\n\n**Porting to Julia:** This would definitely be a worthwhile endeavor and we hope that this can be accomplished in the future, but it remains outside the scope of our current project. We chose to write the library in Python to interface with the popular ML frameworks of PyTorch and Jax. Nevertheless, the design of our library was certainly influenced by Julia and its programming paradigm, and we hope that these ideas can make it back into Julia LinearAlgebra.\n\n**Arnoldi Iterations:** We appreciate the suggestion and wrapping ARPACK. This is something that we have been debating ourselves as there is a clear runtime benefit but at the cost of not modular, not clean and hard to read code for the user. However, for the particular Arnoldi case that you mention, the best way to proceed is simply to wrap ARPACK and the result will still be autograd enabled because of how we define our custom autograd rules.\n\n**Typos:** We will correct the typos you have identified."}]

Author rebuttal:
We thank the reviewers for their thoughtful and strongly supportive feedback. In this general post, we highlight some of the new experiments that we conducted inspired by reviewer comments, and provide some general remarks about CoLA. We also have separate posts individually replying to each reviewer. 

We were happy to see that reviewers share our enthusiasm for CoLA. Linear algebra is a core foundation for machine learning algorithms, where common modeling assumptions give rise to structure that can be exploited for significant computational savings. CoLA will help make researchers more broadly aware of the structure they can exploit, and significantly reduce the bottleneck to implementing methods that exploit structure, as well as easily prototyping various different structures for their problems (as it is often not clear a priori what structure will be most beneficial for a given problem). 

**Impact and significance.** An analogy with PyTorch and reverse mode automatic differentiation is helpful for understanding the significance of CoLA and potential impact. From a narrow point of view, autograd is merely an application of the chain rule, and yet its impact on machine learning research has been almost immeasurable. Autograd obviates the need for deriving backpropagation rules for each model separately, and bespoke autograd rules (such as our own for iterative solvers in CoLA) can be slotted into an existing language of differentiable functions only when necessary, without requiring the whole structure to be constructed anew. Likewise with CoLA, we have developed an approach such that new rules can be slotted into an existing linear algebra ecosystem, and that ecosystem need not be rewritten for each use case. In this respect, the simplicity of CoLA is a strength, and will help enhance its usefulness to the community. Overall, the NeurIPS and broader ML community has found software libraries that simplify the research process as significant and impactful contributions (for example, consider the NeurIPS papers [1, 2, 3, 4, 5]). 

**Methodological contributions.** We also note that, while they are not the primary focus, we propose novel numerical algorithms that are important to our framework, such as doubly stochastic trace and diagonal estimation. When applied to linear operators with sum structure, we prove that this estimator has considerably lower variance than the standard Hutchinson estimator in Appendix B.1. More broadly, we have provided a novel procedure to automatically compute gradients through iterative methods, as well as an automatic procedure to compute diagonals and transposes / adjoints of linear operators (section 3.1 & 3.4) through their matrix-vector product routine.

**Additional experiments.** Inspired by reviewer feedback, we have put a significant effort into providing some new results (see attached pdf for Figures A, B and C):

 - Figure A (left) shows how our doubly stochastic estimator reduces the runtime by orders of magnitude when applied to large sums such as when estimating the variance (diagonal of the covariance) of the PCA application from Figure 2a.
 - Reviewer qUas suggested, for the Bi-Poisson problem in Figure 1(b), that we perform the comparison with a multi-grid solver, a method which is generally faster for solving elliptic PDEs. We have added this comparison in Figure A (right) using the CoLA decomposition rules to split the solve into two multi-grid solves, which like for our previous case in Figure 1 (b) when we were using conjugate gradients, CoLA also accelerates the convergence of multi-grid.
 - Moreover, for the spectral clustering example in Figure 3 we have now incorporated the LOBPCG algorithm into CoLA (before LOBPCG was giving scikit-learn an edge over us). As you can see in Figure B (left), CoLA’s LOBPCG results have improved significantly.
 - Furthermore, on Figure B (right), we have added an example on how low precision can improve runtime for linear regression.
 - Finally, we added on Figure C the runtime and memory consumption of backpropagating through a log determinant. We compare CoLA’s autograd rules against naively backpropagating through the iterative algorithm used to estimate the log determinant. The plot shows the substantial savings in runtime and memory from our approach. Taken together, Figure C and Figure 4 give the quantitative impact of our backprop rules for the two operations (solves and log determinants) needed for training Gaussian processes. As such, these backprop rules were used for the experiments in Figure 1 (a) and Figure 3 (c).

We are thankful for the questions, and would appreciate it if our responses and clarifications can be considered in your final assessment.

_References_

[1] Paszke et al., 2019. PyTorch: An Imperative Style, High-Performance Deep Learning Library. NeurIPS.

[2] Balandat et al., 2020. BoTorch: A Framework for Efficient Monte-Carlo Bayesian Optimization. NeurIPS.

[3] Daxberger et al., 2021. Laplace Redux – Effortless Bayesian Deep Learning. NeurIPS.

[4] Frank et al., 2021. Cockpit: A Practical Debugging Tool for the Training of Deep Neural Networks. NeurIPS.

[5] Pineda et al. 2022. Theseus: A Library for Differentiable Nonlinear Optimization. NeurIPS.

[6] Gardner et al., 2018. GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration. NeurIPS.

[7] Golub et al., 2018. Matrix Computations. 4th Edition. The Johns Hopkins University Press.

Meta-review: This is a well-written paper presenting work on a linear algebra library that automatically exploits various types of structure present in linear operations. It ought to be of interest to users and numerical library implementors alike. Expert reviewers all found the paper well-written and recommend acceptance to varying degrees. I second this recommendation overall.

The feedback during review can be decomposed roughly into (a) feedback on the framework itself (CoLA), and (b) feedback on the scientific contribution of the paper that presents CoLA.

For the framework itself, reviewers largely agree that CoLA seems interesting and useful. Some noted strengths include portability atop Pytorch and JAX, a variety of numerical procedures and operator structures covered, and automatic differentiation.

Regarding the scientific contribution of the paper, I would highlight the discussion between the authors and reviewers qUas, pi1f, and V5p6. These reviewers all contended that on the algorithmic side, the framework does not extend what is already known in numerical linear algebra and could be implemented on a per-application basis. As the discussion went, there is more substantial contribution elsewhere, e.g. on automation (both in kernels and in automatic differentiation) and evaluation on machine learning applications.
Decision: Accept (poster)