[论文解读] Opening the Blackbox: Accelerating Neural Differential Equations by Regularizing Internal Solver Heuristics

Avik Pal, Yingbo Ma+2 more|arXiv (Cornell University)|May 9, 2021

Model Reduction and Neural Networks参考文献 46被引用 4

一句话总结

本文提出了一种新颖的正则化方法，通过利用自适应求解器的内部启发式方法——特别是局部误差估计和刚性估计——来加速神经微分方程（NDEs），而无需增加训练成本。通过使用这些预先计算、低成本的度量作为正则化信号，该方法将预测时间减少了最多1.84倍，训练时间减少了1.45倍，性能优于高阶导数方法，同时保持了模型的准确性。

ABSTRACT

Democratization of machine learning requires architectures that automatically adapt to new problems. Neural Differential Equations (NDEs) have emerged as a popular modeling framework by removing the need for ML practitioners to choose the number of layers in a recurrent model. While we can control the computational cost by choosing the number of layers in standard architectures, in NDEs the number of neural network evaluations for a forward pass can depend on the number of steps of the adaptive ODE solver. But, can we force the NDE to learn the version with the least steps while not increasing the training cost? Current strategies to overcome slow prediction require high order automatic differentiation, leading to significantly higher training time. We describe a novel regularization method that uses the internal cost heuristics of adaptive differential equation solvers combined with discrete adjoint sensitivities to guide the training process towards learning NDEs that are easier to solve. This approach opens up the blackbox numerical analysis behind the differential equation solver's algorithm and directly uses its local error estimates and stiffness heuristics as cheap and accurate cost estimates. We incorporate our method without any change in the underlying NDE framework and show that our method extends beyond Ordinary Differential Equations to accommodate Neural Stochastic Differential Equations. We demonstrate how our approach can halve the prediction time and, unlike other methods which can increase the training time by an order of magnitude, we demonstrate similar reduction in training times. Together this showcases how the knowledge embedded within state-of-the-art equation solvers can be used to enhance machine learning.

研究动机与目标

• 为解决神经微分方程（NDEs）推理过程中计算成本过高的问题，其依赖于自适应ODE求解器的步长调整。
• 在不增加训练成本的前提下减少预测时间，这是NDE部署中的主要瓶颈。
• 开发一种正则化策略，利用求解器内部的统计信息——而非昂贵的高阶导数——来引导模型学习更易求解的动力学行为。
• 将该方法扩展至神经随机微分方程（SDEs），证明其广泛适用性。
• 提供一种即插即用的解决方案，与现有NDE框架兼容，促进社区更广泛地采用。

提出的方法

• 该方法基于自适应ODE求解器在积分过程中内部计算的局部误差估计和刚性启发式方法引入正则化。
• 这些求解器内部的度量被用作廉价且准确的代价代理，以引导训练过程，使NDEs所需的神经网络函数评估次数（NFEs）更少。
• 正则化通过离散伴随敏感度实现，支持对求解器内部逻辑的高效反向传播，而无需修改底层NDE架构。
• 该方法基于Julia语言的SciML生态系统实现，可无缝集成至现有NDE框架中。
• 通过引入噪声感知启发式方法，该方法可自然扩展至神经SDEs，同时保持效率和准确性。
• 该方法避免了高阶自动微分，从而在不牺牲训练效率的前提下，实现了与更昂贵替代方法相当的性能提升。

实验结果

研究问题

• RQ1 是否可以利用求解器内部的启发式方法对NDEs进行正则化，从而在不增加训练成本的前提下减少预测时间？
• RQ2 与基于高阶导数的正则化相比，局部误差估计和刚性估计在推理速度和训练效率方面表现如何？
• RQ3 该正则化策略是否可推广至ODE之外，包括具有对角乘性噪声的神经SDEs？
• RQ4 使用求解器内部统计信息是否能实现更快的收敛速度和更少的函数评估次数，同时保持模型准确性？
• RQ5 是否存在一种方法，可以利用数值分析启发式方法加速NDEs，而无需修改求解器或模型架构？

主要发现

• 所提方法在监督分类和时间序列任务中，平均实现了1.84倍的预测时间加速，将Spiral方程的神经函数评估次数（NFEs）从1083.0减少至676.2。
• 训练时间平均减少了1.45倍，表明该方法未引入NDE中正则化通常伴随的高计算成本。
• 基于局部误差和刚性估计的正则化在性能上可与高阶导数方法相媲美，但无训练时间惩罚。
• 在图像分类任务中，基于误差的正则化更有效；而在时间序列问题中，刚性正则化表现出更强的性能提升，表明正则化类型应根据任务选择。
• 该方法可推广至神经SDEs，且无需修改底层NDE框架即可应用，具备广泛适用性。
• 该方法成功‘打开ODE求解器的黑箱’，通过利用其内部启发式信息引导模型学习，证明了嵌入式数值知识可提升深度学习效率。