Jet Functors and Weil Algebras in Automatic Differentiation: A Geometric Analysis
2510.14342v1
cs.LG, math.DG, stat.ML
2025-10-18
Авторы:
Amandip Sangha
Abstract
We present a geometric formulation of automatic differentiation (AD) using
jet bundles and Weil algebras. Reverse-mode AD emerges as cotangent-pullback,
while Taylor-mode corresponds to evaluation in a Weil algebra. From these
principles, we derive concise statements on correctness, stability, and
complexity: a functorial identity for reverse-mode, algebraic exactness of
higher-order derivatives, and explicit bounds on truncation error. We further
show that tensorized Weil algebras permit one-pass computation of all mixed
derivatives with cost linear in the algebra dimension, avoiding the
combinatorial blow-up of nested JVP/VJP schedules. This framework interprets AD
theory through the lens of differential geometry and offers a foundation for
developing structure-preserving differentiation methods in deep learning and
scientific computing. Code and examples are available at
https://git.nilu.no/geometric-ad/jet-weil-ad.