[论文解读] Demystifying Differentiable Programming: Shift/Reset the Penultimate Backpropagator

Fei Wang, Zheng, Daniel+4 more|arXiv (Cornell University)|Mar 27, 2018

Advanced Neural Network Applications参考文献 59被引用 13

一句话总结

本文提出了一种使用分隔连续性（shift/reset）的新型、简化的反向模式自动微分（AD）公式，通过操作符重载实现高效、表达性强的可微编程，无需辅助数据结构。通过将该方法与多阶段编程结合，作者构建了Lantern——一个高性能深度学习框架，统一了TensorFlow计算图的效率与PyTorch命令式风格的灵活性。

ABSTRACT

Deep learning has seen tremendous success over the past decade in computer vision, machine translation, and gameplay. This success rests in crucial ways on gradient-descent optimization and the ability to learn parameters of a neural network by backpropagating observed errors. However, neural network architectures are growing increasingly sophisticated and diverse, which motivates an emerging quest for even more general forms of differentiable programming, where arbitrary parameterized computations can be trained by gradient descent. In this paper, we take a fresh look at automatic differentiation (AD) techniques, and especially aim to demystify the reverse-mode form of AD that generalizes backpropagation in neural networks. We uncover a tight connection between reverse-mode AD and delimited continuations, which permits implementing reverse-mode AD purely via operator overloading and without any auxiliary data structures. We further show how this formulation of AD can be fruitfully combined with multi-stage programming (staging), leading to a highly efficient implementation that combines the performance benefits of deep learning frameworks based on explicit reified computation graphs (e.g., TensorFlow) with the expressiveness of pure library approaches (e.g., PyTorch).

研究动机与目标

• 通过揭示其与分隔连续性的关联，向编程语言研究者阐明反向模式自动微分（AD）的本质。
• 通过使用shift/reset控制操作符，消除反向模式AD中显式tape或trace数据结构的需求。
• 通过分阶段编程，统一define-then-run框架（如TensorFlow）的性能与define-by-run框架（如PyTorch）的表达性。
• 通过纯粹的局部、基于操作符重载的变换，在任意程序中（包括递归和高阶函数）实现梯度计算。
• 构建一个高性能的深度学习领域特定语言（DSL）Lantern，该语言在运行时重用计算图，同时支持无限制的控制流。

提出的方法

• 将反向模式AD形式化为柯里化传递风格（CPS）变换程序的符号微分，揭示其与分隔连续性的深层联系。
• 使用shift/reset控制操作符实现反向模式AD，实现无需管理辅助tape的纯粹局部、基于操作符重载的微分。
• 将基于连续性的AD公式与多阶段编程（分阶段）结合，实现在运行时生成高效、可重用的计算图。
• 在Scala中使用轻量级模块化分阶段（LMS）框架，无缝组合分阶段与基于连续性的AD，用于DSL构建。
• 将生成的DSL编译为C++/CUDA以实现高性能执行，达到与现有框架相当的推理与训练速度。
• 通过嵌套的shift/reset实现多级CPS变换，正确处理复杂控制流与嵌套抽象。

实验结果

研究问题

• RQ1 是否能仅通过操作符重载与分隔连续性实现反向模式自动微分，而无需显式管理tape？
• RQ2 如何在一个系统中结合define-by-run框架（如PyTorch）的表达性与define-then-run框架（如TensorFlow）的性能？
• RQ3 连续性在支持任意一等计算（包括递归与高阶函数）的梯度计算中起到什么作用？
• RQ4 能否构建一个支持无限制控制流的深度学习框架，同时仍能实现高效的图级优化与编译？
• RQ5 是否存在一种更简单、更基础的反向模式AD公式，避免依赖范畴论或复杂数据结构？

主要发现

• 反向模式AD可通过shift/reset实现纯粹基于操作符重载，无需显式tape或trace数据结构。
• 反向模式AD与分隔连续性的关联，为可微编程提供了清晰、局部且可组合的基础。
• 将基于连续性的AD与多阶段编程结合，可构建同时支持动态控制流与高效图级编译的框架。
• 所构建的框架Lantern在SqueezeNet、ResNet、DeepSpeech2与TreeLSTM等前沿模型上达到具有竞争力的性能，速度与效率与现有框架相当或更优。
• 该方法原生支持一般递归与高阶函数，克服了以往基于显式数据结构的AD系统的关键局限。
• 实现表明，深度学习框架可被理解为两种正交抽象的组合：分阶段与分隔连续性。