今年 1 月に ALBERT に入社した清水です。深層学習プログラマとして自社プロダクト開発をしております。このブログを書くのは始めてなのですが、今日はちょっとプログラミング寄りの記事を。残暑厳しい折ですが、実装の詳細にまで立ち入りつつアツく Yuxin Wu および Kaiming He の考案した手法 Group Normalization（ECCV 2018 採択済）について語ります。Kaiming He 氏は ResNet を筆頭に優れた convolutional neural networks (CNN) の設計で知られていることもあり、みなさんも注目している手法ではないでしょうか？
Chainer v5.0.0b4 を用いて、Chainer に入っている実装と弊社の独自実装での速度比較、また Batch Normalization との速度・精度比較を行いました。その上で、速度差が生じる原因の調査や CUDA を使った高速化の工夫について詳しく記します。ソースコードは MIT ライセンスで GitHub に公開していますので、ぜひご覧ください。
Group Normalization って？
Group Normalization の発想はシンプルです。早速ですが、Group Normalization を端的に表す図（論文 Figure 2 より引用）を見てみましょう。
(N, C, HW) の 3 次元からなる多次元配列が出てきましたが、これは CNN の中間層の出力だと考えると想像しやすいかと思います。バッチサイズ N, チャンネル数 C, 画像の縦幅 H, 画像の横幅 W です。
図に示された Batch Normalization, Layer Normalization, Instance Normalization およびこの記事の本題 Group Normalization の 4 つの手法は、いずれも青色で示された領域ごとに算出される平均・分散によって入力を正規化 (normalize) します。Batch Normalization が各チャンネルごとに平均・分散を求めることは有名ですね。精度・収束性を劇的に向上させる Batch Normalization は今や CNN のデファクトスタンダードですが、しかし
- 画像の解像度が大きくてメモリが不足するなどの理由でバッチサイズが小さくなる場合に、平均・分散の推定が不安定になり学習ができなくなる
- 複数 GPU にまたがって平均・分散を推定することで実質的なバッチサイズを増やすことは可能だが、高価なハードウェアが必要になる上に実装や最適化が複雑
- ビデオの隣接フレームといった相関がある画像をミニバッチとして入力する場合も平均・分散の推定が不安定になる
- 学習時に平均・分散の移動平均を覚えておいて推論時に用いるといった処理が煩雑
- Finetune 時に移動平均をどう扱うべきかよくわからない
といった難点も併せ持っており、いつでも使えるわけではありません。このためミニバッチに依存しない正規化手法が待ち望まれていました。
そのような手法として、全チャンネルにまたがって平均・分散を取る Layer Normalization と、各チャンネル独立に画像の縦横方向についてのみ平均・分散を取る Instance Normalization が考案されました。とはいえ十分にバッチサイズが確保されている条件下での精度は Batch Normalization に比べてかなり劣っており、主流とはなっていません。
そこで登場したのが Group Normalization です。チャンネルを G 個にグルーピングして Layer Normalization と Instance Normalization の中間的な処理を行うことで、画像分類などのタスクで Batch Normalization に匹敵する精度を実現しました。グループ数 G を 32 にした場合がベストだったと論文に述べられていますが、それほど G の値に対して敏感ではないようです。Group Normalization の論文では、Instance Normalization には複数のチャンネルの組み合わせによって表現される特徴を歪めてしまう問題があると考察されています。その対極になるのが Layer Normalization ですが、こちらは大域的すぎて、特に重要ではないチャンネルが画像全体にわたって高い値を示した場合に他の全てのチャンネルが抑制されてしまいます。中間的な Group Normalization は良いとこ取りをできるというわけです。
なんだか素晴らしそうですね。Chainer では v5.0.0b3 から Group Normalization がサポートされているのでお手軽に使えます。しかし、本当に Batch Normalization をドロップインで置き換えて精度低下は起きないのでしょうか？ Batch Normalization と同等の速度で動作するのでしょうか？ この疑問を検証します。
結論から言えば、Batch Normalization を単に Group Normalization に置き換えるだけでは精度がかなり落ちてしまいました。なので Group Normalization を使う場合は精度の確認やパラメータチューニングをきちんとやるべきでしょう。また、Chainer v5.0.0b3 に追加された Group Normalization の実装はあまり効率的ではなく、CNN 全体の実行速度を大きく下げてしまうことがわかりました。この原因や、より効率のいい実装方法についても詳述します。