Backpropagation with Continuation Callbacks: Foundations for Efficient and Expressive Differentiable Programming
Training of deep learning models depends on gradient descent and end-to-end differentiation. Under the slogan of differentiable programming, there is an increasing demand for efficient automatic gradient computation for emerging network architectures that incorporate dynamic control flow, especially in NLP. In this paper we propose an implementation of backpropagation using functions with callbacks, where the forward pass is executed as a sequence of function calls, and the backward pass as a corresponding sequence of function returns. A key realization is that this technique of chaining callbacks is well known in the programming languages community as continuation-passing style (CPS). Any program can be converted to this form using standard techniques, and hence, any program can be mechanically converted to compute gradients. Our approach achieves the same flexibility as other reverse-mode automatic differentiation (AD) techniques, but it can be implemented without any auxiliary data structures besides the function call stack, and it can easily be combined with graph construction and native code generation techniques through forms of multi-stage programming, leading to a highly efficient implementation that combines the performance benefits of define-then-run software frameworks such as TensorFlow with the expressiveness of define-by-run frameworks such as PyTorch.
=== Summary  The fields of high-performance computing and systems have played a large role in the recent developments in deep learning. More recently, important research from the compilers and programming languages perspective has been contributing to deep learning, addressing the lack of usability and flexibility that many frameworks have traded for efficiency. This work falls squarely in that category.  The authors propose a new implementation approach to reverse mode automatic differentiation. The technique relies on operator overloading, but does not rely on a tape (and the accompanying interpreter that walks the tape backwards). Instead, any primal computation takes a callback to the rest of the forward pass which it calls before performing the adjoint computation, making the "tape" implicit in the function call stack.  The authors provide a basic implementation of this idea in Scala, along with a cleaner implementation which relies on delimited continuations (using  the shift/reset oeprators) in order to construct the necessary callbacks automatically.  Lastly, lightweight modular staging (LMS) is used to reduce the overhead of these abstractions (dual numbers, callbacks, etc.) and generate efficient C++ code. Experiments highlight that this approach is competitive with existing frameworks such as PyTorch and TensorFlow.  === Review  In the abstract the authors highlight that the use of continuation passing style (CPS) is a key ingredient to their approach. They claim that this can be implemented without any "auxiliary data structures". This claim is repeated in lines 36-37. However, frameworks that rely on having a single global scope (e.g. TensorFlow) or that use closures (e.g. Myia) don't need these data structures either.  The same confusion appears in section 2: The text suggests that TensorFlow uses graph as in figure 1, where nodes hold values for both the forward and backward pass. This is not the case. TensorFlow augments the forward graph with additional primitives i.e. it creates a larger graph in which there is no distinction between the forward and backward pass anymore. Hence, the graph itself is never used as an auxiliary data structure that stores intermediate variables and there is no "graph iterator which performs a forward pass first, then a backward pass" (line 75).  The remainder of the section does a good job at conveying the main ideas of the paper. In the last two paragraphs the authors compare their approach to Siskind's approach. Although it is true that that transformation is more complex, it must be noted that it operates on lambda calculus, whereas the proposed implementation here relies on (1) compiler support for control delimiters and (2) side effects (i.e. instead of requiring a backpropagator to return partial derivatives to a particular scope it can call `that.d += y.d`). The different restrictions under which the transformations apply makes it misleading to compare them without qualifying this more clearly in the text.  In section 3 lightweight modular staging (LMS) is introduced as an approach to achieving some of the performance gains that libraries such as TensorFlow can achieve by optimizing and compiling the computation graph. I think the text would possible benefit from a small code sample that involves a loop or conditional here, but overall the text is clear.  In the last two sections the authors argue that their approach is applicable to other languages. "Fundamentally" almost anything can be done in any language, but I would have appreciated more insight as to what requirements their approach puts on a language. For example, I think it is far less feasible to implement delimited continuations and LMS in Python, given its dynamically typed, interpreted nature.  The authors argue that their approach benefits more from general compiler transformations, but I believe the same is true for e.g. TensorFlow when using the XLA backend, since it lowers the entire graph to LLVM in that case.  In section 4 a few simple models are implemented using the proposed approach and benchmarked against PyTorch, NumPy and TensorFlow. I do not think it is necessary to show training curves; that they decrease actually says very little (neural networks and SGD are very robust, and the loss will decrease just as well even if there are mistakes in the gradients). The correctness of gradients can just be checked using finite differences by the authors and be taken for granted by the reader.  The runtime benchmarks look promising, but as the authors hint at in the text: The results mainly highlight that if you consider very small models, batch sizes, and datasets, your performance will no longer depend on e.g. your compute kernels, high-level linear algebra optimizations, or your ability to exploit large amounts of parallelism (the things that TensorFlow and PyTorch are optimized for). Instead, your performance will be entirely determined by the overhead of your language and AD implementation, in which case a statically typed language like Scala which will use LMS to compile custom compute kernels will always beat a dynamically typed, interpreted language like Python which will be calling e.g. MKL kernels--the authors failed to mention which BLAS library they linked to--which are optimized for larger workloads on more cores. That said, there is no perfect way to do these benchmarks and for people who care about small workloads in low-resource environments, these results are encouraging. That said, statements such as "Lantern performs comptetitively on cutting edge deep learning applications which push the boundaries of existing frameworks [in] performance" could be qualified, considering that neither the models nor hardware in this benchmark are representative of the common deep learning workloads that existing frameworks are optimized for.  In the related work section I think it would be worthwhile to discuss Siskind and Pearlmutter's more recent work on checkpointing, considering that uses both CPS and control operators in the context of reverse mode AD.  Finally, some higher-level questions that this work raised for me:  The listing in figure 4 suggests that the partial derivative is explicitly initialized to zero during the forward pass. For multidimensional arrays this can lead to the allocation of large amounts of zeros, significantly increasing the peak memory usages. Most existing libraries try to be more clever about this (e.g. Tangent uses a special zero object). Have the authors given this any thought?  Higer-order differentiation. An important goal in AD is usually that the gradient operator is closed under its own operation. Also in ML, this is the case e.g. TensorFlow, PyTorch, and more research-oriented libraries such as Tangent, Myia, and Autograd, all have support for higher-order derivatives, and often have support for mixing forward and reverse mode. It is a shame that this is not dicsused at all in this paper. Can a single shift operator be matched to multiple reset operators?   === Minor comments  223: Parentheses around citations. 267: lead -> led  === Conclusion  There are a few misunderstandings in the paper about the internals of TensorFlow that should be addressed, but they don't fundamentally change the authors claim that their approach can naturally support control flow while still supporting multi-stage programming. Other things I would like to see addressed are higher-order gradients, and the supposed applicability of this approach to e.g. Python (which I find hard to believe).  Overall I believe the paper is well-written though, the research question relevant, the approach sufficiently novel, and the experimental results, although preliminary, show promise. Hence, I would recommend this work for publication with minor edits.