Don't Unroll Adjoint: Differentiating SSA-Form Programs

2018/10/18 by Michael Innes, Innes, Michael · 4 voices · 19 citations
Computer Science · #Algorithms and Data Compression #Embedded Systems Design Techniques #Parallel Computing and Optimization Techniques #cs.PL

paper · pdf · doi:10.48550/arxiv.1810.07951

openalex publication_date 2018/10/18 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

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.

Citations

• Efficiently computing static single assignment form and the control dependence graph
• Julia: A Fresh Approach to Numerical Computing
• Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs

Cited by

• Efficient optimisation of multi-parameter quantum control protocols for strongly-coupled systems
• Training a force field for proteins and small molecules from scratch
• A Common Interface for Automatic Differentiation
• Sparser, Better, Faster, Stronger: Sparsity Detection for Efficient Automatic Differentiation
• NonlinearSolve.jl: High-Performance and Robust Solvers for Systems of Nonlinear Equations in Julia
• An Adjoint-Based Differentiable Physics Framework for Online Parameter Inversion in Closed-Brayton Gas-Cooled Reactor Digital Twins
• Eigenvalue-accelerated LDOS optimization of high-Q optical resonances
• CODE: A global approach to ODE dynamics learning
• Branching Flows: Discrete, Continuous, and Manifold Flow Matching with Splits and Deletions
• Flinch: A Differentiable Framework for Field-Level Inference of Cosmological parameters from curved sky data
• Residual-Informed Learning of Solutions to Algebraic Loops
• Differentiating a Tensor Language
• Locally Purified Maximally Mixed States At Scale: Entanglement Pruning and Symmetries
• BackPACK: Packing more into backprop
• Differentiating through Log-Log Convex Programs
• Generalization vs. Memorization in Autoregressive Deep Learning: Or, Examining Temporal Decay of Gradient Coherence
• Topology Optimization through Differentiable Finite Element Solver
• Differentiable programming [wikipedia]
• Flux (machine-learning framework) [wikipedia]

Discussions

• Differentiating SSA-Form Programs in Julia (2018) [pdf] [hn, 87 points, 16 comments]
• Don't Unroll Adjoint: Differentiating SSA-Form Programs [hn, 2 points, 0 comments]
• Don't Unroll Adjoint: Differentiating SSA-Form Programs [hn, 1 points, 0 comments]
• Don't Unroll Adjoint: Differentiating SSA-Form Programs [hn, 1 points, 0 comments]

Related

• AbstractDifferentiation.jl: Backend-Agnostic Differentiable Programming in Julia
• Denotationally Correct, Purely Functional, Efficient Reverse-mode Automatic Differentiation
• A Common Interface for Automatic Differentiation
• Efficient Implementation of a Higher-Order Language with Built-In AD
• A Common Interface for Automatic Differentiation