[论文解读] Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE

Juntang Zhuang, Nicha C. Dvornek+4 more|arXiv (Cornell University)|Jun 3, 2020

Model Reduction and Neural Networks参考文献 41被引用 34

一句话总结

本文提出 Adaptive Checkpoint Adjoint (ACA) 方法，通过对齐前向与反向轨迹、降低图深度以及实现自适应求解器来改进神经ODE的梯度估计，从而实现更高的精度与效率。

ABSTRACT

Neural ordinary differential equations (NODEs) have recently attracted increasing attention; however, their empirical performance on benchmark tasks (e.g. image classification) are significantly inferior to discrete-layer models. We demonstrate an explanation for their poorer performance is the inaccuracy of existing gradient estimation methods: the adjoint method has numerical errors in reverse-mode integration; the naive method directly back-propagates through ODE solvers, but suffers from a redundantly deep computation graph when searching for the optimal stepsize. We propose the Adaptive Checkpoint Adjoint (ACA) method: in automatic differentiation, ACA applies a trajectory checkpoint strategy which records the forward-mode trajectory as the reverse-mode trajectory to guarantee accuracy; ACA deletes redundant components for shallow computation graphs; and ACA supports adaptive solvers. On image classification tasks, compared with the adjoint and naive method, ACA achieves half the error rate in half the training time; NODE trained with ACA outperforms ResNet in both accuracy and test-retest reliability. On time-series modeling, ACA outperforms competing methods. Finally, in an example of the three-body problem, we show NODE with ACA can incorporate physical knowledge to achieve better accuracy. We provide the PyTorch implementation of ACA: \url{https://github.com/juntang-zhuang/torch-ACA}.

研究动机与目标

• 解释为何在现有方法下，NODE 的梯度估计会产生偏差或效率低下。
• 开发自适应检查点策略以对齐前向与反向轨迹，从而获得更准确的伴随梯度。
• 通过修剪计算图中的冗余组件来降低计算开销。
• 证明 ACA 在图像分类与时间序列任务上实现更低的误差率和更快的训练速度。
• 表明带有 ACA 的 NODE 能够融入物理知识，以提高动力学系统的精度。

提出的方法

• 引入 Adaptive Checkpoint Adjoint (ACA)，在反向模式求导期间记录前向轨迹以确保梯度的准确性。
• 应用轨迹检查点来对齐前向和反向的计算。
• 删除冗余组件以生成更浅的计算图。
• 在 ACA 框架内支持自适应 ODE 求解器。
• 提供 ACA 的 PyTorch 实现以便可重复性。

实验结果

研究问题

• RQ1 相比传统的伴随方法与朴素方法，ACA 是否能提高神经 ODE 梯度的数值精度？
• RQ2 在图像分类任务中，ACA 是否在保持或提升准确性的同时降低训练时间？
• RQ3 相较于其他梯度估计方法，ACA 在时间序列建模任务中的表现如何？
• RQ4 ACA 是否可用于将物理知识融入 NODE 以提高动力学系统的精度？

主要发现

• 在图像分类任务中，ACA 的误差率大约只有伴随法和朴素方法的一半，且训练时间也只有一半。
• 用 ACA 训练的 NODE 在准确性和测试-再测试的可靠性方面均优于 ResNet。
• 在时间序列建模方面，ACA 相对于竞争方法表现更优。
• 在三体问题示例中，带有 ACA 的 NODE 更好地利用物理知识以提高准确性。