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
Simple Operators Composition Operators Base Case D T P C S Pr A 0 0 B A B C D A -1 Eigs(A) Diag(A) Tr(A) exp(A) det(A)
Table 1: Many structures have explicit composition rules to exploit. Here we show the existence of a dispatch rule ( ) that can be used to accelerate a linear algebraic operation for some matrix structure over 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 are used in other routines. In absence of a rule, the operation will fall back to the iterative and dense base case for each operation (shown in ). Columns are basic linear operator types such as D: 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 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 GMRES for linear 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 (Compositional Linear Algebra), 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). Package GPU Support Autograd Non-symmetric Matrices Complex Numbers Randomized Algorithms Composition Rules
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 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 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.
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 LinearOperators 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
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 are matrix vector multiplies, linear solves, and eigendecomposition. Here ϵ denotes error tolerance. For a given operator of size N × N , we denote τ as its MVM cost, s its linear solve cost, E its eigendecomposition cost and κ its condition number. A lower script indicates to which matrix the operation belongs to.
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
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 a novel approach for computing memory-efficient derivatives of iterative methods through MVMs.
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.
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 LinearOperators together. Some frequently used compositional structures are Sum
Each of these compositional LinearOperators 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 MVMs of the composed LinearOperators. For example, MVM Sum = v → i MVM i (v), where MVM i are the MVM routines for the component LinearOperators.
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 LinearOperators 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. 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.
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. ∆ composed with itself on grid of sizes up to N = 1000 2 . We use CG with a multi-grid αSA preconditioner [9] to solve the linear system required in this application. (c) Finding the nullspace of an equivariant MLP of a 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. (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 a random Fourier features (RFFs) approximation [40] to a RBF kernel with J = 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 as 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 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 ⊙ Ae 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 ⊙ Az 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 : For sufficiently large problems, switching from dense to iterative algorithms provides consistent runtime reductions, especially on a GPU, where matrix multiplies can be effectively parallelized. We plot the ratio between the runtime of a linear algebra operation using CoLA 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 bcsstk18. 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 applied Sum structures. We validate empirically this analysis in Figure 5.
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.
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 k=20 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
Grid Size
10 1 10 1 Runtime (sec) SciPy CoLA 10 2 10 3 10 4 Grid Size 10 1 10 0 10 1 Runtime (sec) SciPy CoLA SciPy (JAX) (d) Spectral Clustering (e) Schrödinger Equation (f) Minimal Surface Figure 4: CoLA is easily applied to numerous applications with competitive performance. Here sk: sklearn, GP: GPyTorch and the tuple (N , D) denotes dataset size and dimensionality. (a): Runtime for PCA decomposition on Buzz (437.4K, 77). (b): Linear regression runtime on Song (386.5K, 90), where we run CoLA on both GPU and CPU. (c): Training efficiency (measure in epochs) on exact GP inference on Elevators (14K, 18) and Kin (20K, 8) on GPU. (d): Spectral clustering runtime on a citations graph (cit-HepPh) consisting on 34.5K nodes and 842K edges. sk(L) denotes sklearn's implicitly restarted Lanczos implementation and sk(A) denotes sklearn's LOBPCG with an algebraic multi-graph preconditioner (PyAMG) [5, 26]. CoLA(L) denotes our Lanczos implementation and CoLA(B) our LOBPCG implementation. (e): Runtimes for finding the smallest eigenfunctions expanding grids of a Schrödinger equation with an expanding finite difference grid. (f): Runtimes for solving the minimal surface equation via root finding on expanding grids. Here SciPy utilizes the ARPACK package, a highly-optimized Fortran implementation of the Arnoldi iteration, while SciPy JAX (the SciPy version integrated with JAX) and CoLA utilize python Arnoldi implementations. Appendix D expands on the experimental details.
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).
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 -2z 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.
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.
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.
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.
We now present the linear algebra identities that we use to exploit structure in CoLA.
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
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
We now assume that the matrices in this section are diagonalizable. That is,
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
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
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
In terms of block composition we have
Finally for the annotated operators we have the following rules.
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)
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
and also if it is symmetric and PSD.
Interestingly, the previous core functions allow us to derive multiple rules from the previous ones.
To illustrate, we have that
A and if A is both symmetric and PSD then f
where in both cases we used Eigs (A) = Λ A , V A . Some example functions for PSD matrices are
Which also this rules allow us to define LogDet (A) = Tr (Log (A)).
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.1 Doubly stochastic diagonal and trace estimation Singly Stochastic Trace Estimator Consider the traditional stochastic trace estimator:
with each z j ∼ N (0,
, 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
From the generating function we can derive the mean and variance of this estimator:
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
Taking derivatives we find that,
Assuming bounded moments on A i , then both
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
, 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. For two different linear algebra operations A -1 θ b and log |A θ |, 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's autograd rules against the autograd default of 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.
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:
∂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:
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.
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.
In simplest form, SVRG [23] performs gradient descent with the varianced reduced gradient w ← w -η(g i (w) -g i (w 0 ) + g(w 0 ))
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 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 -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 ∈ 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.
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.
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 = 14K observations and has D = 18 dimensions.
Algorithm 1 Arnoldi iteration 1: Inputs: A, q 0 = ν 0 / ∥ν 0 ∥ where possibly ν 0 ∼ N (0, I), maximum number of iterations T and tolerance ϵ ∈ (0, 1). 2: for j = 0 to T -1 do
3: ν j+1 ← Aq j 4:
for i = 0 to j do 5:
ν j+1 ← ν j+1 -h i,j q i 7:
end for 8:
if h j+1,j < ϵ then
10: stop 11: else 12: h j = R j ν j 6:
if j < T then 8:
end if 10: end for 11: return H, Q = (q 0 | . . . |q T ) 12: function GET_HOUSEHOLDER_VEC(w, k) 13: u i = 0 for i < k and u i = w i for i > k. 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 a subset that consists of N = 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 a subset consisting of N = 430K observations and has D = 77 dimensions. 4. Song. The full name of this UCI dataset is YearPredictionMSD. 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. 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.
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 = 1K 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 1K. 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 10K. 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 .
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 a 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 = 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 a full sum update on 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 (σ 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 10K, 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 = 12K × 12K 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 = 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.
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.
1. Solves. In this experiment we calculated the % runtime reduction when running torch.linalg.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.
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.