Dual-Numbers Reverse AD for Functional Array Languages

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Smeding, Tom, Konarski, Mikołaj, Jones, Simon Peyton, Fitzgibbon, Andrew
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866911061165211648
author Smeding, Tom
Konarski, Mikołaj
Jones, Simon Peyton
Fitzgibbon, Andrew
author_facet Smeding, Tom
Konarski, Mikołaj
Jones, Simon Peyton
Fitzgibbon, Andrew
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2507_12640
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Dual-Numbers Reverse AD for Functional Array Languages
Smeding, Tom
Konarski, Mikołaj
Jones, Simon Peyton
Fitzgibbon, Andrew
Programming Languages
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.
title Dual-Numbers Reverse AD for Functional Array Languages
topic Programming Languages
url https://arxiv.org/abs/2507.12640