Automatic Differentiation for ML-family languages: correctness via logical relations

TL;DR

This paper presents a logical relations technique to prove the correctness of automatic differentiation code transformations in ML-family functional languages, ensuring reliable forward- and reverse-mode AD.

Contribution

It introduces a novel monadic logical relations approach for recursive types and terms, providing a semantic correctness proof for AD transformations in functional languages.

Findings

Correctness proof for dual numbers style AD in ML-family languages

Unified framework for forward- and reverse-mode AD correctness

Semantic proof based on simple denotational semantics using ω-cpos

Abstract

We give a simple, direct and reusable logical relations technique for languages with term and type recursion and partially defined differentiable functions. We demonstrate it by working out the case of Automatic Differentiation (AD) correctness: namely, we present a correctness proof of a dual numbers style AD code transformation for realistic functional languages in the ML-family. We also show how this code transformation provides us with correct forward- and reverse-mode AD. The starting point is to interpret a functional programming language as a suitable freely generated categorical structure. In this setting, by the universal property of the syntactic categorical structure, the dual numbers AD code transformation and the basic $\omega$-cpo semantics arise as structure preserving functors. The proof follows, then, by a novel logical relations argument. The key to much of our contribution is a powerful monadic logical relations technique for term recursion and recursive types. It provides us with a semantic correctness proof based on a simple approach for denotational semantics, making use only of the very basic concrete model of $\omega$-cpos.