Don't Unroll Adjoint: Differentiating SSA-Form Programs

TL;DR

The paper introduces a reverse‑mode AD approach built on Static Single Assignment (SSA) form and implements it in Julia as Zygote, enabling adjoint code that compiles with existing toolchains. By generalising the Wengert list to SSA, it handles control flow, mutation, data structures, and higher‑order functions via pullbacks and alpha/phi mechanisms, producing efficient, compiler‑friendly derivatives. It demonstrates strong integration with the Julia compiler, leveraging type inference for inlining and memory efficiency, and reports competitive performance against tracing‑based AD on a range of tasks. The work argues that SSA‑based AD offers both expressive power and practical performance, suggesting broad applicability to differentiable programming across modern compiled languages and accelerators.

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.