[论文解读] Forward-Mode Automatic Differentiation in Julia

Jarrett Revels, Miles Lubin+1 more|ENLIGHTEN (Jurnal Bimbingan dan Konseling Islam)|Jul 26, 2016

Distributed and Parallel Computing Systems被引用 100

一句话总结

ForwardDiff.jl 在 Julia 中提供前向模式自动微分，使用栈分配的分块对数以实现高效的梯度和高阶导数，与 C++ 性能相当，并与 JuMP 有强集成。

ABSTRACT

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.

研究动机与目标

• 在 Julia 中激发并实现具竞争力性能的前向模式自动微分。
• 引入支持高阶导数和自定义数类型的多维对数表示。
• 通过栈分配 epsilon 向量并开启分块计算来减少前向模式的内存分配。
• 展示性能特征以及与 JuMP 在优化中的实际集成。

提出的方法

• 定义一个多维对数类型 Dual{N,T}，包含输入值和一个栈分配的偏导向量。
• 通过在输入中追加正交的 epsilon 分量来实现向量前向模式，以跟踪方向导数。
• 引入分块模式，将输入分区并分块计算梯度，以在内存带宽与计算成本之间取得平衡。
• 重载 Julia Base 中的 Dual 的方法，使其能够通过常见函数（例如 sin、cos、-）进行自动微分。
• 通过嵌套 Dual 类型（例如 Dual{M, Dual{N,T}}）来支持高阶导数，以在单次计算中获得 M x N 的二阶导数。
• 在 JuMP 中使用以计算前向对反向的海森-向量积，并实现对嵌入到闭式表达式中的用户自定义函数的微分。

实验结果

研究问题

• RQ1 在大规模梯度计算中，Julia 的分块模式前向模式 AD 是否在实践中优于反向模式 AD？
• RQ2 嵌套的 Dual 类型如何在单次求值中高效实现高阶导数？
• RQ3 ForwardDiff 与 JuMP 的集成在优化模型中对用户自定义函数的微分方面有多有效？
• RQ4 不同分块大小在内存带宽与函数评估成本方面的性能权衡是什么？
• RQ5 Julia 中的前向模式自动微分是否可以利用多线程和 SIMD 来进一步提升性能？

主要发现

• 在报道的基准测试中，ForwardDiff 的梯度时间与 C++ 实现相当，分块大小会影响性能。
• 在输入规模 k=12000 的 Rosenbrock 和 Ackley 函数中，ForwardDiff 的时间随某些分块大小降低，体现了内存/带宽的权衡。
• 与 Python 的 autograd（反向模式）相比，在所测试的规模（例如表 2 中的到 k=10000）下，ForwardDiff 显示出显著的速度优势。
• 在使用 4 线程时，多线程的 ForwardDiff 相对于单线程实现约提升 2 倍。
• 分块模式使梯度计算块大小可以更大或更小，从而影响缓存对齐和内存性能，在 JuMP 的前向对反向 Hessian-向量计算中起到关键作用。
• ForwardDiff 通过嵌套（例如 Dual{M,Dual{N,T}}）支持高阶导数，在单次计算中获得 M x N 的二阶导数。