Towards Denotational Semantics of AD for Higher-Order, Recursive, Probabilistic Languages

TL;DR

This paper develops a denotational semantics framework for automatic differentiation in higher-order, recursive, probabilistic languages, proving correctness and extending previous results to more expressive programming models.

Contribution

It extends the correctness theorem of AD to a highly expressive language with recursion, higher-order functions, and probabilistic primitives, providing a solid theoretical foundation.

Findings

AD computes correct derivatives for a broad class of programs

The framework generalizes existing theorems to recursive and probabilistic languages

It offers a denotational proof of recent probabilistic program results

Abstract

Automatic differentiation (AD) aims to compute derivatives of user-defined functions, but in Turing-complete languages, this simple specification does not fully capture AD's behavior: AD sometimes disagrees with the true derivative of a differentiable program, and when AD is applied to non-differentiable or effectful programs, it is unclear what guarantees (if any) hold of the resulting code. We study an expressive differentiable programming language, with piecewise-analytic primitives, higher-order functions, and general recursion. Our main result is that even in this general setting, a version of Lee et al. [2020]'s correctness theorem (originally proven for a first-order language without partiality or recursion) holds: all programs denote so-called $\omega$PAP functions, and AD computes correct intensional derivatives of them. Mazza and Pagani [2021]'s recent theorem, that AD disagrees with the true derivative of a differentiable recursive program at a measure-zero set of inputs, can be derived as a straightforward corollary of this fact. We also apply the framework to study probabilistic programs, and recover a recent result from Mak et al. [2021] via a novel denotational argument.