RNN より高速な Feedforward Sequential Memory Networks (FSMN) を TensorFlow で実装してみた

LSTM や GRU など RNN の一般的なアーキテクチャの弱点としては、DNN や CNN に比べた場合に処理の遅さがあげられます。それは、アーキテクチャからは自明で、LSTM や GRU のような系列の記憶としての隠れ変数を使用する層の場合、あるステップの計算を行う時には事前のステップまでの計算を完了している必要があるため計算の並列化が難しい点、ステップ数が深くなればなるほど逆伝搬が深くなってしまいその分の計算コストが高くなる点があるため、GPU を使った場合も処理速度に限界があります。

Zhang 氏らが提案している Feedforward Sequential Memory Networks: A New Structure to Learn Long-term Dependency（FSMN） は、既存の RNN の構造とは異なり、メモリを使用することによりシーケンシャルな構造を考慮した学習を行う事ができます。

論文中では、スカラーFSMN、ベクトルFSMN、双方向FSMNなどのアーキテクチャを提案していますが、この記事では代表してスカラーFSMNを紹介したいと思います。FSMNの構造はシンプルで、FSMN 層に入力された隠れ変数 をメモリーに記憶しておき、そのメモリーに記憶されている値も考慮して次の層への出力を生成するという形をしています。

スカラーFSMNの更新式は、下記の2つで表すことができます。

[tex:{ \tilde{h}^{l}{t} = \sum^{N}{i} a^{l}{i} h^{l}{t-i} }]

[tex:{ h^{l+1}{t} = f(W^{l} h^{l}{t} + \tilde{W}^{l}{t} \tilde{h}^{l}{t} + b^{l}) }]

上記を見ての通りで、メモリブロックを考慮した値として [tex:{\tilde{h}^{l}{t}}] を生成しますが、スカラーFSMNの場合は[tex:{a^{l}{i}}]という係数を重みとしたメモリの和として定義します。あとはもとの入力とこの値を普通のニューラルネットの重み計算を行うことにより、FSMN層の出力とします。

順伝搬計算の場合、通常の RNN とは違い、事前のステップの入力が必要になりますがこの層内での計算結果を必要としないため、層単位で並列に計算することができます。 また逆伝搬の計算ですが、ステップ数が深くなっても学習に必要な逆伝搬が深くならず、さらに並列に計算することが可能となるため、通常の RNN と比べ計算を高速に処理することができるようになります。 一応論文中では、識別器としての性能も大差ないよと言っています。

というわけで、簡単に紹介したところで、TensorFlow で実装し既存のLSTMと比較してみましょう。比較実験のお題としては、TensorFlow の PTB サンプルを使用します。実験環境としては、g2.2xlarge を使用します。FSMN用のコードは、github にもアップしています。

実験条件としては、TensorFlow の PTB サンプルに付属している reader.py を使用して、バッチサイズを 20 、ステップ数を 50 として、エンベッド層、一層のFSMN層 もしくは LSTM層、出力層の三層構造でそれぞれの隠れユニットは 400 、エポック数は 10 とします。

	TIME	Validation PPL
LSTM	22m11.378s	122.462
FSMN	23m58.838s	148.335

結果を見る限りだと、PPL はともかく残念ながら速度が全然改善してませんね。論文によると4分の3くらいの速度になるようなので、僕の実装の問題なのでしょうが、TensorFlow との相性もあるんでしょうか。コード公開してもらえると再現実験楽でいいんですけどね。

2016/12/11 追記

コードの冗長性を除去して実験したところ、下記のように倍くらいの速度で処理をすることができました。PPLは残念ながら変わらずですが（変わったら変わったで何のこっちゃですが）、処理を見直すのってやはり大事ですね。修正済みのコードは、github にもアップしていますので良ければご活用ください。

	TIME	Validation PPL
FSMN	11m51.884s	148.740