[论文解读] Heavy Ball Neural Ordinary Differential Equations

Hedi Xia, Vai Suliafu+5 more|arXiv (Cornell University)|Oct 10, 2021

Model Reduction and Neural Networks参考文献 56被引用 15

一句话总结

该论文提出了一种基于动量动力学的二阶神经普通微分方程（HBNODEs），通过用二阶常微分方程替代标准的一阶NODE，实现训练和推理的加速。HBNODEs通过确保伴随方程本身也是一个HBNODE，显著减少了函数评估次数（NFEs），并通过具有良好结构的谱特性，提升了长期依赖关系的学习能力，在图像分类、动力学建模和序列任务中，以更少的NFEs和更高的准确率，超越了NODEs及相关模型。

ABSTRACT

We propose heavy ball neural ordinary differential equations (HBNODEs), leveraging the continuous limit of the classical momentum accelerated gradient descent, to improve neural ODEs (NODEs) training and inference. HBNODEs have two properties that imply practical advantages over NODEs: (i) The adjoint state of an HBNODE also satisfies an HBNODE, accelerating both forward and backward ODE solvers, thus significantly reducing the number of function evaluations (NFEs) and improving the utility of the trained models. (ii) The spectrum of HBNODEs is well structured, enabling effective learning of long-term dependencies from complex sequential data. We verify the advantages of HBNODEs over NODEs on benchmark tasks, including image classification, learning complex dynamics, and sequential modeling. Our method requires remarkably fewer forward and backward NFEs, is more accurate, and learns long-term dependencies more effectively than the other ODE-based neural network models. Code is available at \url{https://github.com/hedixia/HeavyBallNODE}.

研究动机与目标

• 解决标准神经ODE（NODEs）在训练和推理过程中函数评估次数（NFEs）过高及收敛缓慢的问题。
• 克服NODEs在学习序列数据中长期依赖关系方面的已知局限性。
• 通过在连续深度模型中引入动量动力学，提升计算效率和模型实用性。
• 通过引入跳跃连接和门控机制等架构改进，缓解训练HBNODEs时可能出现的数值发散问题。

提出的方法

• 用包含阻尼项的二阶HBNODE（HBODE）替代一阶NODE的常微分方程，以在连续时间中建模动量。
• 推导HBNODEs的伴随方程，发现其本身也是一个HBNODE，从而实现更快速的反向常微分方程求解。
• 在HBNODEs中引入跳跃连接和门控机制，构建广义HBNODEs（GHBNODEs），以提升训练稳定性。
• 使用Dormand–Prince-45求解器，对HBODE与标准梯度下降ODE在基准函数上的收敛性进行数值比较。
• 将HBNODEs和GHBNODEs应用于图像分类（CIFAR-10、MNIST）、序列建模（ODE-RNN）以及动力系统（Walker 2D、平面振动系统）。
• 通过时间填充和可学习参数（γ, ξ）提升模型的灵活性与性能。

实验结果

研究问题

• RQ1 用基于动量的二阶常微分方程替代一阶NODEs，是否能显著减少前向和反向传播中的函数评估次数（NFEs）？
• RQ2 HBNODE的伴随方程是否也遵循HBNODE结构，从而实现更快速高效的训练？
• RQ3 由于具有良好结构的谱特性，HBNODEs是否能有效学习序列数据中的长期依赖关系？
• RQ4 在基准任务中，HBNODEs在准确率、NFEs和训练时间方面与NODEs、ANODEs和SONODEs相比表现如何？
• RQ5 GHBNODEs是否能在保持或提升性能的同时，有效缓解训练过程中的不稳定性与数值发散问题？

主要发现

• 与NODEs、ANODEs和SONODEs相比，HBNODEs在前向和反向传播中均显著减少了NFEs，尤其在低误差容差（如10⁻⁵）下效果更明显。
• HBNODE的伴随状态满足另一个HBNODE方程，从而实现更快速的常微分方程求解，显著降低训练和推理阶段的计算成本。
• 在CIFAR-10图像分类任务中，HBNODEs以更少的NFEs和更短的训练时间，实现了比NODEs及其他相关模型更高的测试准确率。
• 由于具有良好结构的谱特性，HBNODEs能有效缓解梯度消失问题，从而在序列数据中学习到长期依赖关系。
• GHBNODEs进一步提升了训练稳定性与性能，在具有挑战性的基准测试中，优于标准HBNODEs及其他基于ODE的模型。
• 在Rosenbrock和Beale函数上的数值实验表明，HBODEs比标准梯度下降ODE收敛更快，达到最小值。