[论文解读] Automatic differentiation in ML: Where we are and where we should be going

Bart van Merriënboer, Olivier Breuleux+2 more|arXiv (Cornell University)|Oct 26, 2018

Evolutionary Algorithms and Applications被引用 39

一句话总结

本文提出了一种用于机器学习中自动微分（AD）的新型基于图的中间表示（IR），可完全支持高阶函数、递归和闭包——这些特性在现有框架中通常缺失。通过源码转换，Myia 将 Python 代码编译为优化后的无磁带 AD 流水线，支持提前编译优化和高阶导数，性能可与 TensorFlow 比肩，同时保留了 PyTorch 的灵活性。

ABSTRACT

We review the current state of automatic differentiation (AD) for array programming in machine learning (ML), including the different approaches such as operator overloading (OO) and source transformation (ST) used for AD, graph-based intermediate representations for programs, and source languages. Based on these insights, we introduce a new graph-based intermediate representation (IR) which specifically aims to efficiently support fully-general AD for array programming. Unlike existing dataflow programming representations in ML frameworks, our IR naturally supports function calls, higher-order functions and recursion, making ML models easier to implement. The ability to represent closures allows us to perform AD using ST without a tape, making the resulting derivative (adjoint) program amenable to ahead-of-time optimization using tools from functional language compilers, and enabling higher-order derivatives. Lastly, we introduce a proof of concept compiler toolchain called Myia which uses a subset of Python as a front end.

研究动机与目标

• 解决现有机器学习框架在支持数组编程通用自动微分方面的局限性，特别是针对递归和高阶函数的支持。
• 设计一种原生支持闭包和函数组合的基于图的中间表示（IR），以支持更富表达力和可组合的机器学习模型。
• 实现无磁带、基于源码转换的自动微分，支持高阶导数，并可通过函数式语言编译器技术进行优化。
• 构建一个概念验证编译工具链（Myia），将 Python 子集编译为高效、优化的导数代码，支持完整的类型和形状推断。
• 通过将高级优化和类型推断技术整合到 AD 流水线中，弥合机器学习框架与函数式语言编译器研究之间的差距。

提出的方法

• 设计一种基于图的中间表示（IR），将程序表示为一等公民值，支持函数调用、闭包和递归。
• 使用源码转换（ST）从 IR 生成伴随程序，通过利用闭包表示避免运行时磁带的需要。
• 使用协程实现类型和形状推断系统，无需类型注解即可推断类型、值和形状。
• 使用函数式编译器工具链中的工具，对导数程序应用标准函数式语言优化（如内联、常量折叠）。
• 通过虚拟机和 GPU 后端将优化后的导数代码编译为高效的机器代码，针对不同输入类型进行专门化处理。
• 将 Python 函数解析为抽象语法树（AST），转换为 IR，通过源码转换应用 AD，并优化生成的导数图。

实验结果

研究问题

• RQ1 如何在支持高阶函数和递归的数组编程框架中，高效且通用地支持自动微分？
• RQ2 能否设计一种基于图的中间表示，原生支持闭包和函数组合，从而在无运行时磁带开销的情况下实现完整的自动微分？
• RQ3 函数式语言编译器优化技术（如 A-正则形式、持久数据结构）在机器学习框架的导数程序中可应用到何种程度？
• RQ4 基于源码转换的自动微分系统能否实现与 TensorFlow 等静态编译框架相当的性能，同时保留 PyTorch 等动态框架的灵活性？
• RQ5 在无需显式类型注解的情况下，如何在 Python 等动态类型语言中高效地执行类型和形状推断？

主要发现

• 所提出的基于图的 IR 原生支持高阶函数、递归和闭包，使得机器学习模型的实现比现有框架更具表达力和可组合性。
• 通过结合闭包表示的源码转换，系统避免了运行时磁带的需要，实现了高效、优化的导数代码生成。
• Myia 生成的导数程序可通过函数式语言编译器技术进行提前编译优化，提升性能并减少开销。
• Myia 中的类型和形状推断无需类型注解即可运行，使用户能够编写动态 Python 代码，同时仍可高效编译和类型检查。
• Myia 原型表明，可在支持递归和高阶函数等特性（这些特性在静态计算图框架中缺失）的同时，实现与 TensorFlow 相当的性能。
• 该系统通过清晰、可组合的表示支持高阶导数，克服了现有框架在导数计算上仅限于一阶或不可组合形式的局限。