asdex: Automatic Sparse Differentiation in JAX

Adrian Hill; Guillaume Dalle

cs.MSOct 8, 2026

asdex: Automatic Sparse Differentiation in JAX

Authors: Adrian Hill, Guillaume Dalle

Organizations: BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany · Machine Learning Group, Technical University of Berlin, Berlin, Germany · LVMT, ENPC, Institut Polytechnique de Paris, Univ Gustave Eiffel, Marne-la-Vallée, France

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.

Figures & tables

Package	Detection	Coloring
sparsejac 1 1 1 https://github.com/mfschubert/sparsejac	None	Distance-1 (from networkx )
sparsediffax 2 2 2 https://github.com/gdalle/sparsediffax	None	Distance-2 & star (from SMC)
jax-nansparse 3 3 3 https://github.com/nardi/jax-nansparse	NaN tracing	None
jax2sympy 4 4 4 https://github.com/johnviljoen/jax2sympy	jaxpr symbolic conversion	None
asdex	jaxpr tracing	Distance-2 & star (based on SMC)

Table 1

Explore similar work

Truncated automatic sparse differentiation for machine learning interatomic potentials
Sep 17, 2026Marcel F. Langer, Adrian Hill, Michele CeriottiAutomatic Differentiation Machine Learning Interatomic Potentials
Laboratory of Computational Science and Modeling (COSMO), EPFL, Lausanne, Switzerland · Machine Learning Group, Technical University of Berlin, Berlin, Germany · BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany
jNO: A JAX Library for Neural Operator and Foundation Model Training
May 11, 2026Leon Armbruster, Rathan Ramesh, Georg Kruse +1PDE Operator Learning Physics-Informed Neural Operators
Fraunhofer Institute for Integrated Systems and Device Technology IISB, Germany
AdamX: Cosine similarity meets gradient descent
Sep 10, 2026Francisco Caldas, Ruben Belo, Cláudia SoaresDeep Learning Optimization Adaptive Gradient Methods
NOVA School of Science of Technology Universidade Nova de Lisboa, Caparica, Portugal

Sep 17, 2026, physics.chem-phJ/K move · Enter open · S save

Truncated automatic sparse differentiation for machine learning interatomic potentials

Marcel F. Langer, Adrian Hill, Michele Ceriotti

Laboratory of Computational Science and Modeling (COSMO), EPFL, Lausanne, Switzerland · Machine Learning Group, Technical University of Berlin, Berlin, Germany · BIFOLD – Berlin Institute for the Foundations of Learning and Data, Berlin, Germany

Automatic Differentiation Machine Learning Interatomic Potentials Computational Materials Science

Open paper AskPDF arXiv

Machine learning interatomic potentials (MLIPs) learn the mapping from atomic positions to potential energy. The forces, the negative gradient of this energy, drive molecular dynamics and are readily obtained using automatic differentiation. Higher-order derivatives, most notably the Hessian, describe collective motion and allow the direct prediction of experimental observables, but are considered computationally inaccessible for large systems. We suggest a solution: in physical systems, interactions decay with distance, and most MLIPs build on this locality through message passing up to a finite receptive field. This implies both sparsity of higher-order derivatives and their decay with distance. This structure can be exploited using automatic sparse differentiation (ASD). We explain how to compute the sparsity pattern for MLIP derivatives and demonstrate that, for multiple foundation MLIPs, ASD computes full Hessians of large porous materials exactly, but with modest speedups at best. The larger gains come from truncated ASD: discarding small, but nonzero, Hessian entries between distant atoms yields order-of-magnitude speedups with negligible impact on predicted observables.

Abstract

Explore similar work

Truncated automatic sparse differentiation for machine learning interatomic potentials

jNO: A JAX Library for Neural Operator and Foundation Model Training

AdamX: Cosine similarity meets gradient descent