論文レビュー: Neural Ordinary Differential Equations
概要および要約
「Neural Ordinary Differential Equations」はRicky T. Q. Chenら3名が2018年に発表した論文であり、2018 NeurIPS Best Papersに選ばれた。以下のように簡単な1階常微分方程式である非自律システムをニューラルネットワークで近似する方法を提案している。
$$ \dfrac{\mathrm{d}y}{\mathrm{d}t} = f(y, t) $$
注目すべきは、ニューラルネットワークが近似（予測）するのは$y$ではなく、変化率$f$であるということだ。この式の両辺を$0$から任意の$T$まで積分すると次のようになる。
$$ y(T) - y(0) = \int\limits_{0}^{T} f(y, t) \mathrm{d}t \implies y(T) = y(0) + \int\limits_{0}^{T} f(y, t) \mathrm{d}t $$
したがって、$f$を正確に近似できさえすれば、任意の$T$に対して$y(T)$を得ることができる。
1 はじめに
下の図のように、等間隔の時間である8つのデータポイントから成る時系列データが与えられているとする。
このような状況で、通常我々がやりたいことは、$t_{9}$、$t_{10}$、$t_{11}$、$\dots$に対するデータ$\mathbf{h}_{9}$、$\mathbf{h}_{10}$、$\mathbf{h}_{11}$、$\dots$を予測することである。これを実現するための簡単に考えられる方法の一つは、残差ニューラルネットワーク、RNN、ノーマライジングフローなどを使用して以下のようにモデリングするものである。
$$ \mathbf{h}_{t+1} = \mathbf{h}_{t} + f(\mathbf{h}_{t}; \theta) \tag{1} $$
ここで$f$はニューラルネットワークであり、$\theta$は$f$の学習すべきパラメータ（重み）である。こういった方法はシンプルで直感的であり、最初に試してみることのできる方法の一つだが、次のような欠点がある。
- データを補間することが難しい。$(1)$のようなモデルでは$t_{3}$と$t_{4}$の間のデータを予測することが難しい。特に任意の時間$t_{3.472}$に対するデータを予測することが簡単でない。
- 時間間隔が均等でないデータに対しては適用が難しい。$(1)$で$\mathbf{h}_{t+1}$を更新するルールは$f$を1回、2回、整数回足していくものであるため、時間間隔が均等でなければ$(1)$の等式は成立しない。
この論文で提案する方法は、離散自律システムである$(1)$を次のように連続な非自律システムに変えて考えるものである。
$$ \dfrac{\mathrm{d}\mathbf{h}}{\mathrm{d}t} = f(\mathbf{h}(t), t; \theta) \tag{2} $$
「非自律システムで外力（あるいは速度）$f$をパラメータ$\theta$を持つ人工ニューラルネットワークで定義した$(2)$」をNeural Ordinary Differential Equations（Neural ODE, NODE）という。
このようなODEを解いてくれるソルバー（すなわち$f$を積分してくれるもの）はすでに多く研究されており、性能の良いさまざまな方法が提案されている。したがってニューラルネットワークを通じて$f$をうまく近似したならば、以下の式を通じて任意の時間$t$について$\mathbf{h}(t)$を予測することができる。両辺を積分すると、
$$ \mathbf{h}(t) = \mathbf{h}(0) + \int\limits_{0}^{t} f(\mathbf{h}(t), s; \theta) \mathrm{d}s \tag{3} $$
右辺は初期値$\mathbf{h}(0)$と$f$が与えられているならばODEソルバーを通して得られる値であるため、論文では次のように表記する。
$$ \mathbf{h}(t_{1}) = \operatorname{ODESolve}(\mathbf{h}(t_{0}), f, t_{0}, t_{1}, \theta) $$
論文ではODEソルバーによってモデルを定義し値を計算することによって得られるメリットを以下のように説明している。
Memory efficiencyメモリ効率:
下記の2章で、逆伝播アルゴリズムを使用せずに損失関数の勾配を計算する方法を紹介する。したがってモデルの関数値を計算するときに中間計算結果を保存する必要がなく、これはメモリ使用量がモデルのサイズに依存せず一定に維持されることを意味する。
Adaptive computation適応計算:
オイラー法以来120年が経過する間に、より正確で効率的なODEソルバーが多く開発されてきた。特に最近のODEソルバーは、誤差のモニタリングおよび高精度達成を目的として実行中に評価戦略evaluation strategyを動的に調整する機能を提供する。
Scalable and invertible normalizing flowsノーマライジングフローの一般化:
連続的なモデルの副次的利点の一つは、変数変換の公式を計算しやすいということであり（4章で説明する）、これは効率の良いノーマライジングフローモデルを構築できることを意味する。
Continuous time-series models連続時系列モデル:
訓練および予測データが等間隔に離散化されている必要があるRNNとは異なり、連続で定義されたNeural ODEは任意の時間$t$に対して予測を行うことができる。この内容は5章で詳しく説明される。
2 ODEソリューションのリバースモード自動微分
Neural ODEに関する核心的なアイデアと原理ははじめにで全て説明された。2章では学習法について扱う。著者はNeural ODEを実装する方法を開発し、GitHubで公開している。直接実装するわけではないので、学習に関する内容が特に気にならなければ飛ばしても構わないと思う。3〜5章はNeural ODEがどのように応用されるかを扱っているため、自分のフィールドでNeural ODEを活用したいならば3〜5章を詳しく読むことをお勧めする。特に4章の内容はディフュージョンモデルの次に生成モデルの主流を引っ張っているフローマッチングにもつながる。
観測値（正解データ）はニューラルネットワークのパラメータ$\theta$に依存しない定数であるため、損失関数を次のように簡単に表現しよう。
$$ L(\mathbf{z}(t_{1})) = L \left( \mathbf{z}(t_{0}) + \int_{t_{0}}^{t_{1}} f(\mathbf{z}(t), t; \theta) \mathrm{d}t \right) = L \left( \operatorname{OSESolve}(\mathbf{z}(t_{0}), f_{\theta}, t_{0}, t_{1}) \right) $$
論文で損失関数の勾配を計算するための方法はadjoint sensitivity methodであり、これは別のODEを時間逆順に積分することだ。その結果だけを見ると次のようになる。
$$ \dfrac{\mathrm{d}L}{\mathrm{d}\theta} = \int\limits_{t_{1}}^{t_{0}} \left( \dfrac{\partial L}{\partial \mathbf{z}(t)} \right)^{\mathsf{T}} \dfrac{\partial f(\mathbf{z}(t), t; \theta)}{\partial \theta} \mathrm{d}t $$
詳細はAppendix B、C、Dで確認できる。(ダウンロードリンク) また下記の例で説明したコードと学習に関連するメソッドは torchdiffeq パッケージに含まれており GitHubで公開されている。
3 監督学習のための残差ネットワークをODEに置き換える
Neural ODEは$(1)$を$(3)$として置き換えたものであると考えれば、残差ネットワークの一般化として見ることができる。
$$ \text{residual: } \mathbf{h}_{t+1} = \mathbf{h}_{t} + f(\mathbf{h}_{t}; \theta) \quad\overset{\text{generalization}}{\implies}\quad \dfrac{\mathrm{d}\mathbf{h}}{\mathrm{d}t} = f(\mathbf{h}(t), t; \theta) $$
したがってResNetを利用した監督学習においてNeural ODEを活用できる。著者たちはMNISTデータセットの判別問題でResNetの部分をそのままNeural ODEで置き換えて性能を比較した。
ODEの初期値問題を数値的に解くために選ばれたソルバーはアダムズメソッドであり、scipy.integrate パッケージで提供される陰的メソッドを使用した。最適化手法であるadjoint sensitivity methodを pytorch.autograd で実装した。
実験結果としてODE-NetはResNetに比べて$1/3$程度のパラメータを持つにもかかわらず、ほぼ同じ性能を示した。$L$はResNetの層数を意味し、ResNetの計算時間と逆伝播に必要なメモリは$L$に比例する。$\tilde{L}$はODEソルバーのnumber of function evaluations (NFE)であり、積分区間をどれだけ分割して計算するかを意味する。ODEソルバーは数値積分を行うので積分区間を分割することになり、NFEはこの区間の数を示す。例えば$[0, 1]$区間を$0.05$の間隔で積分を行うと$\text{NFE} = \dfrac{1}{0.05} = 20$である。論文ではこれをODE-Netの層数と解釈することができると述べている。NFEが大きいほど正確に計算できるが、計算に時間がかかるため、適度な数を選ぶことが重要である。Neural ODEではNFEをモデルが自動で選択する。
Figure 3a、3bではそれぞれ順方向計算におけるNFEに応じた誤差と計算時間を示し、直感に矛盾しない結果と言える。Figure 3cを見ると逆方向計算が順方向計算に比べておおよそ半分のレベルであることがわかる。これは提案する学習方法であるadjoint sensitivity methodが非常に効率的であることを示している。Figure 3dは訓練が進むにつれてNFEが増加することを示しており、これはモデルが訓練されるにつれて徐々に複雑になり、それをよく表現するために精巧になっていると説明している。
4 連続ノーマライジングフロー
$(1)$のような離散的なモデリングはノーマライジングフローnormalizing flow (NF)においても見受けられる。
$$ \mathbf{z}_{1} = f(\mathbf{z}_{0}) $$
$$ \log p(\mathbf{z}_{1}) = \log p(\mathbf{z}_{0}) - \log \left| \det \dfrac{\partial f}{\partial \mathbf{z}_{0}} \right| \tag{a4} $$
簡単にplanar flowを例として挙げると、
$$ \mathbf{z}_{t+1} = \mathbf{z}_{t} + \mathbf{u} \tanh (\mathbf{w}^{\mathsf{T}}\mathbf{z}_{t} + b) $$
$$ \log p(\mathbf{z}_{t+1}) = \log p(\mathbf{z}_{t}) - \log \left| 1 + \mathbf{u}^{\mathsf{T}} \dfrac{\partial \tanh}{\partial \mathbf{z}} \right| $$
ノーマルフローを学習する上で最も重要な考慮点は、$(a4)$の行列式の計算である。これは$\mathbf{z}$の次元や重みの数に対して三乗の計算コストを持つ。興味深いことに離散的な構造$(1)$を連続的構造$(3)$に変えるとそれに関連する計算がさらに簡素化される。
$\textbf{Theorem 1 }\text{(Instantaneous Change of Variables).}$
$\mathbf{z}(t)$を連続確率変数、$p(\mathbf{z}(t))$をそれに対する確率密度関数としよう。また非自律微分方程式が次のように与えられたとしよう。
$$ \dfrac{\mathrm{d} \mathbf{z}}{\mathrm{d}t} = f(\mathbf{z}(t), t) \tag{a5} $$
$f$が変数$\mathbf{z}$に関してリプシッツ連続であり、変数$t$に関して連続であるとしよう。すると、対数確率密度の導関数は次のようになる。
$$ \dfrac{\partial \log p(\mathbf{z}(t))}{\partial t} = -\Tr \left( \dfrac{\partial f}{\partial \mathbf{z}(t)} \right) $$
ここで$\Tr$はトレースである。
$(a4)$の行列式計算が単純なトレース計算に変わった。また標準的な離散ノーマライジングフローでは、今回は$f$が変数分離可能である必要がない。上記で例として挙げたplanar flowを上記$\textbf{Theorem 1}$に合わせて連続的形態に変えれば次のようになる。
$$ \dfrac{\mathrm{d}\mathbf{z}(t)}{\mathrm{d}t} = \mathbf{u} \tanh(\mathbf{w}^{\mathsf{T}}\mathbf{z}(t) + b), \qquad \dfrac{\partial \log p (\mathbf{z}(t))}{\partial t} = -\Tr \left( \mathbf{u} \dfrac{\partial \tanh(\mathbf{w}^{\mathsf{T}}\mathbf{z}(t) + b)}{\partial \mathbf{z}(t)} \right) $$
複数の隠れユニットを線形コストで使用する
行列式は線形ではないが、トレースは線形であるため$\Tr \sum\limits_{n} J_{n} = \sum\limits_{n} \Tr J_{n}$が成立する。したがって$(a5)$の$f$を線形結合で表しても、対数確率密度の導関数は次のように簡単に表現できる。
$$ \dfrac{\mathrm{d} \mathbf{z}(t)}{\mathrm{d}t} = \sum\limits_{n=1}^{M} f_{n}(\mathbf{z}(t)), \qquad \dfrac{\mathrm{d}\log p(\mathbf{z}(t))}{\mathrm{d}t} = -\sum\limits_{n=1}^{M} \Tr \left( \dfrac{\partial f_{n}}{\partial \mathbf{z}(t)} \right) $$
これは隠れ層のノード数$M$に対して計算コストが線形的に増加することを意味する。既存のノーマライジングフローは$\mathcal{O}(M^{3})$の計算コストのために多くの場合、一つのノードを多くの層で使用する構造をしている。
時間依存ダイナミクス
論文ではgatingと呼ばれる方法を使用してフローを以下のように定義する。
$$ \dfrac{\mathrm{d} \mathbf{z}}{\mathrm{d}t} = \sum\limits_{n} \sigma_{n}(t)f_{n}(\mathbf{z}), \qquad \sigma_{n}(t) \in (0 1) $$
すなわち、$f(\mathbf{z}, t)$が変数分離可能であると仮定する。ここで$\sigma_{n}$も$f_{n}$も全て学習する必要があるニューラルネットワークである。論文ではこれを連続ノーマライジングフローcontinuous normalizing flow (CNF)と呼ぶ。[フローマッチング]はこのCNFをより効率的に学習することができる一つの方法である。
4.1 連続ノーマライジングフローを用いた実験
CNFとNFを比較する。CNFは上記で説明した通りに実装され、Adamオプティマイザーを用いて10,000回の反復を行い訓練された。NFは初提案された論文と同様に実装され、RMSpropオプティマイザーを用いて500,000回の反復で訓練された。CNF側の反復回数は顕著に少ない。結果は以下の図から確認できる。
下のFigure 5の上の2行はCNF、最後の行はNFに対する結果として見られる。CNFは滑らかに変換される一方で、NFはそうではなく、Two Moonsデータに対して正しくモデル化ができていない。
5 ジェネレーティブ潜在関数時系列モデル
Neural ODEは非自律システムをニューラルネットワークで近似する方法であるため、時系列データの予測、補間、欠損値生成などにも使用できる。論文が提案する訓練方法は以下の通りである。
時系列データ$\left\{ \mathbf{x}_{i}, t_{i} \right\}_{i=1}^{N}$が与えられているとする。
RNNをエンコーダーとして利用し、潜在変数$\boldsymbol{\mu}$および$\boldsymbol{\sigma}$を抽出する。
$$ (\boldsymbol{\mu}, \log\boldsymbol{\sigma}^{2}) = \operatorname{RNNencoder}\left( \left\{ \mathbf{x}_{i}, t_{i} \right\} \right) $$
正規分布$N(\mathbf{0}, I)$から$\boldsymbol{\epsilon}$を抽出し、$\boldsymbol{\mu}$と$\boldsymbol{\sigma}$を使用して潜在変数の初期値$\mathbf{z}_{t_{0}}$を生成する。
$$ \mathbf{z}_{t_{0}} = \exp(0.5 \log\boldsymbol{\sigma}^{2}) \odot \boldsymbol{\epsilon} + \boldsymbol{\mu} $$
こうすることで、あたかも$\operatorname{RNNencoder}$が時系列データを潜在空間$N(\boldsymbol{\mu}, \diag(\boldsymbol{\sigma}^{2}))$にマッピングするように見せるが、逆伝播が可能である。
潜在変数の初期値$\mathbf{z}_{t_{0}}$とNeural ODEを通じて$\mathbf{z}_{t_{i}}$を計算する。
$$ (\mathbf{z}_{t_{1}}, \mathbf{z}_{t_{2}}, \dots, \mathbf{z}_{t_{N}}) = \operatorname{ODESolve}(\mathbf{z}_{t_{0}}, f, t_{0}, \dots, t_{N}, \theta) $$
$\mathbf{z}_{t_{i}}$をFCNとして定義されたデコーダーに入力し、予測された時系列データ$\mathbf{x}_{i}$を得る。
$$ \mathbf{x}_{t_{i}} = \operatorname{FCNdecoder}(\mathbf{z}_{t_{i}}) $$
ELBOを最大化する。
$$ \text{ELBO} = \sum\limits_{i} \log p(\mathbf{x}_{t_{i}} | \mathbf{z}_{t_{i}}) + \log p(\mathbf{z}_{t_{0}}) - \log q(\mathbf{z}_{t_{0}} | \left\{ \mathbf{x}_{t_{i}}, t_{i} \right\}) $$
5.1 時系列潜在ODE実験
実験に使用したニューラルネットワークは次のとおりである。
- $\operatorname{RNNencoder}$: ノードが25個のRNN
- 潜在空間は4次元。$(\boldsymbol{\mu}), (\boldsymbol{\sigma}) \in \mathbb{R}^{2 \times 2}$
- NeuralODE $f$: ノードが20個の隠れ層を1つ持つMLP
- $\operatorname{FCNdecoder}$: ノードが20個の隠れ層を1つ持つ別のMLP
データセットに関する条件は以下の通りである。
- 2次元螺旋データ1000個を生成。
- 半分は時計回り、残りは反時計回り。
- 各軌道は異なる初期値から始まり等間隔な時間で100個の点をサンプリング。
- 観測値にガウシアンノイズを追加した。
- 学習時には各軌道の100個の点から重複なしで無作為に30/50/100個を抽出して使用した。
以下の表を見るとNeuralODEの性能が明らかに優れていることが分かり、さらにデータが少ない時も多い時に比べて性能が大幅に落ちないということが確認できる。
以下の図はRNNとNeuralODEの予測結果である。NeuralODEの結果の方が滑らかで、学習しなかった外部領域での予測もより優れていることが示されている。
以下の図は潜在変数の軌道を最初の2次元について描いたものである。軌道は時計回りの螺旋と反時計回りの螺旋という2つのクラスタに分かれており、これは提案する方法が潜在変数をうまく表現していることを意味する。
