Composing Automatic Differentiation with Custom Derivatives of Higher-Order Functions

TL;DR

The paper addresses the need for extensible automatic differentiation by allowing programmers to attach and propagate custom derivatives for higher-order functions and closures within a reverse-mode AD framework. It formalizes this capability in an untyped lambda calculus setting, introducing an operator to attach derivatives and rules for extraction, inversion, and composition that preserve closure properties. Through concrete examples, it demonstrates improved numerical stability, flexible handling of primitives, one-sided and implicit differentiation, and space-efficient differentiation of vector operations like map. This approach broadens the design space of differentiable programming, enabling practitioners to encode expert insights directly into the differentiation process and potentially achieve better performance and stability in real-world computations.

Abstract

Recent theoretical work on automatic differentiation (autodiff) has focused on characteristics such as correctness and efficiency while assuming that all derivatives are automatically generated by autodiff using program transformation, with the exception of a fixed set of derivatives for primitive operations. However, in practice this assumption is insufficient: the programmer often needs to provide custom derivatives for composite functions to achieve efficiency and numerical stability. In this work, we start from the untyped lambda calculus with a reverse-mode autodiff operator, extend it with an operator to attach manual derivatives, and demonstrate its utility via several examples.