Efficient and Sound Differentiable Programming in a Functional Array-Processing Language

TL;DR

This paper introduces an efficient and sound automatic differentiation system for a higher-order functional array-processing language, enabling fast gradient and Jacobian computations suitable for machine learning and optimization tasks.

Contribution

It presents a novel AD system that combines source-to-source forward-mode AD with global optimizations in a functional language, achieving performance comparable to reverse mode.

Findings

Gradient computation with forward-mode AD is as efficient as reverse mode.

Jacobian matrices for numerical algorithms can be computed efficiently.

The system supports higher-order functions and array processing.

Abstract

Automatic differentiation (AD) is a technique for computing the derivative of a function represented by a program. This technique is considered as the de-facto standard for computing the differentiation in many machine learning and optimisation software tools. Despite the practicality of this technique, the performance of the differentiated programs, especially for functional languages and in the presence of vectors, is suboptimal. We present an AD system for a higher-order functional array-processing language. The core functional language underlying this system simultaneously supports both source-to-source forward-mode AD and global optimisations such as loop transformations. In combination, gradient computation with forward-mode AD can be as efficient as reverse mode, and the Jacobian matrices required for numerical algorithms such as Gauss-Newton and Levenberg-Marquardt can be efficiently computed.