A Common Interface for Automatic Differentiation

TL;DR

The paper tackles the challenge of choosing among diverse Automatic Differentiation backends in scientific computing by introducing DifferentiationInterface.jl (DI), a common frontend that unifies access to multiple AD packages. It emphasizes a modular design built on backends, operators, and a novel preparation mechanism that amortizes one-time computations, enabling advanced capabilities like sparse AD. Key contributions include a standardized API across backends, support for contexts and sparsity, backend composition and translation, and an accompanying testing framework for benchmarking. This approach lowers the barrier to experiment with different AD solutions, improves performance for complex differentiable workflows, and provides a scalable path for broader adoption in Julia-based scientific computing.

Abstract

For scientific machine learning tasks with a lot of custom code, picking the right Automatic Differentiation (AD) system matters. Our Julia package DifferentiationInterface$.$jl provides a common frontend to a dozen AD backends, unlocking easy comparison and modular development. In particular, its built-in preparation mechanism leverages the strengths of each backend by amortizing one-time computations. This is key to enabling sophisticated features like sparsity handling without putting additional burdens on the user.

Figures (2)

• Figure 1: Comparison of the AD ecosystems in Python and Julia for applications to optimization, differential equations and probabilistic programming
• Figure 2: Impact of preparation on gradient performance of $f : x \mapsto \lVert x \rVert^2$