Title: CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra
Venue: NeurIPS 2023
Accepted: yes
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

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.

1 Introduction

The framework of automatic differentiation has revolutionized machine learning. Although the rules that govern derivatives have long been known, automatically computing derivatives was a 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 as PyTorch, Chainer, Tensorflow, JAX, and others [1, 8, 30, 31, 38, 47] 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-scale linear algebra (e.g. matrix solves, eigenvalue problems, nullspace computations). These ubiquitous operations are at 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-such as diagonal dominance, sparsity, or a low-rank factorization. Given a 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 a general-purpose routine (e.g., the cubic-time Cholesky decomposition). However, exploiting

Base

Case

D

T

P

C

S

Pr

∑

∏

⊗

[

A

0

0

B

]

[

A

B

C

D

]

A

-

1

(

A

)

Diag

(

A

)

Tr(

A

)

exp(

A

)

A

)

structure for faster computation is often an intensive implementation process. Rather than having an object A in code that represents a low-rank-plus-diagonal matrix and simply calling solve ( A , b ) , a practitioner must instead store the low-rank factor F as a matrix, the diagonal d as a vector, and implement the Woodbury formula from scratch. Implementing structure-aware routines in machine learning models is often seen as a major research undertaking. For example, a nontrivial portion of the Gaussian process literature is devoted to deriving specialty inference algorithms for structured kernel matrices [e.g. 7, 11, 19, 25, 29, 46, 52, 53, 24].

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 LinearOperator objects which adhere to the same API as standard dense matrices. For example, a user can call A -1 b or eig ( A ) on any LinearOperator A , and under-the-hood our framework derives a computationally efficient algorithm built from our set of compositional dispatch rules (see Table 1). If little is known about A , the derived algorithm reverts to a general-purpose base case (e.g. Gaussian elimination or GMRESforlinear solves). Conversely, if A is known to be the Kronecker product of a lower triangular matrix and a 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 special-purpose implementations across numerous applications despite relying on only a small number of base LinearOperator types.

Furthermore, our framework offers additional novel functionality that is necessary for ML applications (see Table 2). In particular, we automatically compute gradients, diagonals, transposes and adjoints of linear operators, and we modify classic iterative algorithms to ensure numerical stability in low precision. We also support specialty algorithms, such as SVRG [23] and a novel variation of Hutchinson's diagonal estimator [22], which exploit implicit structure common to matrices in machine learning applications (namely, the ability to express matrices as large-scale sums amenable to stochastic approximations). Moreover, our framework is easily extensible in both directions : a user can implement a new linear operator (i.e. one column in Table 1), or a new linear algebraic operation (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.

We term our framework CoLA ( Co mpositional L inear A lgebra), which we package in a library that supports both PyTorch and JAX. We showcase the extraordinary versatility of CoLA with

a 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 https://github.com/wilson-labs/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 problems [14, 44]; (3) Kronecker-factorizable matrices, which admit efficient spectral decompositions; (4) Toeplitz or circulant matrices, which admit fast matrix-vector products. See Section 3 and Section 4 for applications that use 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, it is possible to design algorithms that exploit their inherent properties (see Section 3.3).

Iterative matrix-free algorithms Unlike direct methods, which typically require dense instantiations 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-such as conjugate gradients, GMRES, Lanczos and Arnoldi iteration-fall under the category of Krylov subspace methods, which iteratively apply MVMs to refine a solution until a desired error tolerance is 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.) or when using accelerated hardware (GPUs or TPUs) designed for efficient parallel MVMs [e.g. 10, 20, 51].

Multiple dispatch Popularized by Julia [6], multiple dispatch is a functional programming paradigm for defining type-specific behaviors. Under this paradigm, a given function (e.g. solve ) can have multiple definitions, each of which are specific to a particular set of input types. A base-case definition solve[LinearOperator] would use a generic matrix-vector solve algorithm (e.g. Gaussian elimination or GMRES), while a type-specific definition (e.g. solve[Sum] , for sums of matrices) would use a special purpose algorithm that makes use of the subclass' structure (e.g. SVRG, see Section 3.3). When a user calls solve ( A , b ) at 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 compositional patterns of 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 to define a specific SumOfLowRankAndDiagonal 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 is often a manual effort. Consider for example the problems of second order/natural gradient optimization, which require matrix solves with (potentially large) Hessian/Fisher matrices. Researchers have proposed tackling these solves with matrix-free methods [33], diagonal approximations [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 SciPy Sparse [50], Spot [49], PyLops [41], or GPyTorch [20], which offer a unified interface for using matrix-free algorithms with any type of structured matrices. A 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

gradients; see Figure 1.) Moreover, these libraries often lack modern features (e.g. GPU acceleration or automatic differentiation) or are specific to certain types of matrices (see Table 2).

3 CoLA: Compositional Linear Algebra

We now discuss all the components that make CoLA. In Section 3.1 we first describe the core MVMbased 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 CoLA exploits structure frequently encountered in ML applications beyond wellknown analytic formulae (e.g. the Woodbury identity). Finally, in Section 3.4 we present CoLA's machine learning-specific 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 is the LinearOperator : a linear function on a finite dimensional vector space, defined by how it acts on vectors via a matrix-vector multiply MVM A : v ↦→ Av . While this function has a matrix representation for a given basis, we do not need to store or compute this matrix to perform a MVM. Avoiding the dense representation of the operator saves memory and often compute.

Some basic examples of LinearOperators are: unstructured Dense matrices, which are represented by a 2-dimensional array and use the standard MVM routine [ Av ] i = ∑ j =1 A ij v 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 are represented by a 1-dimensional array of the diagonal entries and where the MVM is given by [ Diag ( d ) v ] i = d i v i ; Convolution operators, which are represented by a convolutional filter array and where the MVM is given by Conv ( a ) v = a ∗ v ; or JVP operatorsthe Jacobian represented implicitly through an autograd Jacobian Vector Product-represented by a function and an input x and where the MVM is given by Jacobian ( f, x ) v = JVP( f, x , v ) . In CoLA, each of these examples are sub-classes of the LinearOperator superclass.

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 ( e 1 ) , . . . , MVM ( e N ) , on each unit vector e i . We now describe several key operations supported by our framework, some well-established, and others novel to 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 of these algorithms can be applied to any non-symmetric and/or complex linear operator. When LinearOperator s are annotated with additional structure (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

matrix free (and thus memory efficient), amenable to 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 a user must manually define a transposed MVM v ↦→ A ⊺ v for linear operator objects. In contrast, CoLA uses a novel autograd trick to derive the transpose from the core MVM routine. We note that A ⊺ v is the vector-Jacobian product (VJP) of the vector v and the Jacobian ∂ Aw /∂ w . Thus, the function transpose ( A ) returns a LinearOperator object that uses VJP( MVM A , 0 , v ) as its MVM . We extend this idea to Hermitian conjugates, using the fact that A ∗ v = ( A ) ⊺ v = ( A ⊺ v ) .

Other operations In Section 3.3 we outline how to stochastically compute diagonals and traces of operators with MVM s, and in Section 3.4 we discuss a novel approach for computing memory-efficient derivatives of iterative methods through MVM s.

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 a minor implementation detail: as we demonstrate in the next section, it is crucial for the efficiency and compositional power of 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 as 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 -1 = Diag ( a -1 ) without needing to solve a linear system. Similarly, solves with triangular matrices can be inverted efficiently through back substitution, and solves with circulant matrices can be computed efficiently in the Fourier domain Conv ( a ) = F -1 Diag ( F a ) F (where F is the Fourier transform linear operator). We offer more examples in Table 1 (left).

As described in Section 2, we use multiple dispatch to implement these special case methods. For example, we implement the solve[Diagonal] , solve[Triangular] , and solve[Circulant] dispatch rules using the efficient routines described above. If a specific LinearOperator subclass does not have a 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.

The dispatch framework makes it easy to implement one-off rules for the basic LinearOperator sub-classes described in Section 3.1. 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 (e.g. Sparse , Diagonal , Convolution ), our framework provides mechanisms to compose multiple LinearOperator s together. Some frequently used compositional structures are Sum ( ∑ i A i ), Product ( Π i A i ), Kronecker ( A ⊗ B ), KroneckerSum ( A ⊕ B ), BlockDiag [ A , 0; 0 , B ] and Concatenation [ A , B ] . Each of these compositional LinearOperator s are defined by (1) the base LinearOperator objects to be composed, and (2) a corresponding MVM routine, which is typically written in terms of the MVM s of the composed LinearOperators . For example, MVM Sum = v ↦→ ∑ i MVM i ( v ) , where MVM i are the MVM routines for the component LinearOperator s.

Dispatch rules for compositional operators are especially powerful. For example, consider Kronecker products where we have the rule ( A ⊗ B ) -1 = A -1 ⊗ B -1 . Though simple, this rule yields highly efficient routines for numerous structures. For example, suppose we want to solve ( A ⊗ B ⊗ C ) x = b 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 MVM Kron ).

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 = A n ; in general, the runtime is upper-bounded by O ( nτ √ κ n log 1 ϵ ) , where τ 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 ( nτ √ κ log n ϵ ) . 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 dispatch is the ability to write simple special rules for compositions of specific operators. While a 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 a solve[Sum[LowRank, Diagonal]] dispatch rule that is used if the Sum operator is specifically combining a LowRank and a Diagonal linear operator. Coding these rules without multiple dispatch would require specialty defining sub-classes like LowRankPlusDiagonal over the LinearOperator object, increasing complexity and hampering extendibility.

Decoration/annotation operators Finally, we include several decorator types that annotate existing LinearOperator s with additional structure. For example, we define SelfAdjoint (Hermetian/symmetric), Unitary (orthonormal), and PSD (positive semi-definite) operators, each of which wraps an existing LinearOperator object. None of these decorators define a specialty MVM ; 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, 6 decoration linear operators, and roughly 70 specialty dispatch rules for solve , eig , and other operations. (See Table 1 for a short summary and Appendix A for a 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 or 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 (1) define a 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 explicit matrix structures and composition rules for which there are simple analytic formulas easily found in 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 linear operators composed of large summations which are amenable to stochastic algorithms. Below we 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) is 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 + κ/M ) log 1 ϵ ) where κ is the condition number and ϵ is the desired accuracy.

0

25

50

75

100

Size (10 2

)

10

4

10

3

10

2

10

1

10 0

Runtime (sec)

Dense

Iterative

CoLA

0

25

50

75

100

Size (10 4

)

10

2

10

1

10 0

10 1

Runtime (sec)

Multi-grid

Multi-grid + CoLA

0

50

100

150

Size (10 2

)

10 0

10 1

10 2

Runtime (sec)

Dense

Iterative

CoLA

(a) Multi-task GPs (

⊗

)

(b) Bi-Poisson (

Π

)

(c) EMLP









A

0

0

B







0

50

100

10

6

10

4

10

2

10 0

10 2

Convergence Criteria

CoLA (SVRG)

Iterative

0

50

100

150

200

10

6

10

3

10 0

Residual

CoLA (SVRG)

Iterative

0

50

100

150

200

10

4

10

2

10 0

Residual

CoLA (SVRG)

Iterative

(a) PCA (Eig)

(b) GPs RFF (Solve)

(c) Neural-IVP (Solve)

When the condition number and the number of elements in the sum is large, SVRG becomes a desirable alternative even to classical deterministic iterative algorithms such as CG or Lanczos whose runtimes are bounded by O ( √ κ log 1 ϵ ) . 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 or the diagonal of a linear operator. While collecting the diagonal for a dense matrix is a trivial task, it is a costly algorithm for an arbitrary LinearOperator defined only through its MVM -it requires computing Diag ( A ) = ∑ N i =1 e i ⊙ A e i where ⊙ 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 ) = 1 n ∑ n j =1 z j ⊙ A z j where the z j ∈ R N are any randomly sampled probe vectors with covariance I . We extend this randomized estimator to use randomization both in the probes, and random draws from a sum when A = ∑ M i =1 A i :

0

0.25

0.5

0.75

1.0

Runtime

Tre(20K)

Solves

Eig

bcs18(11.9K)

Logdet

Torch

CoLA

CPU

GPU

In Section B.1 we derive the variance of this estimator and we show that it converges faster than the base Hutchinson estimator when applied Sum structures. We validate empirically this analysis in Figure 5.

3.4 Automatic Differentiation and Machine Learning Readiness

Memory efficient auto-differentiation In ML applications, we want to backpropagate through operations like A -1 , Eigs ( A ) , Tr ( A ) , exp( A ) , log det( A ) . To achieve this, in CoLA we define a novel concept of the gradient of a LinearOperator which we detail in Appendix B. For routines like GMRES, SVRG, and Arnoldi, we utilize a 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 in 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 PyTorch and JAX. This compatibility not only makes our framework plug-and-play with existing implemented models, but it also adds GPU/TPU support, differentiating it 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 on hardware such as GPUs. Figure 3 empirically shows the additional impact of hardware accelerators across different datasets and linear algebra operations.

4 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 3 (Figure 1 & Figure 2), the applications presented here have a basic structure (sparse, vector-product, etc) but not a 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

k=5

k=10

PCA Components

10

1

10 0

10 1

Runtime (sec)

CoLA

sk

10 2

10 3

10 4

10 5

Dasetset Size

10

2

10 0

Runtime (sec)

sk

CoLA (GPU)

CoLA (CPU)

0

10

20

30

40

50

Runtime (sec)

0

25

50

75

100

Epochs

CoLA (ele)

GP (ele)

CoLA (kin)

GP (kin)

(a) PCA

(b) Linear Regression

(c) GPs

10 4

10 5

10 6

Edges

10

1

10 0

10 1

Runtime (sec)

sk (B)

sk (L)

CoLA (L)

CoLA (B)

10 2

10 3

10

1

10 1

Runtime (sec)

SciPy

CoLA

10 2

10 3

10 4

10

1

10 0

10 1

Runtime (sec)

SciPy

CoLA

SciPy (JAX)

(d) Spectral Clustering

(e) Schrödinger Equation

(f) Minimal Surface

performs in any application. We compare against several well-known libraries, sometimes providing runtime improvements but other times performing equally. This is remarkable as 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 or SciPy ). Below we describe each of the applications found in Figure 4.

Principal Component Analysis PCA is a 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 ∈ R N × D . 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 a diagonal matrix containing the degree of the nodes and W the weights on the edges between nodes. This problem requires finding the smallest k eigenvectors of L . We run 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 T K -1 y (where y is the vector of observations).

Schrödinger Equation In this problem we characterize the spectrum of an atom or molecule by finding the eigenspectrum of a PDE operator in a Schrodinger equation H ψ = Eψ . After discretizing ψ to a grid, we compute the smallest eigenvalues and eigenvectors of the operator H which for this experiment is non-symmetric as we perform a compactfying transform.

Minimal Surface Here we solve a 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+ z 2 x ) z yy -2 z x z y z xy +(1+ z 2 y ) z xx = 0 and solved by root finding on a discrete grid. Applying Newton-Raphson, we iteratively solve the non-symmetric linear system z ← z -J -1 f ( z ) where J is the Jacobian of the PDE operator.

Bi-Poisson Equation The Bi-Poisson equation ∆ 2 u = ρ is a linear boundary value PDE relevant in continuum mechanics, where ∆ is the Laplacian. When discretized using a grid, the result is a 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 PDEs. One approach for initial value problems requires advancing an ODE on the neural network parameters θ , where ˙ θ = M ( θ ) -1 F ( θ ) where M is an operator defined from Jacobian of the neural network which decomposes as the sum over data points M = 1 N ∑ i M i and where F 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] as shown in Figure 2(c).

Equivariant Neural Network Construction As shown in [16], constructing the equivariant layers of a neural network for a 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 a block diagonal matrix of concatenated Kronecker products and Kronecker sums of sparse matrices. In Figure 1(c) we show the empirical benefits of exploiting this structure.

5 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 a wide variety of practical applications. Algorithms like SVRG and a novel variation of Hutchinson'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-framework 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.

References

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: Largescale machine learning on heterogeneous systems, 2015. URL https://www.tensorflow. org/ . Software available from tensorflow.org.

Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable Second Order Optimization for Deep Learning. Preprint arXiv 2002.09018v2 , 2020.

Walter Edwin Arnoldi. The principle of minimized iterations in the solution of the matrix eigenvalue problem. Quarterly of applied mathematics , 9(1):17-29, 1951.

S Becker and Yann Lecun. Improving the convergence of back-propagation learning with second-order methods. In Proceedings of the 1988 Connectionist Models Summer School, San Mateo , pages 29-37. Morgan Kaufmann, 1989.

Nathan Bell, Luke N. Olson, Jacob Schroder, and Ben Southworth. PyAMG: Algebraic Multigrid Solvers in Python. Journal of Open Source Software , 2023.

Jeff Bezanson, Alan Edelman, Stefan Karpinski, and Viral B. Shah. Julia: A Fresh Approach to Numerical Computing. arXiv preprint arXiv:1411.1607 , 2014.

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.

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 of Python+NumPy programs. SoftwareX , 2018. URL http://github.com/google/jax .

M. Brezina, R. Falgout, S. MacLachlan, T. Manteuffel, S. McCormick, and J. Ruge. Adaptive Smoothed Aggregation ( α SA) Multigrid. SIAM Review , 2005.

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 , 22(1):3457-3462, 2021.

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.

Marco Cuturi. Sinkhorn Distances: Lightspeed Computation of Optimal Transport. Advances in Neural Information Processing Systems (NeurIPS) , 2013.

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.

Timothy A Davis. Direct methods for sparse linear systems . SIAM, 2006.

Yifan Du and Tamer A Zaki. Evolutional Deep Neural Network. Physical Review E , 104(4): 045303, 2021.

Marc 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.

Marc Finzi, Andres Potapczynski, 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.

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.

Jacob Gardner, Geoff Pleiss, Ruihan Wu, Kilian Weinberger, and Andrew Wilson. Product kernel interpolation for scalable gaussian processes. In International Conference on Artificial Intelligence and Statistics , pages 1407-1416. PMLR, 2018.

Jacob R. Gardner, Geoff Pleiss, David Bindel, Kilian Q. Weinberger, and Andrew Gordon Wilson. GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration. Advances in Neural Information Processing Systems (NeurIPS) , 2018.

Gene H Golub and Charles F Van Loan. Matrix Computations . The Johns Hopkins University Press, 2018. Fourth Edition.

Michael F Hutchinson. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics-Simulation and Computation , 18(3):10591076, 1989.

Rie Johnshon and Tong Zhang. Accelerating Stochastic Gradient Descent using Predictive Variance Reduction. Advances in Neural Information Processing Systems (NeurIPS) , 2013.

Sanyam Kapoor, Marc Finzi, Ke Alexander Wang, and Andrew Gordon Gordon Wilson. SKIing on Simplices: Kernel Interpolation on the Permutohedral Lattice for Scalable Gaussian Processes. International Conference on Machine Learning (ICML) , 2021.

Matthias Katzfuss and Joseph Guinness. A general framework for vecchia approximations of gaussian processes. Statistical science , 36(1):124-141, 2021.

Andrew Knyazev. Toward The Optimal Preconditioned Eigensolver: Locally Optimal Block Preconditioned Conjugate Gradient Method. SIAM Journal on Scientific Computing , 2000.

Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Neural Operator: Learning Maps Between Function Spaces. Preprint arXiv 2108.08481v3 , 2021.

Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the Intrinsic Dimension of Objective Landscapes. International Conference on Learning Representations (ICLR) , 2018.

Jackson Loper, David Blei, John P Cunningham, and Liam Paninski. A general linear-time inference method for gaussian processes on one dimension. The Journal of Machine Learning Research , 22(1):10580-10615, 2021.

Dougal Maclaurin. Modeling, inference and optimization with composable differentiable procedures . PhD thesis, School of Engineering and Applied Sciences, Harvard University, 2016.

Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter optimization through reversible learning. In International conference on machine learning , pages 2113-2122. PMLR, 2015.

Wesley J. Maddox, Andres Potapczynski, and Andrew Gordon Wilson. Low-Precision Arithmetic for Fast Gaussian Processes. Conference on Uncertainty in Artificial Intelligence (UAI) , 2022.

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.

James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning , pages 2408-2417. PMLR, 2015.

Per-Gunnar Martinsson and Joel Tropp. Randomized Numerical Linear Algebra: Foundations &Algorithms . arXiv 2002.01387v3, 2020.

Andrew Y. Ng, Michael I. Jordan, and Yair Weiss. On Spectral Clustering: Analysis and an algorithm. Advances in Neural Information Processing Systems (NeurIPS) , 2001.

Eric Nguyen, Karan Goel, Albert Gu, Gordon W. Downs, Preey Shah, Tri Dao, Stephen A. Baccus, and Christopher Ré. S4ND: Modeling Images and Videos as Multidimensional Signals Using State Spaces. Preprint arXiv 2210.06583v2 , 2022.

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.

Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron C. Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018 , pages 3942-3951. AAAI Press, 2018.

Ali Rahimi and Ben Recht. Random Features for Large-Scale Kernel Machines. Advances in Neural Information Processing Systems , 2007.

Matteo Ravasi and Ivan Vasconcelos. PyLops-A linear-operator Python library for scalable algebra and optimization. SoftwareX , 11:100361, 2020. ISSN 2352-7110. doi: https://doi.org/10. 1016/j.softx.2019.100361. URL https://www.sciencedirect.com/science/article/ pii/S2352711019301086 .

Nicolas Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algorithm. Advances in neural information processing systems , 20, 2007.

Youcef Saad and Martin H Schultz. Gmres: A generalized minimal residual algorithm for solving nonsymmetric linear systems. SIAM Journal on scientific and statistical computing , 7 (3):856-869, 1986.

Yousef Saad. Iterative methods for sparse linear systems . SIAM, 2003.

Ohad Shamir. A Stochastic PCA and SVD Algorithm with an Exponential Convergence Rate. arXiv preprint arXiv:1409.2848v5 , 2015.

Edward Snelson and Zoubin Ghahramani. Sparse gaussian processes using pseudo-inputs. Advances in neural information processing systems , 18, 2005.

Seiya Tokui, Kenta Oono, Shohei Hido, and Justin Clayton. Chainer: a next-generation open source framework for deep learning. In NeurIPS Workshop on Machine Learning Systems (LearningSys) , volume 5, pages 1-6, 2015.

Lloyd N. Trefethen and David Bau. Numerical Linear Algebra . SIAM, 1997.

Ewout van den Berg and Michael P. Friedlander. Spot - A Linear-Operator Toolbox. SoftwareX , 2013. URL http://www.cs.ubc.ca/labs/scl/spot/ .

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, C J 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-272, 2020. doi: 10.1038/s41592-019-0686-2.

Ke Wang, Geoff Pleiss, Jacob Gardner, Stephen Tyree, Kilian Q Weinberger, and Andrew Gordon Wilson. Exact gaussian processes on a million data points. Advances in neural information processing systems , 32, 2019.

Andrew Wilson and Hannes Nickisch. Kernel interpolation for scalable structured gaussian processes (kiss-gp). In International conference on machine learning , pages 1775-1784. PMLR, 2015.

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.

Max A Woodbury. Inverting modified matrices . Department of Statistics, Princeton University, 1950.

Peng Xu, Bryan He, Christopher De Sa, Ioannis Mitliagkas, and Chris Re. Accelerated stochastic power iteration. In International Conference on Artificial Intelligence and Statistics , pages 58-67. PMLR, 2018.

Appendix Outline

This Appendix is 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-efficient autograd implementation.

In Appendix C we include pseudo-code on various of the iterative methods incorporated in CoLA and discuss modifications to improve lower precision performance.

In Appendix D we expand on 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 ) -1 = ( B -1 A -1 ) and for Kronecker product we have ( A ⊗ B ) -1 = A -1 ⊗ B -1 . In terms of block compositions we have the following identities:

Finally, for sum we have the Woodbury identity and its variants. Namely, for Woodbury we have

the Kailath variant where

and the rank one update via the Sherman-Morrison formula

Besides the compositional operators, we have some rules for some special operators. For example, for A = Diag ( a ) we have A -1 = Diag ( a -1 ) . Also, if Q is unitary then Q -1 = Q ∗ or if Q is orthonormal then Q -1 = Q ⊺ .

A.1.2 Eigendecomposition

We now assume that the matrices in this section are diagonalizable. That is, Eigs ( A ) = Λ A , V A , where A = V A Λ A V -1 A . In terms of the compositional operators, there is not a general rule for product or sum. However, for the Kronecker product we have Eigs ( A ⊗ B ) = Λ A ⊗ Λ B , V A ⊗ V B and for the Kronecker sum we have Eigs ( A ⊕ B ) = Λ A ⊕ Λ B , V A ⊗ V B . Finally, for block diagonal we have

A.1.3 Diagonal

As a base case, if we need to compute Diag ( A ) for a general matrix A we may compute each diagonal element by e ⊺ i Ae i . Additionally, if A is large enough we switch to randomized estimation Diag ( A ) ≈ ( Z ⊙ AZ ) 1 /N with Z ∼ N (0 , 1) d × N where N is the number of samples used to approximate 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 ⊗ B ) = vec ( Diag ( A ) Diag ( B ) ⊺ ) and for Kronecker sum Diag ( A ⊕ B ) = vec ( Diag ( A ) 1 ⊺ + 1 Diag ( B ) ⊺ ) . Finally, for block composition we have

A.1.4 Transpose / Adjoint

As explained in Section 3.1, as a base case we have an automatic procedure to compute the transpose or adjoint of any operator A via autodiff. However, we also incorporate the following rules. For sum we have ( A + B ) ∗ = A ∗ + B ∗ and ( A + B ) ⊺ = A ⊺ + B ⊺ . For product we have ( AB ) ∗ = B ∗ A ∗ and ( AB ) ⊺ = B ⊺ A ⊺ . For Kronecker product we have ( A ⊗ B ) ∗ = A ∗ ⊗ B ∗ and ( A ⊗ B ) ⊺ = A ⊺ ⊗ B ⊺ . For the Kronecker sum we have ( A ⊕ B ) ∗ = A ∗ ⊕ B ∗ and ( A ⊕ B ) ⊺ = A ⊺ ⊕ B ⊺ . In terms of block composition we have

Finally for the annotated operators we have the following rules. A ∗ = A if A is self-adjoint and A ⊺ = A if A is symmetric.

A.1.5 Pseudo-inverse

As a base case, if we need to compute A + , we may use SVD ( A ) = U , Σ , V and therefore set A + = UΣ + V ∗ , where Σ + inverts the nonzero diagonal scalars. If the size of A is too large, then we may use randomized SVD. Yet, it is uncommon to simply want A + , usually we want to solve a least-squares problem and therefore we can use solvers that are not as expensive to run as SVD. For the compositional operators we have the following identities. For product ( AB ) + = ( A + AB ) + ( ABB + ) + and for Kronecker product we have ( A ⊗ B ) + = A + ⊗ B + . For block diagonal we have

.

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 Q is orthonormal. If A is self-adjoint, then A + = A -1 and also if it is symmetric and PSD.

A.2 Derived Functions

Interestingly, the previous core functions allow us to derive multiple rules from the previous ones. To illustrate, we have that Tr ( A ) = ∑ i Diag ( A ) i . Additionally, if A is PSD we have that f ( A ) = V A f ( Λ A ) V -1 A and if A is both symmetric and PSD then f ( A ) = V A f ( Λ A ) V ⊺ A . where in both cases we used Eigs ( A ) = Λ A , V A . Some example functions for PSD matrices are Sqrt ( A ) = V A Λ 1 / 2 A V -1 A or Log ( A ) = V A log Λ A V -1 A . Which also this rules allow us to define LogDet ( A ) = Tr ( Log ( A )) .

A.3 Other matrix identities

We emphasize that there are a myriad more matrix identities that we do not intentionally include such as Tr ( A + B ) = Tr ( A ) + Tr ( B ) or Tr ( AB ) = Tr ( BA ) when A and B are squared. These

additional cases are not part of our dispatch rules as 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:

with each z j ∼ N ( 0 , I D ) where A is a D × D matrix. When A is itself a sum A = 1 m ∑ m i =1 A i , we can expand the trace as Tr [ Base ]( A ) = 1 mn ∑ n j =1 ∑ m i =1 z ⊺ j A i z j , with probe variables shared across elements of the sum.

Consider the quadratic form Q := z ⊺ Az , which for Gaussian random variables has a cumulant generating function of K Q ( t ) = log E [ e tQ ] = -1 2 log det( I -2 t A ) . From the generating function we can derive the mean and variance of this estimator: E [ Q ] = K ′ Q (0) = Tr( A ) and Var[ Q ] = K ′′ Q (0) = 2Tr( A 2 ) . Since Tr [ Base ]( A ) is a sum of independent random draws of Q , we see:

Doubly Stochastic Trace Estimator For the doubly stochastic estimator, we choose probe variables which are sampled independently for each element of the sum:

Separating out the elements of the sum, we can write the estimator as Tr [ Sum ]( A ) = 1 n ∑ n j =1 R j where R j are independent random samples of the value R = 1 m ∑ m i =1 z ⊺ i A i z i . The cumulant generating function is merely K R ( t ) = ∑ m i =1 K Q i ( t/m ) where Q i = z ⊺ A i z . Taking derivatives we find that,

Assuming bounded moments on A i , then both A = 1 m ∑ i A i and S ( A ) = 1 m ∑ i A 2 i will converge to fixed values as m →∞ . Given that Tr [ Sum ]( A ) = 1 n ∑ n j =1 R j , we can now write the mean and variance of the doubly stochastic estimator:

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 / √ n ) (even when applied to sum structures), whereas the error for Tr [ Sum ] is O (1 / √ nm ) , a significant asymptotic variance reduction.

The related stochastic diagonal estimator

achieves the same O (1 / √ nm ) convergence rate, though we omit this derivation for brevity as it is follows the same steps.

In Figure 5 we empirically how our doubly stochastic diagonal estimator outperforms the standard Hutchinson estimator.

10

2

10 0

10 2

Passes through dataset

10

2

10

1

10 0

Diagonal Estimate Relative Error

Hutchinson

Doubly Stochastic

10 2

10 3

10 4

CG iters

10 0

10 1

10 2

Runtime (sec)

(

A

1

b

)

Unrolled autodiff

CoLA vjp rules

10 2

10 3

10 4

CG iters

10 0

10 1

10 2

Memory (MB)

(

A

1

b

)

Unrolled autodiff

CoLA vjp rules

10 2

10 3

Lanczos iters

10 0

10 1

Runtime (sec)

A

|

Unrolled autodiff

CoLA vjp rules

10 2

10 3

Lanczos iters

10 0

10 1

10 2

Memory (MB)

A

|

Unrolled autodiff

CoLA vjp rules

(a) Backwards pass runtime

(b) Backwards pass memory

B.2 Autograd rules for iterative algorithms

For machine learning applications, we want to 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 an operation like f = CGSolve , CGSolve ( A , b ) = A -1 b we must define a Vector Jacobian Product: VJP ( f, ( A , b ) , v ) = ( v ⊺ ∂f ∂ A , v ⊺ ∂f ∂ b ) . 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: θ = flatten [ A ] . For example for A = Kron ( Diag ( θ 1 ) , Conv ( θ 2 ) ) , flatten [ A ] = [ θ 1 , θ 2 ] . From this perspective, we consider A as a container or tree of its arguments θ , and define v ⊺ ∂f ∂ A := unflatten [ v ⊺ ∂f ∂ θ ] which coincides with the usual definition for dense matrices. Applying to inverses, we can now write a simple VJP :

for v ⊺ ∂f ∂ θ = v ⊺ ( A -1 ) ⊺ ( ∂ θ A θ ) A -1 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.

In Figure 6 we show the practical benefits of our autograd rules. We take gradients of different linear solves A -1 θ b that were derived using conjugate gradients (CG), where each solve required an increasing number of 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 we use as base algorithms in CoLA. Finally, for the third point we expand upon CoLA's strategy for dealing with the different numerical precisions that we support.

C.1 SVRG

In simplest form, SVRG [23] performs gradient descent with the varianced reduced gradient

where g i represents the stochastic gradient evaluated at only a single element or minibatch of the sum, and g ( w 0 ) is the full batch gradient evaluated at the anchor point w 0 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

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 0 , then the corresponding linear algebra solution has been recovered.

While it may seem that we need to take three complete passes through { A i } per SVRG epoch (due to 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 i W -A i W 0 with A i ( W -W 0 ) where appropriate. In all of the Sum structure experiments where we leverage SVRG, the x-axis measures the total number of passes through { A i } i m =1 , two for each epoch for SVRG.

C.2 Iterative methods

In Table 5 we list the different iterative methods (base cases) that we use 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].

C.3 Lower precision linear algebra

The accumulation of round-off error is usually the breaking point of several numerical linear algebra (NLA) routines. As such, it is common to use precisions like float64 or higher, especially when running these routines on a 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 as detrimental when training machine learning models (as we are already

running noisy optimization algorithms) as when solving linear algebra problems (where round-off error can lead us to poor solutions). Thus, it 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 better 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 at 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 are two algorithms that are key for eigendecompositions. The first is Arnoldi (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 , a Householder projector (or Householder reflector) is defined as the following operator R = I -2 uu ∗ . 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 ∈ R 2 and u = e 1 . Hence,

which is exactly the reflection of the vector across the axis generated by e 2 . Most notably, R is unitary RR ∗ = 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, ∥ R ( A + E ) -RA ∥ = ∥ 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 Arnoldi present in CoLA. The implementations are notably different and also it is easy to see how Algorithm 2 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 1 we only have one (line 4-6). Worse, the two for loops in Algorithm 2 require more flops than the only for loop in Algorithm 1. 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 is 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 hyperparameters of different algorithms and the specific choices of algorithms used both for CoLA but also for the alternatives. We run each of the experiments 3 times and compute the mean dropping the first observation (as usually the first run contains some compiling time much is 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 Intel(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://archive.ics.uci.edu/ml/datasets.php . Also, a community repo hosting these UCI benchmarks can be found here: https://github.com/treforevans/uci_ datasets (we have no affiliation).

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 = 14 K observations and has D = 18 dimensions.

̸

Algorithm 1 Arnoldi iteration

Algorithm 2 Householder Arnoldi iteration

Kin40K . The full name of this UCI dataset is Statlog (Shuttle) Data Set . This dataset contains information about NASA shuttle flights and we used a subset that consists of N = 40 K observations and has D = 8 dimensions.

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 a subset consisting of N = 430 K observations and has D = 77 dimensions.

Song . The full name of this UCI dataset is YearPredictionMSD . This dataset consists of N = 386 . 5 K observations and it has D = 90 audio features such as 12 timbre average features and 78 timbre covariance features.

cit-HepPh . This dataset is based on arXiv's HEP-PH (high energy physics phenomenology) citation graph and can be found here: https://snap.stanford.edu/data/cit-HepPh. html . The dataset covers all the citations from January 1993 to 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.

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 1's panels.

(a) The multi-task GP problem exploits the structure of the following Kronecker operator K T ⊗ K X , where K T is a kernel matrix containing the correlation between the tasks and K X is a RBF kernel on the data. For this experiment, we used a synthetic Gaussian dataset where the train data x i ∼ N ( 0 , I D ) which has dimension D = 33 , N = 1 K and we used T = 11 tasks (where the tasks basically set the size of K T ). We used conjugate gradients (CG) as the iterative method, where we set the hyperparameters to a tolerance of 10 -6 and to a maximum number of iterations to 1 K. We used the exact same hyperparameters for CoLA.

(b) For the bi-poisson problem we set up the maximum grid to be N = 1000 2 . Since this PDE problem involves solving a symmetric linear system, we used CG as the iterative method with a tolerance of 10 -11 and a maximum number of iterations of 10 K. The previous

parameters also apply for CoLA. We note that PDE problems are usually solved to higher tolerances as the numerical error compounds as we advance the PDE.

(c) For the EMLP experiment we consider solving the equivariance constraints to find the equivariant linear layers of a graph neural network with 5 nodes. To solve this problem, we need to find the nullspace of a 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 a tolerance of 10 -5 .

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 10 K and the learning rate to 0 . 0008 . We note that a single power iteration roughly contains 43 / 2 = 21 . 5 times more MVMsthan 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 = 43 times less elements per iteration, where 10 comes from the 10 K batch size. Finally, the 2 is explained by noting that SVRG incurs in a full sum update on every epoch.

(b) In this experiment we trained a GP by estimating the covariance RBF kernel with J = 1 K random Fourier features (RFFs). The hyperparameters for the RBF kernel are the following: length scale ( ℓ = 0 . 1 ), output scale ( a = 1 ) and likelihood noise ( σ 2 = 0 . 1 ). Moreover, we used CG as the iterative solver with a tolerance of 10 -8 and 100 as 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 10 K, a 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 a wave equation in two spatial dimensions. At each integrator step, a linear system M ( θ ) ˙ θ = F ( θ ) must be solved to find ˙ θ , for a d = 12 K × 12 K dimensional matrix. While Finzi et al. [17] use conjugate gradients to solve the linear system, we demonstrate the advantages of using SVRG, as M ( θ ) = 1 m ∑ m i =1 M i ( θ ) is a sum over the evaluation at m = 50 K 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 PyTorch 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 a 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 / 100 = 0 . 4 for the GPU column.

Solves . In this experiment we calculated the % runtime reduction when running torch.linalg.solve on the Trefethen N = 20 K 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.

Eigenvalue estimation . In this experiment we calculated the % runtime reduction when running torch.linalg.eigh on the mhd4800b N = 4 . 8 K matrix market sparse operator. In this experiment, Lanczos was run with a tolerance of 10 -9 and a maximum number of iterations equal to 100 .

Log determinant computation . In this experiment we calculated the % runtime reduction when running torch.linalg.logdet on the bcsstk18 N = 11 . 9 K matrix market sparse operator. In this experiment, the stochastic Lanczos 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 a 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 1 K. Additionally, we ran 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 a leading package such as sklearn .

(c) In this experiment we fit a GP with a RBF kernel on two datasets: Elevators and Kin40K . We only used up to 20K observations from Kin40K as that was 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 Lanczos quadrature (SLQ) to compute and optimize the negative log-marginal likelihood (loss function). Both experiments were run on a GPU for 100 iterations using Adam as an optimizer with learning rate of 0 . 1 with the default values of β 1 = 0 . 9 and β 2 = 0 . 999 . Additionally, for both GPyTorch and CoLA, the CG tolerance was set to 10 -4 with a maximum number of CG iterations of 250 and 20 probes were used for SLQ. Note that both CoLA and GPyTorch have similar throughputs, for example GPyTorch runs a 100 iterations on Elevators on 43 seconds whereas CoLA runs a 100 iterations on 49 seconds. When training a GP, we solve a 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 a stopping criteria based on the convergence of all the solves.

(d) In this experiment we run spectral clustering on the cit-HepPh dataset using an embedding size of 8 and also 8 clusters for k-means (with only 1 run of k-means after estimating the embeddings). We compare against sklearn using two different solvers, one based on Lanczos iterations using ARPACK and another using an Algebraic Multi-Grid solver AMG . In this case, CoLA also uses Lanczos iterations with a default tolerance of 10 -6 . We see how sklearn 's AMG solver runs faster than CoLA's but this is mostly the algorithmic constants as they have similar asymptotical 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 = 5 K 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 Arnoldi implementation with 20 iterations where CoLA also uses Arnoldi with the same number of iterations. Surprisingly, CoLA's JAX jitted code has a competitive runtime when compare to SciPy 's runtime using ARPACK .

(f) In this experiment we solve a minimal surface problem on a grid of maximum size of N = 100 2 points. To solve this problem we have to run Netwon-Rhapson where each

inner step involves a 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 as it is done in CoLA). The tolerance for this experiment was 5e-3. We see how CoLA's GMRES implementation is competitive with SciPy (JAX) but it still does not beat ARPACK mostly due to the faster runtime of using a lower level GMRES implementation.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSimilarly 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n**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\n\nLIMITATIONS:\nThe article has no Limitations section but seems compared extensively to other libraries in sec. 3.4.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-13"}

comments: 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.

STRENGTHS:
1. Designing efficient numerical linear algebra libraries is important to the machine learning community.

2. 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.

QUESTIONS:
n/a

LIMITATIONS:
n/a
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\nOverall 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?\n\nLIMITATIONS:\nLimitations 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThe limitations of the work are clear from the presentation and no issues require further acknowledgment, to the best of my knowledge.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n-) 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.\n\nWEAKNESSES:\n-) 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?\n\nQUESTIONS:\n-) 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.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}