SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Andres Potapczynski*1 Marc Finzi*2 Geoff Pleiss3,4 Andrew Gordon Wilson1
$^{1}$ New York University,  $^{2}$ Carnegie Mellon University,  $^{3}$ University of British Columbia,
$^{4}$ Vector Institute
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.
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
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  $\mathbf{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  $\mathbf{F}$  as a matrix, the diagonal  $\mathbf{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  $\mathbf{A}^{-1}\mathbf{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.
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 nonexhaustive 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
Scipy Sparse	X	X	✓	✓	X	X
PyLops	*	*	✓	✓	X	X
GPyTorch	✓	✓	X	X	X	X
CoLA	✓	✓	✓	✓	✓	✓
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 well-known 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  $\mathbf{M}\mathbf{V}_{A}:\mathbf{v}\mapsto \mathbf{A}\mathbf{v}$ . 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  $\left[\mathbf{A}\mathbf{v}\right]_i = \sum_{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  $\left[\operatorname{Diag}(\mathbf{d}) \mathbf{v}\right]_i = d_i v_i$ ; Convolution operators, which are represented by a convolutional filter array and where the MVM is given by  $\operatorname{Conv}(\mathbf{a}) \mathbf{v} = \mathbf{a} * \mathbf{v}$ ; or JVP operators—the Jacobian represented implicitly through an autograd Jacobian Vector Product—represented by a function and an input  $\mathbf{x}$  and where the MVM is given by  $\operatorname{Jacobian}(f, \mathbf{x}) \mathbf{v} = \operatorname{JVP}(f, \mathbf{x}, \mathbf{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  $\mathrm{MVM}(\mathbf{e}_1)$ , ...,  $\mathrm{MVM}(\mathbf{e}_N)$ , on each unit vector  $\mathbf{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
ΠiM A i	ΣiM A i	BlockDiag(A, B)	Kron(A, B)
MVM (τ)	Σi τi	Σi τi	τA + τB	τANB + NARTB
Solve (s)	Σi κiτi log M/ε	(1 + κ/M)τ log 1/ε	sA + sB	sANB + NAsB
Eigs (E)	τ log M/εΠiκi	(1 + κ/M)τ log 1/ε	EA + EB	EA + EB
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  $\epsilon$  denotes error tolerance. For a given operator of size  $N\times N$ , we denote  $\tau$  as its MVM cost,  $s$  its linear solve cost,  $E$  its eigendecomposition cost and  $\kappa$  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  $\mathbf{v} \mapsto \mathbf{A}^{\top}\mathbf{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  $\mathbf{A}^{\top}\mathbf{v}$  is the vector-Jacobian product (VJP) of the vector  $\mathbf{v}$  and the Jacobian  $\partial \mathbf{A}\mathbf{w} / \partial \mathbf{w}$ . Thus, the function transpose(A) returns a LinearOperator object that uses VJP(MVM $_{\mathbf{A}}$ , 0, v) as its MVM. We extend this idea to Hermitian conjugates, using the fact that  $\mathbf{A}^{*}\mathbf{v} = (\overline{\mathbf{A}})^{\top}\mathbf{v} = (\overline{\mathbf{A}^{\top}\overline{\mathbf{v}}})$ .
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.
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  $\mathbf{A} = \operatorname{Diag}(\mathbf{a})$ , then we know that  $\mathbf{A}^{-1} = \operatorname{Diag}(\mathbf{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  $\operatorname{Conv}(\mathbf{a}) = \mathcal{F}^{-1}\operatorname{Diag}(\mathcal{F}\mathbf{a})\mathcal{F}$  (where  $\mathcal{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  $(\sum_{i}\mathbf{A}_{i})$  Product  $(\Pi_i\mathbf{A}_i)$ , Kronecker  $(\mathbf{A}\otimes \mathbf{B})$ , KroneckerSum  $(\mathbf{A}\oplus \mathbf{B})$ , BlockDiag  $[\mathbf{A},0;0,\mathbf{B}]$  and Concatenation  $[\mathbf{A},\mathbf{B}]$ . 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,  $\mathrm{MVM}_{\mathrm{Sum}} = \mathbf{v}\mapsto \sum_{i}\mathrm{MVM}_{i}(\mathbf{v})$ , where  $\mathrm{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  $(\mathbf{A} \otimes \mathbf{B})^{-1} = \mathbf{A}^{-1} \otimes \mathbf{B}^{-1}$ . Though simple, this rule yields highly efficient routines for numerous structures. For example, suppose we want to solve  $(\mathbf{A} \otimes \mathbf{B} \otimes \mathbf{C})\mathbf{x} = \mathbf{b}$  where  $\mathbf{A}$  is dense,  $\mathbf{B}$  is diagonal, and  $\mathbf{C}$  is triangular. From the rules, the solve would be split over the product, using GMRES for  $\mathbf{A}$ , diagonal inversion for  $\mathbf{B}$ , and forward substitution for  $\mathbf{C}$ . This breakdown is much more efficient than the base case (GMRES with  $\mathrm{MVM}_{\mathrm{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  $\mathbf{B} = \mathbf{A}^n$ ; in general, the runtime is upper-bounded by  $O(n\tau \sqrt{\kappa^n}\log \frac{1}{\epsilon})$ , where  $\tau$  is the time for a MVM with  $\mathbf{A}$ ,  $\kappa$  is the condition number of  $\mathbf{A}$ , and  $\epsilon$  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\tau \sqrt{\kappa}\log \frac{n}{\epsilon})$ . 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. 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].