arXiv：cs.LG· Adrian Hill, Guillaume Dalle·· 3 小时前
asdex：JAX 中的自动稀疏微分工具包
asdex: Automatic Sparse Differentiation in JAX
AI 导读
asdex 发布了 JAX 生态首个独立自动稀疏微分（ASD）工具包，通过检测稀疏模式、图着色、压缩微分与解压四步，将 AD 所需次数降到与问题维度无关。它提供 sparse 版 drop-in 替换 jax.jacfwd 和 jax.jacrev，例如含 b 个连续带的带状 Jacobian 只需 b 次 AD。
正文
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 this http URL and this http URL, it provides sparse drop-in replacements for this http URL and this http URL.
| Comments: | 1 table |
| Subjects: | Mathematical Software (cs.MS); Machine Learning (cs.LG); Numerical Analysis (math.NA) |
| Cite as: | arXiv:2610.12336 [cs.MS] |
| (or arXiv:2610.12336v1 [cs.MS] for this version) | |
| https://doi.org/10.48550/arXiv.2610.12336 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Adrian Hill [view email]
[v1]
Thu, 8 Oct 2026 17:13:09 UTC (12 KB)
来源：arXiv：cs.LG · arxiv.org