Source-linked AI summary
Forward-Mode Automatic Differentiation in Julia
Jarrett Revels, Miles Lubin, Theodore Papamarkou
TL;DR
The paper addresses efficient automatic differentiation in a high-level language by presenting ForwardDiff, a Julia package built around JIT compilation and chunk-mode forward differentiation. ForwardDiff supports higher-order and custom-number differentiation, outperforms Python autograd on some large-dimensional gradient problems, and is used within JuMP, although second-order derivatives of user-defined functions are not yet supported there.
Problem
The paper targets the challenge of providing efficient automatic differentiation in a high-level language while supporting higher-order differentiation and custom number types.
Method
ForwardDiff uses Julia JIT recompilation, nested multidimensional dual numbers, and chunk mode with stack-allocated partial vectors for forward-mode differentiation.
Results
ForwardDiff is performance competitive with C++, can outperform Python autograd for gradients at nontrivially large dimensions, and yields 30% speedups for JuMP Hessian-vector benchmarks using chunk mode.
Takeaways & Limitations
ForwardDiff provides a practical AD implementation used in JuMP and adopted by 41 unique GitHub repositories across diverse application areas.
Takeaways & Limitations
JuMP does not yet support second-order derivatives of user-defined functions, and the extended abstract cannot adequately cover some implementation details.
Abstract
from arXiv · showhide
We present ForwardDiff, a Julia package for forward-mode automatic differentiation (AD) featuring performance competitive with low-level languages like C++. Unlike recently developed AD tools in other popular high-level languages such as Python and MATLAB, ForwardDiff takes advantage of just-in-time (JIT) compilation to transparently recompile AD-unaware user code, enabling efficient support for higher-order differentiation and differentiation using custom number types (including complex numbers). For gradient and Jacobian calculations, ForwardDiff provides a variant of vector-forward mode that avoids expensive heap allocation and makes better use of memory bandwidth than traditional vector mode. In our numerical experiments, we demonstrate that for nontrivially large dimensions, ForwardDiff's gradient computations can be faster than a reverse-mode implementation from the Python-based autograd package. We also illustrate how ForwardDiff is used effectively within JuMP, a modeling language for optimization. According to our usage statistics, 41 unique repositories on GitHub depend on ForwardDiff, with users from diverse fields such as astronomy, optimization, finite element analysis, and statistics. This document is an extended abstract that has been accepted for presentation at the AD2016 7th International Conference on Algorithmic Differentiation.
1 Introduction
ForwardDiff is a Julia forward-mode AD package designed to combine performance competitive with C++ and efficient support for advanced differentiation. Its experiments and adoption examples show strong performance and practical use across optimization and other fields.
- ForwardDiff is a Julia package for forward-mode automatic differentiation with performance competitive with low-level languages like C++.
- JIT compilation transparently recompiles AD-unaware user code, supporting higher-order differentiation and custom number types including complex numbers.
- A vector-forward variant avoids expensive heap allocation and improves memory-bandwidth use for gradient and Jacobian calculations.
- For nontrivially large dimensions, ForwardDiff computes gradients faster than reverse-mode Python autograd in numerical experiments.
- 41 unique GitHub repositories depend on ForwardDiff, with users spanning astronomy, optimization, finite element analysis, and statistics.
2 Methodology
ForwardDiff uses multidimensional dual numbers and nested Dual types to propagate directional derivatives, gradients, and higher-order derivatives through Julia code. Its chunk mode reduces heap-allocation costs while retaining tunable gradient computation, and nesting supports custom and complex number types.
- Vector forward mode: Orthogonal epsilon components appended to input vectors track individual directional derivatives for vector forward mode.This enables gradient computation by propagating several directional derivatives together.
- Chunk mode: Chunk mode stack-allocates epsilon vectors and tunes their length, balancing memory use, cache behavior, and function-evaluation cost.For chunk size N and input length k, the gradient requires ⌈k/N⌉ passes through f.
- Dual-number representation: ForwardDiff represents multidimensional dual numbers as Dual{N,T}, storing a value and an N-component stack-allocated partials vector.N sets the epsilon-vector length, while T specifies its element type.
- Type support: Julia method definitions support derivatives of trigonometric functions over nested, complex, and custom Dual-based number types.The listed forms include Complex{Dual{N,T}}, Custom{Dual{N,T}}, and Dual{N,Custom}.
- Scope: The abstract omits a thorough treatment of complex and custom-type differentiation because explaining Julia’s multiple dispatch and JIT compilation requires more exposition.The authors identify this as a limitation of the document’s scope.
- Higher-order differentiation: Nested Dual types enable higher-order vector-mode differentiation, with Dual{M,Dual{N,T}} computing M x N second-order derivatives in one target-function pass.A d-level nesting implements a dth-order hyper-dual number and scales to arbitrary dimensions.
3 Performance Analysis
The performance analysis evaluates ForwardDiff gradient computation across chunk sizes and against reverse-mode autograd. Results show that chunk-size effects reflect hardware and function-evaluation trade-offs, while ForwardDiff can outperform autograd despite less favorable asymptotic scaling.
- Chunk-size effects: Increasing chunk size reduces univariate-function evaluations, with relative gains depending on whether the target contains many or few expensive univariate functions.Ackley and Rosenbrock provide contrasting cases for this effect.
- Chunk-size effects: The optimal chunk size trades off memory bandwidth, memory alignment, cache performance, and function evaluation cost.ForwardDiff Rosenbrock performance worsens from N = 4 to N = 5, while C++ performance for both functions worsens from N = 1 to N = 2.
- ForwardDiff versus autograd: Table 2 compares autograd reverse mode with ForwardDiff forward mode for varying input sizes at chunk size N = 10.The table also reports an experimental multithreaded ForwardDiff implementation.
- ForwardDiff versus autograd: ∼2x speed-up is obtained with 4 threads relative to the single-threaded ForwardDiff implementation.This result comes from the experimental multithreaded implementation.
- ForwardDiff versus autograd: Although forward mode scales quadratically and reverse mode linearly for these functions, autograd is slower when k ≤10000.The comparison concerns functions with linear complexity in input dimension k.
4 ForwardDiff within JuMP
ForwardDiff is integrated into JuMP to differentiate user-defined functions inside optimization models and to accelerate Hessian-related computations. This extends automatic differentiation beyond JuMP’s closed-form algebraic expressions, although second-order derivatives of user-defined functions remain unsupported.
- Hessian computations: JuMP’s forward-over-reverse Hessian-vector products use chunk mode to compute Hessian-matrix products instead of independent Hessian-vector products.This use of chunk mode yielded speedups of 30% on benchmarks presented in [5].
- Integration with JuMP: ForwardDiff enables JuMP to automatically differentiate user-defined functions embedded within closed-form expressions.The paper describes this as the first such capability in an AML to the authors’ knowledge.
- Integration with JuMP: The demonstrated JuMP model registers a Newton-method squareroot function for automatic differentiation inside a nonlinear optimization model.The example constrains the Euclidean norm of two variables through the user-defined function.
- Integration with JuMP: JuMP integrates ForwardDiff gradients into its reverse-mode computations for user-defined functions such as squareroot.This supports user-defined functions embedded in otherwise closed-form optimization expressions.
- Current limitation: JuMP does not yet support second-order derivatives of user-defined functions.The authors describe this functionality as immature and identify reverse-mode taping as a possible improvement.
5 Future Work
Future work targets ForwardDiff’s performance, perturbation handling, and matrix-operation support. Planned directions include SIMD vectorization, compile-time interception of unwanted perturbations, and direct overloading of linear-algebra functions.
- Performance and usability: The authors are investigating SIMD vectorization of derivative propagation to improve ForwardDiff’s performance.Implementation is described as being in preliminary phases.
- Performance and usability: ForwardDiff aims to address perturbation confusion by intercepting unwanted perturbations at compile time using Julia metaprogramming.This is presented as a planned research direction.
- Performance and usability: The authors plan to improve matrix-operation support, including eigenvalue computations, by directly overloading linear-algebra functions.The passage frames this as a usability improvement.