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

TL;DR

CoLA introduces a flexible framework that leverages matrix structure to optimize large-scale linear algebra computations in machine learning, enabling automatic, memory-efficient, and accelerated algorithms across diverse applications.

Contribution

It presents a general, compositional approach for efficient linear algebra computations that integrates with existing machine learning frameworks like JAX and PyTorch.

Findings

Accelerates linear algebra operations in various applications.

Enables automatic differentiation and low precision computation.

Provides a versatile, easy-to-use tool for structured matrix problems.

Abstract

Many areas of machine learning and science involve large linear algebra problems, such as eigendecompositions, solving linear systems, computing matrix exponentials, and trace estimation. The matrices involved often have Kronecker, convolutional, block diagonal, sum, or product structure. In this paper, we propose a simple but general framework for large-scale linear algebra problems in machine learning, named CoLA (Compositional Linear Algebra). By combining a linear operator abstraction with compositional dispatch rules, CoLA automatically constructs memory and runtime efficient numerical algorithms. Moreover, CoLA provides memory efficient automatic differentiation, low precision computation, and GPU acceleration in both JAX and PyTorch, while also accommodating new objects, operations, and rules in downstream packages via multiple dispatch. CoLA can accelerate many algebraic operations, while making it easy to prototype matrix structures and algorithms, providing an appealing drop-in tool for virtually any computational effort that requires linear algebra. We showcase its efficacy across a broad range of applications, including partial differential equations, Gaussian processes, equivariant model construction, and unsupervised learning.