[论文解读] 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 Networks参考文献 6被引用 89

一句话总结

DiffEqFlux.jl 将 DifferentialEquations.jl 的求解器与 Flux 神经网络结合，以在神经网络中嵌入微分方程层，从而实现神经ODE、SDE、DDE 等，并通过伴随反向传播进行训练。

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.

研究动机与目标

• 展示微分方程如何为数据科学中的非线性变换提供结构，并与机器学习互补地发挥作用。
• 展示如何将 DifferentialEquations.jl 的问题集成到 Flux 神经网络中，反之亦然。
• 突出完整的 DifferentialEquations.jl 求解器套件用于神经网络层的优势，包括刚性、时滞和随机性。
• 介绍神经 ODE 及相关模型（神经 SDE、DDE），并讨论用于通过求解器反向传播的伴随方法。

提出的方法

• 描述将 ODE 求解器嵌入为神经网络层的做法，使用返回求解器输出的前向传递。
• 介绍 diffeq_rd，这是一个在 Flux 中集成 ODE 求解的接口，允许通过反向传播进行训练。
• 展示替代的梯度方法：diffeq_fd（前向模式自动微分）和 diffeq_adjoint（伴随灵敏度分析）。
• 展示神经 ODE 和神经 SDE 的示例，如 Lotka–Volterra，以及在各种网络架构中的神经 ODE 层。
• 强调需要一个完整的求解器套件（ODEs、SDEs、DDEs、DAEs）以提高鲁棒性和性能，包括刚性和非刚性情况。
• 讨论通过求解器进行反向传播的挑战，以及自动微分（ForwardDiff.jl、Zygote.jl）和伴随方法如何解决这些挑战。

实验结果

研究问题

• RQ1 如何将神经网络与微分方程求解器融合，以创建可微分的端到端可训练模型？
• RQ2 与求解器内部实现相比，完整的 DifferentialEquations.jl 求解器套件为神经网络层提供了哪些优势？
• RQ3 哪些梯度方法（前向模式、反向模式、伴随灵敏度）在对微分方程求解器进行微分时最有效，它们如何随参数数量扩展？
• RQ4 神经 ODE、SDE、DDE 能否在包括 GPU 部署的标准深度学习流水线中进行实际训练？
• RQ5 各种方程类型（ODEs、DDEs、SDEs）作为神经网络层在训练稳定性和外推方面的表现如何？

主要发现

• DiffEqFlux.jl 能将完整的 DifferentialEquations.jl 求解器套件嵌入为神经网络层，从而实现 ODEs、SDEs、DDEs、刚性方程和基于伴随的训练。
• 该框架支持多种梯度路径（diffeq_rd、diffeq_fd 和 diffeq_adjoint），在最少代码修改的情况下实现通过求解器反向传播。
• 神经 ODE 层可以实现并通过将初始条件和网络保持在 GPU 上来实现 GPU 加速，ODE 求解器在 GPU 上运行。
• 使用完整的求解器套件可以提高在简单积分策略失败的挑战性问题上的鲁棒性，如 ROBER、刚性和非刚性示例所示。
• 利用自动微分（ForwardDiff、ReverseDiff）和即将推出的源到源 AD（Zygote），优化相对于求解器参数的梯度，从而实现可扩展的训练。