[論文レビュー] 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 は PyTorch で実装され、パブリックパッケージとして公開されている： https://jzkay12.github.io/TorchDiffEqPack

実験結果

リサーチクエスチョン

• RQ1 長時間の統合期間にわたり、逆方向時間における精度を損なわせることなく、一定のメモリ使用量を維持できる非線形 ODE 統合手法を設計できるか？
• RQ2 MALI は、従来の手法がメモリ制限や勾配の不正確さにより失敗する ImageNet のような大規模データセットにおいて、非線形 ODE の学習を可能にするか？
• RQ3 MALI は、時系列モデルタスクにおいてアドジョイント法と比較して、勾配の正確性と性能に優れているか？
• RQ4 MALI は、連続的正規化フローおよびその他の連続的生成モデルタスクで、最先端の結果を達成できるか？
• RQ5 逆方向時間における軌道の正確性は、多様な機械学習タスクにおける最終的なモデル性能にどのような影響を与えるか？

主な発見

• MALI は、ImageNet における非線形 ODE の学習を初めて可能にし、チューニングされた ResNet を上回る性能を達成した。
• MALI は、時系列モデルタスクにおいてアドジョイント法を上回り、優れた予測精度を示した。
• MALI は、特に正規化フローに基づく密度推定において、連続的生成モデル分野で新たな最先端の性能を達成した。
• MALI は、統合時間に対して一定のメモリコストを維持する一方で、ACA や単純な手法は線形に増大する。
• MALI は逆方向時間における軌道の高精度を維持し、信頼性が高く正確な勾配推定を可能にした。
• PyPI でのオープンソース実装により、研究者がさまざまな連続的深さモデルに MALI を容易に採用・拡張できるようになった。