Don't Unroll Adjoint: Differentiating SSA-Form Programs

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.

References in corpus (6)

• DyNet: The Dynamic Neural Network Toolkit
• Intel nGraph: An Intermediate Representation, Compiler, and Executor for Deep Learning
• Reversible Architectures for Arbitrarily Deep Residual Neural Networks
• Automatic differentiation in ML: Where we are and where we should be going
• Tangent: Automatic Differentiation Using Source Code Transformation in Python
• Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs

Cited by in corpus (22)

• The ITensor Software Library for Tensor Network Calculations
• Differentiable Programming Tensor Networks
• DiffEqFlux.jl - A Julia Library for Neural Differential Equations
• TensorFlow Eager: A Multi-Stage, Python-Embedded DSL for Machine Learning
• Discovering Sparse Interpretable Dynamics from Partial Observations
• Stochastic Estimation of Dynamical Variables
• Trace formulation for photonic inverse design with incoherent sources
• Variational optimization and data assimilation in chaotic time-delayed systems with automatic-differentiated shadowing sensitivity
• Differentiating a Tensor Language
• Arrhenius.jl: A Differentiable Combustion SimulationPackage
• Dynamic Automatic Differentiation of GPU Broadcast Kernels
• Lyceum: An efficient and scalable ecosystem for robot learning
• Differentiating through Log-Log Convex Programs
• Spline parameterization of neural network controls for deep learning
• Opening the Blackbox: Accelerating Neural Differential Equations by Regularizing Internal Solver Heuristics
• BackPACK: Packing more into backprop
• Differentiate Everything with a Reversible Embeded Domain-Specific Language
• Tuning of Online Feedback Optimization for setpoint tracking in centrifugal compressors
• Mapping the Internet: Modelling Entity Interactions in Complex Heterogeneous Networks
• Self-explaining variational posterior distributions for Gaussian Process models
• Topology Optimization through Differentiable Finite Element Solver
• Error Autocorrelation Objective Function for Improved System Modeling