Efficient CHAD

TL;DR

This paper optimizes the CHAD algorithm for reverse-mode automatic differentiation, achieving correct computational complexity through well-known techniques and formal proof, and extends applicability to parallel functional array programs.

Contribution

It demonstrates how standard optimizations can be applied to CHAD to achieve efficient, composable reverse-mode AD with formal complexity guarantees, including for parallel array computations.

Findings

Optimized CHAD achieves expected reverse-mode AD complexity.

Functional mutable updates reduce log-factors in complexity.

Applicability to parallel array programs with data-parallel derivatives.

Abstract

We show how the basic Combinatory Homomorphic Automatic Differentiation (CHAD) algorithm can be optimised, using well-known methods, to yield a simple, composable, and generally applicable reverse-mode automatic differentiation (AD) technique that has the correct computational complexity that we would expect of reverse-mode AD. Specifically, we show that the standard optimisations of sparse vectors and state-passing style code (as well as defunctionalisation/closure conversion, for higher-order languages) give us a purely functional algorithm that is most of the way to the correct complexity, with (functional) mutable updates taking care of the final log-factors. We provide an Agda formalisation of our complexity proof. Finally, we discuss how the techniques apply to differentiating parallel functional array programs: the key observations are 1) that all required mutability is (commutative, associative) accumulation, which lets us preserve task-parallelism and 2) that we can write down data-parallel derivatives for most data-parallel array primitives.