Conditional Variational Autoencoder

はじめに

　今回は生成モデルの一種であるConditional Variational Autoencoder（CVAE）を紹介する。最初に定式化を行い、その後、PythonフレームワークであるPixyzを用いた実装例を示す。最初の例として、適当な数字を指定するとその手書き文字画像を自動生成するモデルを見る。次に、手書き文字画像を与えるとその数字を予測するモデルを取り上げる。

ベイズの変分推論

　 $N$ 個の観測値のペア $\{x_n,y_n\}_{n=1}^{N}$ を考える。 $X=\{x_1,\cdots,x_N\},Y=\{y_1,\cdots,y_N\}$ とし潜在変数 $z$ を考え、これら3つの同時確率分布 $p(X,Y,z)$ にベイズの定理を適用する。

(1) $\begin{equation*} p(z|X,Y)=\frac{p(Y|X,z)p(z)}{p(Y|X)} \end{equation*}$

式変形の途中で $p(X|z)=p(X)$ を用いた（観測値 $X$ と潜在変数 $z$ は独立である）。事後確率 $p(z|X,Y)$ をベイズの変分推論により求める。 $p(z|X,Y)$ を近似する関数としてパラメータ $\phi$ を持つ関数 $q_{\phi}(z|X,Y)$ を導入し、次のKullback-Leibler Divergenceを最小にすることを考える。

(2) $\begin{equation*} D_{KL}\left[q_{\phi}(z|X,Y)||p(z|X,Y)\right] \equiv \int dz\; q_{\phi}(z|X,Y)\ln{\frac{q_{\phi}(z|X,Y)}{p(z|X,Y)}} \end{equation*}$

右辺を変形すると

(3) $\begin{equation*} D_{KL}\left[q_{\phi}(z|X,Y)||p(z|X,Y)\right]= D_{KL}\left[q_{\phi}(z|X,Y)||p(z)\right]-E_{q_{\phi}(z|X,Y)}\left[\ln{p(Y|X,z)}\right]+\ln{p(Y|X)} \end{equation*}$

となる。ここで、 $E_{q_{\phi}(z|X,Y)}\left[\cdot\right]$ は $q_{\phi}(z|X,Y)$ についての期待値を表す。右辺の第3項は $\phi$ に依存しないので、右辺を $\phi$ について最小にするには次式を最小にすれば良い。

(4) $\begin{equation*} L\left[\phi\right]\equiv D_{KL}\bigl[q_{\phi}(z|X,Y)||p(z)\rbig]-\matchbb{E}_{q_{\phi}(z|X,Y)}\left[\ln{p(Y|X,z)}\right] \end{equation*}$

ところで、 $L\left[\phi\right]$ を用いて式(3)を変形すると

(5) $\begin{equation*} \ln{p(Y|X)}=D_{KL}\left[q_{\phi}(z|X,Y)||p(z|X,Y)\right]-L\left[\phi\right] \end{equation*}$

を得る。Kullback-Leibler Divergenceは常に0以上であるから

(6) $\begin{equation*} \ln{p(Y|X)}\geq -L\left[\phi\right] \end{equation*}$

となる。左辺の量 $p(Y|X)$ はEvidenceと呼ばれる量であり、この対数の下限が $-L\left[\phi\right]$ であることを上式は示している。この下限をEvidence Lower Bound（ELBO）と呼ぶ。すなわち、ELBOを最大することと、 $D_{KL}\left[q_{\phi}(z|X,Y)||p(z|X,Y)\right]$ を最小にすることとは等価である。いずれにせよ、式(4)を最小するような $\phi$ を見つけることが目標になる。

深層生成モデル

　確率分布 $q_{\phi}(z|X,Y)$ を正規分布で表し、そのパラメータである平均 $\mu_{\phi}$ と分散 $\sigma_{\phi}^2$ をネットワークを用いて計算する。

(7) $\begin{equation*} q_{\phi}(z|X,Y)=\mathcal{N}(z|\mu_{\phi},\sigma_{\phi}^2) \end{equation*}$

ここで、 $\mu_{\phi},\sigma_{\phi}$ は $X,Y$ の関数であり、パラメータ $\phi$ はネットワークの各種重みに相当する。ネットワークを用いて表現された生成モデルを深層生成モデルと呼ぶ。ここでは手書き文字画像のデータセットMNISTを用いて、次の二通りの生成モデルを考える。

$X$ を数字、 $Y$ を画像とみなす。このとき、 $p(Y|X,z)$ は数字 $X$ と潜在変数 $z$ から画像 $Y$ を生成するモデルになる。この生成モデルを「生成モデル1」と呼ぶことにする。 $p(Y|X,z)$ としてベルヌーイ分布を仮定する。

(8) $\begin{eqnarray*} p_{\theta}(Y|X,z) &=& \prod_{n=1}^{N}p(y_n|\eta_{\theta}(X,z)) \\ &=& \prod_{n=1}^{N}\prod_{m=1}^{M}{\rm Bern}(y_{n,m}|\eta_{\theta}(X,z)) \end{eqnarray*}$

$m$ は個々の画素を指定する番号である。ベルヌーイ分布 ${\rm Bern}(y|\eta)$ は次式で定義される。

(9) $\begin{equation*} {\rm Bern}(y|\eta)=\eta^y(1-\eta)^{1-y} \end{equation*}$

ここで、 $\eta$ は $\eta\in(0,1)$ を満たす実数である。ベルヌーイ分布は0か1のいずれかを生成する。すなわち、MNIST画像を2値画像とみなすということである（実際はグレイ画像だけど）。パラメータ $\eta_{\theta}$ を、入力を $X,z$ とするネットワークで計算する。 $\theta$ はネットワークの重みを表す。
$X$ を画像、 $Y$ を数字とみなす。このとき、 $p(Y|X,z)$ は画像 $X$ と潜在変数 $z$ から数字 $Y$ を生成するモデルになる。この生成モデルを「生成モデル2」と呼ぶことにする。 $p(Y|X,z)$ としてカテゴリカル分布を仮定する。

(10) $\begin{equation*} p_{\theta}(Y|X,z)=\prod_{n=1}^{N}{\rm Cat}(y_n|\eta_{\theta}(X,z)) \end{equation*}$

パラメータ $\eta_{\theta}$ を、入力を $X,z$ とするネットワークで計算する。 $\theta$ はネットワークの重みを表す。カテゴリカル分布は次式で定義される確率分布である。

(11) $\begin{equation*} {\rm Cat}(y|\eta)=\prod_{k=1}^{K}\pi_k^{s_k} \end{equation*}$

ただし、 $s_k$ は0か1のどちらかの値をとり、かつ $\sum_{k=1}^{K}s_k=1$ である。また、 $\pi_k$ は、 $\pi_k\in(0,1)$ 、かつ $\sum_{k=1}^{K}\pi_k=1$ を満たす実数である。 $y$ は1から $K$ までのいずれかの整数値を取る。

生成モデル1、生成モデル2のどちらのモデルにおいても $p(z)$ として標準正規分布を仮定する。

(12) $\begin{equation*} p(z)=\mathcal{N}(z|0,1) \end{equation*}$

ここまでの様子を図にすると以下のようになる。

結果

　深層生成モデルのためのPythonフレームワークPixyzを用いて実装した。ソースはここにある。最初に、標準正規分布 $p(z)$ から発生させた乱数 $z$ と数字5( $=x$ )のペアから生成した手書き文字画像( $=y$ )を示す（生成モデル1）。

左図は1epochの訓練後、右図は10epochの訓練後の生成画像である。64個の乱数を発生させ、それぞれに対応する画像を示した。次に再構成画像を示す。 $q_{\phi}(z|X,Y)$ のあと $p_{\theta}(Y|X,z)$ により再構成された画像 $Y$ に相当する。

左図下の段は1epochの訓練後の再構成画像、右図下の段は10epochの訓練後の再構成画像である。上の段は左右ともに元画像である。

　次に、画像と乱数を与えて、数字をあてる生成モデル（予測器）の精度を示す。これは生成モデル2の結果である。

各エポックの終わりに $p(z)$ から乱数 $z$ を発生させ（訓練データについては5万個、テストデータについては1万個）、訓練データ・テストデータが持つ画像 $X$ を用いて、カテゴリカル分布 $p_{\theta}(Y|X,z)$ によりラベルを予測し、精度を測定した。特に優れた結果になるわけではないが、生成モデルを予測器として使えるのは純粋に面白いと思う。

Pixyzの特長

　今回の生成モデル1のソースコードを取り上げ、Pixyzの特長を解説する。まず最初に、 $q_{\phi}(z|X,Y)$ を実装した部分である。

# q(z|x,y)
class Inference(Normal):

    def __init__(self):
        super().__init__(var=["z"], cond_var=["x", "y"], name="q")

        self.fc1 = nn.Linear(X_DIM + Y_DIM, H_DIM)
        self.fc2 = nn.Linear(H_DIM, H_DIM)
        self.fc31 = nn.Linear(H_DIM, Z_DIM)
        self.fc32 = nn.Linear(H_DIM, Z_DIM)

    def forward(self, x, y):
        h = F.relu(self.fc1(torch.cat([x, y], 1)))
        h = F.relu(self.fc2(h))
        # scale is variance
        return {"loc": self.fc31(h), "scale": F.softplus(self.fc32(h))}

# q(z|x,y)

class Inference(Normal):

def __init__(self):

super().__init__(var=["z"], cond_var=["x", "y"], name="q")

self.fc1 = nn.Linear(X_DIM + Y_DIM, H_DIM)

self.fc2 = nn.Linear(H_DIM, H_DIM)

self.fc31 = nn.Linear(H_DIM, Z_DIM)

self.fc32 = nn.Linear(H_DIM, Z_DIM)

def forward(self, x, y):

h = F.relu(self.fc1(torch.cat([x, y], 1)))

h = F.relu(self.fc2(h))

# scale is variance

return {"loc": self.fc31(h), "scale": F.softplus(self.fc32(h))}

先に説明したように

(13) $\begin{equation*} q_{\phi}(z|X,Y)=\mathcal{N}(z|\mu_{\phi},\sigma_{\phi}^2) \end{equation*}$

と仮定した。これを表現するため、正規分布クラスNormalを継承している（2行目）。さらに、平均locと分散scaleをネットワークで計算させている（13から16行目）。Normalを継承したクラスは、メソッドforwardでlocとscaleを返さなければならない仕様になっている。この約束により、クラスInferenceは正規分布として働くことができる。

　次に $p_{\theta}(Y|X,z)$ を実装した部分である。コード内で $X$ と $Y$ の表記が逆になっているのは実装上の都合のためである。読み替えて見てほしい。

# p(x|z,y)
class Generator(Bernoulli):

    def __init__(self):
        super().__init__(var=["x"], cond_var=["z", "y"], name="p")

        self.fc1 = nn.Linear(Z_DIM + Y_DIM, H_DIM)
        self.fc2 = nn.Linear(H_DIM, H_DIM)
        self.fc3 = nn.Linear(H_DIM, X_DIM)

    def forward(self, z, y):
        h = F.relu(self.fc1(torch.cat([z, y], 1)))
        h = F.relu(self.fc2(h))
        return {"probs": torch.sigmoid(self.fc3(h))}

# p(x|z,y)

class Generator(Bernoulli):

def __init__(self):

super().__init__(var=["x"], cond_var=["z", "y"], name="p")

self.fc1 = nn.Linear(Z_DIM + Y_DIM, H_DIM)

self.fc2 = nn.Linear(H_DIM, H_DIM)

self.fc3 = nn.Linear(H_DIM, X_DIM)

def forward(self, z, y):

h = F.relu(self.fc1(torch.cat([z, y], 1)))

h = F.relu(self.fc2(h))

return {"probs": torch.sigmoid(self.fc3(h))}

先に見たようにベルヌーイ分布を用いて

(14) $\begin{equation*} p_{\theta}(Y|X,z)=\prod_{n=1}^{N}{\rm Bern}(y_n|\eta_{\theta}) \end{equation*}$

と仮定した。これを表現するため今度はBeroulliクラスを継承している（2行目）。そしてパラメータ $\eta_{\theta}$ を計算するためネットワークを用いている（12から14行目）。probsは $\eta_{\theta}$ に相当する。メソッドforwardでprobsを返すことにより、クラスGeneratorはベルヌーイ分布として働くことができる。

　二つの例で見たように、あらかじめ用意されている確率分布NormalBernoulliを継承することで容易に目的とする確率分布を実装することができる。これら2つの確率分布、 $p_{\theta}(Y|X,z), q_{\phi}(z|X,Y)$ から計算される最小にすべき量

(15) $\begin{equation*} L\left[\phi\right]\equiv D_{KL}\bigl[q_{\phi}(z|X,Y)||p(z)\rbig]-\matchbb{E}_{q_{\phi}(z|X,Y)}\left[\ln{p(Y|X,z)}\right] \end{equation*}$

は以下のようにコーディングされる。

    # p(x|y,z)
    p = net.Generator().to(device)

    # q(z|x,y)
    q = net.Inference().to(device)

    # prior p(z)
    prior = Normal(loc=torch.tensor(0.0), scale=torch.tensor(1.0),
                   var=["z"], features_shape=[net.Z_DIM], name="p_{prior}").to(device)

    loss = (KullbackLeibler(q, prior) - Expectation(q, LogProb(p))).mean()
    model = Model(loss=loss, distributions=[p, q], optimizer=optim.Adam, optimizer_params={"lr": 1e-3})

# p(x|y,z)

p = net.Generator().to(device)

# q(z|x,y)

q = net.Inference().to(device)

# prior p(z)

prior = Normal(loc=torch.tensor(0.0), scale=torch.tensor(1.0),

var=["z"], features_shape=[net.Z_DIM], name="p_{prior}").to(device)

loss = (KullbackLeibler(q, prior) - Expectation(q, LogProb(p))).mean()

model = Model(loss=loss, distributions=[p, q], optimizer=optim.Adam, optimizer_params={"lr": 1e-3})

2行目で $p_{\theta}(Y|X,z)$ を、5行目で $q_{\phi}(z|X,Y)$ を、8行目で $p(z)$ （標準正規分布）を作っている。そして、11行目で式(15)を計算している。KullbackLeiblerExpectationLogProbはいずれもPixyzが用意しているモジュールである。12行目では、勾配降下法を行うためのモデルを作っている。このように、ほぼ理論式のままコーディングできるエレガントさがPixyzの大きな特長である。

まとめ

　今回は生成モデルの1つであるConditional Variational Autoencoder（CVAE）を取り上げ、確率分布のパラメータを計算する過程にネットワークを導入した。このようにネットワークを利用した生成モデルを深層生成モデルと呼ぶ。さらに、MINISTを用いた応用例をPixyzを用いて実装し、その結果を示した。Pixyzを用いると、ほぼ理論式のままコーディングすることができる。深層学習の中でも深層生成モデルは今後も発展する分野であると考えている。
　

Kumada Seiya

仕事であろうとなかろうと勉強し続ける、その結果”中身”を知ったエンジニアになれる

毎月アルゴリズム – 第1回 –

Salesforce Apexプログラミング／アノテーション

Conditional Variational Autoencoder

はじめに

ベイズの変分推論

深層生成モデル

結果

Pixyzの特長

まとめ

Kumada Seiya

最近の記事

オープンソースのLLM（ELYZA-japanese-Llama-2-7b-instruct）

GPT4oを用いた固有表現抽出

三角比を用いたピタゴラスの定理の証明

顔認証（DeepFace vs GPT4o）

時系列解析再訪

πの新しい公式

レイアウト解析

Universal Approximation Theorem

JAX

OpenCVで深層学習

PyTorch C++の導入〜その3〜

テンソルネットワークの入り口

Google Vision APIでOCR

アーカイブ

カテゴリー