注意機構・自由エネルギー原理・ニューラルネットの概念の獲得

モダン・ホップフィールドネットと正準ニューラルネットには，どちらも生物学的妥当性をもつ再帰型ニューラルネットとして提案された，という共通点がある．この投稿では，モダン・ホップフィールドネットと正準ニューラルネットにまたがる共通点を掘り下げていくことによって，注意機構，自由エネルギー原理，そしてニューラルネットにおける概念の獲得という一見異なる容貌をした３つの概念が，隠れた水路を通じてお互いにつながり合うことを明らかにしたい．

1. ホップフィールドネット

1−1. モダン・ホップフィールドネットと注意機構

モダン・ホップフィールドネットは，離散的な状態しか持たない古典的なホップフィールドネットを拡張したモデルで，連続的な状態とそれに対応した状態の更新式を備えている[1]．この新しいホップフィールドネットには，①多くのパターンを連想空間に保存する指数関数的な記憶容量を持ち，②1回の更新でパターンを取得し，③指数関数的に小さな検索誤差を持つ，という特徴がある．i 番目の可視ニューロンの状態を {v}_i(t) としたとき，古典的なホップフィールドネットからモダン・ホップフィールドネットにかけて，ニューロンの状態の更新式は次のように置き換えられる．

(1)   \begin{align*} \begin{gathered} v_i(t+\Delta t) = \operatorname{sign}\left(\sum_{j=1}^{N_f} \xi_{\mu j} v_j(t)\right) \\ \big\Downarrow \\ v_i(t+\Delta t) = \sum_{\mu=1}^{N_h} \xi_{i \mu} \operatorname{softmax}\left(\sum_{j=1}^{N_f} \xi_{\mu j} v_j(t)\right) \end{gathered} \end{align*}

(1)式 の後段の更新式が示すように，ニューロンが連続的な状態をとるように修正された新しい状態の更新式は，トランスフォーマーなどの深層学習モデルに搭載されている注意機構を表現することができる．つまり，モダン・ホップフィールドネットを使ってトランスフォーマーを作れることが明らかになった[2]．論文[1]，ブログ[2]のモダン・ホップフィールドネットの導出はかなり技巧的で人工的な印象があったが，ドミトリー・クロトフ＆ジョン・ホップフィールド「神経生物学と機械学習における大規模連想記憶問題 (2021)」では，あるラグランジアンからごく自然にモダン・ホップフィールドネットが導出できることが示された[3]．こちらのアプローチでは，i 番目の可視ニューロンの状態 {v}_i(t)，\mu 番目の記憶ニューロン（隠れニューロン）の状態 {h}_\mu(t) について

(2)   \begin{eqnarray*} & \displaystyle \left\{\begin{array}{l} \tau_f \dot{v}_i(t) = \sum_{\mu=1}^{N_h} \xi_{i \mu} f\left({h}_\mu(t)\right) - v_i(t) + I_i(t) \\ \tau_h \dot{h}_\mu(t) = \sum_{i=1}^{N_f} \xi_{\mu i} g\left({v}_i(t)\right) - h_\mu(t) \end{array}\right. \end{eqnarray*}

という状態の更新式を考える．この新しいモデルでは記憶ニューロンの活動は外部の観測者からは観測できないということを仮定するため，記憶ニューロンの自由度 {h}_\mu(t) は最終的に積分されてモデルから消去されてしまう．この一見迂遠な操作を経由することによって，ニューロン間の2体相互作用（2体シナプス結合）のみを用いて，生物学的な妥当性がないと言われている3つ以上のニューロン間の相互作用（3体シナプス結合）をモデルに取り込むことができるようになった．その帰結として，モダン・ホップフィールドネットはトランスフォーマーなどの現代的な（かつ非生物学的な）深層学習モデルを内包しながら，一定の生物学的な妥当性をもつ稀有なモデルになった．ここで (2)式 のうち，可視ニューロンについての状態の更新式に注目してみよう．

(3)   \begin{eqnarray*} & \displaystyle \tau_f \dot{v}_i(t) = \sum_{\mu=1}^{N_h} \xi_{i \mu} f\left({h}_\mu(t)\right) - v_i(t) + I_i(t) \end{eqnarray*}

記憶ニューロンのダイナミクスが可視ニューロンのそれと比べて非常に早い状況を仮定すると，これは数学的には \tau_h \rightarrow 0 を考えることに対応するので (2)式 から

(4)   \begin{eqnarray*} & \displaystyle {h}_\mu(t)　=　\sum_{j=1}^{N_f} \xi_{\mu j} g\left({v}_j(t)\right) \end{eqnarray*}

が導かれる．結局，可視ニューロンについての状態の更新式は

(5)   \begin{eqnarray*} & \displaystyle \tau_f \dot{v}_i(t) = \sum_{\mu=1}^{N_h} \xi_{i \mu} f\left(\sum_{j=1}^{N_f} \xi_{\mu j} g\left({v}_j(t)\right)\right) - v_i(t) + I_i(t) \end{eqnarray*}

となる．加えて，ここで \dot{v}_i(t) = \frac{v_i(t+\Delta t) - v_i(t)}{dt} の関係に注意すると

(6)   \begin{eqnarray*} & \displaystyle v_i(t+\Delta t) = v_i(t) + \frac{dt}{\tau_f}\left[ \sum_{\mu=1}^{N_h} \xi_{i \mu} f\left(\sum_{j=1}^{N_f} \xi_{\mu j} g\left({v}_j(t)\right)\right) - v_i(t) + I_i(t) \right] \end{eqnarray*}

となるので，d t=\tau_f のとき

(7)   \begin{eqnarray*} & \displaystyle v_i(t+\Delta t) = \sum_{\mu=1}^{N_h} \xi_{i \mu} f\left(\sum_{j=1}^{N_f} \xi_{\mu j} g\left({v}_j(t)\right)\right) + I_i(t) \end{eqnarray*}

となる．(7)式 は，特に，活性化関数 f(x)=\operatorname{softmax}(x)，g(x)=x かつ I_i(t) \rightarrow 0 のとき

(8)   \begin{eqnarray*} & \displaystyle v_i(t+\Delta t) = \sum_{\mu=1}^{N_h} \xi_{i \mu} \operatorname{softmax}\left(\sum_{j=1}^{N_f} \xi_{\mu j} {v}_j(t)\right) \end{eqnarray*}

に帰着する．この更新式はトランスフォーマーの核心とも言える注意機構の更新式と厳密に一致する[3]．

2. 正準ニューラルネット

2−1. 正準ニューラルネットの拡張

能動的推論を自然に導くニューラルネット，すなわち正準ニューラルネットは，推論（神経活動）と学習（シナプス可塑性）に関わるコスト関数が能動的推論の目的関数（変分ベイズ自由エネルギー）と厳密な対応関係を持つ，というある種の深遠さを感じさせる印象的な性質を備えている．正準ニューラルネットの状態の更新式は，磯村拓也・カール・フリストン「コスト関数の性質の解明のためのニューラルネットのリバースエンジニアリング (2020)」によって初めて科学の表舞台に引き揚げられた．正準ニューラルネットの中間層の可視ニューロンの状態を {x}_i(t)，出力層の可視ニューロンの状態を {y}_i(t) とすると，それぞれの状態の更新式は

(9)   \begin{eqnarray*} & \displaystyle \left\{\begin{array}{l} \dot{x}_i(t) = -\operatorname{sigmoid}^{-1}\left({x}_i(t)\right) + \sum_{j=1}^{N_x} k_{i j} x_j(t-\Delta t) + I_i(t) \\ \dot{y}_i(t) = -\operatorname{sigmoid}^{-1}\left({y}_i(t)\right) + \sum_{j=1}^{N_x} v_{i j} x_j(t-\Delta t) \end{array}\right. \\ \end{eqnarray*}

であらわせる[4]．この更新式の平衡状態，すなわち \dot{x}_i(t)=0 かつ \dot{y}_i(t)=0 の場合を考えると

(10)   \begin{eqnarray*} & \displaystyle \left\{\begin{array}{l} x_i(t) = \operatorname{sigmoid}\left(\sum_{j=1}^{N_x} k_{i j} x_j(t-\Delta t) + I_i(t)\right) \\ y_i(t) = \operatorname{sigmoid}\left(\sum_{j=1}^{N_x} v_{i j} x_j(t-\Delta t)\right) \end{array}\right. \end{eqnarray*}

となる．さらに重要な性質として，論文[5]の中で，（完全クラス定理の帰結として）正準ニューラルネットのコスト関数を最小化するプロセスが（ある事前分布の下で）変分ベイズ推論を実行することと等価であること，が示された．以下では，活性化関数を \operatorname{sigmoid} 関数に固定せず，より一般の非線形関数 f(x)，g(x) とした場合について考察する．

(11)   \begin{eqnarray*} & \displaystyle \left\{\begin{array}{l} \dot{x}_i(t) = -f^{-1}\left({x}_i(t)\right) + \sum_{j=1}^{N_x} k_{i j} x_j(t-\Delta t) + I_i(t) \\ \dot{y}_i(t) = -g^{-1}\left({y}_i(t)\right) + \sum_{j=1}^{N_x} v_{i j} x_j(t-\Delta t) \end{array}\right. \end{eqnarray*}

この更新式についても平衡状態，すなわち \dot{x}_i(t)=0 かつ \dot{y}_i(t)=0 の場合を考えると

(12)   \begin{eqnarray*} & \displaystyle \left\{\begin{array}{l} x_i(t) = f\left(\sum_{j=1}^{N_x} k_{i j} x_j(t-\Delta t) + I_i(t)\right) \\ y_i(t) = g\left(\sum_{j=1}^{N_x} v_{i j} x_j(t-\Delta t)\right) \end{array}\right. \end{eqnarray*}

となる．さらに x_i(t-\Delta t)=x_i(t)-\dot{x}_i(t) \Delta t+\mathcal{O}\left(\Delta t^2\right) の関係があることに注意すると

(13)   \begin{eqnarray*} & \displaystyle \left\{\begin{array}{l} x_i(t) = f\left(\sum_{j=1}^{N_x} k_{i j} x_j(t-\Delta t) + I_i(t)\right) \\ y_i(t) = g\left(\sum_{j=1}^{N_x} v_{i j} x_j(t)\right) \end{array}\right. \end{eqnarray*}

と書き直せる．結局，この更新式は活性化関数 f(x)=\operatorname{softmax}(x)，g(x)=x かつ v_{i j}=k_{i j}，I_i(t) \rightarrow 0 のとき

(14)   \begin{eqnarray*} & \displaystyle y_i(t+\Delta t) = \sum_{\mu=1}^{N_x} k_{i \mu} \operatorname{softmax}\left(\sum_{j=1}^{N_x} k_{\mu j} x_j(t)\right) \end{eqnarray*}

となる．こうして正準ニューラルネットでも注意機構の更新式を構成できることが示された．ここで，この正準ニューラルネットの変種についてコスト関数を導出しておこう．状態の更新式とコスト関数の間には以下の関係

(15)   \begin{align*} \displaystyle \dot{x}_i(t) &= -\frac{d}{d t}\frac{\partial L_{x_i}}{\partial {x}_i} \\ \dot{y}_i(t) &= -\frac{d}{d t}\frac{\partial L_{y_i}}{\partial {y}_i} \end{align*}

があるので，中間層，出力層，それぞれのニューロンについてのコスト関数は一般に

(16)   \begin{align*} \displaystyle L_{x_i} &= \int_{0}^{t} \left\{ \left[\int f^{-1}\left({x}_i(\tau)\right) d{x}_i\right] - x_i(\tau) \sum_{j=1}^{N_x} k_{i j} x_j(\tau-\Delta \tau) - x_i(\tau)I(\tau) \right\} d\tau & \\ L_{y_i} &= \int_{0}^{t} \left\{ \left[\int g^{-1}\left({y}_i(\tau)\right) d{y}_i\right] - y_i(\tau) \sum_{j=1}^{N_x} v_{i j} x_j(\tau-\Delta \tau) \right\} d\tau & \end{align*}

と書ける．論文[4]，論文[5]で検討された正準ニューラルネット (9)式（活性化関数 f(x)=\operatorname{sigmoid}(x)，g(x)=\operatorname{sigmoid}(x)）のコスト関数は

(17)   \begin{align*} \displaystyle \left[\int f^{-1}\left({x}_i(\tau)\right) d{x}_i\right] &= x_i(t)\ln{x_i(t)} + \left(1-x_i(t)\right)\ln{\left(1-x_i(t)\right)} & \\ \left[\int g^{-1}\left({y}_i(\tau)\right) d{y}_i\right] &= y_i(t)\ln{y_i(t)} + \left(1-y_i(t)\right)\ln{\left(1-y_i(t)\right)} & \end{align*}

となり，本稿で検討する注意機構を導く正準ニューラルネットの変種 (11)式（活性化関数 f(x)=\operatorname{softmax}(x)，g(x)=x）のコスト関数は

(18)   \begin{align*} \displaystyle \left[\int f^{-1}\left({x}_i(\tau)\right) d{x}_i\right] &= x_i(t)\left(\ln{x_i(t)} - 1\right) + x_i(t)\ln{Z} & \\ \left[\int g^{-1}\left({y}_i(\tau)\right) d{y}_i\right] &= \frac{1}{2} {y}_i^{2}(t) & \end{align*}

となる．ただし，Z は \operatorname{softmax}(x)\equiv\frac{e^{x}}{Z} の分配関数 Z=\int e^{x} dx を表している．

2−2. 正準ニューラルネットと注意機構

2-1. で求めた注意機構の更新式 (14)式 が，トランスフォーマーに搭載されている注意機構

(19)   \begin{eqnarray*} & \displaystyle \operatorname{Attention}\left(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\right) = \operatorname{softmax}\left(\boldsymbol{Q}\boldsymbol {K}^{\top}\right) \boldsymbol{V} \end{eqnarray*}

を再現することを確かめよう．正準ニューラルネットの中間層と出力層のニューロンの数は等しい（N_{x}=N_{y}）という条件のもとで

(20)   \begin{align*} \displaystyle \boldsymbol{Q}^{\mathrm{new}} &= \left(y_1(t+\Delta t), \cdots, y_{N_{y}}(t+\Delta t)\right) & \\ \boldsymbol{Q} &= \left(x_1(t), \cdots, x_{N_{x}}(t)\right) & \\ \boldsymbol{K} &= \left[k_{i j}\right] & \\ \boldsymbol{V} &= \boldsymbol{K} \boldsymbol{W}_V & \end{align*}

と置くと，(14)式 は次のように書き直せる．

(21)   \begin{align*} \displaystyle \begin{gathered} y_i(t+\Delta t) = \sum_{\mu=1}^{N_x} k_{i \mu} \operatorname{softmax}\left(\sum_{j=1}^{N_x} k_{\mu j} x_j(t)\right) \\ \big\Downarrow \\ \boldsymbol{Q}^{\text {new }} = \operatorname{softmax}\left(\boldsymbol{Q} \boldsymbol{K}^{\top}\right) \boldsymbol{K} \end{gathered} \end{align*}

この (21)式 の意味を，一部が欠損した『ザ・シンプソンズ』の登場人物の画像から，元の画像はどの登場人物だったのかを当てるタスク[2]を例にとって説明してみよう．まず，正準ニューラルネットへの感覚入力が，中間層のニューロンの状態を変化させ，クエリ行列 \boldsymbol{Q} を作る（図中では，一部が欠損した『ザ・シンプソンズ』の登場人物の画像に対応する）．そして，過去の学習で記憶した「それぞれの訓練データに対応するニューロンの状態のパターン」はキー行列 \boldsymbol{K} に記録される（図中では，『ザ・シンプソンズ』の登場人物の画像を並べた行列に対応する）．クエリ行列 \boldsymbol{Q} とキー行列 \boldsymbol{K} の積をとり，さらにその値を活性化関数 \operatorname{softmax}(x) によって全体の和が1になるように非線形変換することによって，正準ニューラルネットへの感覚入力と過去の学習で記憶したそれぞれの訓練データとの類似度 \operatorname{softmax}\left(\boldsymbol{Q} \boldsymbol{K}^{\top}\right) が評価される．最後に，この類似度と訓練データを記憶したキー行列に重みをつけたバリュー行列 \boldsymbol{V}=\boldsymbol{K} \boldsymbol{W}_V の積をとることで，正準ニューラルネットが過去に学習した訓練データのうち，どの訓練データに注意を向けるべきか，という情報 \operatorname{Attention}\left(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\right) が出力される．

したがって，\operatorname{Attention}\left(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\right) = \boldsymbol{Q}^{\mathrm{new}} \boldsymbol{W}_V とすれば (21)式 より

(22)   \begin{align*} \displaystyle \operatorname{Attention}\left(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}\right) &= \boldsymbol{Q}^{\mathrm{new}} \boldsymbol{W}_V & \\ &= \operatorname{softmax}\left(\boldsymbol{Q}\boldsymbol {K}^{\top}\right) \boldsymbol{V} & \end{align*}

が得られる．こうして正準ニューラルネットでも注意機構を実装できることが確かめられた．

3. 注意機構と自由エネルギー原理

3−1. 自由エネルギー原理に従うトランスフォーマー

自由エネルギー原理に従う正準ニューラルネット (9)式 の活性化関数はそれぞれ以下のようになっていた．

(23)   \begin{eqnarray*} & \displaystyle \left\{\begin{array}{l} f(x_f)=\operatorname{sigmoid}(x_f)=\frac{1}{1+e^{-x_f}} \\ g(x_f)=\operatorname{sigmoid}(x_f)=\frac{1}{1+e^{-x_f}} \\ \end{array}\right. \end{eqnarray*}

自己注意機構を搭載した正準ニューラルネット (11)式 の活性化関数はそれぞれ以下のようになっていた．

(24)   \begin{eqnarray*} & \displaystyle \left\{\begin{array}{l} f(x_a)=\operatorname{softmax}(x_a)=\frac{e^{x_a}}{Z} \\ g(x_a)=x_a \\ \end{array}\right. \end{eqnarray*}

これらの活性化関数は次の変数変換によって互いに移り変わる．したがって，正準ニューラルネットとトランスフォーマーはある特別な”座標変換”のもとで等価なモデルとみなすことができる．

(25)   \begin{align*} \displaystyle f(x):\ x_f &=\ln\frac{e^{x_a}}{Z-e^{x_a}}& \\ g(x):\ x_f &=\ln\frac{x_a}{1-x_a}& \end{align*}

参考文献

  1. Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlović, Geir Kjetil Sandve, Victor Greiff, David Kreil, Michael Kopp, Günter Klambauer, Johannes Brandstetter and Sepp Hochreiter; Hopfield Networks is All You Need. International Conference on Learning Representations (ICLR) 2021.
  2. Hopfield Networks is All You Need, Blog post.
  3. Dmitry Krotov, John Hopfield; Large Associative Memory Problem in Neurobiology and Machine Learning. International Conference on Learning Representations (ICLR) 2021.
  4. Takuya Isomura, Hideaki Shimazaki and Karl Friston; Canonical neural networks perform active inference. Commun Biol 2022; 5, 55.
  5. Takuya Isomura, Karl Friston; Reverse-Engineering Neural Networks to Characterize Their Cost Functions. Neural Comput 2020; 32 (11): 2085–2121.

コメントを残す