VAEの誤差関数にMSEを使うと学習がうまく進まない。

Question

やったこと

Chainerユーザーです。Chainerを使ってVAEを実装しました。参考にしたURLは

・Variational Autoencoder徹底解説
・AutoEncoder, VAE, CVAEの比較

・PyTorch＋Google ColabでVariational Auto Encoderをやってみた

などです。実装したコードのコアになる部分は以下の通りです。

class VAE(chainer.Chain):

    def __init__(self, n_in, n_latent, n_h, act_func=F.tanh):
        super(VAE, self).__init__()
        self.act_func = act_func
        with self.init_scope():
            # encoder
            self.le1        = L.Linear(n_in, n_h)
            self.le2        = L.Linear(n_h,  n_h)
            self.le3_mu     = L.Linear(n_h,  n_latent)
            self.le3_ln_var = L.Linear(n_h,  n_latent)

            # decoder
            self.ld1 = L.Linear(n_latent, n_h)
            self.ld2 = L.Linear(n_h,      n_h)
            self.ld3 = L.Linear(n_h,      n_in)

    def __call__(self, x, sigmoid=True):
        return self.decode(self.encode(x)[0], sigmoid)

    def encode(self, x):
        h1 = self.act_func(self.le1(x))
        h2 = self.act_func(self.le2(h1))
        mu = self.le3_mu(h2)
        ln_var = self.le3_ln_var(h2) 
        return mu, ln_var

    def decode(self, z, sigmoid=True):
        h1 = self.act_func(self.ld1(z))
        h2 = self.act_func(self.ld2(h1))
        h3 = self.ld3(h2)
        if sigmoid:
            return F.sigmoid(h3)
        else:
            return h3

    def get_loss_func(self, C=1.0, k=1):
        def lf(x):
            mu, ln_var = self.encode(x)
            batchsize = len(mu.data)
            # reconstruction error
            rec_loss = 0
            for l in six.moves.range(k):
                z = F.gaussian(mu, ln_var)
                z.name = "z"
                rec_loss += F.bernoulli_nll(x, self.decode(z, sigmoid=False)) / (k * batchsize)
            self.rec_loss = rec_loss
            self.rec_loss.name = "reconstruction error"
            self.latent_loss = C * gaussian_kl_divergence(mu, ln_var) / batchsize
            self..name = "latent loss"
            self.loss = self.rec_loss + self.latent_loss
            self.loss.name = "loss"
            return self.loss
        return lf

rec_lossは再構成誤差、すなわち入力と出力がどの程度等しいかを表していて、latent_lossの方は特徴量空間における分布が正規分布からどれくらいことなるを表す誤差だと認識しています。

MNISTで実験してみた結果、
１．lossが減少していく
２．再構成がちゃんと行われる（input画像が3なら、output画像も3になっている)
３．特徴量空間でランダムサンプリングを行った結果、ちゃんと数字が出力される。
などが確かめられました。

疑問

ところで、疑問なのですが、rec_lossは再構成誤差なので、素朴には平均二乗誤差をつかうのが自然だと思われます。そこでrec_lossの部分を

rec_loss += F.mean_squared_error(x, self.decode(z)) / k

と書き換え、ほかの条件は全部そのままで（他の部分は一切書き換えずに）、実験すると

1. rec_lossが2epoch目以降、全く減少しない。
2. 再構成が行われない。適当なinputを与えても、意味のないoutput画像が得られる。（ちなみにinput画像の種類によらずoutput画像は一定のようです）

などの結果が得られて、学習がうまくいっていっていないようです。

これはなぜでしょうか。

追記1

bernoulli_nllはデフォルトではすべて合計する一方で、mean_squared_errorは二乗誤差をバッチとピクセルの両方で平均するので、再構成項が過小評価されてしまっているかもしれないと考えました。

MNSITは28×28の画像なので、MSEを用いる際には28×28×F.mean_squared_error(x, decode(z))とすれば良いと思い、試してみましたが結果は変わりませんでした。

追記2

chainerでMSEを使ったVAEの実装を行っているコードを見つけました。
（https://github.com/maguro27/VAE-CIFAR10_chainer/blob/master/VAE_CIFAR10.ipynb）
なぜ、このコードでは動いて、私の上のコードでは学習がうまくいかないのでしょうか。

追記3

optimizerの問題ではないかとの指摘を受けて、Adam,AdaDelta,SGDで試しましたが結果は変わらず…

解決策？

passerbyさんに言われた通り、

rec_loss += F.mean_squared_error(x, self.decode(z)) / k 

を

rec_loss += F.mean(F.sum((x - self.decode(z)) ** 2, axis=1))

に変えるとうまくいきます。でもなぜだ…？

Comments

[tiitoi]

回答依頼をいただき申し訳ないですが、Chainer は使ったことがないので、実際に動かして検証することができず、回答はできないです。
ただ理論的には、reconstruction error は MSE でも問題ないはずです。
実際、Keras の VAE のサンプルコードでは、MSE が使われています。
https://github.com/keras-team/keras/blob/master/examples/variational_autoencoder.py

損失関数の形状が変わると適切な最適化手法というのも変わるので、いくつか最適化手法を変更してみたりはどうでしょうか (adam、SGD、adadelta など)

[physics303]

ありがとうございます。
私もpytorchでMSEを採用したVAEを確認しました。
https://gist.github.com/koshian2/2098e2261d673c818f6bdc51fa485e86
なので、特に何も工夫することなくMSEを使えるかと私も思っているのですが…

Adam, AdaDelta, SGDで試してみましたが結果に変化はありませんでした。

[passerby]

rec_loss += F.mean(F.sum((x - self.decode(z)) ** 2, axis=1))はどうですか？

[physics303]

コメントありがとうございます。やってみます。

[physics303]

MNISTの場合についてですが、うまくいきました！！

rec_loss += F.mean_squared_error(x, self.decode(z)) / k

がうまくいかずに、

rec_loss += F.mean(F.sum((x - self.decode(z)) ** 2, axis=1))

でうまくいくのはどういう理屈なのでしょう。どちらも同じことをやっている印象なのですが…

Answer 1

passerbyさんに言われた通り、

rec_loss += F.mean_squared_error(x, self.decode(z)) / k 

を

rec_loss += F.mean(F.sum((x - self.decode(z)) ** 2, axis=1))

に変えるとうまくいきます。でもなぜだ…？