Don't Unroll Adjoint: Differentiating SSA-Form Programs

TL;DR

This paper introduces a source code transformation approach for reverse-mode AD using SSA form, enabling efficient differentiation of complex programs and integrating with existing compilers, demonstrated through the Julia-based tool Zygote.

Contribution

It presents a novel SSA-based source transformation method for reverse-mode AD that supports advanced language features and is implemented in the Julia tool Zygote.

Findings

Supports control flow, recursion, and data structures in AD

Produces highly efficient differentiated code with existing compilers

Improves usability and performance of AD tools

Abstract

This paper presents reverse-mode algorithmic differentiation (AD) based on source code transformation, in particular of the Static Single Assignment (SSA) form used by modern compilers. The approach can support control flow, nesting, mutation, recursion, data structures, higher-order functions, and other language constructs, and the output is given to an existing compiler to produce highly efficient differentiated code. Our implementation is a new AD tool for the Julia language, called Zygote, which presents high-level dynamic semantics while transparently compiling adjoint code under the hood. We discuss the benefits of this approach to both the usability and performance of AD tools.