asdex：JAX 中的自动稀疏微分

arXiv:2610.12336v1 公告类型：交叉收录 摘要：科学计算和机器学习中的许多任务都需要计算函数的雅可比矩阵或海森矩阵。自动微分（AD）能够以机器精度计算这些导数，但显式构建一个稠密的 $m \times n$ 雅可比矩阵需要执行 $n$ 次前向模式或 $m$ 次反向模式的 AD，每次计算一列或一行。对于一大类函数，每个输出仅依赖少量输入，因此其导数矩阵是稀疏的。自动稀疏微分（ASD）通过四个步骤利用这一结构：检测与输入无关的稀疏模式；通过图着色，将可以共享一次 AD 计算的列或行分组；执行压缩微分，每种颜色进行一次 AD 计算，得到压缩的导数矩阵；最后将其解压缩，恢复原始稀疏模式。颜色数量以及相应的 AD 计算次数通常与问题维度无关：例如，具有 $b$ 条连续带的带状雅可比矩阵，无论规模多大，都只需要 $b$ 种颜色。asdex 是广受欢迎的 JAX 生态系统中首个独立的 ASD 工具包。它通过 asdex.jacobian 和 asdex.hessian，为 jax.jacobian 和 jax.hessian 提供可直接替换的稀疏版本。