Don't Unroll Adjoint: Differentiating SSA-Form Programs

2018/10/18 by Michael Innes, Innes, Michael · 4 voices · 62 citations
Computer Science · #Algorithms and Data Compression #Automatic differentiation #Code (set theory) #Compiler #Computation #Computer science #Control flow #Embedded Systems Design Techniques #Nesting (process) #Parallel Computing and Optimization Techniques #Parallel computing #Programming language #Recursion (computer science) #Semantics (computer science) #Theoretical computer science #Transformation (genetics) #cs.PL

paper · pdf · doi:10.48550/arxiv.1810.07951

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2018/10/18 · arxiv created 2019/03/09 · arxiv updated 2019/03/12 · openalex created_date 2025/10/10 · openalex updated_date 2026/08/05

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
• Lyceum: An efficient and scalable ecosystem for robot learning
• Automated Optimization of Laser Fields for Quantum State Manipulation
• Universal Differential Equations for Scientific Machine Learning of Node-Wise Battery Dynamics in Smart Grids
• New Time Integrators and Capabilities in SUNDIALS Versions 6.2.0-7.4.0
• Differentiable molecular simulation can learn all the parameters in a coarse-grained force field for proteins
• Quantum Circuits as a Dynamical Resource to Learn Nonequilibrium Long-Range Order
• QuantumToolbox.jl: An efficient Julia framework for simulating open quantum systems
• Fast and General Automatic Differentiation for Finite-State Methods
• Stochastic estimation of dynamical variables
• ad-trait: A Fast and Flexible Automatic Differentiation Library in Rust
• Inverse design of multiresonance filters via quasi-normal mode theory
• Deriving the Gradients of Some Popular Optimal Transport Algorithms
• 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