[論文レビュー] Proximal Implicit ODE Solvers for Accelerating Learning Neural ODEs

Justin N. Baker, Hedi Xia+8 more|arXiv (Cornell University)|Apr 19, 2022

Model Reduction and Neural Networks被引用数 6

ひとこと要約

本稿では、implicit time integrationとproximal最適化を組み合わせることで、剛性ODE系の安定かつ効率的な解法を可能にする、proximal implicit ODE solverを提案する。この手法により、神経ODEの学習が加速され、ベンチマークタスク（グラフニューラルネットワークや正規化流れを含む）において、DOPRI5/8などのexplicit solverと比較して、前向き/逆向きのNFEを最大10倍まで削減し、優れた数値安定性を達成する。

ABSTRACT

Learning neural ODEs often requires solving very stiff ODE systems, primarily using explicit adaptive step size ODE solvers. These solvers are computationally expensive, requiring the use of tiny step sizes for numerical stability and accuracy guarantees. This paper considers learning neural ODEs using implicit ODE solvers of different orders leveraging proximal operators. The proximal implicit solver consists of inner-outer iterations: the inner iterations approximate each implicit update step using a fast optimization algorithm, and the outer iterations solve the ODE system over time. The proximal implicit ODE solver guarantees superiority over explicit solvers in numerical stability and computational efficiency. We validate the advantages of proximal implicit solvers over existing popular neural ODE solvers on various challenging benchmark tasks, including learning continuous-depth graph neural networks and continuous normalizing flows.

研究の動機と目的

• explicit solverにおける剛性と小さなステップサイズによる、神経ODE学習の計算ボトルネック（前向き・逆向きのNFEが高くなること）を解消すること。
• 内側と外側の反復を分離することで、高次元ODEに対する標準的なimplicit solverの非効率性を、proximal最適化を活用して克服すること。
• 剛性神経ODE系において、unconditionalなエネルギー安定性と収束性を保証するフレームワークの開発。
• 連続的深さを持つグラフニューラルネットワークや連続的正規化流れといった、挑戦的なベンチマークタスクにおける優れた性能の実証。

提案手法

• implicit ODE solver（後退オイラー、クランクニコルソン、BDF2–4）を、implicit更新ステップを処理するためのproximal演算子を用いて凸最適化問題に定式化する。
• 解法プロセスを内側・外側の反復に分解する：内側反復は、速やかな最適化（例：FR法）によりproximal部分問題を解き、外側反復は時間ステップを進める。
• 任意のステップサイズにおいてエネルギー減少を保証し、定常点への収束を保証するため、proximal定式化による数値安定性を証明する。
• 逆伝播に随伴法を用い、逆向きの時間においても同じ安定な時間積分器を再利用することで、バックワードNFEを削減する。
• 自動微分を維持しながら、メモリ効率を保ちつつ、proximal solverを神経ODE学習パイプラインに統合する。
• 固有値比が無限大であるような系（例：拡散モデル（GRANDなど）や正規化流れ）に、本手法を適用する。

実験結果

リサーチクエスチョン

• RQ1 Proximal implicit solversは、剛性ODE系において、数値精度を維持したまま、前向き・逆向きのNFEを削減できるか？
• RQ2 異なるステップサイズや誤差許容範囲の下で、proximal定式化のエネルギー安定性は、explicit solverと比べてどの程度優れているか？
• RQ3 グラフニューラルネットワークや正規化流れにおいて、計算効率の観点から、proximal implicit solversは、explicit adaptive solvers（例：DOPRI5, DOPRI8）をどの程度上回れるか？
• RQ4 内側・外側反復構造は、深層学習応用における高次元で剛性のあるODEのスケーラブルかつ安定な解法を可能にするか？
• RQ5 BDF2とBDF4などの異なるimplicitスキームは、proximal ODE solversにおける収束性と最終的な解の精度に、どの程度の影響を与えるか？

主な発見

• CoauthorCSグラフノード分類タスクにおいて、GRANDを用いて、DOPRI5およびDOPRI8と比較して、前向き・逆向きのNFEを最大10倍まで削減した。
• 1次元の拡散方程式において、proximal BDF4スキームはステップサイズ1/2000で最終ステップ誤差1.15e-6を達成し、Crank-NicolsonやBDF2を上回る精度を示した。
• proximal部分問題の強制的性質のおかげで、ODE系が非凸であっても、unconditionalなエネルギー安定性と定常点への収束を保証する。
• GRANDモデルでは、誤差許容範囲を1e-3から1e-6に低下させた場合、explicit solverではNFEが急激に増加したが、proximal implicit solverは優れた安定性のおかげで低NFEを維持した。
• 内側最適化ソルバ（FR法）は、許容誤差5e-9に設定した場合、1ステップあたり10～20回の反復で収束し、実用的な効率性を確保した。
• proximal定式化により、通常のimplicit solverでは不安定または実用的でない高次精度implicitスキーム（例：BDF4）を、神経ODEに安定して統合可能となった。