いものやま。

雑多な知識の寄せ集め

強化学習用のニューラルネットワークをSwiftで書いてみた。（その4）

技術 AI 強化学習 ニューラルネットワーク HME Swift

昨日は強化学習用のニューラルネットワークの実装を行った。

今日はHMEの実装を行うために、同様にHMEの計算を行列で表現していく。

強化学習用のHMEの計算

ここでは、HMEへの入力を \boldsymbol{x} \in \mathbb{R}^{n}、出力を y \in \mathbb{R} とする。
また、エキスパートネットワーク i の出力を y_i \in \mathbb{R}、パラメータを \boldsymbol{\theta}_i とし、それに対応するゲートネットワークの出力を g_i \in [0, 1]、重みを \boldsymbol{w}_i \in \mathbb{R}^{n} とする。

行列で表現しない計算

ニューラルネットワークのときと同様に、まずは行列で表現しない計算から。

  1. エキスパートネットワークの出力 y_i とその出力のパラメータに関する勾配 \frac{\partial y_i}{\partial \boldsymbol{\theta}_i} を求める。
  2. ゲートネットワークの出力 g_i とHMEの出力の重みに関する勾配 \frac{\partial y}{\partial \boldsymbol{w}_i} を求める：
    1. s_i = \boldsymbol{w}_i^{\mathrm{T}} \boldsymbol{x}
    2. g_i = \frac{\mathrm{exp}(s_i)}{\sum_j \mathrm{exp}(s_j)}
    3. \frac{\partial y}{\partial \boldsymbol{w}_i} = \sum_j y_j \cdot \frac{\partial g_j}{\partial s_i} \cdot \boldsymbol{x}
      ただし、 i = j ならば \frac{\partial g_j}{\partial s_i} = g_j (1 - g_j)、 i \neq j ならば \frac{\partial g_j}{\partial s_i} = - g_j^{2}。
  3. HMEの出力 y を求める：
    1. y = \sum_i g_i y_i
  4. HMEの出力のエキスパートネットワークのパラメータに関する勾配 \frac{\partial y}{\partial \boldsymbol{\theta}_i} を求める：
    1. \frac{\partial y}{\partial \boldsymbol{\theta}_i} = g_i \frac{\partial y_i}{\partial \boldsymbol{\theta}_i}

行列で表現した計算

まず、エキスパートネットワークの各出力をまとめて、ベクトル \boldsymbol{y} \in \mathbb{R}^{m}（ただし、 m はエキスパートネットワークの数）で表現することにする。

そして、ゲートネットワークの重みを W \in \mathbb{R}^{m \times n} で次のように表すことにする：

 {
W = \left( \begin{array}{c}
\boldsymbol{w}_1 {}^{\mathrm{T}} \\
\vdots \\
\boldsymbol{w}_m {}^{\mathrm{T}}
\end{array} \right)
}

さらに、次のような行列 G \in \mathbb{R}^{m \times m} を考える：

 {
\begin{align}
G &= \left( \begin{array}{ccc}
\frac{\partial g_1}{\partial s_1} & \cdots & \frac{\partial g_m}{\partial s_1} \\
\vdots & \ddots & \vdots \\
\frac{\partial g_1}{\partial s_m} & \cdots & \frac{\partial g_m}{\partial s_m}
\end{array} \right) \\
&= \left( \begin{array}{ccc}
g_1 (1 - g_1) & \cdots & - g_m^2 \\
\vdots & \ddots & \vdots \\
- g_1^2 & \cdots & g_m (1 - g_m)
\end{array} \right) \\
&= (\boldsymbol{g} \otimes \boldsymbol{1}) \odot (I - \boldsymbol{g} \otimes \boldsymbol{1})
\end{align}
}

ただし、 I \in \mathbb{R}^{m \times m} は単位行列、 \boldsymbol{1} \in \mathbb{R}^{m} はすべての要素が1のベクトル、 \otimes はベクトルの直積（外積）、 \odot はアダマール積（要素ごとの積）。

すると、上の計算は次のように書き直すことが出来る：

  1. エキスパートネットワークの出力 \boldsymbol{y} とその出力のパラメータに関する勾配 \frac{\partial y_i}{\partial \boldsymbol{\theta}_i} を求める。
  2. ゲートネットワークの出力 \boldsymbol{g} とHMEの出力の重みに関する勾配 \frac{\partial y}{\partial W} を求める：
    1. \boldsymbol{s} = W \boldsymbol{x}
    2. \boldsymbol{t} = \mathrm{exp}(\boldsymbol{s})
    3. \boldsymbol{g} = \frac{1}{\|\boldsymbol{t}\|_1} \boldsymbol{t} （※L1ノルムで正規化）
    4. \frac{\partial y}{\partial W} = (G \boldsymbol{y}) \otimes \boldsymbol{x}
  3. HMEの出力 y を求める：
    1. y = \boldsymbol{g}^{\mathrm{T}} \boldsymbol{y}
  4. HMEの出力のエキスパートネットワークのパラメータに関する勾配 \frac{\partial y}{\partial \boldsymbol{\theta}_i} を求める：
    1. \frac{\partial y}{\partial \boldsymbol{\theta}_i} = g_i \frac{\partial y_i}{\partial \boldsymbol{\theta}_i}

かなりトリッキーな感じもするけど、地道に追えば、元の計算と同等になっていることが分かると思う。

なお、途中、 \boldsymbol{t} をL1ノルムで正規化しているけど、これは指数関数の値域が (0, \infty) であることから、 \| \boldsymbol{t} \|_1 = \sum_i \left|t_i\right| = \sum_i t_i = \sum_i \mathrm{exp}(s_i) となることを利用している。

あと、実際に実装する場合には、 \boldsymbol{s} の一番大きい値を見つけて、その値を \boldsymbol{s} のすべての要素から一律に引く必要がある。
というのも、指数関数は急激に大きくなるので、下手するとNaNになってしまって計算が破綻することがあるから。
ソフトマックス関数は各要素から定数を引いても値は変わらないという性質があるので、一番大きい値を一律に引いておくことで、指数関数の値が (0, 1] に収まるようになる。
（強化学習とニューラルネットワークを組合せてみた。（その12） - いものやま。も参照）

今日はここまで！