【技術】合成データアルゴリズム VAE

技術

2023.05.08

URL copied

Acompany無料メルマガに
登録しませんか？

プライバシーテックに関する情報を配信！

はじめに

Deep Learningベースの生成モデルにおいて代表的なアルゴリズムといえば、Generative Adversarial Networks(GAN) [1] 、Variational Auto-Encoder(VAE)[2]、Flow[3] などがありますがその中で VAE について解説します。

VAE論文[2]ではSGVB(Stochastic Gradient Variational Bayes)という推定量を導入して、それを用いた認識モデルの最適化手法としてAEVB(Auto-Encoding Variational Bayes)を提案しています。その応用としてVAE(Variational Auto-Encoder)を例示しています。そのためまずはSGVB推定量の解説をします。オートエンコーダーとの関係についても適宜補足します。

潜在変数モデル

SGVB推定量を定義するにあたって潜在変数モデルの説明からはじめます。この記事で潜在変数モデルとは、観測可能変数を表す確率変数$x$と潜在変数を表す確率変数$z$およびそれらの同時確率分布$p(x, z)$からなる確率モデルのこととします。またパラメータ$\theta$を持つ場合は$p(x, z) = p_{\theta}(x, z)$と書くことにします。

潜在変数モデルでは観測データに最も適合するパラメータ$\theta$を求めることや、事後分布$p_{\theta}(z | x)$を近似する確率分布$q_{\varphi}(z | x)$を求めるのが応用上重要な関心事になります。

💡 オートエンコーダとの関係１

オートエンコーダの観点からは観測可能変数$x$を入出力データだと見なします。例えば画像のオートエンコーダを考えたい場合は$x$は画像データになります。そして潜在変数$z$は入出力データを変換して得られるコードだと見なします。

上記のように見ると事後分布を近似する$q_{\varphi}(z | x)$は確率的なエンコーダのモデルと見なせます。なぜなら入力データ$x$が与えられるとコード$z$の確率分布が得られるためです。同様にコード$z$から出力データ$x$の確率分布が得られるため、$p_{\theta}(x | z)$は確率的なデコーダのモデルだと見なせます。

通常のオートエンコーダと違う点はエンコーダ・デコーダの出力が決まったひとつの値ではなく確率分布になることのみです。VAEの実装でエンコーダ・デコーダの出力をなぜ確率的にするのかと疑問に思ったことがある方もいるかもしませんが、裏側には以下で説明するような理論的背景があります。

SGVBは以下のような状況を想定しています。

周辺分布 $p_{\theta}(x) = \int p_{\theta}(x, z) dz$の計算が困難である。したがって事後分布 $p_{\theta}(z | x) = p_{\theta}(x, z) / p_{\theta}(x)$などの計算も困難である。
データ量が多くてバッチ最適化のコストがかかりすぎる。

上記のような状況ではEMアルゴリズムや平均場近似の変分ベイズ法などの、古典的な確率モデルのパラメータ最適化手法は使えません。SGVB推定量はこのような場合にも効率よく事後分布の推定をおこなうための道具だと言えます。

SGVB推定量とAEVBアルゴリズム

観測データ$X = \{x_1, \cdots, x_N\}$が得られたとします。我々の目的は観測データと最も適合するパラメータを求めることです。モデルの適合度を測る指標として尤度という量があります。尤度は

$$ \begin{align} p_{\theta}(X) = \prod_{i = 1}^N p_{\theta}(x_i) \end{align} $$

によって定義される量です。またこれの対数を取った対数尤度

$$ \begin{align} \log p_{\theta}(X) = \sum_{i = 1}^N \log p_{\theta}(x_i) \end{align} $$

もよく利用されます。しかし周辺分布の計算が困難な場合には尤度を直接計算することはできません。そこで潜在変数モデルでよく用いられるEvidence Lower Boundという量を導入します。

📖 Evidence Lower Bound

確率変数$(x,z)$の同時分布を$p(x, z)$をする。また$p(x) = \int p(x, z) dz$ をその周辺分布とする。このとき任意の確率分布$q(z)$について

$$ \begin{align*} &D_{KL}(q(z) || p(z | x)) \\ &= \int q(z) \log \frac{q(z)}{p(z | x)} dz \\ &= \int q(z) \log \frac{q(z)p(x)}{p(x, z)} dz \\ &= \int q(z)(\log q(z) + \log p(x) – \log p(x, z))dz \\ &= \log p(x) -\mathbb{E}_{q(z)}[\log p(x, z) – \log q(z)] \end{align*} $$

が成立する。ただし$D_{KL}$はカルバック・ライブラー情報量である。$D_{KL}$は$0$以上なので

$$ \begin{align*} \log p(x) \geq \mathbb{E}_{q(z)}[\log p(x, z) – \log q(z)] \end{align*} $$

が成立する。上式の右辺をELBO(Evidence Lower Bound)という。

ELBOは対数尤度の下限なのでELBOが大きいほどモデルによく適合していると見なされます。したがって以下では直接尤度の評価をする代わりに ELBO を最大化することを考えます。ここで$\mathcal{L}(\theta, \varphi; x)$を確率分布 $p_{\theta}(x)$の$q(z) = q_{\varphi}(z|x)$のときのELBOとして定義します。つまり

$$ \begin{align} \mathcal{L}(\theta, \varphi; x) &= \mathbb{E}_{q_{\varphi}(z | x)}[\log p_{\theta}(x, z) – \log q_{\varphi}(z | x)] \end{align} $$

となります。また観測データ$X=\{x_n\}_{n = 1}^N$に対して

$$ \begin{align} \mathcal{L}(\theta, \varphi; X) &= \sum_{n = 1}^N \mathcal{L}(\theta, \varphi; x_n) \end{align} $$

と定めます。(3)式を少し変形すると

$$ \begin{align*} \mathcal{L}(\theta, \varphi; x) &= \mathbb{E}_{q_{\varphi}(z | x)}[\log p_{\theta}(x, z) – \log q_{\varphi}(z | x)] \\ &= \mathbb{E}_{q_{\varphi}(z | x)}[\log p_{\theta}(x | z) + \log p_{\theta}(z) – \log q_{\varphi}(z | x)] \\ &= \mathbb{E}_{q_{\varphi}(z | x)}[\log p_{\theta}(x | z)] – \mathbb{E}_{q_{\varphi}(z | x)}[\log \frac{q_{\varphi}(z | x)}{p_{\theta}(z)}] \\ &= \mathbb{E}_{q_{\varphi}(z | x)}[\log p_{\theta}(x | z)] – D_{KL}(q_{\varphi}(z | x) || p_{\theta}(z)) \end{align*} $$

よって

$$ \begin{align} \mathcal{L}(\theta, \varphi; x) = \mathbb{E}_{q_{\varphi}(z | x)}[\log p_{\theta}(x | z)] – D_{KL}(q_{\varphi}(z | x) || p_{\theta}(z)) \end{align} $$

とも書くことができます。

💡 オートエンコーダとの関係２
オートエンコーダの観点からは(5)式の右辺第一項は入力$x$に対する再構成誤差、第二項は正則化項と解釈することができます。

$\mathcal{L(\theta, \varphi; x)}$の最大化が目的なので第一項については$p_{\theta}(x | z)$を$1$に近づけることが目的になります。また積分範囲を見ると$z$の確率分布は$q_{\varphi}(z | x)$です。ここで「オートエンコーダとの関係１」を思い出すと、$q_{\varphi}(z | x)$は確率的エンコーダで$p_{\theta}(x | z)$は確率的デコーダと捉えられるのでした。すると「$p_{\theta}(x | z)$を$1$に近づける」＝「$x$をエンコードした$z$をデコードすると再び$x$になる確率を$1$に近づける」ということになるので第一項を再構成誤差と解釈できることがわかります。

第二項については$D_{KL}(q_{\varphi}(z | x) || p_{\theta}(z))$が非負値を取るので$\mathcal{L}$の最大化のためには$D_{KL}$は$0$に近づくのが望ましいです。そのため事前分布$p_{\theta}(z)$と事後分布$q_{\varphi}(z|x)$の差はカルバック・ライブラー距離にしたがって近づくようになります。この意味で正則化項として機能することが期待できます。

一旦(3)式に戻りましょう。ここには$q_{\varphi}(z | x)$による積分が出てきているので確率分布が$q_{\varphi}(z | x)$になるようなサンプリングを$\{z^{(l)}\}_{l = 1}^L$として

$$ \begin{align} \tilde{\mathcal{L}}^A(\theta, \varphi; x) = \frac{1}{L} \sum_{l = 1}^L \big( \log p_{\theta}(x, z^{(l)}) – \log q_{\varphi}(z^{(l)}|x) \big) \end{align} $$

とすればこれはELBOの推定量になります。また(5)式を用いると

$$ \begin{align} \tilde{\mathcal{L}}^B(\theta, \varphi; x) = – D_{KL}(q_{\varphi}(z | x) || p_{\theta}(z))+\frac{1}{L} \sum_{l = 1}^L \log p_{\theta}(x | z^{(l)}) \end{align} $$

となりこれもELBOの推定量になります。これらの(6)式および(7)式をSGVB推定量といいます。またバッチサイズ$M$のミニバッチ$X^M = \{ x^{(i)} \}_{i = 1}^M$に対して

$$ \begin{align} \tilde{\mathcal{L}}^M(\theta, \varphi; X^M) &= \frac{N}{M} \sum_{i = 1}^M \tilde{\mathcal{L}}(\theta, \varphi; x^{(i)}) \end{align} $$

とすると$\mathcal{L}(\theta, \varphi; X)$の推定量になりこれもSGVB推定量と呼びます。ただし$\tilde{\mathcal{L}}$は$\tilde{\mathcal{L}}^A$か$\tilde{\mathcal{L}}^B$とします。バッチサイズ$M$が十分大きい場合はデータポイント毎の潜在変数$z$のサンプリング数は$L = 1$とすることが多いです。ここで定義したSGVB推定量を用いた以下のような学習プロセスをAEVBアルゴリズムと呼びます。

📖 AEVBアルゴリズム

$\theta,\varphi$を初期化する.
$\theta, \varphi$が収束するまで以下を繰り返す.
1. ミニバッチ$X^M$をデータ$X$からサンプリングする.
2. $\nabla_{\theta, \varphi} \tilde{\mathcal{L}^M}(\theta, \varphi; X^M)$に基づいて$\theta, \varphi$を改善する.
$\theta, \varphi$を出力する.

Variational Auto-Encoder

いよいよ上で導入したSGVB推定量およびAEVBアルゴリズムを用いてVAEを構成します。

潜在変数の事前分布は$p_{\theta}(z) = \mathcal{N}(z ; 0, I)$とします。ただし$\mathcal{N}(z;0, I)$は潜在変数$z$の次元$J$と同じ次元の標準正規分布です。事後分布を近似する確率モデルは

$$ q_{\varphi}(z | x) = \mathcal{N}(z; \mu_{\varphi}(x), \sigma_{\varphi}^2(x)) $$

とします。ただしこれらは一例であって必ずしも正規分布を使う必要はないです。また$\mu_{\varphi}(x),\sigma_{\varphi}(x)$はパラメータ$\varphi$付きの関数であって$\mu_{\varphi}(x)$は潜在変数$z$の平均を表す$J$次元ベクトル、$\sigma_{\varphi}(x)$は標準偏差を表す$J \times J$の対角行列です。これらはニューラルネットワークでの実装を想定しています。正規分布の標準化から$z$のサンプリングは

$$ z = \mu_{\varphi}(x) + \sigma_{\varphi}(x) \varepsilon,~\varepsilon \sim \mathcal{N}(0, I) $$

とすればよいです。

💡 Reparametrization Trick
VAE論文[2]では$z^{(i)} \sim q_{\varphi}(z | x^{(i)})$のサンプリングを微分可能な関数$g_{\varphi}$とパラメータを持たない乱数$\epsilon \sim p(\epsilon)$によって$z^{(i)} = g_{\varphi}(\epsilon, x^{(i)})$と表現することをreparametrization trickと呼んでいます。

本来サンプリングのようなランダムな操作は微分不可能などころか数学的な関数ですらないですが、ランダム性をパラメータを持たない部分に押し付けることで学習が可能な仕組みを実現しています。

SGVB推定量を$L = 1$として計算すると

$$ \begin{align*} &\tilde{\mathcal{L}}^B(\theta, \varphi; x^{(i)}) \\ &= -D_{KL}(q_{\varphi}(z | x^{(i)}) || p_{\theta}(z)) + \log p_{\theta}(x^{(i)} | z^{(i)}) \\ &= \frac{1}{2} \sum_{j = 1}^J(1 + \log \sigma_{\varphi}^2(x^{(i)})_j – \mu_{\varphi}^2(x^{(i)})_j – \sigma_{\varphi}^2(x^{(i)})_j) + \log p_{\theta}(x^{(i)} | z^{(i)}) \\ &~~~~~~~~(z^{(i)} = \mu_{\varphi}(x^{(i)}) + \sigma_{\varphi}(x^{(i)}) \varepsilon,~\varepsilon \sim \mathcal{N}(0, I)) \end{align*} $$

となります。再構成誤差の部分は、入出力データ$x^{(i)}$がバイナリのときは負のバイナリクロスエントロピー

$$ \begin{align*} \log p_{\theta}(x^{(i)} | z^{(i)}) &\sim x^{(i)} \log y^{(i)} + (1 – x^{(i)}) \log (1 – y^{(i)}),~y^{(i)} = h_{\theta}(z^{(i)}) \end{align*} $$

連続値のときは

$$ \begin{align*} \log p_{\theta}(x^{(i)} | z^{(i)}) &\sim \log \mathcal{N}(x^{(i)}; \mu_{\theta}(z^{(i)}), \sigma_{\theta}^2(z^{(i)})) \\ \end{align*} $$

によって表現します。もちろんこれらにも他の確率モデルを用いてもよいです。エンコーダと同様に$h_{\theta}(z),\mu_{\theta}(z),\varphi_{\theta}(z)$はニューラルネットワークでの実装を想定しています。

💡 画像データの再構成誤差
画像サイズを$S$として再構成データ$y = (y_s)_{s = 1}^S$は$[0,1]^S$に値をとります。これをベルヌーイ分布$P_y$として確率分布だと思うことができます。つまり$\{0, 1\}^S$に値をとる確率分布であって$s$番目の要素が$1$になる確率が$y_s$であるような確率分布と見なします。仮に$x$がバイナリ(すなわち$[0, 1]^S$ではなく$\{0, 1\}^S$に値を取る)だとすると$\log P_{y}(x) = \sum_{s \in S} x_s \log y_s$となります。これを$x$が$[0,1]^S$に値をとるようなバイナリでない場合にも拡張することで、画像データの再構成誤差にバイナリクロスエントロピーを利用することがあります。

モデルのアーキテクチャを図示すると以下のようになります。割とシンプルな構成ですね。このニューラルネットワークに対して負のSGVB推定量を損失関数として学習させればVAEが実現できます。VAEを試してみたい方は[7]のような実装例があります。