asdex: Automatic Sparse Differentiation in JAX

2026-10-08

AI Summary

This study addresses the inefficiency of computing sparse Jacobian and Hessian matrices via dense automatic differentiation (AD) in JAX. To this end, this work proposes the first standalone automatic sparse differentiation (ASD) toolkit within the JAX ecosystem. By integrating sparsity pattern detection, graph coloring algorithms, and compression-decompression techniques, the proposed method reduces the required number of AD passes to a constant independent of problem dimensionality. Furthermore, the toolkit serves as a seamless drop-in replacement for `jax.jacobian` and `jax.hessian`, significantly lowering computational complexity while maintaining full interface compatibility. Ultimately, this framework provides efficient infrastructure for large-scale sparse optimization and scientific computing applications.

Abstract

Many tasks in scientific computing and machine learning require the Jacobian or Hessian matrix of a function. Automatic differentiation (AD) computes these derivatives to machine precision, but materializing a dense $m \times n$ Jacobian requires $n$ forward-mode or $m$ reverse-mode AD passes, one per column or row. For a large class of functions, each output depends on only a few inputs, making the derivative matrix sparse. Automatic sparse differentiation (ASD) exploits this structure in four steps: detection of the input-agnostic sparsity pattern, coloring of a graph to group columns or rows that can share an AD pass, compressed differentiation to compute a compressed derivative matrix with one AD pass per color, and finally decompression into the original sparsity pattern. The number of colors, and hence of AD passes, is often independent of the problem dimension: a banded Jacobian with $b$ contiguous bands, for instance, only ever requires $b$ colors, regardless of its size. asdex offers the first standalone ASD toolkit in the popular JAX ecosystem. With asdex.jacobian and asdex.hessian, it provides sparse drop-in replacements for jax.jacobian and jax.hessian.

Problem

• Automatic Differentiation
• Sparse Jacobian
• Hessian Matrix
• JAX
• Scientific Computing

Innovation

• Automatic Sparse Differentiation
• JAX
• Graph Coloring
• Jacobian
• Hessian