Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces asdex, a novel Automatic Sparse Differentiation (ASD) toolkit for JAX, addressing the computational inefficiency of dense Jacobian and Hessian computation for functions with sparse derivatives. By leveraging sparsity patterns, graph coloring, and compressed differentiation, asdex significantly reduces the number of required automatic differentiation passes. The key result is the provision of drop-in sparse Jacobian and Hessian functions that offer substantial performance gains for relevant scientific computing and machine learning tasks.
Exploit sparse derivative structure to slash AD passes for Jacobians and Hessians, offering drop-in replacements for JAX users.
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.