Jet Functors and Weil Algebras in Automatic Differentiation: A Geometric Analysis
This work offers a foundational geometric abstraction for computational differentiation systems, potentially improving efficiency and transparency in machine learning and scientific computing.
The authors tackled the problem of providing a unified, coordinate-free formulation for automatic differentiation (AD) by using jet functors and Weil algebras, resulting in an implementation that computes all mixed derivatives in a single forward pass with cost linear in the algebra dimension, achieving algebraically exact and numerically stable differentiation.
We present a differential-geometric formulation of automatic differentiation (AD) based on jet functors and Weil algebras. In this framework, forward- and reverse-mode differentiation arise naturally as pushforward and cotangent pullback, while higher-order differentiation corresponds to evaluation in a Weil algebra. This construction provides a unified, coordinate-free view of derivative propagation and clarifies the algebraic structure underlying AD. All results are realized in modern JAX code, where the Weil-mode formulation computes all mixed derivatives in a single forward pass with cost linear in the algebra dimension. The resulting implementation achieves algebraically exact and numerically stable differentiation with predictable scaling, demonstrating that geometric abstraction can yield more efficient and transparent computational differentiation systems. Code is available at https://git.nilu.no/geometric-ad/jet-weil-ad