[Paper Review] Don't Unroll Adjoint: Differentiating SSA-Form Programs

Michael Innes|arXiv (Cornell University)|Oct 18, 2018

Parallel Computing and Optimization Techniques22 references61 citations

TL;DR

The paper introduces a reverse-mode AD approach using SSA-form programs, implemented in Julia as Zygote, enabling efficient differentiation of complex language features and interactions with compilers.

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.

Motivation & Objective

• Motivate the need for expressive, high-performance AD beyond tracing-based methods.
• Propose SSA-form based AD that integrates with existing compilers to produce optimized derivative code.
• Demonstrate the Zygote implementation and its interaction with Julia's compiler and runtime.
• Show how this approach supports language features like control flow, mutation, and higher-order derivatives.

Proposed method

• Represent differentiation via a J operator and pullbacks to propagate gradients through SSA-form IR.
• Use reversed basic blocks and phi nodes to handle control flow and data dependencies in the adjoint.
• Store and resolve alpha values on a stack to manage SSA iteration definitions during reverse differentiation.
• Handle complex data structures and mutation with corresponding adjoint pullbacks for cons cells, boxes, and closures.
• Leverage Julia's type inference and compiler optimizations to inline pullbacks and produce efficient code.
• Position the approach as a bridge between high-level differentiable programming and traditional compiler optimization.

Experimental results

Research questions

• RQ1 Can reverse-mode AD over SSA-form IR accommodate full language features including control flow, mutation, and higher-order functions?
• RQ2 How does differentiating SSA-form programs interact with a host compiler to produce efficient differentiated code?
• RQ3 What is the performance and usability trade-off compared with tracing-based AD systems?
• RQ4 How can adjoint code be generated to preserve correct dataflow across branches and loops?
• RQ5 Is this approach scalable to real-world models and compiler backends (e.g., LLVM, TPUs)?

Key findings

• Zygote provides a working SSA-form AD tool for Julia that yields highly efficient differentiated code.
• The adjoint construction preserves correct dataflow through reversed CFGs using phi nodes and alpha tracking.
• Pullbacks can be composed to support higher-order derivatives and complex language features including mutation and data structures.
• Experiments show competitive performance with hand-written derivatives in several cases and favorable benchmarks against tracing approaches.
• The approach enables leveraging existing compiler optimizations and future accelerators without sacrificing expressiveness.
• SSA-based differentiation can be integrated with modern compiler stacks to enable differentiable programming across languages.