Dual-Numbers Reverse AD for Functional Array Languages

TL;DR

This paper enhances reverse-mode automatic differentiation for array languages by introducing a first-class multidimensional array support with minimal performance overhead, using a novel vectorisation transformation and symbolic interpretation.

Contribution

It presents a new approach combining vectorisation, lifting, and symbolic interpretation to efficiently support multidimensional arrays in dual-numbers reverse AD.

Findings

Achieves low-overhead support for multidimensional arrays in reverse AD.

Transforms higher-order array programs into first-order programs for AD.

Enables dual-array lifting with minimal modifications.

Abstract

The standard dual-numbers construction works well for forward-mode automatic differentiation (AD) and is attractive due to its simplicity; recently, it also has been adapted to reverse-mode AD, but practical performance, especially on array programs, leaves a lot to be desired. In this paper we introduce first-class support for multidimensional arrays in dual-numbers reverse-mode AD with little to no performance overhead. The algorithm consists of three loosely-coupled components: a semantics-preserving vectorisation code transformation (the bulk-operation transform or BOT), a fairly straightforward lifting of the basic dual-numbers reverse AD algorithm to a mostly first-order array language, and symbolic interpretation to achieve an end-to-end compilation pipeline. Unfortunately, we lose some of the nice generalisable aspects of dual-numbers AD in the process, most importantly support for higher-order code. We do support some higher-order array combinators, but only a carefully-chosen set: 'build' (elementwise array construction), 'gather' and 'scatter'. In return, the BOT can eliminate the essential (for AD) higher-orderness of the input program, meaning that AD gets essentially presented with a first-order program. This allows the naive trick of lifting dual numbers to "dual arrays" to work without much modification.