Composing Automatic Differentiation with Custom Derivatives of Higher-Order Functions

TL;DR

This paper introduces a formal framework extending reverse-mode automatic differentiation with an operator for custom derivatives, enabling programmers to specify derivatives for composite functions to improve efficiency and stability.

Contribution

It presents a lambda calculus-based formalization that incorporates manual derivatives into autodiff, addressing practical needs beyond automatic generation.

Findings

Demonstrates the extended calculus with practical examples

Shows improved efficiency and stability with custom derivatives

Provides a foundation for further research in autodiff customization

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.