[Paper Review] DiffEqFlux.jl - A Julia Library for Neural Differential Equations

Christopher Rackauckas, Mike Innes+4 more|arXiv (Cornell University)|Feb 6, 2019

Model Reduction and Neural Networks6 references89 citations

TL;DR

DiffEqFlux.jl combines DifferentialEquations.jl solvers with Flux neural networks to embed differential equation layers in neural nets, enabling neural ODEs, SDEs, DDEs, and more with adjoint backpropagation.

ABSTRACT

DiffEqFlux.jl is a library for fusing neural networks and differential equations. In this work we describe differential equations from the viewpoint of data science and discuss the complementary nature between machine learning models and differential equations. We demonstrate the ability to incorporate DifferentialEquations.jl-defined differential equation problems into a Flux-defined neural network, and vice versa. The advantages of being able to use the entire DifferentialEquations.jl suite for this purpose is demonstrated by counter examples where simple integration strategies fail, but the sophisticated integration strategies provided by the DifferentialEquations.jl library succeed. This is followed by a demonstration of delay differential equations and stochastic differential equations inside of neural networks. We show high-level functionality for defining neural ordinary differential equations (neural networks embedded into the differential equation) and describe the extra models in the Flux model zoo which includes neural stochastic differential equations. We conclude by discussing the various adjoint methods used for backpropogation of the differential equation solvers. DiffEqFlux.jl is an important contribution to the area, as it allows the full weight of the differential equation solvers developed from decades of research in the scientific computing field to be readily applied to the challenges posed by machine learning and data science.

Motivation & Objective

• Demonstrate how differential equations provide structure for nonlinear transforms in data science and their complementary role with machine learning.
• Show how to integrate DifferentialEquations.jl problems into Flux neural networks and vice versa.
• Highlight the advantages of a full DifferentialEquations.jl solver suite for neural network layers, including stiffness, delays, and stochasticity.
• Present neural ODEs and related models (neural SDEs, DDEs) and discuss adjoint methods for backpropagation through solvers.

Proposed method

• Describe embedding an ODE solver as a neural network layer using a forward pass that returns the solver output.
• Introduce diffeq_rd, an interface that integrates ODE solving within Flux and allows training via backpropagation.
• Show alternative gradient methods: diffeq_fd (forward-mode AD) and diffeq_adjoint (adjoint sensitivity analysis).
• Demonstrate neural ODEs and neural SDEs with examples like Lotka–Volterra and neural ODE layers in various network architectures.
• Emphasize the need for a full solver suite (ODEs, SDEs, DDEs, DAEs) for robustness and performance, including stiff and non-stiff cases.
• Discuss backpropagation challenges through solvers and how automatic differentiation (ForwardDiff.jl, Zygote.jl) and adjoint methods address them.

Experimental results

Research questions

• RQ1 How can neural networks and differential equation solvers be fused to create differentiable, end-to-end trainable models?
• RQ2 What advantages does a full DifferentialEquations.jl solver suite provide for neural network layers compared to solver-internal implementations?
• RQ3 Which gradient methodologies (forward-mode, reverse-mode, adjoint sensitivity) are most efficient for differentiating through differential equation solvers, and how do they scale with parameter count?
• RQ4 Can neural ODEs, SDEs, and DDEs be practically trained within standard deep learning pipelines, including GPU deployment?
• RQ5 How do various equation types (ODEs, DDEs, SDEs) behave as neural network layers in terms of training stability and extrapolation?

Key findings

• DiffEqFlux.jl enables embedding a full DifferentialEquations.jl solver suite as neural network layers, allowing ODEs, SDEs, DDEs, stiff equations, and adjoint-based training.
• The framework supports multiple gradient paths (diffeq_rd, diffeq_fd, and diffeq_adjoint) to backpropagate through solvers with minimal code changes.
• Neural ODE layers can be implemented and GPU-accelerated by keeping the initial conditions and networks on the GPU, with the ODE solver operating on the GPU.
• Using a full solver suite improves robustness for challenging problems where simple integration strategies fail, as shown with stiff and non-stiff examples (e.g., ROBER, stiff cases).
• Automatic differentiation (ForwardDiff, ReverseDiff) and forthcoming source-to-source AD (Zygote) are leveraged to optimize gradients with respect to solver parameters, enabling scalable training.