[論文レビュー] 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) 手法を導入し、前向きと後向きの軌道を揃えることで Neural ODEs における勾配推定を改善し、グラフの深さを削減し、適応解法を可能にすることで、より高い精度と効率を達成します。

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 の勾配推定がなぜ偏りが生じるのか、または非効率であるのかを説明する。
• 正確な adjoint 勾配のために、前方と逆向きの軌道を揃える適応的なチェックポイント戦略を開発する。
• 計算グラフ内の冗長な要素を削減して計算オーバーヘッドを削減する。
• ACA が画像分類および時系列タスクでより低い誤差率とより高速な学習を達成することを示す。
• ACA を用いた NODE が物理的知識を組み込み、動力学系における精度を向上させることを示す。

提案手法

• 逆モード微分中に前方の軌道を記録して勾配の正確性を確保する Adaptive Checkpoint Adjoint (ACA) を導入する。
• 前方と逆方の計算を揃えるために軌道チェックポイントを適用する。
• 冗長な要素を削除して計算グラフを浅くする。
• ACA フレームワーク内で適応的な ODE ソルバーをサポートする。
• 再現性のための ACA の PyTorch 実装を提供する。

実験結果

リサーチクエスチョン

• RQ1 従来の adjoint 法や単純法と比べて、ACA は Neural ODEs の勾配の数値精度を向上させることができるか？
• RQ2 画像分類タスクで精度を維持または向上させつつ、ACA は訓練時間を短縮できるか？
• RQ3 時系列モデリングタスクにおける ACA の性能は、他の勾配推定手法と比べてどうか？
• RQ4 ACA を用いて NODE に物理知識を組み込み、動力学系の精度を向上させることができるか？

主な発見

• ACA は画像分類タスクで、adjoint 法および naive 法と比較して誤差率を約半分、学習時間を半分で達成する。
• NODE が ACA で訓練された場合、精度と再現性の信頼性の両方で ResNet を上回る。
• 時系列モデリングでは競合手法と比較して ACA の性能が優れている。
• 三体問題の例では、ACA を用いた NODE が物理的知識をより効果的に活用して精度を高める。