Source-linked AI summary
A Differentiable Programming System to Bridge Machine Learning and Scientific Computing
Mike Innes, Alan Edelman, Keno Fischer, Chris Rackauckas, Elliot Saba, Viral B Shah, Will Tebbutt
TL;DR
Machine learning and scientific computing increasingly require shared computational infrastructure, but their interaction is limited by the difficulty of differentiating complex scientific programs and accessing broad software ecosystems. The paper presents Zygote, a Julia Differentiable Programming system that differentiates general programs and existing packages. It demonstrates shared applicability across machine learning and scientific computing, including comparable performance with existing machine-learning frameworks and examples in neural SDEs and quantum machine learning.
Problem
Machine learning and scientific computing are increasingly interconnected, creating a need for shared computational infrastructure that supports both domains.
Method
Zygote extends Julia with source-to-source automatic differentiation for existing packages and general program structures, including state-based control flow and user-defined types.
Results
Zygote matches existing machine-learning frameworks for deep learning on CPUs, GPUs, and TPUs and in reinforcement learning, while also supporting neural SDEs and quantum machine learning.
Takeaways & Limitations
Differentiable Programming can serve as common infrastructure for composing machine-learning methods with scientific-computing programs.
Takeaways & Limitations
Zygote can currently generate code that pessimizes compiler assumptions and causes slower execution in some cases.
Abstract
from arXiv · showhide
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.
1 Introduction
Machine learning and scientific computing share deep infrastructure and increasingly benefit from interaction through Differentiable Programming. The paper presents a Julia-based system that differentiates general programs and applies it across machine-learning and scientific-computing examples.
- 1 Introduction: Neural ODEs reduce residual-network memory and computational costs while accommodating irregularly sampled time series.The formulation replaces multiple ResNet layers with one ODE solved in fewer adaptive-solver steps and is not restricted to a grid.
- 1 Introduction: Scientific computing can benefit from machine learning through surrogate modeling, adjoint sensitivity analysis, inverse problems, and probabilistic programming.These applications respectively approximate expensive simulations, accelerate adjoint calculations, differentiate through simulators, and support richer composable statistical models.
- 1 Introduction: Differentiable Programming can provide shared infrastructure connecting machine learning and scientific computing.The paper frames the disciplines as historically linked through numerical linear algebra and as capable of benefiting from stronger interaction.
- 1 Introduction: Zygote extends Julia with source-to-source automatic differentiation that works directly on existing packages, user-defined types, state-based control flow, and scalar operations.The system is designed to differentiate programs without requiring models to use special data types or extensive refactoring.
- 1.1 A simple sin example: Differentiate Programs not Formulas: The introductory Taylor-series example deliberately uses loops, conditionals, printing, native function calls, and a convergence-dependent number of terms.Forward- and reverse-mode differentiation are demonstrated at x = 1.0, with the derivative matching cos(1.0).
2 Implementation
Zygote uses source-to-source transformation and a user-extensible differential operator to differentiate general Julia programs while preserving control flow and supporting composable code. Its design combines recursive chain-rule application with multiple dispatch, enabling derivatives through custom types and existing packages.
- Motivation: Tracing-based AD can unroll control flow and require recompilation for new inputs, limiting its fit for dynamic scientific programs.Scientific workloads also use scalar operations, adaptive algorithms, user-defined data structures, and nonlinear models.
- Source-to-source transformation: Zygote generates one derivative function from source code that handles all input values, preserves control flow, and can be compiled, optimized, and reused.The source-to-source transformation targets Julia’s dynamic high-level language and existing scientific and machine-learning ecosystem.
- Differential operator: The differential operator J returns a first-class function and supports recursive chain-rule implementation, allowing higher-order derivatives through repeated application.Combining forward and backward operations also enables reuse of values computed during the forward pass.
- Extensibility: An extensible ∂ function intercepts automatically generated transformations through Julia’s multiple dispatch, allowing custom behavior at any stack level.The generated implementation recurses to ∂ rather than J, providing a user-extensible fallback mechanism.
- Extensibility: Two primitive definitions for addition and multiplication suffice for composed real-valued programs, while other functions and types fall back to the AD transform.Zygote differentiates operations involving the user-defined Measurement type without requiring the package to know about the AD system.
- Composability: Custom gradients use the same mechanism as system-provided gradients, are compiler-co-optimized, and can be targeted through multiple dispatch.Because Julia’s language and package ecosystem are largely pure Julia, most functions and types are automatically supported.
3 ∂P in Practice
Zygote applies differentiable programming across deep learning and scientific-computing workflows, including control, rendering, finance, quantum simulation, and stochastic differential equations. These examples combine source-to-source automatic differentiation with existing Julia packages and domain-specific computation.
- 3.1 Deep Learning: Zygote differentiates ordinary Julia model code, including recurrent networks, without requiring special model data types or AD-aware layer definitions.The same model can execute on CPU, GPU, and Google TPU architectures with little to no change.
- 3.1 Deep Learning: Zygote’s source-to-source transformation adds little runtime overhead beyond computing the backwards pass, although compiler pessimization can currently slow some generated code.The authors describe this as an ongoing limitation while working toward true zero overhead.
- 3.1 Deep Learning: 568.8 ns per adjoint definition is the reported average overhead for Zygote across stacked LSTM benchmarks.The benchmark varies batch size and the number of stacked LSTM operations to estimate fixed per-operation overhead.
- 3.2 Differentiating a Trebuchet: Differentiating through a trebuchet simulator trains a neural network for inverse control, producing constant-time aiming that is 100× faster than direct optimization.The network maps target distance and wind conditions to simulator parameters, then backpropagates through the ODE-based simulation.
- 3.3 Computer Vision: Differentiable rendering optimizes a point light’s position by comparing rendered output with a reference image and updating the source using gradients.The loss function takes the light source as input, renders the scene, and compares the result with the target.
- 3.4 Financial Derivatives: Automatic differentiation composes through financial bootstrapping and Newton iterations, enabling higher-order derivatives such as key rate durations.The programmer need not explicitly manage the several differentiation operations involved in differentiating through the solver.
4 Conclusion
The paper presents differentiable programming as shared infrastructure for machine learning and scientific computing, demonstrating applications across both domains.
- Zygote differentiates programs across machine learning and scientific computing using the same technology.
- On the machine-learning side, Zygote matches existing frameworks in deep learning on CPUs, GPUs, and TPUs and in reinforcement learning.
- In scientific computing, the paper demonstrates neural stochastic differential equations and quantum machine learning.
- The system is open source and presented as a basis for shared infrastructure across the two disciplines.