[Paper Review] Demystifying Differentiable Programming: Shift/Reset the Penultimate Backpropagator

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

Advanced Neural Network Applications59 references13 citations

TL;DR

This paper proposes a novel, simplified formulation of reverse-mode automatic differentiation (AD) using delimited continuations (shift/reset), enabling efficient, expressive differentiable programming via operator overloading without auxiliary data structures. By combining this with multi-stage programming, the authors build Lantern—a high-performance deep learning framework that unifies the efficiency of TensorFlow’s computation graphs with the flexibility of PyTorch’s imperative style.

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).

Motivation & Objective

• To demystify reverse-mode automatic differentiation (AD) for programming language researchers by revealing its connection to delimited continuations.
• To eliminate the need for explicit tape or trace data structures in reverse-mode AD by using shift/reset control operators.
• To unify the performance of define-then-run frameworks (like TensorFlow) with the expressivity of define-by-run frameworks (like PyTorch) via staging and continuations.
• To enable gradient computation in arbitrary programs, including those with recursion and higher-order functions, using a purely local, operator-overloading-based transformation.
• To build a high-performance deep learning DSL, Lantern, that reifies computation graphs at runtime while supporting unrestricted control flow.

Proposed method

• Formalize reverse-mode AD as symbolic differentiation of CPS-transformed programs, revealing a deep connection to delimited continuations.
• Implement reverse-mode AD using shift/reset control operators, enabling purely local, operator-overloading-based differentiation without managing auxiliary tapes.
• Integrate the continuation-based AD formulation with multi-stage programming (staging) to generate efficient, reified computation graphs at runtime.
• Use the Lightweight Modular Staging (LMS) framework in Scala to seamlessly compose staging and continuation-based AD for DSL construction.
• Compile the resulting DSL to C++/CUDA for high-performance execution, achieving competitive inference and training speeds.
• Encode multi-level CPS transformations via nested shift/reset, enabling correct handling of complex control flow and nested abstractions.

Experimental results

Research questions

• RQ1 Can reverse-mode automatic differentiation be expressed purely through operator overloading using delimited continuations, without explicit tape management?
• RQ2 How can the expressivity of define-by-run frameworks (e.g., PyTorch) be combined with the performance of define-then-run frameworks (e.g., TensorFlow) in a single system?
• RQ3 What is the role of continuations in enabling gradient computation for arbitrary, first-class computations, including recursive and higher-order functions?
• RQ4 Can a deep learning framework be built that supports unrestricted control flow while still enabling efficient graph-level optimizations and compilation?
• RQ5 Is there a simpler, more foundational formulation of reverse-mode AD that avoids reliance on category theory or complex data structures?

Key findings

• Reverse-mode AD can be implemented purely via operator overloading using shift/reset, eliminating the need for explicit tape or trace data structures.
• The connection between reverse-mode AD and delimited continuations provides a clean, local, and composable foundation for differentiable programming.
• Combining continuation-based AD with multi-stage programming enables a framework that supports both dynamic control flow and efficient graph-level compilation.
• The resulting framework, Lantern, achieves competitive performance on state-of-the-art models like SqueezeNet, ResNet, DeepSpeech2, and TreeLSTM, matching or exceeding existing frameworks in speed and efficiency.
• The approach supports general recursion and higher-order functions natively, overcoming a key limitation of prior AD systems based on explicit data structures.
• The implementation demonstrates that deep learning frameworks can be understood as the composition of two orthogonal abstractions: staging and delimited continuations.