[论文解读] ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs

Amir Gholami, Kurt Keutzer+1 more|arXiv (Cornell University)|Feb 27, 2019

Model Reduction and Neural Networks参考文献 26被引用 77

一句话总结

ANODE 引入 discretize-then-optimize adjoint 框架并结合 checkpointing，以在神经 ODEs 中计算无条件准确的梯度，将内存从 O(LN_t) 降至 O(L)+O(N_t)，并避免以往 reverse-time 方法的不稳定性。

ABSTRACT

Residual neural networks can be viewed as the forward Euler discretization of an Ordinary Differential Equation (ODE) with a unit time step. This has recently motivated researchers to explore other discretization approaches and train ODE based networks. However, an important challenge of neural ODEs is their prohibitive memory cost during gradient backpropogation. Recently a method proposed in [8], claimed that this memory overhead can be reduced from O(LN_t), where N_t is the number of time steps, down to O(L) by solving forward ODE backwards in time, where L is the depth of the network. However, we will show that this approach may lead to several problems: (i) it may be numerically unstable for ReLU/non-ReLU activations and general convolution operators, and (ii) the proposed optimize-then-discretize approach may lead to divergent training due to inconsistent gradients for small time step sizes. We discuss the underlying problems, and to address them we propose ANODE, an Adjoint based Neural ODE framework which avoids the numerical instability related problems noted above, and provides unconditionally accurate gradients. ANODE has a memory footprint of O(L) + O(N_t), with the same computational cost as reversing ODE solve. We furthermore, discuss a memory efficient algorithm which can further reduce this footprint with a trade-off of additional computational cost. We show results on Cifar-10/100 datasets using ResNet and SqueezeNext neural networks.

研究动机与目标

• 动机并分析训练神经 ODEs 时的内存挑战，识别 reverse-time 梯度方法的失败之处。
• 提出 ANODE，一种基于 checkpointing 的 adjoint 框架，能够实现无条件准确的梯度。
• 在 CIFAR-10/100 上使用基于 ResNet/SqueezeNext 的 ODE 块，演示内存效率和稳定性。

提出的方法

• 将神经网络建模为带残差块的 ODE，并使用前向欧拉法或其他离散化。
• 证明将正向 ODE 反转以计算梯度会导致数值和一致性问题。
• 采用 discretize-then-optimize 的梯度计算，并结合 checkpointing 仅存储必要的激活。
• 定义基于 DTO 的 adjoint 计算，复用存储的正向轨道以计算梯度。
• 提供一种内存管理方案，在可选对数级 checkpointing 下将存储从 O(LN_t) 降至 O(L)+O(N_t)。
• 在 CIFAR-10/100 上使用 ResNet 与 SqueezeNext 主干网络、Euler 与 RK2 方案对 ANODE 进行实验。

实验结果

研究问题

• RQ1 在内存受限的训练下，神经 ODE 能否获得无条件准确的梯度？
• RQ2 带有 checkpointing 的 discretize-then-optimize adjoint 框架是否能够为通用神经网络模块（包括 ReLU 激活与卷积）提供稳定且准确的梯度？
• RQ3 与先前的 reverse-time 方法相比，ANODE 的内存与计算权衡如何？
• RQ4 离散化选择（Euler、RK2）如何影响神经 ODE 的训练稳定性和准确性？

主要发现

• 逆向时间（回代）ODE 求解在数值上可能不稳定，并且会对一般网络产生不正确的梯度。
• 先优化再离散化（OTD）的梯度可能与离散化的正向求解不一致，导致在小时间步时训练发散。
• 采用 discretize-then-optimize 与 checkpointing 的 ANODE 能获得无条件准确的梯度并实现稳定训练。
• ANODE 将内存从 O(LN_t) 降至 O(L)+O(N_t)，前向计算成本与对立方法相同。
• 使用离散化 DTO 配合 checkpointing 可实现高效的梯度计算，同时不约束权重范数以保证稳定性。
• 在 CIFAR-10/100 上使用 ResNet/SqueezeNext ODE 块的实验表明，ANODE 实现了稳定收敛，且性能优于 reverse-time 基线。