[Paper Review] A Differentiable Programming System to Bridge Machine Learning and Scientific Computing

Mike Innes, Alan Edelman+5 more|arXiv (Cornell University)|Jul 17, 2019

Neural Networks and Applications129 citations

TL;DR

The paper presents ∂P, a differentiable programming system built in Julia that enables differentiating through general Julia programs, integrates with ML and scientific computing packages, and demonstrates diverse applications from deep learning to differentiable physics and finance.

ABSTRACT

Scientific computing is increasingly incorporating the advancements in machine learning and the ability to work with large amounts of data. At the same time, machine learning models are becoming increasingly sophisticated and exhibit many features often seen in scientific computing, stressing the capabilities of machine learning frameworks. Just as the disciplines of scientific computing and machine learning have shared common underlying infrastructure in the form of numerical linear algebra, we now have the opportunity to further share new computational infrastructure, and thus ideas, in the form of Differentiable Programming. We describe Zygote, a Differentiable Programming system that is able to take gradients of general program structures. We implement this system in the Julia programming language. Our system supports almost all language constructs (control flow, recursion, mutation, etc.) and compiles high-performance code without requiring any user intervention or refactoring to stage computations. This enables an expressive programming model for deep learning, but more importantly, it enables us to incorporate a large ecosystem of libraries in our models in a straightforward way. We discuss our approach to automatic differentiation, including its support for advanced techniques such as mixed-mode, complex and checkpointed differentiation, and present several examples of differentiating programs.

Motivation & Objective

• Motivate a unified differentiable programming approach that spans ML and scientific computing.
• Provide a Julia-based system that can differentiate through general programs with minimal user intervention.
• Demonstrate practical applications mixing ML with existing scientific packages (e.g., ray tracing, quantum simulations, finance).
• Show performance characteristics and extensibility with custom data types and complex control flow.

Proposed method

• Use a source-to-source automatic differentiation approach (Zygote) to generate gradients from Julia code.
• Define a general differential operator J that wraps function outputs with Jacobian-vector products for backpropagation.
• Provide a user-extensible partial derivative interface that allows custom gradients via multiple dispatch.
• Demonstrate differentiating through mixed-language and complex data types without requiring primitive changes to packages.

Experimental results

Research questions

• RQ1 Can a differentiable programming system differentiate through general Julia programs with control flow, mutation, and complex data types?
• RQ2 How well does a differentiable programming approach integrate with existing ML and scientific computing ecosystems in Julia?
• RQ3 What are the performance implications of differentiating large, real-world Julia programs on CPUs, GPUs, and TPUs?
• RQ4 Can differentiable programming enable applications like differentiable ray tracing, quantum machine learning, and neural SDEs while preserving ecosystem composability?

Key findings

• Zygote enables low-overhead automatic differentiation via source-to-source transformation with information reuse across forward and backward passes.
• The per-adjoint operation overhead is around 568.8 ns on average for stacked LSTMs, competitive with PyTorch’s overhead.
• Differentiating through diverse domains (ray tracing, quantum simulation, neural SDEs, finance) is feasible and productive using ∂P in Julia.
• Custom gradients can be easily defined for new types via the partial derivative interface without modifying the underlying AD system.
• The approach supports training with CPUs, GPUs, and TPUs, matching performance of existing ML frameworks in deep learning tasks.