[论文解读] Efficient Implementation of a Higher-Order Language with Built-In AD

Jeffrey Mark Siskind, Barak A. Pearlmutter|arXiv (Cornell University)|Nov 10, 2016

Numerical Methods and Algorithms参考文献 6被引用 7

一句话总结

本文通过运行时反射嵌入自动微分（AD）算子，并利用多态流分析将这些算子迁移至编译时，提出了一套高性能、一流的自动微分（AD）系统，该系统运行于动态、函数式语言中。结果在数值性能上可与手工优化的Fortran代码媲美，甚至在嵌套导数计算中超越基于预处理器的AD工具。

ABSTRACT

We show that Automatic Differentiation (AD) operators can be provided in a dynamic language without sacrificing numeric performance. To achieve this, general forward and reverse AD functions are added to a simple high-level dynamic language, and support for them is included in an aggressive optimizing compiler. Novel technical mechanisms are discussed, which have the ability to migrate the AD transformations from run-time to compile-time. The resulting system, although only a research prototype, exhibits startlingly good performance. In fact, despite the potential inefficiencies entailed by support of a functional-programming language and a first-class AD operator, performance is competitive with the fastest available preprocessor-based Fortran AD systems. On benchmarks involving nested use of the AD operators, it can even dramatically exceed their performance.

研究动机与目标

• 弥合高级动态语言与低级数值语言（如Fortran）在自动微分（AD）方面的性能差距。
• 在不牺牲数值效率的前提下，实现在函数式编程语言中的一流、嵌套和高阶AD。
• 通过使用激进的静态分析，将AD转换从运行时迁移至编译时，彻底消除运行时开销。
• 证明具有内置AD的函数式语言可以实现与传统基于预处理器的AD系统相当或更优的性能。
• 为现代编译器和高级数值编程提供一个实用且高效的AD基础。

提出的方法

• 引入一组新颖的高阶原语（例如 j*、primal、tangent、bundle），通过源到源转换表达前向和反向AD。
• 利用运行时反射，在（derivative f）求值期间动态将函数转换为其AD等价形式（例如 f → f_forward）。
• 使用全程序、过程间多态流分析，消除运行时反射及所有非数值支撑代码。
• 将AD转换从运行时迁移至编译时，生成高度优化的数值代码，其与手工优化的Fortran代码无异。
• 利用语言结构作为具有数值和AD基础的急切lambda演算，其与高性能编译器中的中间表示同构。
• 采用激进的优化编译，移除所有运行时开销，仅保留浮点数运算。

实验结果

研究问题

• RQ1 动态、函数式语言是否能原生支持一流、嵌套和高阶AD，且性能可与手工优化的Fortran媲美或超越？
• RQ2 是否能通过静态分析系统性地消除AD的运行时反射，从而实现Fortran级别的性能？
• RQ3 是否可行实现一个高级AD系统，兼具函数式编程的表达性与低级数值代码的效率？
• RQ4 多态流分析是否能有效识别并移除AD转换代码中的非数值运行时支撑代码？
• RQ5 将AD转换从运行时迁移至编译时是否能带来可测量的性能提升，特别是在嵌套或高阶导数场景中？

主要发现

• 该系统在数值性能上与最快的基于预处理器的Fortran AD工具（如Tapenade和ADOL-C）相当。
• 在涉及嵌套AD应用的基准测试中，该系统显著优于传统的基于预处理器的AD系统，包括C++和Fortran中的实现。
• 在使用前向向量AD的反向传播示例（backprop-Fv）中，Chicken的基线运行时间慢了1626.73倍，而该系统的优化版本实现了归一化运行时间1626.73，表明优化有效且可测量。
• 在粒子和鞍点示例中，该系统的性能分别比某些实现的基线快14.97倍和8.09倍，显示出强大的可扩展性。
• 多态流分析成功消除了所有运行时反射和非数值支撑代码，生成的数值代码具有与Fortran相当的效率。
• 原型证明，高级语言中的一流、嵌套AD既具表达性又高效，挑战了高级语言特性必然导致性能下降的假设。