rarilureloの日記

筋肉が学んだことを書きます. 機械学習とか.

脳のように非同期学習を行うニューラルネットワークの実装 with keras tensorflow backend

はじめに

タイトルには脳の非同期学習というようにまるで脳が非同期的に学習をしているかのように書きましたが, そこんところ実際はどうなっているかよくわかりません.
自転車を漕ぎながら考えごとをしたり, サッカーでドリブルしながらシュートかパスか考えたり, 少なくとも思考と運動の処理自体は並列非同期であるよう私自身は思います.
まぁよくわからないんですが, とりあえずニューラルネットワークでは非同期でも学習できたよ! というのが今回紹介する論文です.

Decoupled Neural Interfaces using Synthetic Gradients

目次

どんなことをしているか

まず言葉の意味から行くと, Decoupledとは分離という意味です.
なにを分離するのかというと層の依存関係を分離します.

ニューラルネットワークでは, 当たり前なのですが, 次の層に渡る前に前の層で値が計算される必要があります.
これをForward Lockingといいます.
順伝搬方向に並列化つまり分離して非同期に処理を行うことができないということです.
また, 層の更新も誤差の伝播を待つ必要があり, Update Lockingといいます.
誤差自体の計算も計算するには上の層での計算を待たなければならないのでこれをBackward Lockingといいます.

これら３つの依存関係のうち, Backward Lockingを分離したのが論文内でDNI(Decoupled Neural Interfaces)と呼ばれる層の結合関係です(実はForward Lockingを分離したモデルも提案している)

DNI

論文中にかなりざっくりとした近似式で書かれていますが, 勾配の計算を上の層全て使うのではなく現在の層が出した値で近似してしまえば上からの誤差なんて待たずにすむやんって感じです.
そして, 近似といえばニューラルネットワークです.
ニューラルネットワークを学習するための勾配を生成するニューラルネットワーク(論文内でM)を構築してやることでこの近似を成り立たせBackward Lockingを分離します.
ちなみにM自体の学習には１つ上のMが出した誤差を使って学習するので, １個上の層の計算を待つだけで更新ができます

この学習の過程はDeepMindのアニメーションがすごいわかりやすいです.

実験

実験は通常のFeedForwardネットワークに加え, rnn対して行っています.
特にrnnは通常であればtruncateする部分でMが出した誤差を使うことで収束性能の向上や, DNC(Differentiable Neural Computers)の前身であるNTM(Neural Turing Machine)で行われていた記号列の記憶タスクで良い結果を出しています.

特におもしろいのがMultiNetworkSystemに関する実験で, 一風変わったタスクを解いています.
流れてくる数字をrnnで構成されたネットワークAで受取り, T回に一回, 受け取った数字の内奇数の個数を出力する, かつ, ネットワークBに隠れ層を渡します.
渡されたBはT^2回に一回流れてきた数字の内「３」の個数を出力します.
従来の誤差逆伝搬法ではT^2回に一回しか誤差が計算されず, かといってAのタスクだけで更新するとBに必要な情報が取れなくなります.
そこでBに順伝搬する時にMを使うことで収束を速め, かつ精度を維持しています.
単純な入出力だけでなく異なるタスク間で通信を行うネットワーク構造は, 入力に対して考えを巡らせる人間のようでおもしろいです.

詳しい実験結果

以前輪読会で発表したときにまとめた資料があるのでそれを観ると上の説明ではよくわからなかった部分もわかるかもしれません, よりわからなくなるかもしれません.


追実験

正直DeepMindから出ているとはいえ, 実際の勾配を使わずに学習できるのか信じがたいので再実験をすることにしました.
一方向のネットワークを実装し, MNISTの分類タスクに適用しました.
https://github.com/rarilurelo/tensorflow-synthetic_gradient
このページにあるようにちゃんと学習できています(少し収束が遅いですが).

実装

勾配に手を加える実装は通常のフレームワークを使用していると難しいところです.
ほとんどのフレームワークは勾配計算を自動で行ってくれてしまうからです.

幸いtensorflowにはtf.gradientsという関数が準備されていて勾配の操作を行うことができます.
この関数は引数にobjective functionと勾配を計算したいパラメタを入力することでそのパラメタに対する勾配を自動で計算してくれます.
これに加え, 勾配の初期値をgrad_ys引数を使って設定することができます.


まずLayerを順伝搬の分とM(cDNI)の分で定義します.

# Layer1
layer1 = Sequential()
layer1.add(Dense(256, input_dim=784))
layer1.add(BatchNormalization(mode=2))
layer1.add(Activation('relu'))
# Layer2
layer2 = Sequential()
layer2.add(Dense(256, input_dim=256))
layer2.add(BatchNormalization(mode=2))
layer2.add(Activation('relu'))
# Layer3
layer3 = Sequential()
layer3.add(Dense(256, input_dim=256))
layer3.add(BatchNormalization(mode=2))
layer3.add(Activation('relu'))
# Layer4
layer4 = Sequential()
layer4.add(Dense(10, activation='softmax', input_dim=256))

# cDNI1 belongs to Layer2, so it accepts Layer1's output and emmits grad_y of Layer1
cDNI1 = Sequential()
cDNI1.add(Dense(1024, input_dim=256+10))
cDNI1.add(BatchNormalization(mode=2))
cDNI1.add(Activation('relu'))
cDNI1.add(Dense(1024))
cDNI1.add(BatchNormalization(mode=2))
cDNI1.add(Activation('relu'))
cDNI1.add(Dense(256, weights=[np.zeros(shape=[1024, 256])], bias=False))
# cDNI2 belongs Layer3
cDNI2 = Sequential()
cDNI2.add(Dense(1024, input_dim=256+10))
cDNI2.add(BatchNormalization(mode=2))
cDNI2.add(Activation('relu'))
cDNI2.add(Dense(1024))
cDNI2.add(BatchNormalization(mode=2))
cDNI2.add(Activation('relu'))
cDNI2.add(Dense(256, weights=[np.zeros(shape=[1024, 256])], bias=False))
# cDNI3 belongs Layer4
cDNI3 = Sequential()
cDNI3.add(Dense(1024, input_dim=256+10))
cDNI3.add(BatchNormalization(mode=2))
cDNI3.add(Activation('relu'))
cDNI3.add(Dense(1024))
cDNI3.add(BatchNormalization(mode=2))
cDNI3.add(Activation('relu'))
cDNI3.add(Dense(256, weights=[np.zeros(shape=[1024, 256])], bias=False))


Layer1の出した値をLayer2に伝搬して同時にMで勾配を算出します.

# layer1の算出した値y_l1をlayer2の入力x_l2に代入
x_l2 = y_l1

# layer2の順伝搬
y_l2 = layer2(x_l2)

# y_l2とlabelをlayer2のMに入れて勾配を算出
p_gy_l2 = cDNI2(K.concatenate((y_l2, labels), axis=1))

# layer1のMのためにlayer2のMが出した勾配をlayer1に逆伝搬
gy_l1 = tf.gradients(y_l2, y_l1, grad_ys=p_gy_l2)[0]

# layer1のMのlossfunction
loss_dni1 = K.mean(K.sum((p_gy_l1-gy_l1)**2, 1))

# layer1のMのパラメタへの勾配を計算
grad_trainable_weights_dni1 = tf.gradients(loss_dni1, cDNI1.trainable_weights)

# layer2のMが算出した勾配を使ってlayer2のパラメタへの勾配を計算
grad_trainable_weights_l2 = tf.gradients(y_l2, layer2.trainable_weights, grad_ys=p_gy_l2)

この操作を各層に渡って行って勾配を計算して更新します.

この時注意するのが全層での勾配の計算が終わってから層を更新する必要があるということです.
tensorflowでは計算グラフに対してtensorをfeedしていくのですが, その計算順序は固定されずノードに値が来たら計算するだけなのでlayerの更新が行われてから勾配計算が行われたりすると不都合です.
そのためtf.control_dependenciesを使って計算依存性をつけなければなりません.

with tf.control_dependencies(gparams):
    updates = optimizer.get_updates(params, gparams)

またtensorflowでは非同期連続でデータを与えることができません(QUEUEがどうなっているのかは知りません).
なので結局非同期更新はできないのですが, このコードで結果をシミュレーションできます.

実装の全体はhttps://github.com/rarilurelo/tensorflow-synthetic_gradientに公開しています.

まとめ

Neural Networkの非同期学習を実装しました.
こうして実際に動いているところをみると勾配をわざわざ微分して求めることはすこし大袈裟なのではないかとすら思えます.
神経科学的妥当性も気になるところです.

非同期にすることで学習を並列に動かし学習を劇的に速くできる可能性があります.
しかし現状のフレームワークでdata feedを並列非同期で行えるのはchainerを改造するぐらいしか思いつきません.(chainerはそもそもdata feedという概念がない)
マルチスレッドでlayerにGPUを割り当てたらできそうな気がするので時間があったら実装してみたいです.
今後Neural Networkの学習がどんな発展をするかまったく予測はできませんが, 前提とされている計算順序なども壊されてくるのならより柔軟な操作が必要になるのでしょう.