# CoLA: Exploiting Compositional Structure for Automatic and Efficient Numerical Linear Algebra

arXiv: https://arxiv.org/abs/2309.03060 · HTML: https://mlsanity.com/paper/2309.03060/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 61 references that are in the map.
- First posted on arXiv 6 Sep 2023.

## Builds on (most cited first, 24 of 24)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration](https://mlsanity.com/paper/1809.11165/) (2018)
- [Hungry Hungry Hippos: Towards Language Modeling with State Space Models](https://mlsanity.com/paper/2212.14052/) (2022)
- [Neural Operator: Learning Maps Between Function Spaces](https://mlsanity.com/paper/2108.08481/) (2021)
- [Measuring the Intrinsic Dimension of Objective Landscapes](https://mlsanity.com/paper/1804.08838/) (2018)
- [Gradient-based Hyperparameter Optimization through Reversible Learning](https://mlsanity.com/paper/1502.03492/) (2015)
- [A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups](https://mlsanity.com/paper/2104.09459/) (2021)
- [Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP)](https://mlsanity.com/paper/1503.01057/) (2015)
- [Exact Gaussian Processes on a Million Data Points](https://mlsanity.com/paper/1903.08114/) (2019)
- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](https://mlsanity.com/paper/2004.11127/) (2020)
- [Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations](https://mlsanity.com/paper/1903.05895/) (2019)
- [A general framework for Vecchia approximations of Gaussian processes](https://mlsanity.com/paper/1708.06302/) (2017)
- [Evolutional Deep Neural Network](https://mlsanity.com/paper/2103.09959/) (2021)
- [S4ND: Modeling Images and Videos as Multidimensional Signals Using State Spaces](https://mlsanity.com/paper/2210.06583/) (2022)
- [Product Kernel Interpolation for Scalable Gaussian Processes](https://mlsanity.com/paper/1802.08903/) (2018)
- [Accelerated Stochastic Power Iteration](https://mlsanity.com/paper/1707.02670/) (2017)
- [A Stochastic PCA and SVD Algorithm with an Exponential Convergence Rate](https://mlsanity.com/paper/1409.2848/) (2014)
- [A Stable and Scalable Method for Solving Initial Value PDEs with Neural Networks](https://mlsanity.com/paper/2304.14994/) (2023)
- [Linear-time inference for Gaussian Processes on one dimension](https://mlsanity.com/paper/2003.05554/) (2020)
- [Low-Precision Arithmetic for Fast Gaussian Processes](https://mlsanity.com/paper/2207.06856/) (2022)
- [SKIing on Simplices: Kernel Interpolation on the Permutohedral Lattice for Scalable Gaussian Processes](https://mlsanity.com/paper/2106.06695/) (2021)

## Built on by (most cited first, 7 of 7)

- [Hidden Breakthroughs in Language Model Training](https://mlsanity.com/paper/2506.15872/) (2025)
- [Searching for Efficient Linear Layers over a Continuous Space of Structured Matrices](https://mlsanity.com/paper/2410.02117/) (2024)
- [Compute-Optimal LLMs Provably Generalize Better With Scale](https://mlsanity.com/paper/2504.15208/) (2025)
- [Perspectives on the State and Future of Deep Learning - 2023](https://mlsanity.com/paper/2312.09323/) (2023)
- [Hamiltonian Monte Carlo Inference of Marginalized Linear Mixed-Effects Models](https://mlsanity.com/paper/2410.24079/) (2024)
- [Diffusing Differentiable Representations](https://mlsanity.com/paper/2412.06981/) (2024)
- [Spectral Estimation with Free Decompression](https://mlsanity.com/paper/2506.11994/) (2025)

## Abstract (arXiv)

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.