🔥

LTCsにおける順伝搬【解説】

に公開
AI
memory
rnn
ode
lnn
idea

はじめに

この記事では前回紹介した LTCs について、順伝搬をコンピュータ上ではどのように離散化して処理するのか、いくつかの手法を紹介します。LNN を構成するニューロンは LTCs により状態が記述されますが、外部からの入力を層の数だけ順伝搬をさせていく必要があります。このとき常微分方程式をそのまま解いて状態を求めるわけにはいかないですから ODE ソルバーを考えます。
多くの場合、この ODE ソルバーの効率が処速度やメモリ消費量などを決定づけます。

  • LTCs の基本式
    \frac{dx(t)}{dt} = -\left[\frac{1}{\tau} + f(x(t), I(t), t; \theta)\right]x(t) + f(x(t), I(t), t; \theta)A

前回の記事で紹介したように LTCs は Neural ODE に基づく常微分方程式であるため、ここでは簡単に Neural ODE での解法を考えてみます。

https://zenn.dev/yryromrk/articles/71f60cf8cd0efb

  • Neural ODE [Chen et al., 2018]
    \frac{dx(t)}{dt} = f(x(t), I(t), t; \theta)

そもそも初期値のある常微分方程式（ODE）の問題は、Neural ODE の場合であれば厳密解は以下のように表現されます。

\mathbf{x}(t_1) = \mathbf{x}(t_0) + \int_{t_0}^{t_1} f(\mathbf{x}(t), t, \theta) dt

しかし、このままではこの積分を計算することはできません。コンピュータ上で連続時間を扱うことは出来ないためです。そこでこの常微分方程式を解く手法として大きく数値的ソルバーと解析的ソルバーの二つの手法が存在します。

数値的ソルバー

数値的ソルバーは常微分方程式の連続を離散ステップにより近似的に離散化して扱います。この離散ステップの取り方にいくつかの手法が提案されています。

まず基本的解法となる Explicit Euler（陽解法）, Implicit Euler（陰解法）について紹介します。

Explicit Euler（陽解法）

陽解法とは次の状態 x_{i+1} を表現する式に x_{i+1} 自体が含まれない、単純な式で表現されたものです。計算コストは小さいですが、安定性には懸念があり、急激な変化を如実に反映してしまう傾向があります。

x_{i+1} = x_i + \Delta t \cdot f(x_i)

Implicit Euler（陰解法）

一方で陰解法は次の状態 x_{i+1} を表現する式に x_{i+1} も含み、自明でないような式変形や追加の数値計算を必要とするものです。陽解法に対して計算コストが大きい一方で、安定性は高くなります。

x_{i+1} = x_i + \Delta t \cdot f(x_{i+1})

この二種類の解法を合わせたものが、2020 年に Hasani らが提案したFused Eulerです。現在の LTCs の近似的解法にはこの手法を用います。
具体的には\frac{dx}{dt} = f(x) とすると、以下のように近似します。

x_{i+1} \approx x_i + \Delta t \cdot f(x_i, \text{ nonlinear part}) + \Delta t \cdot f(x_{i+1}, \text{ linear part})

Note:
線形なものは f(x_{i+1}) に置き換え、非線形なものは f(x_i) にすることが多いようです。これは、安定性に寄与する線形部分を陰解法的に、計算コストの大きい非線形部分を陽解法的に扱うことで、安定性と計算速度の両立を図るためです。

この解法を実際に LTCs (Liquid Time-constant Networks) に適用すると、次の式が得られます。

x_{i+1} = (1 - \Delta t \cdot \frac{1}{\tau_i}) x_i + \Delta t (\frac{1}{\tau_i} \odot f(x_i, \text{Input}_i))

活性化関数 f は任意の非線形関数をとることができます。入力系列の長さを L, 離散化ステップ数を T とすると、計算量は O(L \times T) となり、同じニューロン・セル数の LSTM と同程度の計算量です。

Runge-Kutta Method (ルンゲ＝クッタ法)

この手法は ODE の数値的近似手法としてよく用いられますが、LTCs に対しては不適切であるとされています。なぜなら、LTCs は硬い（stiff な）微分方程式であり、変化が非常に速い部分と遅い部分が混在していることから、指数関数的にステップ幅を小さくする必要がある RK 法には不適切であるためです (Press et al. 2007)。

ただし、積分の評価を改良することで、Euler 法よりも精度の良い計算を行うことができます。

2 次の Runge-Kutta 法

Euler 法ではステップの始点 t_i での傾きを使うのに対し、積分区間の中点 t_i + \frac{\Delta t}{2} で評価された傾きを用いることで精度の向上が期待できます。

まず、中点での y の値を Euler 法で近似します。

k_1 = \Delta t \cdot f(t_i, y_i)
k_2 = \Delta t \cdot f(t_i + \frac{\Delta t}{2}, y_i + \frac{k_1}{2})

この k_2 を使って y_{i+1} を計算します。

y_{i+1} = y_i + k_2

4 次の Runge-Kutta 法

より広く使われている高精度な手法です。4 つの評価点（始点、中点 2 つ、終点）での傾きを計算し、それらを重み付き平均して更新量を求めます。

k_1 = \Delta t \cdot f(t_i, y_i)
k_2 = \Delta t \cdot f(t_i + \frac{\Delta t}{2}, y_i + \frac{k_1}{2})
k_3 = \Delta t \cdot f(t_i + \frac{\Delta t}{2}, y_i + \frac{k_2}{2})
k_4 = \Delta t \cdot f(t_i + \Delta t, y_i + k_3)
y_{i+1} = y_i + \frac{1}{6}(k_1 + 2k_2 + 2k_3 + k_4)

解析的ソルバー

Closed-form Continuous-time Neural Networks（CfC）[R.Hasani et al., 2021]

ODE ソルバーを用いると、連続時間モデルにおける推論や訓練の速度が他のモデルと比較して遅くなるという課題があります。そこで、閉形式（Closed-form）の解を用いるアプローチが考えられました。閉形式というのは加減乗除や初等関数の合成関数による解の表し方です。つまり、複雑な微積分を線形的に表現しようというような解法になります。

具体的には、前述した LTCs の式をもとに、x(t) を以下のように閉形式で近似できます。

x(t) \approx (x(t_0) - \text{bias})e^{-t/\tau} \cdot \text{sigmoid}(\text{Input}(t) \cdot W) + \text{bias}

この方法では ODE ソルバーを必要としないため、従来の手法と比べて計算量が圧倒的に少なくなります。
ただし、この手法の欠点として閉形式が存在する方程式が限られており、少しの状態の変化で表せなくなる可能性が存在することと、本来の LTCs を用いる良さである微分方程式の説明可能性が失われるといったものが存在します。

ほかにも、モデルの計算を部分的に省略することで高速化を図る Sparse Flows や、初期値問題（IVP）に基づく解くべき微分方程式に合わせて専用のソルバーを自動生成する Hypersolvers といった、より効率的な ODE ソルバーの研究も進められています。

まとめ

LTC ニューロンで構成される LNN での順伝搬の際に用いられる常微分方程式の解法の紹介でした。これまで様々な手法が提案されていますが、結局精度と効率や説明可能性はトレードオフであるため、場合に応じて適用する解法を切り替えることが重要になるでしょう。
次回は逆伝搬について解説したいと思います！

参考文献

[1] T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud, "Neural ordinary differential equations," Advances in neural information processing systems, vol. 31, 2018.
[2] R. Hasani, M. Lechner, A. Amini, D. Rus, and R. Grosu, "Liquid time-constant networks," Proceedings of the AAAI Conference on Artificial Intelligence, vol. 35, no. 9, pp. 7657–7666, 2020.
[3] R. Hasani, M. Lechner, A. Amini, L. Liebenwein, A. Ray, M. Tschaikowski, G. Teschl, and D. Rus, "Closed-form continuous-time neural models," Nature Machine Intelligence, vol. 4, pp. 992–1003, 2021.

Discussion