SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 optimizing 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 + \kappa / M)\log \frac{1}{\epsilon})$  where  $\kappa$  is the condition number and  $\epsilon$  is the desired accuracy.
(a) Multi-task GPs  $(\otimes)$
(b) Bi-Poisson (II)
(c) EMLP  $\left(\left[ \begin{array}{cc}\mathbf{A} & 0\\ 0 & \mathbf{B} \end{array} \right]\right)$
Figure 1: Empirically, our composition rules yield the best runtimes across applications consisting of linear operators with different structures (more application details in Section 4). We plot mean runtime (over 3 repetitions) for different methods (dense, iterative and ours (CoLA)) against the size of the linear operator. (a) Computing solves on a multi-task GP problem [7] for a linear operator having Kronecker structure  $\mathbf{K}_T \otimes \mathbf{K}_X$ , where  $\mathbf{K}_T$  is a kernel matrix containing the correlation between the tasks and  $\mathbf{K}_X$  is a RBF kernel on the data. For this experiment we used a synthetic Gaussian dataset which has dimension  $D = 33$ ,  $N = 1\mathrm{K}$  and we used  $T = 11$  tasks. (b) Computing solves on the 2-dimensional Bi-Poisson PDE problem for the composition of the Laplacian operator  $\Delta$  composed with itself on grid of sizes up to  $N = 1000^2$ . We use CG with a multi-grid  $\alpha$  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  $15\mathrm{K}$ . See Appendix D for further details.
(a) PCA (Eig)
Figure 2: CoLA exploits the sum structure of linear operators through stochastic routines. (a) Eigenvalue convergence criteria against number of MVMs for computing the first principal component on Buzz ( $N = 430\mathrm{K}$ ,  $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 = 1\mathrm{K}$  features on Elevators ( $N = 12.5\mathrm{K}$ ,  $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.
(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\left(\sqrt{\kappa}\log \frac{1}{\epsilon}\right)$ . 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  $\mathrm{Diag}(\mathbf{A}) = \sum_{i=1}^{N} e_i \odot \mathbf{A} e_i$  where  $\odot$  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]  $\overline{\mathrm{Diag}}(\mathbf{A}) = \frac{1}{n} \sum_{j=1}^{n} z_j \odot \mathbf{A} z_j$  where the  $z_j \in \mathbb{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  $\mathbf{A} = \sum_{i=1}^{M} \mathbf{A}_i$ :
$$
\overline {{\operatorname {D i a g}}} \left(\sum_ {i = 1} ^ {M} \mathbf {A} _ {i}\right) := \sum_ {i j} z _ {i j} \odot \mathbf {A} _ {i} z _ {i j}.
$$
Figure 3: 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 bcstk18. 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.
3.4 Automatic Differentiation and Machine Learning Readiness
Memory efficient auto-differentiation In ML applications, we want to backpropagate through operations like  $\mathbf{A}^{-1}$ , Eigs(A),  $\operatorname{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 backprogating 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
(a) PCA
(b) Linear Regression
(c) GPs
(d) Spectral Clustering
(e) Schrodinger 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 sklearnn's implicitly restarted Lanczos implementation and sk(A) denotes sklearnn's LOBPCG with an algebraic multi-graph preconditioner (PyAMG) [5, 26]. CoLA(L) denotes our Lanzos 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  $\mathbf{X} \in \mathbb{R}^{N \times 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  $\mathbf{L} = \mathbf{D} - \mathbf{W}$  where  $\mathbf{D}$  denotes a diagonal matrix containing the degree of the nodes and  $\mathbf{W}$  the weights on the edges between nodes. This problem requires finding the smallest  $k$  eigenvectors of  $\mathbf{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 |\mathbf{K}|$  and of a quadratic term  $\mathbf{y}^T\mathbf{K}^{-1}\mathbf{y}$  (where  $\mathbf{y}$  is the vector of observations).
Schrodinger 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  $\mathbf{H}\psi = E\psi$ . After discretizing  $\psi$  to a grid, we compute the smallest eigenvalues and eigenvectors of the operator  $\mathbf{H}$  which for this experiment is non-symmetric as we perform a compactifying 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_x^2)z_{yy} - 2z_xz_yz_{xy} + (1 + z_y^2)z_{xx} = 0$  and solved by root finding on a discrete grid. Applying Newton-Raphson, we iteratively solve the non-symmetric linear system  $z \gets z - \mathbf{J}^{-1}f(z)$  where  $\mathbf{J}$  is the Jacobian of the PDE operator.
Bi-Poisson Equation The Bi-Poisson equation  $\Delta^2 u = \rho$  is a linear boundary value PDE relevant in continuum mechanics, where  $\Delta$  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  $\theta$ , where  $\dot{\theta} = \mathbf{M}(\theta)^{-1}F(\theta)$  where  $\mathbf{M}$  is an operator defined from Jacobian of the neural network which decomposes as the sum over data points  $\mathbf{M} = \frac{1}{N}\sum_{i}\mathbf{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  $\mathbf{C}\mathbf{v} = \mathbf{0}$ , where the constraint matrix  $\mathbf{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 work is supported by NSF Award 1922658, NSF CAREER IIS-2145492, BigHat Biosciences, Capital One, and an Amazon Research Award.
[1] Martin 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 Mane, 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 Xiaogiang Zheng. TensorFlow: Large-scale machine learning on heterogeneous systems, 2015. URL https://www.tensorflow.org/. Software available from tensorflow.org.
[2] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable Second Order Optimization for Deep Learning. Preprint arXiv 2002.09018v2, 2020.
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  $(\mathbf{AB})^{-1} = (\mathbf{B}^{-1}\mathbf{A}^{-1})$  and for Kronecker product we have  $(\mathbf{A}\otimes \mathbf{B})^{-1} = \mathbf{A}^{-1}\otimes \mathbf{B}^{-1}$ . In terms of block compositions we have the following identities:
$$
\begin{array}{l} \left[ \begin{array}{l l} \mathbf {A} & \mathbf {0} \\ \mathbf {0} & \mathbf {D} \end{array} \right] ^ {- 1} = \left[ \begin{array}{c c} \mathbf {A} ^ {- 1} & \mathbf {0} \\ \mathbf {0} & \mathbf {D} ^ {- 1} \end{array} \right] \quad \text {a n d} \quad \left[ \begin{array}{l l} \mathbf {A} & \mathbf {B} \\ \mathbf {0} & \mathbf {D} \end{array} \right] ^ {- 1} = \left[ \begin{array}{c c} \mathbf {A} ^ {- 1} & - \mathbf {A} ^ {- 1} \mathbf {B D} ^ {- 1} \\ \mathbf {0} & \mathbf {D} ^ {- 1} \end{array} \right] \\ \left[ \begin{array}{l l} \mathbf {A} & \mathbf {B} \\ \mathbf {C} & \mathbf {D} \end{array} \right] ^ {- 1} = \left[ \begin{array}{l l} \mathbf {I} & - \mathbf {A} ^ {- 1} \mathbf {B} \\ \mathbf {0} & \mathbf {I} \end{array} \right] \left[ \begin{array}{l l} \mathbf {A} & \mathbf {0} \\ \mathbf {0} & \mathbf {D} - \mathbf {C A} ^ {- 1} \mathbf {B} \end{array} \right] ^ {- 1} \left[ \begin{array}{l l} \mathbf {I} & \mathbf {0} \\ - \mathbf {C A} ^ {- 1} & \mathbf {I} \end{array} \right] \\ \end{array}
$$
Finally, for sum we have the Woodbury identity and its variants. Namely, for Woodbury we have
$$
\left(\mathbf {A} + \mathbf {U B V}\right) ^ {- 1} = \mathbf {A} ^ {- 1} - \mathbf {A} ^ {- 1} \mathbf {U} \left(\mathbf {B} ^ {- 1} + \mathbf {V A} ^ {- 1} \mathbf {U}\right) ^ {- 1} \mathbf {V A} ^ {- 1},
$$
the Kailath variant where
$$
\left(\mathbf {A} + \mathbf {B C}\right) ^ {- 1} = \mathbf {A} ^ {- 1} - \mathbf {A} ^ {- 1} \mathbf {B} \left(\mathbf {I} + \mathbf {C A} ^ {- 1} \mathbf {B}\right) \mathbf {C A} ^ {- 1}
$$
and the rank one update via the Sherman-Morrison formula
$$
\left(\mathbf {A} + \mathbf {b c} ^ {\intercal}\right) ^ {- 1} = \mathbf {A} ^ {- 1} - \frac {1}{1 + \mathbf {c} ^ {\intercal} \mathbf {A b}} \mathbf {A} ^ {- 1} \mathbf {b c} ^ {\intercal} \mathbf {A} ^ {- 1}.
$$
Besides the compositional operators, we have some rules for some special operators. For example, for  $\mathbf{A} = \operatorname{Diag}(\mathbf{a})$  we have  $\mathbf{A}^{-1} = \operatorname{Diag}(\mathbf{a}^{-1})$ . Also, if  $\mathbf{Q}$  is unitary then  $\mathbf{Q}^{-1} = \mathbf{Q}^{*}$  or if  $\mathbf{Q}$  is orthonormal then  $\mathbf{Q}^{-1} = \mathbf{Q}^{\intercal}$ .
A.1.2 Eigendecomposition
We now assume that the matrices in this section are diagonalizable. That is,  $\operatorname{Eigs}(\mathbf{A}) = \boldsymbol{\Lambda}_{\mathbf{A}}$ ,  $\mathbf{V}_{\mathbf{A}}$  where  $\mathbf{A} = \mathbf{V}_{\mathbf{A}} \boldsymbol{\Lambda}_{\mathbf{A}} \mathbf{V}_{\mathbf{A}}^{-1}$ . In terms of the compositional operators, there is not a general rule for product or sum. However, for the Kronecker product we have  $\operatorname{Eigs}(\mathbf{A} \otimes \mathbf{B}) = \boldsymbol{\Lambda}_{\mathbf{A}} \otimes \boldsymbol{\Lambda}_{\mathbf{B}}$ ,  $\mathbf{V}_{\mathbf{A}} \otimes \mathbf{V}_{\mathbf{B}}$  and for the Kronecker sum we have  $\operatorname{Eigs}(\mathbf{A} \oplus \mathbf{B}) = \boldsymbol{\Lambda}_{\mathbf{A}} \oplus \boldsymbol{\Lambda}_{\mathbf{B}}$ ,  $\mathbf{V}_{\mathbf{A}} \otimes \mathbf{V}_{\mathbf{B}}$ . Finally, for block diagonal we have
$$
\operatorname {E i g s} \left(\left[ \begin{array}{c c} \mathbf {A} & \mathbf {0} \\ \mathbf {0} & \mathbf {D} \end{array} \right]\right) = \left[ \begin{array}{c c} \mathbf {\Lambda_ {A}} & \mathbf {0} \\ \mathbf {0} & \mathbf {\Lambda_ {D}} \end{array} \right], \left[ \begin{array}{c c} \mathbf {V _ {A}} & \mathbf {0} \\ \mathbf {0} & \mathbf {V _ {D}} \end{array} \right].
$$
A.1.3 Diagonal
As a base case, if we need to compute  $\operatorname{Diag}(\mathbf{A})$  for a general matrix  $\mathbf{A}$  we may compute each diagonal element by  $\mathbf{e}_i^{\intercal}\mathbf{A}\mathbf{e}_i$ . Additionally, if  $\mathbf{A}$  is large enough we switch to randomized estimation  $\operatorname{Diag}(\mathbf{A}) \approx (\mathbf{Z} \odot \mathbf{A}\mathbf{Z})\mathbf{1}/N$  with  $\mathbf{Z} \sim \mathcal{N}(0,1)^{d \times N}$  where  $N$  is the number of samples used to approximate the diagonal. In terms of compositional operators, we have that for sum  $\operatorname{Diag}(\mathbf{A} + \mathbf{B}) = \operatorname{Diag}(\mathbf{A}) + \operatorname{Diag}(\mathbf{B})$ . For Kronecker product we have  $\operatorname{Diag}(\mathbf{A} \otimes \mathbf{B}) = \operatorname{vec}\bigl(\operatorname{Diag}(\mathbf{A})\operatorname{Diag}(\mathbf{B})^{\intercal}\bigr)$  and for Kronecker sum  $\operatorname{Diag}(\mathbf{A} \oplus \mathbf{B}) = \operatorname{vec}\bigl(\operatorname{Diag}(\mathbf{A})\mathbf{1}^{\intercal} + \mathbf{1}\operatorname{Diag}(\mathbf{B})^{\intercal}\bigr)$ . Finally, for block composition we have
$$
\operatorname {D i a g} \left(\left[ \begin{array}{c c} \mathbf {A} & \mathbf {B} \\ \mathbf {C} & \mathbf {D} \end{array} \right]\right) = [ \operatorname {D i a g} (\mathbf {A}), \operatorname {D i a g} (\mathbf {D}) ].
$$
A.1.4 Transpose / Adjoint