[论文解读] MALI: A memory efficient and reverse accurate integrator for Neural ODEs

Juntang Zhuang, Nicha C. Dvornek+2 more|arXiv (Cornell University)|May 3, 2021

Model Reduction and Neural Networks参考文献 38被引用 7

一句话总结

MALI 是一种内存高效且反向精确的神经 ODE 积分器，利用异步跃迁法（ALF）求解器实现积分时间内的恒定内存使用，与传统方法随时间增长的内存消耗不同。它通过在反向时间保持轨迹保真度，实现精确的梯度估计，在 ImageNet 训练、时间序列建模和连续生成建模任务中表现优于现有方法。

ABSTRACT

Neural ordinary differential equations (Neural ODEs) are a new family of deep-learning models with continuous depth. However, the numerical estimation of the gradient in the continuous case is not well solved: existing implementations of the adjoint method suffer from inaccuracy in reverse-time trajectory, while the naive method and the adaptive checkpoint adjoint method (ACA) have a memory cost that grows with integration time. In this project, based on the asynchronous leapfrog (ALF) solver, we propose the Memory-efficient ALF Integrator (MALI), which has a constant memory cost w.r.t integration time similar to the adjoint method, and guarantees accuracy in reverse-time trajectory (hence accuracy in gradient estimation). We validate MALI in various tasks: on image recognition tasks, to our knowledge, MALI is the first to enable feasible training of a Neural ODE on ImageNet and outperform a well-tuned ResNet, while existing methods fail due to either heavy memory burden or inaccuracy; for time series modeling, MALI significantly outperforms the adjoint method; and for continuous generative models, MALI achieves new state-of-the-art performance. We provide a pypi package: https://jzkay12.github.io/TorchDiffEqPack

研究动机与目标

• 解决现有神经 ODE 训练方法中高内存消耗和反向时间轨迹估计不准确的问题。
• 开发一种积分器，无论积分时间长短，均保持恒定内存使用，与伴随方法的效率相当。
• 确保反向时间积分的数值准确性，从而提高梯度估计的可靠性。
• 实现大规模数据集（如 ImageNet）上神经 ODE 的可行训练，这些数据集此前因内存或精度限制而使先前方法失效。
• 在连续生成建模和时间序列建模任务中实现最先进性能。

提出的方法

• MALI 基于异步跃迁法（ALF）求解器构建，可实现稳定且精确的时间积分，且内存开销极低。
• 该方法仅存储反向时间积分所需的中间状态，确保内存成本不随积分时长增加而上升。
• MALI 通过在伴随计算期间保持一致的时间步长同步，确保反向时间轨迹的准确性。
• 其检查点策略兼具内存效率和数值稳定性，避免反向传播中误差累积。
• 该积分器支持自适应时间步长，同时保持精度，适用于复杂动力学系统。
• MALI 已在 PyTorch 中实现，并作为公开包发布：https://jzkay12.github.io/TorchDiffEqPack

实验结果

研究问题

• RQ1 能否设计一种神经 ODE 积分器，在长时间积分过程中保持恒定内存使用，同时不牺牲反向时间精度？
• RQ2 MALI 是否能使神经 ODE 在 ImageNet 等大规模数据集上实现可行训练，而此前方法因内存或梯度不准确而失败？
• RQ3 在时间序列建模任务中，MALI 与伴随方法相比，在梯度精度和性能方面表现如何？
• RQ4 MALI 是否能在连续归一化流及其他连续生成建模任务中实现最先进结果？
• RQ5 反向时间轨迹精度对各类机器学习任务最终模型性能的影响如何？

主要发现

• MALI 首次实现了在 ImageNet 上对神经 ODE 的可行训练，性能超越了经过仔细调优的 ResNet。
• MALI 在时间序列建模任务中优于伴随方法，展现出更优的预测精度。
• MALI 在连续生成建模任务中达到新的最先进性能，尤其在基于归一化流的密度估计方面表现突出。
• 与 ACA 和朴素方法线性增长内存消耗不同，MALI 的内存成本不随积分时间变化。
• MALI 确保了反向时间轨迹的高精度，从而实现可靠且精确的梯度估计。
• 其在 PyPI 上的开源实现使研究人员能够轻松采用并扩展 MALI 以支持各类连续深度模型。