Source-linked AI summary
A Review of automatic differentiation and its efficient implementation
Charles C. Margossian
TL;DR
Computing derivatives is increasingly difficult for complex statistical models and algorithms, motivating methods that automate differentiation without sacrificing efficiency. The paper reviews automatic differentiation techniques and mathematical strategies, finding that implementation choices and specialized derivatives can substantially affect runtime and memory use. It also identifies higher-order differentiation and specialized routines as open problems.
Problem
As probabilistic models and algorithms become more complex, derivative computation becomes formidable, while naive automatic-differentiation implementations can be slow and memory-intensive.
Method
The paper reviews automatic-differentiation modes, implementation schemes, computational optimizations, and mathematical techniques for differentiating complex algorithms.
Results
30% reduction in runtime was reported when intermediate expressions used reverse-mode differentiation within an overall forward-mode framework.
Takeaways & Limitations
Choosing differentiation modes and implementation techniques according to the problem can improve runtime, memory usage, and applicability.
Takeaways & Limitations
Efficient higher-order differentiation remains an open question, with recursive automatic differentiation reported to produce inefficient and sometimes numerically unstable code.
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 paper reviews automatic differentiation for mathematical functions and algorithms, emphasizing its role in computational statistics and its implementation considerations.
- The review covers forward and reverse-mode chain rules, reduced expression graphs, and super nodes.
- Its applications include Hamiltonian Monte Carlo sampling in computational statistics.
- Automatic differentiation differentiates mathematical functions and algorithms.
Introduction
Derivatives are central to optimization, neural-network training, and Bayesian inference, but increasing model and algorithmic complexity makes their computation difficult. The paper reviews derivative-calculation methods and implementation techniques for improving automatic differentiation efficiency.
- Derivatives support objective optimization, neural-network training, and Bayesian posterior inference.
- Finite and symbolic differentiation can perform poorly on complex functions, motivating automatic differentiation.
- Automatic differentiation can still produce inefficient code when implemented naively, so the paper reviews computational schemes and optimization strategies.
- The paper considers accuracy, runtime, memory usage, applicability, implementation effort, and readability when evaluating differentiation approaches.
How automatic differentiation works
Automatic differentiation decomposes a programmed function into elementary operations and applies the chain rule through its expression graph. Forward mode propagates directional derivatives from inputs, while reverse mode propagates adjoints from outputs; their efficiency depends on the input-output dimensions and implementation strategy.
- Core mechanism: A target function is represented as a composition of elementary maps whose Jacobians multiply according to the chain rule.This decomposition allows Jacobian–vector and transposed-Jacobian–vector products to be evaluated sequentially rather than forming the full Jacobian explicitly.
- Forward mode: Forward-mode AD evaluates the function while propagating a directional derivative from an input direction through each intermediate variable.Choosing a direction with one nonzero input component computes one Jacobian column, so a full m×n Jacobian requires n forward sweeps.
- Reverse mode: Reverse-mode AD first evaluates the function, then propagates adjoints backward from an output to the input variables.For a one-dimensional output, one reverse sweep computes the full gradient, although intermediate values must be available during the reverse pass.
- Choosing a mode: Reverse mode is preferable when n >> m, whereas forward mode is preferable when n ≤m because the required sweep counts and storage overhead differ.Reverse mode requires m sweeps for all Jacobian rows, while forward mode requires n sweeps for all columns; forward mode can also avoid reverse-mode expression-graph storage.
- Choosing a mode: Intermediate functions can use a different AD mode from the overall computation, and intermediate reverse mode reduced runtime by 30% versus uncached standard forward AD in a fluid-mechanics case.When forward AD already used caching, the additional gain was less than 10%.
Computational implementation
Efficient AD implementations balance applicability, runtime, and memory through choices such as operator overloading, tapes, checkpointing, and expression templates. These techniques introduce trade-offs: broader program support and lower peak memory can require additional implementation effort, storage, or recomputation.
- Naive AD can produce prohibitively slow code and excessive memory use, motivating careful implementation choices.
- Operator overloading: Operator overloading represents differentiable variables as objects carrying values and differential components, with overloaded operators propagating derivatives.Forward mode stores derivatives with respect to an input; reverse mode stores adjoints with respect to an output.
- Memory management: Reverse-mode AD requires persistent storage of the expression graph and intermediate values because the reverse trace follows the forward evaluation.A tape can be reused across sweeps when the graph and intermediate values remain unchanged, while retaping updates the graph when control flow changes.
- Checkpointing: Checkpointing lowers peak memory by recording only selected portions of the computation and recomputing forward traces between checkpoints.Recompute-all repeatedly reruns partial traces from earlier checkpoints, trading reduced memory for multiple forward sweeps.
- Expression templates: Expression templates can improve performance by delaying evaluation and allowing compilers to combine expression-graph nodes and eliminate redundancies.Stan Math took 40% less time on the second implementation, while careful user-written code can achieve similar results with more coding effort.
Mathematical implementation
Efficient AD combines computational techniques that reduce expression-graph size with mathematical methods tailored to complex algorithms. These choices can improve runtime and memory use, but may increase coding effort and reduce clarity.
- Computational techniques: Operator overloading handles most program statements, while memory management requires care.
- Computational techniques: Retaping is fast when differentiating the same expression graph repeatedly, but requires a new tape when the graph changes.
- Code optimization: Eliminating redundant operations reduces both runtime and the AD expression graph, while intermediate variables can increase ordinary memory usage but improve AD memory usage.
- Code optimization: The optimized 2 × 2 matrix-exponential implementation is about a third faster because it produces an expression graph with fewer nodes.The benchmark evaluates exp(A) and 16 partial derivatives on 1000 randomly generated matrices.
- Trade-offs: Manual optimization can require more code and reduce clarity, while expression templates may fail to remove repeated expressions.The authors identify a trade-off between performance and readability, illustrated by the variable ad sinh half delta.
- Specialized mathematical methods: Using the implicit function theorem for algebraic solvers is orders of magnitude faster than standard AD, with analytical Jy providing additional gains as state dimension increases.The study compares numerical algebraic-solver implementations and measures runtime; super nodes provide specialized Jacobian computation with smaller graphs.
Open and practical problems
AD packages and optimization strategies remain incomplete because techniques can conflict, depend on problem structure, or require substantial implementation effort. Higher-order differentiation is especially difficult, costly, and still an open area for improvement.
- Package limitations: No single AD package comprehensively implements all reviewed optimization techniques, partly because methods may be incompatible or offer problem-specific benefits.
- Experimental evidence: The algebraic-solver experiment uses 100 repetitions and reports uncertainty through sample standard deviation, while the compared solvers differ in their Jacobian strategy.The implicit-function-theorem approach is contrasted with standard AD and with analytical versus AD-computed Jy.
- Package limitations: Retaping is well suited to computing one Jacobian when the expression graph is fixed, but is not helpful when the graph changes across evaluation points.
- Specialized methods: Specialized differentiation methods work for some ODE and algebraic-equation problems but are not generally transferable to PDEs or field-specific analytical solutions.The coding burden becomes more severe when developers must determine which specialized routines are worth implementing.
- Higher-order differentiation: Higher-order derivatives are less common and can be prohibitively expensive, contributing to practitioners’ reluctance to deploy them.RHMC requires second- and third-order derivatives, whereas first-order HMC has become widely used for statistical modeling.
- Higher-order differentiation: Recursive application of AD for higher-order derivatives may be inefficient and numerically unstable, and whether it is optimal remains an open question.Alternatives include univariate Taylor series and explicitly constructed higher-order differential operators.
Conclusion
AD supports computational statistics and machine learning applications, while users can improve efficiency by understanding package capabilities and mathematical optimization techniques. Package selection should match the application, differentiation mode, computational techniques, and available function library.
- Applications: AD supports applications including NUTS, Stan, and PyMC3 in computational statistics and machine learning.
- Practical use: Users can optimize AD code by applying mathematical techniques or adding expression-template routines and custom derivative methods.
- Package selection: Package choice should consider programming language, motivating applications, compatible computational techniques, and the library of available functions.Stan Math is optimized for reverse-mode AD and log-posterior gradients, while its forward mode is less optimized and tested.
Appendix: computer experiments
The appendix describes two computer experiments and specifies their hardware, compiler, libraries, and wall-clock runtime measurement procedure.
- Experimental setup: The appendix complements the performance tests described in the main text.
- Experimental setup: The experiments use a 2015 MacBook Pro with a 2.7 GHz Intel Core i5 processor and 8 GB of memory.
- Experimental setup: The compiler is clang++ version 4.2.1.
- Experimental setup: The libraries are Eigen 3.3.3, Boost 1.66.0, and Stan Math 2.17-development9.
- Measurement: Runtime is measured as wall-clock time while evaluating the target function and calculating all sensitivities of interest.The measurement uses std::chrono system_clock::now().
Differentiating the 2 × 2 Matrix Exponential
The 2 × 2 matrices are generated by sampling four elements uniformly between 1 and 10, ensuring that Δ in Equation 11 is real.
- Four matrix elements are sampled from U(1, 10), with lower and upper bounds of 1 and 10.The implementation uses std::uniform_real_distribution<T> unif(1, 10).
- The conservative sampling range ensures that Δ in Equation 11 has no imaginary part.
Differentiating a numerical algebraic solver
The paper tests differentiation of a numerical algebraic solver using a pharmacometric steady-state problem. It extends the example to multiple patients and computes sensitivities for patient-specific and population parameters.
- Problem setup: The test problem computes steady states for patients undergoing medical treatment, using a standard pharmacometrics problem.The scientific background is presented only as an overview.
- Problem setup: A one-compartment pharmacokinetic model describes first-order drug absorption from the gut and diffusion into a central compartment.y1 denotes gut drug mass, y2 central-compartment drug mass, and k1 and k2 diffusion rates.
- Numerical solver: The analytically solvable ODE system is represented by an evolution operator g(y0, δt) that advances the patient state over a time interval.The operator takes an initial condition and a time interval as inputs.
- Numerical solver: Equilibrium is defined by the patient reaching the same state after repeated dosing cycles with interval δt.Each cycle instantaneously adds drug mass m to the gut before the state is evolved.
- Differentiation: The demonstration solves the algebraic equation numerically and computes sensitivities for k1 and k2.For n patients, the construction yields 2n states and 2n + 2 parameters requiring sensitivities.