Scientific computing is increasingly incorporating the advancements in machine learning to allow for data-driven physicsinformed modeling approaches. However, re-targeting existing scientific computing workloads to machine learning frameworks is both costly and limiting, as scientific simulations tend to use the full feature set of a general purpose programming language. In this manuscript we develop an infrastructure for incorporating deep learning into existing scientific computing code through Differentiable Programming (@P ). We describe a @P system that is able to take gradients of full Julia programs, making Automatic Differentiation a first class language feature and compatibility with deep learning pervasive. Our system utilizes the one-language nature of Julia package development to augment the existing package ecosystem with deep learning, supporting almost all language constructs (control flow, recursion, mutation, etc.) while generating highperformance code without requiring any user intervention or refactoring to stage computations. We showcase several examples of physics-informed learning which directly utilizes this extension to existing simulation code: neural surrogate models, machine learning on simulated quantum hardware, and data-driven stochastic dynamical model discovery with neural stochastic differential equations.
A casual practitioner might think that scientific
computing and machine learning are different scientific disciplines.
Modern machine learning has made its mark through
breakthroughs in neural networks. Their applicability towards
solving a large class of difficult problems in computer science has
led to the design of new hardware and software to process
extreme amounts of labelled training data, while simultaneously
deploying trained models in devices. Scientific computing,
in contrast, a discipline that is perhaps as old as computing
itself, tends to use a broader set of modelling techniques
arising out of the underlying physical phenomena. Compared to
the typical machine learning researcher, many computational
scientists works with smaller volumes of data but with more
computational complexity and range. However, recent results
like Physics-Informed Neural Networks (PINNs) suggest that
data-efficient machine learning for scientific applications can
be found at the intersection of the methods
Differentiable Programming (@P ) has the potential to be
the lingua franca that can further unite the worlds of
scientific computing and machine learning. Here we present
a @P system which allows for embedding deep neural
networks into arbitrary existing scientific simulations, enabling
these packages to automatically build surrogates for
MLacceleration and learn missing functions from data. Previous
work has shown that differentiable programming systems
in domain specific languages for image processing can
allow for machine learning integration into the domain
applications in a programmer-friendly manner
Our system can be directly used on existing Julia packages, handling user-defined types, general control flow constructs, and plentiful scalar operations through source-tosource mixed-mode automatic differentiation, composing the reverse-mode capabilities of Zygote.jl and Tracker.jl with ForwardDiff.jl for the "best of both worlds" approach. In this paper we briefly describe how we achieve our goals for a @P system and showcase its ability to solve problems which mix machine learning and pre-existing scientific simulation packages.
We start out with a very simple example to differentiate sin(x) written as a program through its Taylor series: sin x = x x3 + x5
: : : : 3! 5!
Note that the number of terms will not be fixed, but will depend on x through a numerical convergence criterion.
To run, install Julia v1.1 or higher, and install the Zygote.jl and ForwardDiff.jl packages with: using Pkg Pkg.add("Zygote") Pkg.add("ForwardDiff") using Zygote, ForwardDiff function s(x) t = 0.0 sign = -1.0 for i in 1:19 if isodd(i) newterm = x^i/factorial(i) abs(newterm)<1e-8 && return t sign = -sign t += sign * newterm end end return t end
While the Taylor series for sine could have been written more compactly in Julia, for purposes of illustrating more complex programs, we purposefully used a loop, a conditional, and function calls to isodd and factorial, which are native Julia implementations. AD just works, and that is the powerful part of the Julia approach. Let’s compute the gradient at x = 1.0 and check whether it matches cos(1.0): julia> ForwardDiff.derivative(s, 1.0) 0.540302303791887 julia> Zygote.gradient(s, 1.0) (0.5403023037918872,) julia> cos(1.0) 0.5403023058681398
Recent progress in tooling for automatic differentiation (AD)
has been driven primarily by the machine learning
community. Many state of the art reverse-mode AD tools such as
Tracker.jl
This choice has been driven largely by the fact that, as
the JAX authors put it, “ML workloads often consist of
large, accelerable, pure-and-statically-composed (PSC)
operations”
However, this assumption does not hold for many scientific inverse problems, or even the cutting edge of ML research. Instead, these problems require a @P system capable of: 1. Low overhead, independent of the size of the executed operation
3. Complete, efficient support for user defined data types
5. Composability with existing code unaware of @P
Particularly, scientific programs tend to have adaptive algorithms, whose control flow depends on error estimates and thus the current state of the simulation, numerous scalar operations, define large nonlinear models using every term individually or implementing specialized numerical linear algebra routines, and pervasive use of user-defined data structures to describe model components, which require efficient memory handling (stack-allocation) in order for the problem to be computationally feasible.
To take these kinds of problems, Zygote does not utilize
the standard methodology and instead generates a derivative
function directly from the original source which is able to
handle all input values. This is called a source-to-source
transformation, a methodology with a long history (Baydin et al.
2018) going back at least to the ADIFOR source-to-source
AD program for FORTRAN 77
One of the primary design decisions of a @P system is how these capabilities should be exposed to the user. One convenient way to do so is using a differential operator J that function J(f . g)(x) a, da = J(f)(x) b, db = J(g)(a) b, z -> da(db(z)) end operates on first class functions and once again returns a first class function (by returning a function we automatically obtain higher order derivatives, through repeated application of J ). There are several valid choices for this differential operator, but a convenient choice is
J (f ) := x ! (f (x); Jf (x)z); i.e. J (f )(x) returns the value of f at x, as well as a function which evaluates the jacobian-vector product between Jf (x) and some vector of sensitivities z. From this primitive we can define the gradient of a scalar function g : Rn ! R which is written as:
rg(x) := [J (g)(x)]2 (1) ([]2 selects the second value of the tuple, 1 = @z=@z is the initial sensitivity).
This choice of differential operator is convenient for several reasons: (1) The computation of the forward pass often computes values that can be re-used for the computation of the backwards pass. By combining the two operations, it is easy to re-use this work. (2) It can be used to recursively implement the chain rule (see figure 1).
This second property also suggests the implementation
strategy: hard code the operation of J on a set of
primitive f ’s and let the AD system generate the rest by repeated
application of the chain rule transform. This same general
approach has been implemented in many systems
However, to achieve our extensibility and composability goals, we implement a slight twist on this scheme. We define a fully user extensible function @ that provides a default fallback as follows
@(f )(args:::) = J (f )(args:::);
where the implementation that is generated automatically
by J recurses to @ rather than J and can thus easily be
intercepted using Julia’s standard multiple dispatch system
at any level of the stack. For example, we might make the
following definitions:
@(f )(:: typeof(+))(a :: IntOrFloat;b :: IntOrFloat) =
@(f )(:: typeof( ))(a :: IntOrFloat;b :: IntOrFloat) =
a + b; z ! (z; z)
a b; z ! (z b; a z)
i.e. declaring how to compute the partial derivative of +
and for two integer or float-valued numbers, but
simultaneously leaving unconstrained the same for other functions or
other types of values (which will thus fall back to applying
the AD transform). With these two definitions, any program
that is ultimately just a composition of ‘+‘, and ‘*‘ operations
of real numbers will work. We show a simple example in
figure 2. Here, we used the user-defined Measurement type
from the Measurements.jl package
knowledge of primitives on new types. By default we
provide implementations of the differentiable operator for
many common scalar mathematical and linear algebra
operations, written with a scalar LLVM backend and
BLASlike linear algebra operations. This means that even when
Julia builds an array type to target TPUs
indirect through @, there is no difference between userdefined custom gradients and those provided by the system. They are written using the same mechanism, are cooptimized by the compiler and can be finely targeted using Julia’s multiple dispatch mechanism.
Since Julia solves the two language problem, its Base, standard library, and package ecosystem are almost entirely pure Julia. Thus, since our @P system does not require primitives to handle new types, this means that almost all functions and types defined throughout the language are automatically supported by Zygote, and users can easily accelerate specific functions as they deem necessary.
Model-based reinforcement learning has advantages over
model-agnostic methods, given that an effective agent must
approximate the dynamics of its environment
Zygote can be used for control problems, incorporating the
model into backpropagation with one call to gradient. We
pick trebuchet dynamics as a motivating example. Instead of
aiming at a single target, we optimize a neural surrogate that
can aim it given any target. The neural net takes two inputs,
the target distance in metres and the current wind speed. The
network outputs trebuchet settings (the mass of the
counterweight and the angle of release) that get fed into a simulator
which solves an ODE and calculates the achieved distance.
We compare to our target and backpropagate through the
entire chain to adjust the weights of the network. Our dataset is
a randomly chosen set of targets and wind speeds. An agent
that aims a trebuchet to a given target can thus be trained in a
few minutes on a laptop CPU, resulting in a network which
is a surrogate to the inverse problem that allows for aiming
in constant-time. Given that the forward-pass of this network
is about 100 faster than performing the full optimization
on the trebuchet system (Figure 3), this surrogate technique
gives the ability to decouple real-time model-based control
from the simulation cost through a pre-trained neural
network surrogate to the inverse. We present the code for this
and other common reinforcement learning examples such as
the cartpole and inverted pendulum
On the one hand, a promising potential application and
research direction for Noisy Intermediate-Scale Quantum
(NISQ) technology
One such state of the art simulator is the Yao.jl
A subtle application is to perform traditional AD of the
quantum simulator itself. As a simple example of this
capability, we consider a Variational Quantum Eigensolver (VQE)
1 H = 4 4 2
X hi;ji
3 ix jx + iy jy + i j 5 z z
We use a standard differentiable variational quantum circuit composed of layers (2 in our example) of (parameterized) rotations and CNOT entanglers with randomly initialized rotation angles. Figure 4 shows the result of the minimization process as performed using gradients provided by Zygote.
Neural latent differential equations
Here we demonstrate automatic construction of Langevin equations from data using neural stochastic differential equations (SDE). Typically, Langevin equation models can be written in the form:
dXt = f (Xt)dt + g(Xt)dWt; where f : Rn ! Rn and g : Rn m ! Rn with Wt as the m-dimensional Wiener process.
To automatically learn such a physical model, we can
replace the drift function f and the diffusion function g
with neural networks and train these against time series
data by repeatedly solving the differential equation 10
times and using the average gradient of the cost function
against the parameters of the neural network. Our simulator
utilizes a high strong order adaptive integration provided
by DifferentialEquations.jl
The analytical formula for the adjoint of the strong solution
of a SDE is difficult to efficiently calculate due to the lack of
classical differentiability of the solution. However, Zygote
still manages to calculate a useful derivative for
optimization with respect to single solutions by treating the Brownian
process as fixed and applying forward-mode automatic
differentiation, showcasing Zygote’s ability to efficiently optimize
its AD through mixed-mode approaches
Acknowledgments This work would not have been possible without the work of many people. We are particularly indebted to our users that contributed the examples in this paper. Thanks to Roger Luo and JinGuo Liu for the quantum experiments. Thanks also to Zenna Tavares, Jesse Bettencourt and Lyndon White for being early adopters of Zygote and its underlying technology, and shaping its development. Much credit is also due to the core Julia language and compiler team for supporting Zygote’s development, including Jameson Nash, Jeff Bezanson and others. Thanks also to James Bradbury for helpful discussions on this paper.