Adversarial variational Bayesを混合正規分布で試す

# はじめに

今回は、adversarial variational Bayesを使って、混合正規分布のパラメータ推定をおこないます。

adversarial variational Bayesについては、以下の文献を参考にしてください。
* [Lars Mescheder, Sebastian Nowozin, Andreas Geiger. Adversarial Variational Bayes: Unifying Variational Autoencoders and Generative Adversarial Networks](https://arxiv.org/abs/1701.04722). arXiv preprint arXiv:1701.04722, 2017.
* [Ferenc Huszár. Variational Inference using Implicit Distributions. arXiv:1702.08235](https://arxiv.org/abs/1702.08235). arXiv preprint arXiv:1702.08235, 2017.

# 観測データの生成

データの次元は簡単のために1次元とします。スカラ量です。また、クラスタ数は$K=3$とします。
つまり、一変数正規分布が3つ混合されているとします。

各クラスタの平均を$\mu\_k$とします。また、クラスタの混合比を$\theta\_k$とします。（$K=1,2,3$）
簡単のために、各クラスタの標準偏差は1としてしまいます。

課題は、$\mu\_k$と$\theta\_k$を、観測データ$x\_1, \ldots, x\_N$から推定することです。

観測データを生成する分布として、以下のような混合正規分布を考えます。
* $\mu\_k \sim \mathcal{N}(0, 20)$, for $k = 1,2,3$. （$\boldsymbol{\mu}$の事前分布）
* $z\_i \sim \mbox{Categorical}(\boldsymbol{\theta}=(\theta_1,\theta_2,\theta_3))$, for $i=1, \ldots, N$.
* $x\_i \sim \mathcal{N}(\mu_{z_i}, 1)$, for $i=1, \ldots, N$.

コードの中では、最初に$\mu\_1, \mu\_2, \mu\_3$を上記の正規分布$\mathcal{N}(0, 20)$からdrawしています。
また、$\theta\_1, \theta\_2, \theta\_3$のほうも、適当に乱数で設定しています。
こうしてコードの最初に設定された値が、正解の値となります。

今回使った乱数のシードでは、正解は
$\boldsymbol{\mu} = (-21.71261207, 19.94690893, 5.65956996)$
$\boldsymbol{\theta} = (0.33707583, 0.24218941, 0.42073475)$
と設定されました。
では、これらを推定する方法を考えましょう。

# 推定

例によってELBOを書いてみます。

```math
\log p(\boldsymbol{x}) 
\geq \int q(\boldsymbol{\mu}) \log \frac{p(\boldsymbol{\mu})}{q(\boldsymbol{\mu})} d\boldsymbol{\mu}
+ \sum_{\boldsymbol{z}} q(\boldsymbol{z}|\boldsymbol{x}) \log \frac{p(\boldsymbol{z})}{q(\boldsymbol{z}|\boldsymbol{x})}
+ \int q(\boldsymbol{\mu})
\sum_{\boldsymbol{z}} q(\boldsymbol{z}|\boldsymbol{x})
\log p(\boldsymbol{x}|\boldsymbol{z},\boldsymbol{\mu}) d\boldsymbol{\mu}
```

初めに述べておくと、今回の問題は普通の混合正規分布の推定ですので、
わざわざAVBを使う必要はありませんし、そもそもNNを使う必要すらありません。
しかし、あえてNNを使えるところには使ってみます。

## クラスタの混合比の推定
$\theta\_1,\theta\_2,\theta\_3$の推定、上の式で言えば$p(z=1), p(z=2), p(z=3)$の推定については、
上のELBOを最大化することでそのまま行うことにします。

## Responsibilityの推定
次に、各観測データ$x\_i$について、そのresponsibility $q(z\_i=k | x\_i)$を推定する必要があります。
今回は、AVBではなく、単純にVAE的な枠組みで推定します。
観測データ$x\_i$を入力とするNNの出力としてカテゴリカル分布のパラメータを得て、
そのカテゴリカル分布から$z\_i$をdrawするように$q(z\_i=k | x_i)$をモデル化します。

## クラスタの平均の推定
最後に、データ集合全体に対して$q(\boldsymbol{\mu})$を推定する必要があります。
今回はこちらをAVBで推定します。
分布$q(\boldsymbol{\mu})$は、データ集合全体に対応しています。
そこで、各クラスタに対応する一次元のembeddingを考え、
このembeddingとノイズとを入力とするNNから$\mu\_1,\mu\_2,\mu\_3$をdrawするよう、モデル化します。
そして、discriminatorは、この$q(\boldsymbol{\mu})$と事前分布$p(\boldsymbol{\mu})$との
密度比を求めるために用います。その際、事前分布$p(\boldsymbol{\mu})$については、
データの生成に使われた$\mathcal{N}(0, 20)$が既知と仮定し、ここからサンプルを得ます。


# 実装

```python:avb_gmm.py
import torch
import torch.nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
import numpy as np
import matplotlib.pyplot as plt
import os

if not os.path.exists('out/'):
    os.makedirs('out/')

torch.manual_seed(123)
np.random.seed(123)

mb_size = 100
eps_size = 100
z_dim = 1
embed_dim = 1
eps_dim = 2
X_dim = 1
h_dim = 8
cnt = 0
lr = 0.0001
sgd_lr = 0.01

def log(x):
    return torch.log(x + 1e-10)

weight_init_dict = {'kn':torch.nn.init.kaiming_normal,
                    'ku':torch.nn.init.kaiming_uniform,
                    'xn':torch.nn.init.xavier_normal,
                    'xu':torch.nn.init.xavier_uniform}
weight_init = 'xn'

# number of clusters
K = 3

# Encoder: q(mu|k, eps)
Q = torch.nn.Sequential(
    torch.nn.Linear(embed_dim + eps_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, z_dim)
)

for l in Q:
    if type(l) == torch.nn.ReLU:
        weight_init_dict[weight_init](prev_l.weight)
    prev_l = l

# Encoder: q(z|X)
Qz = torch.nn.Sequential(
    torch.nn.Linear(X_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, K),
    torch.nn.Softmax()
)

for l in Qz:
    if type(l) == torch.nn.ReLU:
        weight_init_dict[weight_init](prev_l.weight)
    prev_l = l

# Discriminator: T(z)
T = torch.nn.Sequential(
    torch.nn.Linear(z_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, h_dim),
    torch.nn.ReLU(),
    torch.nn.Linear(h_dim, 1)
)

for l in T:
    if type(l) == torch.nn.ReLU:
        weight_init_dict[weight_init](prev_l.weight)
    prev_l = l

def reset_grad():
    Q.zero_grad()
    Qz.zero_grad()
    T.zero_grad()

true_scale = 20.0
true_mu = np.random.normal(loc=0.0, scale=true_scale, size=(K))
print(true_mu)
true_cluster_prob = np.ones(K) + np.random.random(K) * 3
true_cluster_prob /= true_cluster_prob.sum()
print(true_cluster_prob)

def sample_X(size):
    z = np.random.choice(K, size, p=true_cluster_prob)
    X = np.random.normal(loc=true_mu[z], scale=1.0).astype(np.float32)
    X = Variable(torch.from_numpy(X))
    return X

cluster_prob = Variable(torch.randn(K), requires_grad=True)
cluster_embedding = Variable(torch.randn(K, embed_dim), requires_grad=True)

Q_solver = optim.Adam(list(Q.parameters()) + [cluster_embedding], lr=lr)
Qz_solver = optim.Adam(list(Qz.parameters()) + [cluster_prob], lr=lr)
T_solver = optim.SGD(T.parameters(), lr=sgd_lr, momentum=0.9)

for it in range(200000):

    # Discriminator
    eps = Variable(torch.randn(eps_size, K, eps_dim))
    mu = Variable(torch.randn(eps_size, K, z_dim) * true_scale)

    mu_sample = Q(torch.cat([cluster_embedding.unsqueeze(0).repeat(eps_size, 1, 1), eps], 2))
    T_q = F.sigmoid(T(mu_sample))
    T_prior = F.sigmoid(T(mu))

    T_loss = - torch.mean(log(T_q) + log(1.0 - T_prior))

    T_loss.backward()
    T_solver.step()
    reset_grad()

    # Encoder
    X = sample_X((mb_size, 1))
    eps = Variable(torch.randn(K, eps_dim))

    mu_sample = Q(torch.cat([cluster_embedding, eps], 1))
    resp = Qz(X)
    T_sample = T(mu_sample)
    disc = torch.mean(- T_sample)

    temp = torch.pow(X.expand(mb_size, K) - mu_sample.squeeze(), 2)
    temp = log(F.softmax(cluster_prob)) - log(resp) - 0.5 * temp
    loglike = torch.mean(torch.bmm(resp.view(mb_size, 1, K),
                                   temp.view(mb_size, K, 1)))

    elbo = - (disc + loglike)

    elbo.backward()
    Q_solver.step()
    Qz_solver.step()
    reset_grad()

    # Print and plot every now and then
    if it % 1000 == 0:
        print('Iter-{} ; ELBO: {:.4} ; T_loss: {:.4} ;'
              .format(it, -elbo.data[0], -T_loss.data[0]), end=' ')
        print('clst {}'.format(F.softmax(cluster_prob).data.squeeze().numpy()))

        eps = Variable(torch.randn(1000, K, eps_dim))
        mu_sample = Q(torch.cat([cluster_embedding.unsqueeze(0).repeat(1000, 1, 1), eps], 2))
        mu_sample = mu_sample.squeeze().data.numpy()

        plt.figure(figsize=(12, 4))
        for k in range(K):
            n, bins, patches = plt.hist(mu_sample[:, k], 50)
        plt.xlim(-25, 25)
        plt.ylim(0, 80)
        plt.savefig('out/mu_{}_{}.png'.format(weight_init, str(cnt).zfill(4)),
                    bbox_inches='tight')
        plt.clf()

        cnt += 1
```


# 実験結果

やはりNNの重みの初期化をどのようにおこなうかで結果が変わりました。
今回の課題については、kaimingもxavierも、normalでは全然うまくいきませんでした。

## うまくいった場合

一番うまくいったkaiming_uniformによる初期化の場合の結果を示します。

まず、クラスタの混合比$\theta\_1,\theta\_2,\theta\_3$の推定値が、
サイズ100のミニバッチを1,000個見るたびにどのように変化するか、グラフで示します。
![ku.png](https://qiita-image-store.s3.amazonaws.com/0/136528/a5713448-2065-ce41-8363-79937307b5dc.png)
縦軸が混合比で、点線で示されているのが3つのクラスタの混合比の正解です。
横軸の1がミニバッチ1,000個に対応します。
ミニバッチを170,000個程度見た付近（170の目盛付近）以降は、
かなり正解に近い値に落ち着いています。
微妙に振動している様子が興味深いですが、なぜこうなるのかは分かりません。

次に、クラスタの平均$\mu\_1,\mu\_2,\mu\_3$の推定値が、
サイズ100のミニバッチを1,000個見るたびにどのように変化するか、animated gifで示します。
上のグラフの横軸の1が、このanimated gifの一コマに対応しています。
今回は$q(\boldsymbol{\mu})$の推定をAVBで行っているので、
AVBがうまくいったのかどうかは、この結果を見ると分かります。
![mu_ku.gif](https://qiita-image-store.s3.amazonaws.com/0/136528/1fc47e03-eb14-13af-debe-3adda3109c58.gif)
正解は$\boldsymbol{\mu} = (-21.71261207, 19.94690893, 5.65956996)$で、
おおよそこれら3箇所にピークを持った値を$q(\boldsymbol{\mu})$が生成できていると分かります。

ELBOの変化とT_lossの変化も示しておきます。
横軸はクラス混合比のグラフと同じ意味です。
![ku_ELBO.png](https://qiita-image-store.s3.amazonaws.com/0/136528/685daf05-243a-16c0-5da2-6481099b50f2.png)
![ku_Tloss.png](https://qiita-image-store.s3.amazonaws.com/0/136528/80b6fb32-9f49-d505-04b9-2baf1dc87ec2.png)


## うまくいかなかった場合

一番うまくいかなかったkaiming_normalによる初期化の場合の結果を示します。
上のケースに対応する二つの図は、以下の通りです。
![kn.png](https://qiita-image-store.s3.amazonaws.com/0/136528/088d7e57-f563-0b99-ac8a-0d627e5e4ded.png)
![mu_kn.gif](https://qiita-image-store.s3.amazonaws.com/0/136528/1e39496f-a93b-8d95-2732-02e8bfc4f4eb.gif)
クラスタの混合比の推定が、全く正解に収束していないことが分かります。
また、クラスタの平均は、位置は合っています。ただ、分布がデルタ関数のようになっています。
デルタ関数のようになることと、クラスタの混合比の推定がうまくいかないことに、
関係があるのかどうかは、分かりません。

T_lossの変化を示しておきます。ある時点以降、全く変化しなくなっています。
![kn_Tloss.png](https://qiita-image-store.s3.amazonaws.com/0/136528/0fbb734a-f450-01f1-5135-4f86efa3ba5b.png)


# まとめ

懲りずにAVBを試してみました。
この記事を準備するなかで、過去の2つの記事の間違いを見つけました。
AVBについては、もう少し分かりやすい例で使ってみてから、自分の研究に使いたいと思います。


# 2017/11/27の追記
~~クラスタの平均~~Embeddingとクラスタの混合比をNNのパラメータと別に最適化するよう、変更してみました。
学習率も別の値にできるので調整したところ、0.001とNNのパラメータより大きく設定すると良いことが分かりました。

この改良の過程で、上のコードではNN以外のパラメータの勾配をゼロにリセットしていないことに気づきました（汗）。
すると、上の結果ではクラスタ混合比が振動していましたが、下図のようにきれいに収束するようになりました。
![kukuku.png](https://qiita-image-store.s3.amazonaws.com/0/136528/624f8518-2e42-089b-cadd-b9b7d0568b3b.png)
このように、早い段階で収束しています。奇妙な結果は疑ってかかるほうがいいですね・・・。
クラスタ平均の変化については、上の結果とほとんど同じような感じでした。
![mu_kukuku.gif](https://qiita-image-store.s3.amazonaws.com/0/136528/af41203e-c9ea-9f72-a884-7d2542380974.gif)

改良後のコードは私のgistに載せておきます。

* [Adversarial variational Bayes for univariate Gaussian mixture models](https://gist.github.com/tomonari-masada/48c10d01c98eddf27762dbf3a7a0b633)

なお、重みの初期化については、やはりkaiming_uniformが一番良さそうでした。上のコードは

```shell-session
% python avb_gmm.py ku ku ku
```
のように、Q, Qz, Tの重み初期化法を指定して動かせます。ku=kaiming_uniformで、これがベストだと思います。