Forward-Mode Automatic Differentiation in Julia

Introduction

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.

Methodology

ForwardDiff implements a Julia representation of a multidimensional dual number, whose behavior on scalar functions is defined as:

where $\epsilon_{i}\epsilon_{j}=0$ for all indices $i$ and $j$. Storing additional $\epsilon$ components allows for a vector forward-mode implementation of the sort developed by Kahn and Barton . In our formulation, orthogonal $\epsilon$ components are appended to input vector components to track their individual directional derivatives:

Vector forward mode enables the calculation of entire gradients in a single pass of the program defining $f$, but at the cost of additional memory and operations. Specifically, every dual number must allocate an $\epsilon$ vector of equal size to the input vector, and the number of operations required for derivative propagation scales linearly with the input dimension. In practice, especially in memory-managed languages like Julia, the cost of rapidly allocating and deallocating large $\epsilon$ vectors on the heap can lead to slowdowns that practically outweigh the advantage of fewer passes through $f$.

ForwardDiff’s implementation works around this pitfall by stack-allocating the $\epsilon$ vectors, as well as permitting their size to be tunable at runtime relative to the input dimension and performance characteristics of the target function. We call ForwardDiff’s strategy chunk mode, since it allows us to compute the gradient in bigger or smaller chunks of the input vector. The $\epsilon$ vector length is then the chunk size of the computation. For a chunk size $N$ and an input vector of length $k$, it takes $\lceil\frac{k}{N}\rceil$ passes through $f$ to compute $\nabla f(\vec{x})$. For example, it takes two passes through $f$ to evaluate the gradient at a vector of length $k=4$ and chunk size $N=2$:

ForwardDiff implements a multidimensional dual number as the type Dual{N,T}, where the type parameter N denotes the length of the $\epsilon$ vector and the type parameter T denotes the element type (e.g. Dual{2,Float64} has two Float64 $\epsilon$ components). This type has two fields: value, which stores the $x$ component, and partials, which stores the stack-allocated $\epsilon$ vector. It’s straightforward to overload base Julia methods on the Dual type; here’s an example using sin, cos, and - (univariate negation):

These method definitions are all that is required to support the following features:

$n^{th}$-order derivative of sin or cos (through nesting Dual types)

derivative of complex sin or cos via types of the form Complex{Dual{N,T}}

derivative of sin or cos over custom types, e.g. Custom{Dual{N,T}} or Dual{N,Custom}

We unfortunately do not have room in this abstract to adequately cover the latter two items; a proper discussion would require a more thorough exposition of Julia’s multiple dispatch and JIT-compilation facilities.

Instead, we discuss how instances of the Dual type can be nested to enable the use of vector-mode AD for higher-order derivatives. For example, the type Dual{M,Dual{N,T}} can be used to compute M x N $2^{nd}$-order derivatives. As a simple demonstration of the scalar case, we use an instance of the type Dual{1,Dual{1,Float64}} to take the second derivative of sin at the Julia prompt (The notation $\epsilon$[d,k] is used to denote the $k^{th}$ partial nested at level $d$):

Algebraically, the above example is equivalent to the use of hyper-dual numbers described by Fike and Alonso . In fact, a Dual instance with $d$ levels of nesting implements a $d^{th}$-order hyper-dual number, with the added advantage of scaling to arbitrary dimensions. For example, an instance of Dual{M,Dual{N,Dual{L,T}}} can be used to take M x N x L third-order derivatives in a single pass of the target function.

Performance Analysis

In this section, we present timing results for gradient calculations of the Rosenbrock (4) and Ackley (5) functions. Recalling (1) and the discussion in Section 2, increasing chunk size reduces the number of evaluations of the univariate functions within $f$. We choose Ackley and Rosenbrock as our target functions in order to provide a contrast between the relative gains of increasing chunk sizes when the target function contains many and few expensive univariate functions, respectively.

Table 1 shows evaluation times for calculating gradients of our two target functions using ForwardDiff versus a naive equivalent C++ implementation. Various chunk sizes were tested, while the input size was fixed at $12000$ elements. For the sake of simplicity, ForwardDiff’s Dual{N,T} type was translated into a hardcoded C++ class for each $N$.

Table 1 helps illustrate that the optimal chunk size for a given problem is a result of a trade-off between memory bandwidth, memory alignment, cache performance, and function evaluation cost. For example, note that ForwardDiff’s $\nabla(\text{Rosenbrock})$ performance worsens when going from $N=4$ to $N=5$, and that the C++ implementation’s performance with both functions worsens when going from $N=1$ to $N=2$. The former observation is likely due to the large memory bandwidth cost relative to the cost of the arithmetic operations, while the latter observation is likely due to poor alignment of input vector (since each 2-dimensional dual number is essentially a struct of three double values - one for the instance value, and two for the partial components).

Table 2 compares the gradient computation time of the reverse-mode implementation of the Python-based autograd package versus the forward-mode implementation of ForwardDiff for varying input sizes. We also include results obtained using our experimental multithreaded implementation, which show a $\sim$2x speed-up using 4 threads compared to our single-threaded implementation.

Both functions have linear complexity in the input dimension $k$; therefore reverse mode, which requires $O(1)$ passes through each function, scales linearly, while our forward mode, which requires $O(k)$ passes through each function (with fixed $N$), scales quadratically. The results in Table 2 agree with this complexity analysis. Nevertheless, there is a huge performance gap between these two implementations such that autograd is slower on these examples when $k\leq 10000$, despite reverse mode being a superior algorithm in principle for computing gradients.

The code used to generate the timings in this section can be found at https://github.com/JuliaDiff/ForwardDiff.jl/tree/jr/benchmarks/benchmark. Julia benchmarks were run using Julia version 0.5.0-dev+3200, C++ benchmarks were compiled with clang-600.0.57 using -O2, and Python benchmarks were run using Python version 2.7.9.

ForwardDiff within JuMP

Effective use of ForwardDiff has brought improvements to JuMP , a domain-specific language for optimization embedded in Julia where users provide closed-form algebraic expressions using a specialized syntax. JuMP, and similar commercial tools like AMPL , compute derivatives of user models as input to nonlinear optimization solvers, which is quite different from ForwardDiff’s original target case of differentiating general user-defined code.

JuMP computes sparse Hessians by using the graph coloring approach of , which requires computing a small number of Hessian-vector products in order to recover the full Hessian. JuMP’s forward-over-reverse mode implementation for Hessian-vector products makes use of ForwardDiff’s chunk mode, essentially computing Hessian-matrix products instead of independent Hessian-vector products. This use of chunk mode yielded speedups of 30% on benchmarks presented in (under review). The results in include this speedup but are not accompanied by a discussion of the methodology of chunk mode.

On the user-facing side, ForwardDiff has enabled JuMP to be the first AML to our knowledge which performs automatic differentiation of user-defined functions embedded within closed-form expressions. We reproduce an example from illustrating a user-defined square root function within a JuMP optimization model:

JuMP computes gradients of squareroot with ForwardDiff which are then integrated within the reverse-mode computations of JuMP. We do not yet support $2^{nd}$-order derivatives of user-defined functions. While this functionality is immature and leaves room for improvement (we could attempt to tape the user-defined functions and calculate their gradients in reverse mode), it already creates a new and useful way for JuMP users to seamlessly interact with AD when small parts of their model cannot easily be expressed in closed algebraic form.

Future Work

We are currently investigating several avenues of research that could improve ForwardDiff’s performance and usability. We are in the preliminary phases of implementing SIMD vectorization of derivative propagation. We intend to address perturbation confusion by intercepting unwanted pertubations at compile time using Julia’s metaprogramming capabilities. Finally, we wish to improve our support for matrix operations such as eigenvalue computations by directly overloading linear algebra functions, a technique which has already seen use in .

References