[論文レビュー] Symplectic Adjoint Method for Exact Gradient of Neural ODE with Minimal Memory

Takashi Matsubara, Yuto Miyatake+1 more|arXiv (Cornell University)|Feb 19, 2021

Model Reduction and Neural Networks参考文献 36被引用数 5

ひとこと要約

本稿では、シンプレクティック積分法を用いて随伴系を解くことで、ニューラルODEにおける正確な勾配を計算する新しい手法、シンプレクティック随伴法を提案する。バックプロパゲーションをニューラルネットワークの各使用回数に限定し、シンプレクティック積分を活用することで、ステップ数とネットワークサイズに比例するメモリ消費量で正確な勾配を達成する。バックプロパゲーションやチェックポイント法よりも顕著に低く、標準的な随伴法よりも速度と丸め誤差に対するロバスト性に優れる。

ABSTRACT

A neural network model of a differential equation, namely neural ODE, has enabled the learning of continuous-time dynamical systems and probabilistic distributions with high accuracy. The neural ODE uses the same network repeatedly during a numerical integration. The memory consumption of the backpropagation algorithm is proportional to the number of uses times the network size. This is true even if a checkpointing scheme divides the computation graph into sub-graphs. Otherwise, the adjoint method obtains a gradient by a numerical integration backward in time. Although this method consumes memory only for a single network use, it requires high computational cost to suppress numerical errors. This study proposes the symplectic adjoint method, which is an adjoint method solved by a symplectic integrator. The symplectic adjoint method obtains the exact gradient (up to rounding error) with memory proportional to the number of uses plus the network size. The experimental results demonstrate that the symplectic adjoint method consumes much less memory than the naive backpropagation algorithm and checkpointing schemes, performs faster than the adjoint method, and is more robust to rounding errors.

研究の動機と目的

• ニューラルODEの学習におけるバックプロパゲーションおよびチェックポイント方式の高いメモリ消費量を解消すること。
• ニューラルODEにおける勾配計算の計算コストを低減し、数値誤差に対するロバスト性を向上させること。
• 小さなステップサイズを必要とせず、離散時間での正確な勾配計算を可能にすること。
• 随伴法と同等のメモリ効率を実現しながら、高い精度と速度を維持すること。
• 物理的PDEのような剛性または高次系においても、強力でスケーラブルな学習ソリューションを提供すること。

提案手法

• 本手法は、時間の逆方向に随伴系を解くためにシンプレクティック積分法を用い、幾何的構造を保存することで、丸め誤差を除き正確な勾配を達成する。
• バックプロパゲーションは、計算グラフ全体ではなく、ニューラルネットワーク $ f $ の各個々の使用に適用されるため、メモリ使用量が削減される。
• 適切なチェックポイントを用いて中間状態 $ x_n $ および内部段階値 $ X_{n,i} $ を保存し、メモリオーバーヘッドを最小限に抑える。
• 本手法は任意のルンゲ＝クッタ法と互換性があり、明示的および陰的積分法の両方をサポートする。
• シンプレクティック構造のおかげで、後向き統合中の数値的安定性が保たれ、誤差の蓄積が低減される。
• 本手法は、随伴統合に小さなステップサイズを必要としないため、標準的な随伴法よりも高速に計算が可能である。

実験結果

リサーチクエスチョン

• RQ1 シンプレクティック積分法を用いて、最小限のメモリオーバーヘッドでニューラルODEにおける正確な勾配を計算できるか？
• RQ2 提案手法のメモリ消費量は、バックプロパゲーション、チェックポイント、および標準的随伴法と比べてどの程度か？
• RQ3 シンプレクティック随伴法は、勾配の正確性を維持しながら、標準的随伴法よりも高速に計算できるか？
• RQ4 従来のバックプロパゲーションおよび随伴アプローチと比較して、丸め誤差に対してどの程度ロバストか？
• RQ5 本手法は、KdV方程式やCahn–Hilliard方程式のような剛性または高次系を効果的に処理できるか？

主な発見

• シンプレクティック随伴法は、前向き統合と同じステップサイズで正確な勾配（丸め誤差を除く）を達成するが、標準的随伴法とは異なり、誤差を制御するために小さなステップサイズを必要としない。
• メモリ消費量は $ O(MN + s) $ であり、$ M $ は成分数、$ N $ は時間ステップ数、$ s $ は内部段階数を表す。これはバックプロパゲーションの $ O(MNsL) $ より顕著に低く、随伴法と同等の水準である。
• KdV方程式における実験では、79.8 MiB のメモリで MSE が $ 1.61 \pm 4.00 \times 10^{-3} $ に達し、1イテレーションあたり 162 ms と、随伴法（276 ms/itr）よりも高速に性能を発揮した。
• Cahn–Hilliard系では、80.3 MiB のメモリで MSE が $ 5.47 \pm 1.46 \times 10^{-6} $ に達し、剛性系においてもロバストで効率的であることが示された。
• 本手法は、長大なグラフにわたる誤差の蓄積を避けるため、各統合段階ごとに勾配を計算する点で、バックプロパゲーションおよびチェックポイント方式よりも丸め誤差に対してよりロバストである。
• 本手法は、数値誤差を抑制するために小さなステップサイズを必要としないため、実用的に標準的随伴法よりも高速に動作する。