asdex: Automatic Sparse Differentiation in JAX

arXiv:2610.12336 · cs.MS, cs.LG, cs.NA, math.NA · Submitted 2026-10-08 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "asdex: Automatic Sparse Differentiation in JAX".

Jane: The gist: Automatic sparse differentiation (ASD) exploits input-dependent sparsity in derivative matrices to compute them efficiently with fewer automatic differentiation passes than standard methods.

Tom: First, who's behind it and why it matters.

Title and authors: Tom: We're starting with the title and the authors of this paper, "asdex: Automatic Sparse Differentiation in JAX." It’s clear from that name what they’re trying to achieve—making differentiation sparse automatically within the JAX framework.

Jane: The authors are Adrian Hill and Guillaume Dalle, and they’ve been working on this idea for a while now since the problem started getting more pressing for machine learning research.

Lu: They are really focusing on high-level languages like JAX because that’s where most of the new machine learning prototyping is happening right now.

Meng: So, what does this title actually mean in plain English? It means they’re building a system to figure out where the derivatives are zero so they don't waste time calculating those zero entries.

Tom: Exactly. They want to bypass the standard method where you just compute everything dense, and instead use input-dependent sparsity—meaning the derivative matrix is sparse for specific inputs—to save both memory and calculation time.

Jane: It’s about exploiting that input-dependent structure so you get more efficient matrix storage and a faster way to fill in those sparse matrices.

Lalam: So, it’s essentially trying to combine smart pattern detection with clever matrix filling techniques to get the best of both worlds for these derivative calculations.

The paper's summary: Tom: Now, let’s look at what the paper actually summarizes in "asdex: Automatic Sparse Differentiation in JAX." They lay out a four-step process for how this automatic sparse differentiation works.

Jane: First, they have sparsity detection—that’s figuring out which inputs actually influence which outputs. They do this by performing abstract interpretation on a jaxpr7, which is JAX’s intermediate representation of functions.

Lu: That abstract interpretation step is key because it propagates index sets through each primitive to determine the pattern, and importantly, it yields a conservative pattern that never misses any non-zero entries.

Meng: So they aren't just guessing sparsity; they are using JAX’s own structure to find where the dependencies lie before doing any heavy matrix math.

Tom: After detection comes graph coloring. They take the detected pattern and turn it into a graph coloring problem, which they solve approximately using standard greedy algorithms.

Jane: For Jacobians, this involves a distance-two coloring of a bipartite row-column graph to partition columns or rows into structurally orthogonal groups that have no overlapping non-zeros <ref:2610.12336#pg1>.

Lalam: And for Hessians, they use something called star coloring to exploit the symmetry of the matrix further to reduce the number of colors needed.

The paper's improvements: Tom: The paper highlights a few major improvements over just using standard JAX functions like `jax.jacobian` or `jax.hessian`. They show that asdex is designed with several distinct, composable components mirroring the actual stages of ASD.

Jane: One big improvement is the separation between preparation and evaluation. Preparation involves that one-time detection and coloring work which depends only on input shapes but not the actual values.

Lu: That means in an iterative solver or a training loop, you only have to do that expensive setup once, and then you just run the cheap compressed differentiation pass repeatedly.

Meng: That amortization of the preparation step is what makes it practical for real-world applications where you're running lots of optimization steps.

Tom: Another improvement is how they handle second-order detection. They don’t need a separate second-order detector; they get Hessian patterns by applying the same first-order detector to `jax.grad(f)`.

Jane: That’s smart because it keeps the workflow simpler and avoids having to maintain a whole separate system just for the Hessian patterns.

Conclusion: Tom: So, wrapping up this discussion on "asdex: Automatic Sparse Differentiation in JAX," we see a tool that lets us leverage JAX’s speed while bypassing the memory bottleneck of dense derivative matrices by computing them with one AD pass per column or row.

Jane: It offers a way to compute sparse Jacobians and Hessians, which is crucial for applications like second-order optimization algorithms and sensitivity analysis where those matrices are often massive.

Lu: The real impact is that it combines more efficient matrix storage with a faster way to fill it, which addresses both the memory bottleneck and the high computational cost of standard methods.

Meng: For practitioners, this means you can keep using JAX’s performance and JIT compilation but get a much better handle on those large systems without running into immediate out-of-memory errors.

Lalam: Culturally, this kind of tooling makes it easier for researchers to focus on the model dynamics rather than worrying about the underlying hardware limitations of the differentiation process itself.

Tom: So, that’s what we have on "asdex: Automatic Sparse Differentiation in JAX." It’s a method designed to make derivative calculations more efficient by exploiting input-dependent sparsity.

Jane: It gives us a pathway to use JAX for bigger problems without immediately hitting those dense matrix limitations.

Lu: We'll be watching how this pattern detection and coloring approach develops, because it sets a strong foundation for other sparse derivative techniques in high-level systems.

Meng: From an engineering standpoint, the ability to amortize the setup cost is really what makes this viable for large-scale scientific computing tasks.

Lalam: It shows how AI tools can help us build infrastructure that handles complex mathematical structures more intelligently than just brute force computation.

Adrian Hill, Guillaume Dalle

BIFOLD – Berlin Institute for the Foundations of Learning and Data · Machine Learning Group, Technical University of Berlin · LVMT, ENPC, Institut Polytechnique de Paris

cs.MS, cs.LG, cs.NA, math.NA

Submitted: 2026-10-08

Updated: 2026-10-09

Comments: 1 table

Code: https://github.com/mfschubert/sparsejac2https:

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 90/100

The gist: The gist: Automatic sparse differentiation (ASD) exploits input-dependent sparsity in derivative matrices to compute them efficiently with fewer automatic differentiation passes than standard methods.

Key concepts

Sparsity Detection: This step uses abstract interpretation on JAX's internal representation (jaxpr7) to determine which inputs influence which outputs. It propagates index sets through the function, creating a conservative pattern of input-output dependencies. This process is crucial because it identifies the structure of sparsity before any computation begins.
Graph Coloring: The detected sparsity pattern is modeled as a graph coloring problem. For Jacobians, this involves distance-2 coloring to partition inputs into groups that share an AD pass without overlapping nonzeros. This grouping significantly reduces the number of required differentiation passes, which is the core efficiency gain.
Compressed Differentiation: Instead of computing the full dense matrix, asdex computes a compressed derivative matrix using parallelized Jacobian-vector products against a seed matrix derived from the coloring. This allows for one AD pass per color group instead of one pass per dimension, drastically lowering computational cost.
Decompression: The final step takes the compressed derivative information and scatters it back into the original, desired sparse format. This ensures that the output is a usable sparse Jacobian or Hessian matrix matching the original input structure, completing the process.

Summary

The gist: Automatic sparse differentiation (ASD) exploits input-dependent sparsity in derivative matrices to compute them efficiently with fewer automatic differentiation passes than standard methods.

Statement of need

Sparse derivative matrices arise in nonlinear systems of equations, second-order optimization algorithms, sensitivity analysis, and many other applications JAX’s built-in jacfwd, jacrev, and hessian functions materialize dense derivative matrices by performing one AD pass per input or output dimension This standard approach quickly hits two main limitations: a memory bottleneck, if the dense matrix is too large to store, and a high computational cost Computing a sparse Jacobian (Curtis et al., 1974) or Hessian (Coleman & Moré, 1984; Powell & Toint, 1979) with ASD alleviates both of these hurdles by combining more efficient matrix storage with a faster way to fill it

State of the field

ASD initially matured in low-level programming languages such as Fortran and C++, but we restrict our review to high-level languages, which machine learning research favors for the rapid prototyping workflows they enable A general-purpose ASD implementation in a high-level, open-source language recently appeared in Julia (Bezanson et al., 2017), through the combination of DifferentiationInterface.jl (Dalle & Hill, 2026), SparseConnectivityTracer.jl (Hill & Dalle, 2025) for sparsity pattern detection, and SparseMatrixColorings.jl (Montoison et al., 2025) for coloring No general-purpose equivalent exists for JAX, although there are a few prototypes

Software design

Core features of asdex mirror the main stages of ASD as distinct, composable components

((Sparsity detection)

asdex performs abstract interpretation on a jaxpr7 (short for JAX expression), JAX’s intermediate representation of a function Abstract interpretation propagates index sets through each jaxpr primitive to determine which inputs influence which outputs, yielding a conservative pattern that may contain false positives but never misses a nonzero Working at the jaxpr level, the detector reuses JAX’s own program representation and naturally handles functions built from arbitrary JAX primitives It also makes second-order detection free: because jax.grad(f) produces an ordinary JAX function with a jaxpr of its own, and the Hessian of a scalar-valued function f is the Jacobian of its gradient, asdex obtains Hessian patterns by applying the same first-order detector to jax.grad(f), with no separate second-order detector to implement or maintain

((Coloring)

The detected pattern gives rise to a graph coloring problem, which asdex solves approximately using standard greedy algorithms For Jacobians, a distance-2 coloring of a bipartite row-column graph partitions columns (forward mode) or rows (reverse mode) into structurally orthogonal groups (with no overlapping nonzeros) (Gebremedhin et al., 2005) By default, asdex picks whichever mode requires fewer colors For Hessians, a star coloring exploits the symmetry of the matrix to further reduce the number of colors (Gebremedhin et al., 2007) In both cases, the resulting groups are encoded in a seed matrix

((Compressed differentiation and decompression)

The compressed matrix stems from parallelized (jax.vmap-ed8) Jacobian-vector or vector-Jacobian products, evaluated against the seed matrix All nonzero coefficients are obtained in this way and then efficiently scattered back into the original sparse format required for the Jacobian or Hessian

((Performance and interface)

Preparation is a central design choice separates one-time preparation (detection and coloring, which depend on input shapes but not values) from repeated evaluation (compressed differentiation and decompression) of the sparse derivative In an iterative solver or a training loop, preparation is amortized, and only the cheap compressed AD pass is repeated The API deliberately stays close to JAX’s own and matches its semantics Where a dense Jacobian is written as jax.jacobian(f), asdex.jacobian(f, x) differs in only two aspects: it takes a sample input x for preparation and returns a function computing a sparse matrix (defaulting to JAX’s BCOO format)

Research impact statement

asdex was inspired by JAX issue 10329, a feature request for sparse Jacobian and Hessian support, which has remained open since 2019 As for asdex itself, it has already been applied to accelerate the computation of Hessians in machine learning interatomic potentials (Langer et al., 2026) and has been integrated into splineax10, a sparsity-focused extension of the popular linear and least-squares library lineax (Rader et al., 2023)

AI usage disclosure

Generative AI tools (Anthropic’s Claude Opus 4.5 to 5.5 and Fable 5 models) were used in developing the asdex software and writing this manuscript All AI-generated code and text were reviewed and verified by the authors

Acknowledgements

Adrian Hill gratefully acknowledges funding from the German Federal Ministry for Research, Technology and Space under the grant BIFOLD26B We thank Alexis Montoison for his work on the Julia ASD ecosystem, Marcel Langer and Denis Korolev for their feedback on asdex, as well as Jonathan Brodrick, John Viljoen and Johanna Hafner for productive discussions around JAX

References

Baydin et al. (2018) is referenced in the summary Dalle & Hill (2026) are cited for the common interface Gebremedhin et al. (2005) and Gebremedhin et al. (2007) are cited for coloring methods Curtis et al. (1974) is referenced regarding sparse Jacobian estimation Coleman & Moré (1984; Powell & Toint, 1979) are cited for Hessian estimation Bezanson et al. (2017) is referenced regarding Julia Montoison et al. (2025) are cited for sparse matrix coloring and bicoloring Hill & Dalle (Hill & Dalle, 2025) are cited for sparsity detection Pal et al. (2026) is referenced for NonlinearSolve.jl Rader et al. (2023) is referenced regarding Lineax Liu et al. (Langer et al., 2026) are cited for JAX-AMG Pundir et al. (Pundir et al.

Improvements for AI systems

textbfAsdex Integration into JAX Ecosystem for Sparse Derivatives: asdex.jacobian and asdex.hessian support provides sparse drop-in replacements for jax.jacobian and jax.hessian. This allows ML practitioners to leverage JAX’s performance, JIT compilation, and accelerator support while bypassing the memory bottleneck associated with materializing dense derivative matrices in standard AD methods by computing them with one AD pass per column or row.
textbfEnhanced Second-Order Optimization via Hessian Sparsity: asdex exploits structural sparsity through coloring of a graph to group columns or rows that can share an AD pass, enabling the computation of sparse Hessians efficiently. This is particularly useful for second-order optimization algorithms in scientific computing, where the number of AD passes scales with matrix dimension in standard methods, whereas asdex's approach means the number of colors, and hence of AD passes, is often independent of the problem dimension.
textbfAccelerated Sensitivity Analysis: The toolkit supports applications like sensitivity analysis by providing a framework to compute sparse derivative matrices more efficiently. By combining more efficient matrix storage with a faster way to fill it, asdex alleviates both the memory bottleneck and high computational cost that limit standard JAX AD approaches for large nonlinear systems of equations.

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.