Logical Relations for Partial Features and Automatic Differentiation Correctness

TL;DR

This paper introduces a categorical semantics-based logical relations technique to prove the correctness of automatic differentiation methods in ML-like languages, especially handling partial features and recursive types.

Contribution

It develops a new open logical relations framework for reasoning about differentiable partial functions, enabling straightforward correctness proofs for AD implementations.

Findings

Logical relations technique applies to recursive types and partial functions.

Provides a simple proof of correctness for forward- and reverse-mode AD.

Framework is grounded in categorical semantics for semantic rigor.

Abstract

We present a simple technique for semantic, open logical relations arguments about languages with recursive types, which, as we show, follows from a principled foundation in categorical semantics. We demonstrate how it can be used to give a very straightforward proof of correctness of practical forward- and reverse-mode dual numbers style automatic differentiation (AD) on ML-family languages. The key idea is to combine it with a suitable open logical relations technique for reasoning about differentiable partial functions (a suitable lifting of the partiality monad to logical relations), which we introduce.