Sparser, Better, Faster, Stronger: Sparsity Detection for Efficient Automatic Differentiation

TL;DR

This paper introduces a highly efficient, automatic sparsity detection method for Automatic Differentiation that significantly accelerates the computation of Jacobians and Hessians in machine learning applications, enabling scalable and faster differentiation.

Contribution

It presents a novel, operator overloading-based sparsity detection system that is fully automatic, compatible with existing code, and capable of large-scale, high-speed differentiation without prior sparsity knowledge.

Findings

Achieves up to 1000x speed-up in Jacobian and Hessian computations.

Outperforms standard AD methods in one-off calculations without amortization.

Supports both local and global sparsity patterns with flexible index sets.

Abstract

From implicit differentiation to probabilistic modeling, Jacobian and Hessian matrices have many potential use cases in Machine Learning (ML), but they are viewed as computationally prohibitive. Fortunately, these matrices often exhibit sparsity, which can be leveraged to speed up the process of Automatic Differentiation (AD). This paper presents advances in sparsity detection, previously the performance bottleneck of Automatic Sparse Differentiation (ASD). Our implementation of sparsity detection is based on operator overloading, able to detect both local and global sparsity patterns, and supports flexible index set representations. It is fully automatic and requires no modification of user code, making it compatible with existing ML codebases. Most importantly, it is highly performant, unlocking Jacobians and Hessians at scales where they were considered too expensive to compute. On real-world problems from scientific ML, graph neural networks and optimization, we show significant speed-ups of up to three orders of magnitude. Notably, using our sparsity detection system, ASD outperforms standard AD for one-off computations, without amortization of either sparsity detection or matrix coloring.