変分推論（VI）、確率的勾配変分推論（SGVI/SGVB）、変分自己符号化器（VAE）の統合解説

参考文献：

  • VIの解説：PRML 第10章。
  • SGVI論文：Auto-Encoding Variational Bayes -- Kingma。
  • VAE解説1：Variational Autoencodersのチュートリアル -- CARL DOERSCH。
  • VAE解説2：スタンフォード大学 CS236: Deep Generative Models。

関数と汎関数

本記事では、「変分」という名称の由来について理解します。

関数と汎関数の違い

  • 関数：値から値への写像。
  • 汎関数：関数から値への写像。

典型的な汎関数の例：エントロピーの式

[H[p] = \int p(x) \ln p(x) \text{d}x ]

  • 関数の微分：入力値が微小変化したとき、出力値の変化。
  • 汎関数の微分：入力関数が微小変化したとき、出力値の変化。

汎関数の極値

  • 汎関数の極値を求める：すべての可能な関数を探索し、最大化または最小化する関数を見つける。
  • 変分法：汎関数の極値を研究する方法。
  • VIでは、潜在変数モデルの推論問題を汎関数の極値問題に変換するため、「変分」と呼ばれる。

変分推論（VI）

本節では近似推論手法の一つである変分推論について説明します。

問題設定：事後分布と証拠の推定

記号の定義

  • \(Z\)：潜在変数（モデルパラメータを含む）。
  • \(X\)：観測データ（明示的変数）。
  • \(P(Z)\)：潜在変数の事前分布（既知）。
  • \(P(X\mid Z)\)：潜在変数の尤度（既知）。
  • \(P(X,Z)\)：同時分布（\(P(X,Z)=P(X\mid Z)P(Z)\)で計算可能）。
  • \(P(Z\mid X)\)：潜在変数の事後分布（求めたい）。
  • \(P(X)\)：周辺証拠（求めたい）。

可算と不可算

可算（Tractable） 代表例：隠れマルコフモデル（HMM）の特徴：

  • 有限個の隠れ状態を持つ。
  • 模型構造が単純。

モデルパラメータが既知の場合、動的計画法に基づく正確推論が可能：

  • \(P(X)\)の計算：前向き後ろ向きアルゴリズム。
  • \(P(Z\mid X)\)の計算：ヴィタービアルゴリズム。

不可算（Intractable） 多くの場合：

  • 連続確率変数の積分を含む。
  • 模型構造が複雑。

この場合、正確推論は通常不可能。主な困難は積分：\(P(X) = \int P(X,Z)\text{d}Z\).

そのため、近似推論手法が必要です。以下ではその一つである変分推論を紹介します。

ELBOの導出

ベイズ公式の両辺に\(\log\)を適用：

[\log P(X) = \log \frac{P(X,Z)}{P(Z|X)} =\log P(X,Z) -\log P(Z\mid X) ] \(\phi\)でパラメトライズされた確率分布\(q_\phi(Z)\)を導入：

[\begin{aligned} \log P(X) &= \log P(X,Z) -\log P(Z\mid X) \ &= \log \frac{P(X,Z)}{q_\phi(Z)} - \log \frac{P(Z\mid X)}{q_\phi(Z)} \end{aligned} ] 両辺を分布\(q_\phi(Z)\)で平均化：

[\int_Z \log P(X) q_\phi(Z) \text{d} Z = \int_Z q_\phi(Z) \log \frac{P(X,Z)}{q_\phi(Z)}\text{d} Z - \int_Z q_\phi(Z) \log \frac{P(Z\mid X)}{q_\phi(Z)} \text{d}Z ] 左辺の\(\log P(X)\)は\(Z\)に依存しないので、\(\int_Z q_\phi(Z) \text{d}Z =1\)より：

[\log P(X) = \int_Z q_\phi(Z) \log \frac{P(X,Z)}{q_\phi(Z)}\text{d} Z - \int_Z q_\phi(Z) \log \frac{P(Z\mid X)}{q_\phi(Z)} \text{d}Z ] 右辺の第一項を期待値形式、第二項をKLダイバージェンス形式に書き換え：

[\begin{aligned} \log P(X)&= \mathbb E_{Z\sim q_\phi} [\log P(X,Z) - \log q_\phi(Z)] + KL(q_\phi(Z) \mid\mid P(Z\mid X)) \end{aligned} ] 以下の性質に注意：

  • \(KL \ge 0\)より、\(\log P(X) \ge\) 第一項（Evidenceの下界：ELBO）。
  • \(q_\phi(Z) = P(Z\mid X)\)のとき、\(KL=0\)で\(\log P(X) = ELBO\)。
  • KEY1：\(\arg \max_{q_\phi}ELBO\)を求めることで、\(ELBO \to \log P(X)\)かつ\(KL \to 0\)、\(q_\phi(Z) \to P(Z\mid X)\)。
  • KEY2：変分推論の本質は、潜在変数モデルの推論問題をELBOの最適化問題に変換すること。

REMARK1：ELBOを汎関数\(\mathcal L(q_\phi)\)と見なし、\(\arg \max_{q_\phi} \mathcal L(q_\phi)\)を求める（汎関数の極値点を求めるため「変分」）。

REMARK2：変分法は本来近似法ではなく、すべての関数を考慮すれば正確推論が可能。しかし実際には関数の範囲を仮定（例：二次型、平均場分解）するため、近似解になる。

REMARK3：平均場理論に基づく\(q(z)\)の仮定を用いると、座標上升法で最適解を求めることが可能（PRML 第10章参照）。

単一サンプルの場合

  • \(X=\{x^{(1)},...,x^{(n)}\}\)：観測変数。 各サンプルが独立同分布に従うと仮定すると、\(P(X)\)は積形式に分解：

[\log P(X) = \log \prod P(x^{(i)}) = \sum \log P(x^{(i)}) ] \(\phi^{(i)}\)でパラメトライズされた分布\(q_{\phi^{(i)}}(Z\mid x^{(i)})\)を導入し、\(\log P(x^{(i)})\)もELBO+KL形式に：

[\begin{aligned} \log P(x^{(i)}) &= \int_Z q_{\phi^{(i)}}(Z\mid x^{(i)}) \log \frac{P(x^{(i)},Z)}{q_{\phi^{(i)}}(Z\mid x^{(i)})}\text{d} Z - \int_Z q_{\phi^{(i)}}(Z\mid x^{(i)}) \log \frac{P(Z\mid x^{(i)})}{q_{\phi^{(i)}}(Z\mid x^{(i)})} \text{d}Z \ &= \mathbb E_{Z\sim q_{\phi^{(i)}}} [\log P(x^{(i)},Z) - \log q_{\phi^{(i)}}(Z\mid x^{(i)})] + KL(q_{\phi^{(i)}}(Z\mid x^{(i)}) \mid\mid P(Z\mid x^{(i)})) \ &= ELBO +KL \end{aligned} ] NOTE：\(q_\phi\)のパラメータ\(\phi\)は\(\phi^{(i)}\)と表記しているが、各観測変数ごとに事後分布\(P(Z\mid x^{(i)})\)が異なるため。以降のSGVI節では省略。

SGVI

前節では潜在変数の推論をELBOの最適化問題に変換しました。本節ではELBOの勾配を推定する方法について説明します。

  • SGVI（確率的勾配変分推論）は、勾配に基づく最適化法でELBOを最大化。
  • 勾配を求めて\(q_\phi(z)\)のパラメータを更新：\(\phi = \phi + \alpha \nabla_\phi ELBO\).

スコア関数勾配推定子

\(ELBO\)の\(\phi\)に関する勾配は\(\nabla_\phi \mathbb E_{z\sim q_\phi(z)}[f(z)]\)の形。以下に一般化された形式の勾配を分析：

[\nabla_\phi \mathbb E_{z\sim q_\phi(z)}[f(z)] = \nabla_\phi \int q_\phi(z) f(z) \text{d}z ] 勾配を積分内に移動し、\(\nabla \log q_\phi = \frac{\nabla q_\phi}{q_\phi}\)より：

[\begin{aligned} \nabla_\phi \mathbb E &= \int \nabla_\phi q_\phi(z) f(z) \text{d}z \ &= \int q_\phi(z) \nabla_\phi \log q_\phi(z) f(z) \text{d}z \ &= \mathbb E_{z\sim q_\phi(z)}[f(z) \nabla_\phi \log q_\phi(z)] \end{aligned} ] この期待値はMCサンプリングで推定可能：

[\mathbb E_{z\sim q_\phi(z)}[f(z) \nabla_\phi \log q_\phi(z)] \approx \frac{1}{L} \sum_{l=1}^L f(z^{(l)}) \nabla_\phi \log q_\phi(z^{(l)}) ]

  • \(z^{(l)} \sim q_\phi(z)\).
  • この推定法はスコア関数勾配推定子と呼ばれ、\(\nabla_\phi \log q_\phi(z)\)をスコア、\(f(z)\)をコストと呼ぶ。
  • 問題点：分散が大きい。多くの場合使用不可で、ELBOの直接推定には不向き。

NOTE：\(q_\phi(z)\)はtractable（サンプリング・計算が容易）と仮定している。

BTW：強化学習のREINFORCEアルゴリズムでもこの推定子が使われるため、高分散アルゴリズムとして知られる。

再パラメータ化トリック

サンプリングプロセス\(z \sim q_\phi(z\mid x)\)は通常次の2段階に分解される：

  1. ノイズ分布\(p(\epsilon)\)から\(\epsilon\)をサンプリング。
  2. \(\phi\)でパラメトライズされた関数\(g_\phi(,)\)を用いて\(z = g_\phi(\epsilon, x)\).

例えば正規分布\(z \sim \mathcal N(\mu, \sigma)\)からのサンプリング：

  1. \(\epsilon \sim \mathcal N(0,1)\).
  2. \(z = \mu + \epsilon \sigma\).

これが再パラメータ化トリック。

SGVI推定子

再パラメータ化トリックを使用すると、期待値は次のように表現可能：

[ \mathbb E_{z\sim q_\phi(z\mid x^{(i)})}[f(z)]=\mathbb E_{\epsilon \sim p(\epsilon)}[f(g_\phi(\epsilon, x^{(i)}))] \approx \frac{1}{L} \sum_{l=1}^L f(g_\phi(\epsilon^{(l)},x^{(i)})) ] \(f=ELBO\)とし勾配を求める：

[\nabla_\phi ELBO=\nabla_\phi \mathcal L(\phi,x^{(i)}) = \nabla_\phi \bigl(\frac{1}{L}\sum_{l=1}^L \log P(x^{(i)},g_\phi(\epsilon^{(l)},x^{(i)})) - \log q_\phi(g_\phi(\epsilon^{(l)},x^{(i)})\mid x^{(i)}) \bigr) ]

  • \(\epsilon^{(l)} \sim p(\epsilon)\).

NOTE：SGVI推定子は通常スコア関数勾配推定子よりも分散が小さい。

変分自己符号化器（VAE）

本節ではVI手法が潜在変数生成モデルの推論・学習にどのように応用されるかを説明します。

サンプル分布

我々はデータセットと類似したサンプルを生成したい。例えば、顔画像データセットをモデルに与えると、モデルがその分布を学習して新たな顔画像を生成できるようにしたい。

  • サンプルの仮定：データセットは未知分布\(P_{gt}(X)\)からサンプリングされたと考える。
  • 理解：顔画像\(X\)なら\(P_{gt}(X)\)が大きくなり、ノイズ画像\(X\)なら\(P_{gt}(X)\)が小さくなる。
  • 目的：サンプリング可能なモデル\(P(X)\)を構築し、\(P_{gt}\)に近づける。
  • 学習方法：パラメータ化されたモデル\(P_\theta(X)\)のパラメータ\(\theta\)を学習し、観測データ\(X\)の確率を最大化（最大尤度推定）。

潜在変数生成モデル

\(P_{gt}\)を近似するために潜在変数生成モデルを考慮。

例：数字の生成を想定した場合：

  1. 数字\(z \in {0,1,...,9}\)を決定。
  2. \(P(X\mid z;\theta)\)に基づいて対応する画像を生成。

ここで\(z\)が潜在変数。

困難：潜在変数の決定と観測変数との関係の定義が複雑：

  • 例えば、生成する数字に円があるか（0,8,9）やその数を決定。
  • さらに、傾きやフォントなどの複雑な関係を考慮する必要がある。

以下のように対応：

  1. 潜在変数\(z\)に正規事前分布\(z \sim \mathcal N(0,I)\)を仮定。
  2. 模型を正規分布\(P(X\mid z;\theta)=\mathcal N(X\mid \mu_\theta(z),\sigma_\theta(z))\)と仮定し、パラメータを\(z\)の関数とする。
  3. パラメータ\((\mu_\theta(z), \sigma_\theta(z))\)はニューラルネットワークでモデル化。

理由：

  • KEY1：単純な正規分布を十分な非線形変換で任意の分布を表現可能にする。
  • KEY2：ニューラルネットワークは任意の非線形関係をモデル化できるため、任意の分布を表現可能。
  • 例：正規分布に非線形変換を施す。左が正規分布、右が\(g(z) = \frac{z}{10} + \frac{z}{||z||}\)。

これにより、真のデータ分布\(P_{gt}(X)\)を近似するモデル\(P(X;\theta)\)を得る：

[P(X;\theta) = \int P(X,z;\theta)\text{d}z =\int P(X\mid z;\theta) P(z) \text{d}z =\int \mathcal N(X\mid \mu_\theta,\sigma_\theta) \mathcal N(0,I)\text{d}z ]

  • 模型の表現能力は十分。
  • \(z\)は観測不能な潜在変数なので、すべての可能性を考慮（積分）。
  • 推論（マージナル証拠\(P(X;\theta)\)の計算）と学習（パラメータ\(\arg \max_\theta P(X;\theta)\)の最適化）は困難。

モデルの近似推論と学習

変分推論の近似法を採用。

前節の議論より、モデルの\(\log\)証拠はELBO+KL形式に：

[\begin{aligned} \log P(x^{(i)};\theta) &= \mathbb E_{z\sim q_{\phi^{(i)}}} [\log P(x^{(i)},z;\theta) - \log q_{\phi^{(i)}}(z)] + KL(q_{\phi^{(i)}}(z) \mid\mid P(z\mid x^{(i)};\theta)) \end{aligned} ] 違い：モデルパラメータ\(\theta\)を含む。

NOTE：ELBOの最大化は、近似分布\(q\)のパラメータ\(\phi\)とモデルパラメータ\(\theta\)の両方を考慮する必要がある。 NOTE：ELBOの最大化には2つの役割がある：

  1. 推論：KLダイバージェンスを最小化し、\(q\)を真の事後分布に近づける；
  2. 学習：証拠を最大化し、モデル\(P_\theta\)を真のデータ分布\(P_{gt}\)に近づける。

SGVI推定子を用いてELBOの勾配を推定：

[\nabla \mathcal L(x^{(i)},\phi^{(i)},\theta)=\nabla \bigl(\frac{1}{L}\sum_{l=1}^L \log P(x^{(i)},g_{\phi^{(i)}}(\epsilon^{(l)},x^{(i)});\theta) - \log q_{\phi^{(i)}}(g_{\phi^{(i)}}(\epsilon^{(l)},x^{(i)})) \bigr) ]

  • \(g_\phi(\epsilon^{(l)},x^{(i)})\)は再パラメータ化後の\(z\)。
  • \(\epsilon \sim p(\epsilon)\).
  • 勾配演算子\(\nabla = \{\nabla_\phi, \nabla_\theta\}\).

これにより、モデルの学習プロセスは以下の通り：

Amortized Inference

  • 上記学習プロセスの欠点：各\(x\)に対して十分な\(q_{\phi^*}\)を推論する必要がある。
  • KEY：事後分布の近似 = 事後分布のパラメータを学習する（VIによる視点）。

Amortized Inference：\(x^i\)から\(\phi^*\)を学習するパラメータ化関数\(f_\lambda: x^i \to \phi^*\)を導入。各\(x^i\)に対して、事後分布\(q(z;f_\lambda(x^i))\)を取得。

学習方法は再パラメータ化+勾配上昇で同じ。

これにより、モデル学習プロセスは以下のようになる：

NOTE：Amortized Inference節まででEncoder-Decoder構造の原型が現れる：\(x^i \to q(z;f_\lambda(x^i)) \to z \to P(X\mid z;\theta) \to \bar x^i\).

Objectiveの解析

VAEの学習はELBOの最大化を目指す。本節ではELBOの形式からその実際の意味を考察。

簡略版ELBO式：

[\begin{aligned} ELBO &= \int q_\phi(Z)\log P(X,Z) - q_\phi(Z)\log q_\phi(Z) \text{d}Z \ &= \int q_\phi(Z)\log P(X\mid Z) + q_\phi(Z) \log P(Z) - q_\phi(Z)\log q_\phi(Z) \text{d}Z \ &=\int q_\phi(Z)\log P(X\mid Z) + q_\phi(Z) \log \frac{P(Z)}{q_\phi(Z)}\text{d}Z \ &= \int q_\phi(Z)\log P(X\mid Z) \text{d}Z + \int q_\phi(Z) \log \frac{P(Z)}{q_\phi(Z)}\text{d}Z \ &=\mathbb E_{q_\phi}[\log P(X\mid Z)] - KL(q_\phi(Z)\mid\mid P(Z)) \end{aligned} ] ELBOの2つの項の意味：

  1. 第一項：Loss項。生成したい\(X\)とデータセットができるだけ似ていることを示す（最大尤度）。
  2. 第二項：正則化項。事後分布が既知の事前分布\(P(Z)\)に近づくことを示す。\(P(Z)\)から\(z\)をランダムに抽出し、意味のある結果を生成できることを望む。

VAE構造例

以下は、全結合ニューラルネットワークを用いたVAEの構造図。左図は再パラメータ化トリックを用いないもの、右図は用いるもの。

要点：

  • ネットワーク出力：Decoderは潜在変数\(z\)から\(\bar X\)を再構成；Encoderは\(X\)から\(q_\phi(z\mid X)\)を出力（\(X\)を潜在変数に圧縮）。
  • 再パラメータ化と自動微分：再パラメータ化トリックにより、自動微分の逆伝播時にサンプリングを経ずに済む。
  • Objectiveの選択：ELBOに厳密に限定しなくてもよい：
  1. Loss項：再構成誤差を反映。問題に応じて柔軟に選択可能。
  2. 正則項：KLダイバージェンス以外にも他の散度を用いることも可能。
  • 生成について：Decoder部分が潜在変数生成モデル。\(\mathcal N(0,I)\)から\(z\)をランダムに抽出し、Decoderに入力することでサンプルを生成可能。

参考コード

以下はCNNを用いたVAEの実装例。

Encoder

class Encoder(nn.Module):
    def __init__(self, input_channels, latent_dim):
        super(Encoder, self).__init__()
        # 入力は(batch_size, input_channels, 28, 28)（MNISTデータの形状）
        self.conv1 = nn.Conv2d(input_channels, 32, 4, 2) # (28 - 4) / 2 + 1 = 13
        self.conv2 = nn.Conv2d(32, 64, 4, 2, 1) # (13 - 4 + 1) / 2 + 1 = 6
        self.conv3 = nn.Conv2d(64, 128, 4, 2) # (6 - 4) / 2 + 1 = 2

        self.fc_mu = nn.Linear(128 * 2 * 2, latent_dim)
        self.fc_logstd = nn.Linear(128 * 2 * 2, latent_dim)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.relu(self.conv3(x))
        x = torch.flatten(x, start_dim=1)
        mu = self.fc_mu(x)
        log_std = self.fc_logstd(x)
        return mu, log_std

Decoder

class Decoder(nn.Module):
    def __init__(self, output_channels, latent_size):
        super(Decoder, self).__init__()

        self.fc1 = nn.Linear(latent_size, 512)
        self.deconv1 = nn.ConvTranspose2d(512, 64, 5, 2)
        self.deconv2 = nn.ConvTranspose2d(64, 32, 5, 2)
        self.deconv3 = nn.ConvTranspose2d(32, output_channels, 4, 2)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = x.unsqueeze(-1).unsqueeze(-1) # num -> [[num]]
        x = F.relu(self.deconv1(x))
        x = F.relu(self.deconv2(x))
        reconstruction = F.sigmoid(self.deconv3(x)) # 出力範囲(0, 1)
        return reconstruction

VAE

class VAE(nn.Module):
    def __init__(self, input_channels, latent_dim):
        super(VAE, self).__init__()
        self.encoder = Encoder(input_channels, latent_dim)
        self.decoder = Decoder(input_channels, latent_dim)

    def forward(self, x):
        mu, log_std = self.encoder(x)
        z = self.reparameterize(mu, log_std)
        reconstruction = self.decoder(z)
        return reconstruction, mu, log_std

    def reparameterize(self, mu, log_std):
        std = log_std.exp()
        eps = torch.randn_like(std)
        return mu + std * eps

Loss

def vae_loss(x, reconstruction, mu, log_std):
    # 画像生成ではBCEが効果的
    rec_loss = F.binary_cross_entropy(reconstruction, x, reduction='sum')
    # qとN(0, I)のKLダイバージェンスの簡略式
    kl_loss = -0.5 * torch.sum(1 + 2 * log_std - mu.pow(2) - (2*log_std).exp())
    return rec_loss + kl_loss

タグ: 変分推論 確率的勾配変分推論 変分自己符号化器

5月29日 13:49 投稿

ホットタグ

©2026 異端開発室