A memory-efficient neural ODE framework based on high-level adjoint differentiation

要約

ニューラル常微分方程式 (ニューラル ODE) は、動的システムと深層学習の橋渡しをする新しいネットワーク アーキテクチャとして登場しました。
ただし、バニラ ニューラル ODE の連続随伴法で得られた勾配は、逆に正確ではありません。
他のアプローチは、深い計算グラフによる過度のメモリ要件、または時間積分スキームの選択肢の制限のいずれかに悩まされ、大規模で複雑な動的システムへの適用を妨げます。
メモリ効率と柔軟性を損なうことなく正確な勾配を実現するために、高レベルの離散随伴アルゴリズム微分に基づく新しいニューラル ODE フレームワーク PNODE を提案します。
離散随伴時間インテグレータと、これらのインテグレータに合わせた高度なチェックポイント戦略を活用することで、PNODE は、勾配を一貫して正確に計算しながら、メモリと計算コストのバランスを提供できます。
当社は、最も一般的に使用されているポータブルでスケーラブルな科学計算ライブラリの 1 つである PyTorch および PETSc に基づいたオープンソース実装を提供します。
画像分類と連続正規化フロー問題に関する広範な数値実験を通じてパフォーマンスを実証します。
他の逆精度の方法と比較した場合、PNODE が最高のメモリ効率を達成することを示します。
画像分類問題に関しては、PNODE はバニラ ニューラル ODE よりも最大 2 倍高速で、既存の最高の逆精度手法よりも最大 2.3 倍高速です。
また、PNODE により、スティッフな動的システムに必要な陰的な時間積分法の使用が可能になることも示します。

要約(オリジナル)

Neural ordinary differential equations (neural ODEs) have emerged as a novel network architecture that bridges dynamical systems and deep learning. However, the gradient obtained with the continuous adjoint method in the vanilla neural ODE is not reverse-accurate. Other approaches suffer either from an excessive memory requirement due to deep computational graphs or from limited choices for the time integration scheme, hampering their application to large-scale complex dynamical systems. To achieve accurate gradients without compromising memory efficiency and flexibility, we present a new neural ODE framework, PNODE, based on high-level discrete adjoint algorithmic differentiation. By leveraging discrete adjoint time integrators and advanced checkpointing strategies tailored for these integrators, PNODE can provide a balance between memory and computational costs, while computing the gradients consistently and accurately. We provide an open-source implementation based on PyTorch and PETSc, one of the most commonly used portable, scalable scientific computing libraries. We demonstrate the performance through extensive numerical experiments on image classification and continuous normalizing flow problems. We show that PNODE achieves the highest memory efficiency when compared with other reverse-accurate methods. On the image classification problems, PNODE is up to two times faster than the vanilla neural ODE and up to 2.3 times faster than the best existing reverse-accurate method. We also show that PNODE enables the use of the implicit time integration methods that are needed for stiff dynamical systems.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google