Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Jet Functors and Weil Algebras in Automatic Differentiation: A Geometric Analysis

Amandip Sangha

TL;DR

This work places automatic differentiation in a differential-geometric framework using jet and Weil functors, unifying forward-, reverse-, and higher-order differentiation as pushforward, cotangent pullback, and Weil-algebra evaluation. The approach yields a coordinate-free, algebraically exact propagation of derivatives, with a single forward pass computing all mixed higher-order terms at cost linear in the Weil algebra dimension $\dim W$. Theoretical results are complemented by a practical JAX implementation that demonstrates predictable scaling, numerical stability, and substantial efficiency gains over nested first-order schedules. The framework clarifies the algebraic structure of derivative propagation and points toward structure-preserving differentiable programming with broad applicability in optimization and learning.

Abstract

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

Jet Functors and Weil Algebras in Automatic Differentiation: A Geometric Analysis

TL;DR

This work places automatic differentiation in a differential-geometric framework using jet and Weil functors, unifying forward-, reverse-, and higher-order differentiation as pushforward, cotangent pullback, and Weil-algebra evaluation. The approach yields a coordinate-free, algebraically exact propagation of derivatives, with a single forward pass computing all mixed higher-order terms at cost linear in the Weil algebra dimension . Theoretical results are complemented by a practical JAX implementation that demonstrates predictable scaling, numerical stability, and substantial efficiency gains over nested first-order schedules. The framework clarifies the algebraic structure of derivative propagation and points toward structure-preserving differentiable programming with broad applicability in optimization and learning.

Abstract

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
Paper Structure (20 sections, 7 theorems, 129 equations)

This paper contains 20 sections, 7 theorems, 129 equations.

Table of Contents

  1. Introduction
  2. Related Work
  3. Geometric Foundations
  4. Smooth manifolds and differentials
  5. Pushforward, pullback, and functoriality
  6. Automatic differentiation: forward and reverse.
  7. Forward mode (pushforward/JVP).
  8. Reverse mode (pullback/VJP).
  9. Chain rule in both modes.
  10. Jets, Weil algebras, and functorial AD
  11. Weil algebras and Weil functors
  12. AD primitives and their geometric meaning
  13. Core Theorems of the Geometric AD Framework
  14. Complexity and Stability Analysis
  15. Multi-Directional and Cross-Partial Differentiation
  16. ...and 5 more sections

Key Result

Theorem 1

Let $f\!: M \to N$ be a smooth map between manifolds and $\ell\!: N \to \mathbb{R}$ a smooth scalar loss function. Then reverse-mode differentiation (backpropagation) satisfies that is, the gradient of the composed loss is obtained by pulling back the covector $d\ell_{f(x)} \in T^*_{f(x)}N$ along the cotangent map $(d f_x)^*\!: T^*_{f(x)}N \to T^*_xM$. More generally, for any Weil algebra $W = \m

Theorems & Definitions (24)

  • Definition 1: Primitive operation
  • Definition 2: Lifted operation
  • Definition 3: Coefficients per scalar
  • Theorem 1: Backpropagation as cotangent pullback and its Weil extension
  • proof
  • Example 1: Illustration of the cotangent pullback and its Weil extension
  • Theorem 2: Exactness of Weil-mode evaluation
  • proof
  • Example 2: Illustration of Theorem \ref{['thm:weil-exactness']}
  • Corollary 1: Coefficient growth envelope
  • ...and 14 more