🤖 AI Summary
Automatic differentiation (AD) in deep learning and scientific computing suffers from poor structure preservation and inefficient high-order derivative computation. Method: This paper establishes a geometric framework grounded in jet bundles and Weil algebras: reverse-mode AD is interpreted as cotangent pullback, while higher-order Taylor expansions correspond to algebraic evaluation over Weil algebras. Contribution/Results: We introduce tensorized Weil algebras, enabling simultaneous computation of all mixed partial derivatives and circumventing combinatorial explosion from nested Jacobian-vector or vector-Jacobian products. For the first time, correctness and numerical stability of AD are rigorously guaranteed via functorial identities and algebraic exactness. The algorithm exhibits linear complexity in the algebraic dimension and provides explicit bounds on truncation error, thereby establishing a unified theoretical foundation for structure-preserving differential methods.
📝 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.