VAEにはbatch normalizationを入れない

VAEを実装してみようとして調べていた際、batch normalizationは入れなくてよいのかということを気になったのでちょっと調べてみた。

とりあえずpytorch examplesのVAEの実装には入っていない

examples/main.py at 6c51ca5a614cfdbdcd4e8c3e70321c5f6defb177 · pytorch/examples · GitHub

ググってみるとこちらを見つけた

neural networks - Distorted validation loss when using batch normalization in convolutional autoencoder - Cross Validated

元はtensorflowのチュートリアルのようで以下のように書かれている

it's common practice to avoid using batch normalization when training VAEs, since the additional stochasticity due to using mini-batches may aggravate instability on top of the stochasticity from sampling.

どうやらVAEにはbatch normalizationを入れないのが通例の模様。
"VAEの潜在変数のサンプリング時のノイズ"と"mini batchに対してbatch normalizationをかけることによる入力データの変化"が合わさると不安定になってしまう、ということらしい。
(batch normalizationはmini batch全体の平均・分散によって入力データを補正するため、同じ入力データでもepoch毎に異なる値になってしまうことを指していると考えられる)

また、論文も少し探してみたところ関連しそうなものとして以下を見つけた。

https://arxiv.org/pdf/1602.02282.pdf

これは複数の潜在変数を導入したVAEの提案のようで、潜在変数の層を最大で5層まで作って実験している。 提案手法を理解するまでは読んでないものの↓のような実験結果が載っていた。

論文より引用 f:id:y-kamiya:20191204225656p:plain

MNISTのデータセットに対していくつかのパターンで実験した際のlossの値を比較している。
importance samples...と書いてある3つ目のグラフはおいておくとして、左のグラフのVAEとVAE(BN)を見ると層数が2まではほぼ同じ値となっている。層がそれより多くなるとBNを入れた方が値がよくなっている。

ということで普通のVAEの場合はBNは入れなくても良さそうに見える。

ちなみに上記の論文について詳しく書いてあるスライドも見つけた

(DL hacks輪読) How to Train Deep Variational Autoencoders and Probabili…