[论文解读] Proximal Implicit ODE Solvers for Accelerating Learning Neural ODEs

Justin N. Baker, Hedi Xia+8 more|arXiv (Cornell University)|Apr 19, 2022

Model Reduction and Neural Networks被引用 6

一句话总结

本论文提出了一类近端隐式ODE求解器，通过结合隐式时间积分与近端优化，加速了神经ODE的训练过程，实现了对刚性ODE系统的稳定且高效的求解。该方法在包括图神经网络和归一化流在内的基准任务上，相较于DOPRI5/8等显式求解器，实现了更优的数值稳定性以及前向/反向NFE（函数评估次数）减少高达10倍。

ABSTRACT

Learning neural ODEs often requires solving very stiff ODE systems, primarily using explicit adaptive step size ODE solvers. These solvers are computationally expensive, requiring the use of tiny step sizes for numerical stability and accuracy guarantees. This paper considers learning neural ODEs using implicit ODE solvers of different orders leveraging proximal operators. The proximal implicit solver consists of inner-outer iterations: the inner iterations approximate each implicit update step using a fast optimization algorithm, and the outer iterations solve the ODE system over time. The proximal implicit ODE solver guarantees superiority over explicit solvers in numerical stability and computational efficiency. We validate the advantages of proximal implicit solvers over existing popular neural ODE solvers on various challenging benchmark tasks, including learning continuous-depth graph neural networks and continuous normalizing flows.

研究动机与目标

• 解决由于显式求解器在刚性系统中步长过小导致的神经ODE训练计算瓶颈问题，即前向和反向NFE过高。
• 通过利用近端优化将内层与外层迭代解耦，克服标准隐式求解器在高维ODE中效率低下的问题。
• 构建一个框架，确保刚性神经ODE系统具有无条件能量稳定性和收敛性。
• 在连续深度图神经网络和连续归一化流等具有挑战性的基准任务上展示优越性能。

提出的方法

• 利用近端算子将隐式ODE求解器（后向欧拉法、克兰克-尼科尔森法、BDF2–4）建模为凸优化问题，以处理隐式更新步骤。
• 将求解过程分解为内层-外层迭代：内层迭代通过快速优化方法（如FR方法）求解近端子问题，外层迭代推进时间步长。
• 通过证明近端公式能保证能量衰减并使解收敛至驻点，确保数值稳定性，且对任意时间步长均成立。
• 在反向传播中使用伴随法，通过在反向时间中重用相同的稳定时间积分器，显著降低反向NFE。
• 将近端求解器集成至神经ODE训练流水线中，结合自动微分，保持内存效率。
• 将该方法应用于扩散模型（如GRAND）和归一化流等刚性系统，其特征值比为无穷大。

实验结果

研究问题

• RQ1 近端隐式求解器是否能在保持刚性系统数值精度的前提下，减少神经ODE训练中的前向和反向NFE？
• RQ2 在不同步长和误差容限下，近端公式的能量稳定性与显式求解器相比如何？
• RQ3 在图神经网络和归一化流上，近端隐式求解器在计算效率方面相较于显式自适应求解器（如DOPRI5、DOPRI8）能提升多少？
• RQ4 内层-外层迭代结构是否能够实现高维、刚性ODE在深度学习应用中的可扩展且稳定的求解？
• RQ5 不同隐式格式（如BDF2与BDF4）对近端ODE求解器的收敛性和最终解的精度有何影响？

主要发现

• 在使用GRAND的CoauthorCS图节点分类任务中，近端隐式求解器相较于DOPRI5和DOPRI8，将前向和反向NFE减少了高达10倍。
• 在一维扩散方程上，近端BDF4格式在步长为1/2000时达到1.15e-6的最终步长误差，精度优于Crank-Nicolson和BDF2。
• 由于近端子问题具有强制性（coercive nature），该方法即使在ODE系统非凸时，也能保证无条件能量稳定性和收敛至驻点。
• 对于GRAND模型，当误差容限从1e-3降低至1e-6时，显式求解器的NFE迅速增加，而近端隐式求解器由于更强的稳定性，始终保持低NFE。
• 当容差设为5e-9时，内层优化求解器（FR方法）在每个时间步内通常在10–20次迭代内收敛，确保了实际效率。
• 近端公式使得高阶隐式格式（如BDF4）在神经ODE中实现稳定集成，而这些格式在标准隐式求解器中通常不稳定或不切实际。