全4957文字
PR
この記事は日経Robotics 有料購読者向けの過去記事ですが
日経Robotics デジタル版(電子版)』のサービス開始を記念して、特別に誰でも閲覧できるようにしています。
本記事はロボットとAI技術の専門誌『日経Robotics』のデジタル版です
本記事はロボットとAI技術の専門誌『日経Robotics』のデジタル版です

 ニューラルネットワークは有限回の変換の組み合わせで構成され、例えば10層のニューラルネットワークは入力を10回変換した結果として表される。

 これに対し、入力に連続時間の変換を適用し、例えば3.6回分変換といった変換を扱えるNeural ODE1)が提案された。Neural ODEはNeurIPS 2018のベストペーパーに選ばれている。これについて解説する。

 ニューラルネットワークの層のいくつかは変換した値を入力に足しこんでいく形をとる。

\[h_{t+1} = h_t + f(h_t; \theta_t)\]

 この$f$は1層以上から成るニューラルネットワークで表現される。現在の画像認識器の主流となっているResNetやLeaky Unitを使ったRNN、生成モデルの正規化フローなどがこの形をとる。この式は連続な関数を時間$t$についてオイラー法で離散化した時の更新式とみなせる。このステップ幅を極限まで小さくしていくと、次のような常微分方程式(ODE:ordinal differential equation)が得られる。

\[\frac{dh(t)}{dt} = f(h(t), t; \theta)\]

著者の岡野原大輔氏
著者の岡野原大輔氏

 この場合、関数$f$は入力が$h(t)$,時刻$t$の時の微分を与えているとみなせる。入力を時刻$0$における値$h(0)$とし、上記の常微分方程式に従い、時刻$T$における値$h(T)$を出力とした関数を考える。このような常微分方程式を使った関数を層として利用したニューラルネットワークをNeural ODE1)とよぶ(図1)。

多くの利点があるNeural ODE

 このNeural ODEは通常のニューラルネットワークと比べて多くのメリットがある。

図1 Neural ODEと通常のResNetの比較
[画像のクリックで拡大表示]
図1 Neural ODEと通常のResNetの比較
[画像のクリックで拡大表示]
図1 Neural ODEと通常のResNetの比較
通常のResNetは変換の離散的なシーケンスとなるが(左側)、Neural ODEは状態を連続的に変換するベクトル場とみなせる(右側)。丸いプロットは評価点を表す。(図:University of Toronto)

(1)メモリ効率が良い。ResNetは各層毎に異なるダイナミクスを表しているとみなせるが、もし隣り合う層が似たダイナミクスを表す場合はそれらの層は1つにまとめることでパラメータ数を減らせる。Neural ODEは自然に似た層を融合した形で表現できる。RNNのようにパラメータを共有した時間変化する変換を適用していると考えても良い。

(2)速度と精度のトレードオフをとれる。Neural ODEはODEの解を求める際に、通常のRunge-Kutta法などを使うが、その際に評価点数を変えることができる。精度が必要な場合は評価点数を増やし、必要でない場合は評価変数を減らせばよい。

(3)可逆関数であり逆関数を計算できる。ODEで時間を$T$から$0$に向かってたどれば出力から入力を計算できる。後述するがパラメータについての勾配を求める際に時間方向を逆にたどるが、その際に計算の途中結果(活性値)はその場で計算して復元できるため、保存しておく必要がない。

(4)任意の変換を使っても確率密度を効率的に求められる。正規化フローを使った場合、確率密度の計算には隠れ層のユニット数が$n$の時$\mathrm{O}(n^3)$時間必要であり、そのヤコビアンの行列式が効率的に求まるような計算クラスにすると$\mathrm{O}(n)$時間で求まる。Neural ODEは任意の変換を用いて$\mathrm{O}(n)$時間で求められる。

(5)ODEで表されるような問題などは直接扱え、その場合、外挿などの問題を解ける。

誤差逆伝播法を連続ダイナミクスに拡張

 Neural ODEの実現で鍵となるのが、誤差逆伝搬法の連続ダイナミクスへの拡張であり以下に説明していく。ODEの解$z(t_1)$を入力とした損失関数$L$は次のように表される。

\[L(\mathbf{z}(t_1)) = L\left ( \mathbf{z}(t_0) + \int_{t_0}^{t_1}f(\mathbf{z}(t), t, \theta) dt \right ) \]

この損失関数の入力やパラメータについての勾配($\frac{dL}{d\mathbf{z}}$,$\frac{dL}{d\theta}$)を求めることが目標である。まず、各時刻で損失の隠れ状態$\mathbf{z}(t)$についての勾配$\mathbf{a}(t)= \frac{\partial L}{\partial\mathbf{z}(t)}$を求める。この値$\mathbf{a}$は随伴変数とよぶ。この随伴変数は次のようなダイナミクスを持つことが示せる。

\[\frac{d \mathbf{a}(t)}{dt} = -\mathbf{a}(t)^T \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}}\]

 これは、誤差逆伝搬法の連続ダイナミクスへの一般化とみなすことができる(誤差逆伝搬法では$\mathbf{a}$が誤差$\mathbf{e}$に対応し、誤差に各層のヤコビアン$\frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \mathbf{z}}$をかけて逆方向に伝搬する)。このODEを$t=T$から$t=0$に向かって逆向きに解く。この際、途中の$\mathbf{z}(t)$の値も必要となるが、それも同時にODEで求められ、保存しておく必要はない。

 この随伴変数が求まれば、損失のパラメータについての勾配は次のように求められる。

\[\frac{dL}{d\theta} = - \int_{t_1}^{t_0} \mathbf{a}(t)^T \frac{\partial f(\mathbf{z}(t), t, \theta)}{\partial \theta}\]

ベクトルとヤコビアンの積$\mathbf{a}^T \frac{\partial f}{\partial\mathbf{z}}$は、出力の勾配に$\mathbf{a}$を設定した上で自動微分を適用すれば、ヤコビアンを明示的に計算せずに前向き計算$f$の倍程度の計算量で効率的に求めることができる(vjp:matrix-free vector-Jacobian productとよばれる)。

 これら随伴変数と損失関数の入力、パラメータについての勾配は同時にODEを解くことで計算できる。このようにNeural ODEでは勾配を伝搬する際に途中結果を保存しておく必要がない。

連続正規化フロー

 Neural ODEは生成モデルとしても非常に有効である。正規化フローの回2)でも説明したように入力$\mathbf{z}_0$を可逆な関数$f$で変換して$\mathbf{z}_1=f(\mathbf{z}_0)$が得られた時、その対数尤度$\log p(\mathbf{z}_1)$は元の尤度$\log p(\mathbf{z}_0)$を使って、次のように表される。

\[\log p(\mathbf{z}_1) = \log p(\mathbf{z}_0) - \log \left | \det \frac{\partial f}{\partial \mathbf{z}_0} \right |\]

このヤコビアンの行列式は、$\mathbf{z}_0$の空間での体積が、$\mathbf{z}_1$の空間でどのくらい拡大したかを表しており、拡大した分、密度を小さくしている($\log$をとっているので、割り算ではなく引き算となっている)。生成モデルを学習する際はこの計算された対数尤度を最大化するようなパラメータを求めればよい。

 この行列式の計算量はノード数が$n$の時、$\mathrm{O}(n^3)$と大きいため、RealNVPやGlowなどでは効率的に行列式が求まるような変換を扱う手法が提案されていた。一方、連続ダイナミクスを使った場合はこれを$\mathrm{O}(n)$で求めることができる。これについて以下に説明していく。

 変数$\mathbf{z}(t)$は$p(\mathbf{z}(t))$を確率密度とする確率変数とする。このとき、確率密度の時間当たりの変化量は次のように表せられる。

\[\frac{\log p(\mathbf{z}(t))}{\partial t} = - \mathrm{tr} \left ( \frac{df}{d\mathbf{z}(t)} \right )\]

さきほどの正規化フローの時は行列式$\det$だったのに対し、連続ダイナミクスを使用した場合はトレース$\mathrm{tr}$であることに注意してほしい。この計算ではヤコビアン行列はそれぞれ独立に計算するため$\mathrm{O}(n^2)$時間が必要である。

 FFJORD(Free-form Jacobian of Reversible Dynamics)3)は任意の変換を利用しても効率的に確率密度を計算できる手法を提案した。はじめに任意の行列のトレースは、ノイズベクトル$\mathbf{\epsilon}\sim p(\mathbf{\epsilon})$を使って次のように求められる。

\[\mathrm{tr}(A) = \mathrm{E}_{p(\mathbf{\epsilon})} \left [\mathbf{\epsilon}^T A \mathbf{\epsilon} \right ]\]

そしてベクトルとヤコビアン間の積はvjpを使うことでまとめて$\mathrm{O}(n)$時間で評価できる。$p(\epsilon)$としては正規分布や、ラデマッハ分布などを利用できる。

 このFFJORDは任意の変換を利用でき、尤度を正確に求められ、高速にサンプリングすることができる。これに対し、VAEは変分下限しか得られず厳密な尤度は計算できない。GANは尤度を計算できず、入力$\mathbf{x}$から潜在変数$\mathbf{z}$への変換は別の符号化器を学習する必要があり、それは容易ではない。正規化フローは尤度が求められ符号化も可逆変換なので容易だが、行列式が高速に求められる限定された変換しか利用できず表現力は限定的である。

 このFFJORDを使った生成モデルはGANに比べると学習が容易であり、尤度はVAEなど他の手法を超える性能が得られたと報告されている。

今後はライブラリでもサポートの見込み

本記事はロボットとAI技術の専門誌『日経Robotics』のデジタル版です
本記事はロボットとAI技術の専門誌『日経Robotics』のデジタル版です

 このNeural ODEやFFJORDの問題として、複雑なダイナミクスを学習した場合、ODEを解く際に多くの評価点が必要になってしまう問題がある。これに対しては次元を拡張した空間上でODEを解くことで、ダイナミクスを単純化し、評価点数を小さくできることが報告されている4)

 

 今はまだ既存のODEソルバーを利用しているが今後は特化したより効率的な手法が登場してくるだろう。また、今後はディープラーニングフレームワークやライブラリなどでこの連続ダイナミクスをサポートし、他の層と同じように簡単に使えるようになると考えられる。

1)R. T. Q. Chen, et al.,“Neural Ordinary Differential Equations,” NeurIPS 2018
2)岡野原、「Glow:可逆な生成モデル、GANより安定的に学習できる尤度ベースの手法」、『日経Robotics』、2018年9月号、pp.36-37.
3)W. Grathwohl, et al., “FFJORD: Free-Form Continuous Dynamics for Scalable Reversible Generative Models,” ICLR 2019
4)E. Dupont, et al., “Augmented Neural ODEs
岡野原 大輔(おかのはら・だいすけ)
Preferred Networks 取締役副社長
岡野原 大輔(おかのはら・だいすけ) 2006年にPreferred Infrastructureを共同創業。2010年、東京大学大学院博士課程修了。博士(情報理工学)。未踏ソフト創造事業スーパークリエータ認定。東京大学総長賞。