[论文解读] TorchDyn: A Neural Differential Equations Library

Michael Poli, Stefano Massaroli+3 more|arXiv (Cornell University)|Sep 20, 2020

Model Reduction and Neural Networks参考文献 23被引用 14

一句话总结

TorchDyn 是一个基于 PyTorch 的库，提供了一套模块化、可组合且用户友好的框架，用于基于神经微分方程（NDEs）的连续深度模型。它通过可重用组件，实现了对多种 NDE 变体（如神经 ODE、哈密顿网络和拉格朗日网络、连续归一化流）的高效实现，支持伴随梯度和积分损失函数，并通过教程和基准测试加速研究进程。

ABSTRACT

Continuous-depth learning has recently emerged as a novel perspective on deep learning, improving performance in tasks related to dynamical systems and density estimation. Core to these approaches is the neural differential equation, whose forward passes are the solutions of an initial value problem parametrized by a neural network. Unlocking the full potential of continuous-depth models requires a different set of software tools, due to peculiar differences compared to standard discrete neural networks, e.g inference must be carried out via numerical solvers. We introduce TorchDyn, a PyTorch library dedicated to continuous-depth learning, designed to elevate neural differential equations to be as accessible as regular plug-and-play deep learning primitives. This objective is achieved by identifying and subdividing different variants into common essential components, which can be combined and freely repurposed to obtain complex compositional architectures. TorchDyn further offers step-by-step tutorials and benchmarks designed to guide researchers and contributors.

研究动机与目标

• 为深度学习中的连续深度学习解决缺乏可访问、模块化且可组合的软件工具的问题。
• 减少现有神经微分方程（NDEs）实现中复杂的样板代码和冗长的继承链，这些因素阻碍了可复现性和可扩展性。
• 提供统一的、与 PyTorch 兼容的接口，支持多种 NDE 变体，包括哈密顿、拉格朗日和归一化流模型。
• 通过提供逐步教程、基准测试以及与 PyTorch 及其生态系统的完全兼容性，加速研究和入门过程。
• 支持高级训练机制，如伴随梯度和积分损失函数，以提升内存效率和模型灵活性。

提出的方法

• 引入 NeuralDE 和 DEFunc 类作为核心抽象，封装 NDE 模型及其向量场，实现与 PyTorch 的无缝集成。
• 采用专用层（如 DepthCat、GalLinear 和 GalConv2d）处理随深度变化的参数，无需自定义模块定义。
• 提供针对基于能量的模型（如哈密顿、拉格朗日模型）的专用包装器，将其与 DEFunc 集成，并支持混合架构。
• 通过 torchdiffeq 和 torchsde 中的可组合组件和数值求解器，支持连续归一化流（CNFs）和 FFJORD 类模型。
• 通过扩展的伴随方法实现积分损失函数的梯度计算，在反向传播期间实现 O(1) 内存使用。
• 与 PyTorch-Lightning 集成以简化训练循环和日志记录，同时保持独立运行能力，便于底层控制。

实验结果

研究问题

• RQ1 如何在深度学习框架中以模块化、可组合且可扩展的方式实现神经微分方程？
• RQ2 需要哪些架构和软件抽象，才能简化如 ANODE、FFJORD 和哈密顿网络等多样化连续深度模型的实现？
• RQ3 如何在神经 ODE 库中高效支持积分损失函数和伴随梯度？
• RQ4 统一的库在多大程度上可以减少样板代码并提升神经微分方程研究的可复现性？
• RQ5 库如何促进研究人员和实践者对新兴 NDE 变体的采用和实验？

主要发现

• TorchDyn 将复杂的 NDE 组件成功抽象为可重用、可组合的原语（如 NeuralDE、DEFunc 和深度可变层），显著降低了实现复杂度。
• 该库支持一系列前沿模型，包括 ANODE、哈密顿和拉格朗日网络、稳定的归一化流以及基于图的 ODE 模型，如图 3 所示。
• TorchDyn 通过伴随梯度实现高效训练，内存使用为 O(1)，并支持日益广泛应用于现代 NDE 研究的积分损失函数。
• 该库的模块化设计消除了以往实现中常见的复杂继承链，例如原始 FFJORD 代码库中的情况。
• TorchDyn 包含全面的教程和基准测试，引导用户实现 FFJORD、哈密顿网络和神经图 ODE 等模型。
• 通过与 PyTorch 集成并利用 torchdiffeq 和 torchsde，TorchDyn 确保了在 GPU 上的高性能，并支持所有组件的批量操作。