[论文解读] Symplectic Adjoint Method for Exact Gradient of Neural ODE with Minimal Memory

Takashi Matsubara, Yuto Miyatake+1 more|arXiv (Cornell University)|Feb 19, 2021

Model Reduction and Neural Networks参考文献 36被引用 5

一句话总结

本文提出了一种新型的辛伴随方法，通过使用辛积分器求解伴随系统，实现神经ODE中精确梯度的计算。该方法仅在每次神经网络 $ f $ 使用时应用反向传播，并结合辛积分，使内存消耗与步数加网络规模成正比——显著低于反向传播或检查点法，同时在速度和对舍入误差的鲁棒性方面优于标准伴随方法。

ABSTRACT

A neural network model of a differential equation, namely neural ODE, has enabled the learning of continuous-time dynamical systems and probabilistic distributions with high accuracy. The neural ODE uses the same network repeatedly during a numerical integration. The memory consumption of the backpropagation algorithm is proportional to the number of uses times the network size. This is true even if a checkpointing scheme divides the computation graph into sub-graphs. Otherwise, the adjoint method obtains a gradient by a numerical integration backward in time. Although this method consumes memory only for a single network use, it requires high computational cost to suppress numerical errors. This study proposes the symplectic adjoint method, which is an adjoint method solved by a symplectic integrator. The symplectic adjoint method obtains the exact gradient (up to rounding error) with memory proportional to the number of uses plus the network size. The experimental results demonstrate that the symplectic adjoint method consumes much less memory than the naive backpropagation algorithm and checkpointing schemes, performs faster than the adjoint method, and is more robust to rounding errors.

研究动机与目标

• 解决训练神经ODE时反向传播与检查点方案的高内存消耗问题。
• 降低神经ODE梯度计算的计算成本，并提高对数值误差的鲁棒性。
• 在离散时间下使用辛积分器实现精确梯度计算，而无需采用更小的步长。
• 在内存效率上达到与伴随方法相当的水平，同时保持高精度与高速度。
• 为刚性或高阶系统（如物理PDE）提供一种鲁棒且可扩展的神经ODE训练方案。

提出的方法

• 该方法使用辛积分器反向求解伴随系统，保持几何结构，从而在舍入误差范围内实现精确梯度。
• 反向传播仅应用于神经网络 $ f $ 的每次独立使用，而非整个计算图，从而降低内存占用。
• 使用适当的检查点存储中间状态 $ x_n $ 和内部阶段值 $ X_{n,i} $，最小化内存开销。
• 该方法兼容任意龙格-库塔方法，支持显式与隐式积分器。
• 辛结构确保了数值稳定性，并减少了反向积分过程中的误差累积。
• 该方法避免了在伴随积分中使用小步长，从而在计算速度上优于标准伴随方法。

实验结果

研究问题

• RQ1 能否使用辛积分器在神经ODE中实现精确梯度计算，且内存开销极低？
• RQ2 所提方法的内存消耗与反向传播、检查点法及标准伴随方法相比如何？
• RQ3 辛伴随方法在保持梯度精度的同时，是否比标准伴随方法计算更快？
• RQ4 与传统反向传播和伴随方法相比，该方法对舍入误差的鲁棒性如何？
• RQ5 该方法能否有效处理刚性或高阶系统（如KdV方程和Cahn–Hilliard方程）？

主要发现

• 辛伴随方法使用与前向积分相同的步长即可实现精确梯度（舍入误差范围内），而标准伴随方法需使用更小步长以控制误差。
• 内存消耗为 $ O(MN + s) $，其中 $ M $ 为分量数量，$ N $ 为时间步数，$ s $ 为内部阶段数——显著低于反向传播的 $ O(MNsL) $，且与伴随方法相当。
• 在KdV方程实验中，该方法实现了 $ 1.61 \pm 4.00 \times 10^{-3} $ 的MSE，仅消耗79.8 MiB内存，计算速度优于伴随方法（162 ms/itr vs. 276 ms/itr）。
• 在Cahn–Hilliard系统中，该方法实现了 $ 5.47 \pm 1.46 \times 10^{-6} $ 的MSE，内存消耗为80.3 MiB，耗时568 ms/itr，表明其在刚性系统上具有优异的鲁棒性与效率。
• 由于梯度是按积分阶段计算而非在长计算图上累积，该方法对舍入误差的鲁棒性优于反向传播与检查点方案。
• 实践中，该方法比标准伴随方法更快，因为它避免了为抑制数值误差而必须使用小步长。