Automatic Differentiation with Higher Infinitesimals, or Computational Smooth Infinitesimal Analysis in Weil Algebra

Abstract

We propose an algorithm to compute the $C^\infty$-ring structure of arbitrary Weil algebra. It allows us to do some analysis with higher infinitesimals numerically and symbolically. To that end, we first give a brief description of the (Forward-mode) automatic differentiation (AD) in terms of $C^\infty$-rings. The notion of a $C^\infty$-ring was introduced by Lawvere [10] and used as the fundamental building block of smooth infinitesimal analysis and synthetic differential geometry [11]. We argue that interpreting AD in terms of $C^\infty$-rings gives us a unifying theoretical framework and modular ways to express multivariate partial derivatives. In particular, we can “package” higher-order Forward-mode AD as a Weil algebra, and take tensor products to compose them to achieve multivariate higher-order AD. The algorithms in the present paper can also be used for a pedagogical purpose in learning and studying smooth infinitesimal analysis as well.

This work was supported by the Research Institute for Mathematical Sciences, an International Joint Usage/Research Center located in Kyoto University.

A Succinct Multivariate Lazy Tower AD

For completeness, we include the referential implementation of the Tower-mode AD in Haskell, which can be used in Algorithm 2. The method presented here is a mixture of Lazy Multivariate Tower [13] and nested Sparse Tower [9]. For details, we refer readers to the related paper by the author in RIMS Kôkyûroku [7].

The idea is simple: we represent each partial derivative as a path in a tree of finite width and infinite heights. A path goes down if the function is differentiated by the 0^th variable. It goes right if there will be no further differentiation w.r.t. 0^th variable, but differentiations w.r.t. remaining variable can take place. This intuition is depicted by the following illustration of the ternary case:

This can be seen as a special kind of infinite trie (or prefix-tree) of alphabets $\partial _{x_i}$, with available letter eventually decreasing.

This can be implemented by a (co-)inductive type as follows:

A tree can have an infinite height. Since Haskell is a lazy language, this won’t eat up the memory and only necessary information will be gradually allocated. Since making everything lazy can introduce a huge space leak, we force each coefficient when their corresponding data constructors are reduced to weak head normal form, as expressed by field strictness annotation .

Then a lifting operation for univariate function is given by:

Here, we use type-level constraint to represent to a subclass of smooth functions, e.g. for elementary functions. Constraint of form $\forall x k.\ \texttt {c}\ x => \texttt {c}\ (\texttt {STower}\ k\ x)$ is an example of so-called Quantified Constraints. This requires to be implemented for any succinct Tower AD, provided that their coefficient type, say , is also an instance of . This constraint is used recursively when one implements an actual implementation of instance . For example, instance (for elementary floating point operations) can be written as follows:

In this way, we can implement Tower AD for a class of smooth function closed under differentiation, just by specifying an original and their first derivatives.

More general n-ary case of lifting operator is obtained in just the same way: