Source-linked AI summary
A Review of automatic differentiation and its efficient implementation
Charles C. Margossian
TL;DR
Automatic differentiation is useful for derivative calculations, but naive implementations can be prohibitively slow and memory-intensive. This review surveys strategies for efficient implementation and differentiating complex algorithms, and reports that the implicit function theorem was orders of magnitude faster than standard AD in an algebraic-solver experiment.
Problem
Naive automatic differentiation can produce prohibitively slow code and excess memory use.
Source: Computational implementation · p. 10Method
The review surveys computational schemes for efficient automatic differentiation and strategies for optimizing differentiation of complex mathematical algorithms.
Source: Introduction · p. 3Results
In the algebraic-solver experiment, the implicit function theorem was orders of magnitude faster than standard AD.
Source: Differentiating a numerical algebraic… · p. 22Takeaways & Limitations
Users selecting an AD tool should compare the applications that motivated its development with their own goals.
Source: Conclusion · p. 27The review only briefly discusses alternatives for higher-order derivatives, and whether recursively applying AD is optimal remains an open question.
Source: Higher-order differentiation · p. 26
Abstract
from arXiv · showhide
Derivatives play a critical role in computational statistics, examples being Bayesian inference using Hamiltonian Monte Carlo sampling and the training of neural networks. Automatic differentiation is a powerful tool to automate the calculation of derivatives and is preferable to more traditional methods, especially when differentiating complex algorithms and mathematical functions. The implementation of automatic differentiation however requires some care to insure efficiency. Modern differentiation packages deploy a broad range of computational techniques to improve applicability, run time, and memory management. Among these techniques are operation overloading, region based memory, and expression templates. There also exist several mathematical techniques which can yield high performance gains when applied to complex algorithms. For example, semi-analytical derivatives can reduce by orders of magnitude the runtime required to numerically solve and differentiate an algebraic equation. Open problems include the extension of current packages to provide more specialized routines, and efficient methods to perform higher-order differentiation.
Graphical table of content
The review considers computational and mathematical approaches to implementing automatic differentiation efficiently. It covers implementation techniques and differentiation of mathematical functions and algorithms, with applications including Hamiltonian Monte Carlo sampling.
Computational implementation
- The review surveys source transformation, operator overloading, checkpointing, region-based memory, and expression templates for computational implementation.Source: Computational implementation · p. 2
Mathematical implementation
- Hamiltonian Monte Carlo sampling is identified as an application of automatic differentiation in computational statistics.Source: Mathematical implementation · p. 2
- Automatic differentiation has been applied across computational statistics, and the review considers chain rules, reduced expression graphs, and super nodes for differentiating functions and algorithms.Source: Mathematical implementation · p. 2
Introduction
Derivatives are important across optimization, machine learning, and Bayesian computation, but growing model and algorithm complexity makes them difficult to calculate. This review examines how to implement automatic differentiation efficiently and how to assess its performance.
Introduction
- Derivatives support optimization, neural-network training, and Bayesian inference, while complex problems can make their calculation difficult, time-consuming, and error-prone.The passage cites gradient descent and Newton’s method, Hamiltonian Monte Carlo, and variational inference as examples.Source: Introduction · p. 2
- Finite and symbolic differentiation are alternatives to AD, which escapes many of their limitations on complex functions.Several programming-language packages implement AD, including as a black-box capability in PyTorch and Stan.Source: Introduction · pp. 2–3
- The review covers AD fundamentals, implementation schemes, strategies for differentiating complex algorithms, and practical and open problems.It emphasizes that naive AD implementations can produce inefficient code, so AD does not fully automate differentiation in practice.Source: Introduction · p. 3
- The review weighs accuracy, run time, memory use, applicability, and code readability, but excludes compile time because programs are typically compiled once and reused many times.The authors consider accuracy most important, while noting that it rarely distinguishes AD implementation schemes that are exact up to arithmetic precision.Source: Introduction · p. 3
How automatic differentiation works
Automatic differentiation obtains derivatives by applying the chain rule to a program decomposed into elementary operations. Forward and reverse sweeps propagate Jacobian-vector or transposed-Jacobian-vector products, with mode choice depending on input/output dimensions and memory costs.
How automatic differentiation works
- Automatic differentiation decomposes composite functions into elementary operations and applies the chain rule to their Jacobians to compute derivatives.An expression graph represents intermediate variables and their dependencies; derivatives can be propagated operation by operation rather than by explicitly computing every Jacobian element.Source: How automatic differentiation works · pp. 4, 6
Forward mode
- A forward sweep evaluates f and propagates a directional derivative from inputs to outputs, costing at most 2.5 × OPS(f); n sweeps recover the full Jacobian.Initializing the tangent to 1 for one input and 0 for the others computes that Jacobian column as values and directed derivatives are evaluated together.Source: Forward mode · pp. 6–7
Reverse mode
- A reverse sweep computes J^T · w at a cost of at most 4 × OPS(f); for a scalar output, one sweep gives the gradient.It follows a forward evaluation trace with a reverse pass from the output to the input variables, requiring intermediate values to be stored; lazy evaluation can limit work to expressions needed for the gradient.Source: Reverse mode · pp. 7, 9
Forward or reverse mode?
- For an n-input, m-output Jacobian, reverse mode needs m sweeps and is preferable when n>m, while forward mode is preferable when n≤m (Baydin et al., 2018).Reverse mode is especially applicable to many-input, few-output problems; forward mode can also be more efficient when input and output counts are comparable because it has less expression-graph storage overhead.Source: Forward or reverse mode? · pp. 9–10, Table 4 · p. 10
- 30% runtime reduction versus standard forward AD without caching, but less than 10% gain when forward AD already uses caching, was reported for intermediate reverse mode (Phipps & Pawlowski, 2012).The study applied reverse mode to intermediate expressions within an overall forward-mode framework for fluid mechanics.Source: Forward or reverse mode? · p. 10
Computational implementation
Efficient automatic differentiation requires implementation choices that manage runtime and memory costs. This section reviews source transformation, operator overloading, memory-management strategies, checkpointing, and expression templates, whose suitability depends on the problem and computational tradeoffs.
Computational implementation
- Implementation benchmarks are indirect and cover limited problem classes, with no standard suite spanning the diversity of automatic-differentiation applications.A computationally naive implementation can be prohibitively slow and memory-intensive; compared C++ libraries commonly use operator overloading rather than source transformation, while other techniques suit different problems.Source: Computational implementation · pp. 10–11
Source transformation
- Source transformation differentiates code at preprocessing time, but its compile-time information limits handling of typical C++ features such as loops, templates, and object-oriented constructs.It works well for Fortran or C and is advocated for reverse-mode AD, but implementing such a package demands considerable effort; operator overloading suits typical C++ features.Source: Source transformation · p. 11
Operator overloading
- Operator overloading represents expression-graph variables with objects holding values and differential components, with forward components as tangents and reverse components as adjoints.Only root variables requiring sensitivities and intermediate variables depending on them need differentiable objects; the tangent or cotangent initialization must be appropriate to the mode.Source: Operator overloading · pp. 11–12
Memory management
- Reverse-mode operator overloading needs persistent storage of the expression graph and intermediate values, with dynamic memory requirements when control flow includes loops or conditionals.Forward mode computes values and differentials in one sweep using twice the memory of a function evaluation, and its variable objects can be discarded after their dependent nodes are evaluated.Source: Memory management · p. 12
Retaping
- Retaping reuses an expression-graph tape across sweeps when the graph remains unchanged, but changed control flow requires a new tape.The graph can be reused across evaluation points even when intermediate values must be recomputed; ADOL-C and CppAD implement retaping.Source: Retaping · p. 12
Checkpoints
- Checkpointing reduces peak memory by recording only selected parts of a reverse-mode trace and recomputing forward traces, trading memory for runtime.Store-all retains intermediate values at checkpoints, while snapshots offer intermediate memory-use options; no placement is generally optimal, though binomial partitioning is optimal for sequential steps of comparable cost. Checkpoints can also keep segment tapes local across cores, as in Stan Math map-reduce functionality.Source: Checkpoints · pp. 13–14
Region Based Memory
- Region-based memory management let RAD run up to twice as fast as ADOL-C when differentiating quality mesh functions of varying complexity.It allocates derivative memory element-wise on a custom stack, then frees the entire stack at once instead of freeing each element individually.Source: Region Based Memory · p. 15
Expression templates
- Adept was 2.6–9 times faster and 1.3–7.7 times more memory-efficient in specific fluid-dynamics tests; Stan Math took 40% less time with a revised log-sum-exp implementation.Adept achieved optimal speed with both tested implementations, and both packages performed similarly on the revised algorithm. Matching expression-template benefits through careful function writing requires more user effort.Source: Expression templates · p. 16
Mathematical implementation
AD efficiency can be improved by reducing expression-graph size, from reusing intermediate calculations to custom super nodes. In the solver case study, the implicit-function-theorem approach is orders of magnitude faster than standard AD.
Mathematical implementation
- Operator overloading handles most programming statements but requires care with memory; retaping is fast for repeated expressions only when their graph is conserved.Source: Mathematical implementation · p. 17
Reducing the number of operations in a function
- The optimized 2 × 2 matrix-exponential implementation runs about a third faster by reducing redundant operations and graph nodes, but costs coding effort and clarity.Intermediate variables reduce runtime and memory use with AD, while expression templates may not remove repeated computations that the compiler fails to identify.Source: Reducing the number of operations in a… · pp. 18–20, Table 7 · p. 20
Optimizing with super nodes
- Super nodes treat iterative components as stand-alone operators with specialized Jacobian methods, avoiding large graphs that can cause precision errors, excessive memory use, and slow computation.The number of elementary functions in the expanded graph can vary with the numerical solver’s steps.Source: Optimizing with super nodes · p. 20
Differentiating a numerical algebraic solver
- The implicit-function-theorem solver is orders of magnitude faster than standard AD, including Stan Math’s built-in solver; analytic Jy adds a minor gain that grows with state count.The comparison uses an analytical Jy for the custom solver, while Stan Math’s built-in solver computes Jy with AD.Source: Differentiating a numerical algebraic… · p. 22
Open and practical problems
Automatic differentiation packages face practical trade-offs in combining optimization techniques and supporting specialized routines, while extending them can demand substantial expertise. Higher-order differentiation remains more difficult and potentially expensive than first-order differentiation, with the best approach still an open question.
General and specialized code
- No package implements every optimization technique; compatibility, development effort, and impact shape choices, and retaping suits one Jacobian but not changing graphs evaluated at many points.Custom differentiation can be solver-agnostic for ODEs and algebraic equations, but PDE and field-specific analytical solutions can make implementation burdensome.Source: General and specialized code · pp. 22, 25
- The analytical “super node” solver is approximately two to three times faster than the built-in solver using AD to compute Jy when both use the implicit function theorem.The comparison uses an analytical expression for Jy in the super node solver and AD for Jy in the built-in solver.Source: Figure 5 · p. 24
Extensibility
- Extensibility can let users add features, but extending a package takes time, coding expertise, and package-specific knowledge; custom derivative interfaces and better documentation are proposed remedies.PyTorch offers custom reverse-mode derivatives, and CasADi offers custom derivative functions, though its feature is experimental and requires technical understanding.Source: Extensibility · p. 25
Higher-order differentiation
- Higher-order derivatives are harder and can be prohibitively expensive, limiting their use; RHMC needs second- and third-order derivatives, yet is virtually never used while first-order methods are popular.RHMC can overcome geometric pathologies that first-order sampling methods cannot.Source: Higher-order differentiation · pp. 25–26
- Hessian-vector products can be computed with a forward sweep followed by a reverse sweep at cost linear in the complexity of the function, avoiding construction of the full Hessian.The operation can be repeated with basis vectors to obtain the full Hessian; the expression graph can be stored and reused.Source: Higher-order differentiation · p. 26
- Recursive AD can be inefficient and numerically unstable, prompting alternatives based on univariate Taylor series or explicit higher-order differential operators.The differential-operator framework derives second- and third-order operators, but requires differential-geometry expertise and remains a prototype in Nomad.Source: Higher-order differentiation · pp. 26–27
Conclusion
AD supports major applications in computational statistics and machine learning, while users can improve performance by understanding its techniques and choosing tools suited to their goals. Tool choice depends on language, application fit, available functions, extensibility, and advanced computational features.
Conclusion
- AD supports NUTS, the foundation of Stan and PyMC3, variational inference, and neural networks in TensorFlow.The cited examples are NUTS, Stan, PyMC3 (Salvatier, Wiecki, & Fonnesbeck, 2016), variational inference, and TensorFlow (Abadi et al., 2016).Source: Conclusion · p. 27
- Basic understanding of AD can help users optimize code, while advanced users can add expression templates or custom derivative methods.AD often acts as a black box within statistical or modeling software, but users may exploit mathematical techniques or edit package source code to improve performance.Source: Conclusion · p. 27
- Choose AD tools by language and application fit, then assess their computational techniques, function libraries, extensibility, and support for features such as parallelization and GPUs.Language constraints may be relaxed by links between higher- and lower-level languages, such as RCpp for using C++ inside R. Built-in functions can save effort and often provide some optimality; extensibility may matter for esoteric problems.Source: Conclusion · p. 27
- Stan Math optimizes flexible reverse AD for log-posterior gradients, but its forward mode is less optimized and well tested.It also provides mathematical functions for linear algebra, numerical calculations, and probability densities.Source: Conclusion · p. 27
Appendix: computer experiments
The appendix details the setup for two computer experiments complementing the mathematical implementations section. It specifies the test environment and how runtime was measured.
Appendix: computer experiments
- The experiments ran on a 2015 MacBook Pro with a 2.7 GHz Intel Core i5 and 8 GB of memory.The memory was 1867 MHz DDR3.Source: Appendix: computer experiments · p. 28
- The test environment used clang++ 4.2.1, Eigen 3.3.3, Boost 1.66.0, and Stan Math 2.17-develop9.The compiler and library specifications match those used for Stan Math unit tests.Source: Appendix: computer experiments · p. 28
- Runtime was measured as wall-clock time before and after evaluating the target function and calculating all sensitivities of interest.Timing used C++ std::chrono::system_clock::now(); the code is available on GitHub.Source: Appendix: computer experiments · p. 28
Differentiating the 2 × 2 Matrix Exponential
The section describes generating 2 × 2 matrices from four uniform random elements in U(1, 10) to ensure that ∆ in Equation 11 is real.
Differentiating the 2 × 2 Matrix Exponential
- Four elements sampled uniformly from U(1, 10) form each 2 × 2 matrix, keeping ∆ in Equation 11 real.The passage characterizes this choice as conservative and specifies bounds of 1 and 10.Source: Differentiating the 2 × 2 Matrix… · p. 28
Differentiating a numerical algebraic solver
The section demonstrates differentiating a numerical algebraic solver using a pharmacometric steady-state problem, computing sensitivities for model parameters and extending the setup to multiple patients.
Differentiating a numerical algebraic solver
- The test case is a pharmacometrics steady-state problem for an orally administered drug in a one-compartment model with first-order absorption from the gut.The ODE describes drug diffusion and clearance over time.Source: Differentiating a numerical algebraic… · p. 28
- A dosing cycle evolves the patient's state through an operator g(y0, δt), and equilibrium is defined by an algebraic condition relating states before and after a dose.The dose adds mass m to the gut at the start of each cycle.Source: Differentiating a numerical algebraic… · p. 29
- Although the steady-state algebraic equation can be solved analytically, the example solves it numerically and computes sensitivities for k1 and k2.The numerical solution is used for demonstration.Source: Differentiating a numerical algebraic… · p. 29
- For n patients, the system has 2n states and 2n + 2 parameters requiring sensitivities, with patient-specific coefficients scaled by random factors from 0.7 to 1.3.Each patient's coefficients satisfy ki = φik, where k is a population parameter and φi is uniformly sampled between 0.7 and 1.3.Source: Differentiating a numerical algebraic… · p. 29