ベイズ推定では、事後分布 $p(\bm{z}|\bm{x})$ を求めることが中心的な課題です。しかし、多くの実用的なモデルでは事後分布を解析的に計算することが困難です。変分推論(Variational Inference)は、この問題を最適化問題に帰着させるアプローチです。
本記事では、変分推論の中でも最も基本的な「平均場近似(Mean-Field Approximation)」について、理論の導出からPythonでの実装までを解説します。
本記事の内容
- 変分推論の基本的なアイデア
- ELBOの導出
- 平均場近似と座標上昇法
- Pythonでの実装例
前提知識
この記事を読む前に、以下の概念を理解しておくと読みやすくなります。
- ベイズ推定 — 事後分布の考え方
- KLダイバージェンスの定義
変分推論とは
ベイズ推定では、観測データ $\bm{x}$ から潜在変数 $\bm{z}$ の事後分布を求めます。
$$ \begin{equation} p(\bm{z}|\bm{x}) = \frac{p(\bm{x}|\bm{z})p(\bm{z})}{p(\bm{x})} \end{equation} $$
しかし、分母の周辺尤度 $p(\bm{x}) = \int p(\bm{x}|\bm{z})p(\bm{z}) d\bm{z}$ は多くの場合、解析的に計算できません。
変分推論のアイデアは、計算しやすい近似分布 $q(\bm{z})$ を用いて、真の事後分布 $p(\bm{z}|\bm{x})$ を近似することです。「近さ」の尺度としてKLダイバージェンスを使います。
$$ \begin{equation} D_{\text{KL}}(q(\bm{z}) \| p(\bm{z}|\bm{x})) = \int q(\bm{z}) \ln \frac{q(\bm{z})}{p(\bm{z}|\bm{x})} d\bm{z} \end{equation} $$
このKLダイバージェンスを最小化する $q(\bm{z})$ を求めることが、変分推論の目標です。
ELBOの導出
KLダイバージェンスを直接最小化することは、$p(\bm{z}|\bm{x})$ がわからないため困難です。そこで、等価な最適化問題に書き換えます。
KLダイバージェンスを展開すると、
$$ \begin{align} D_{\text{KL}}(q \| p) &= \int q(\bm{z}) \ln \frac{q(\bm{z})}{p(\bm{z}|\bm{x})} d\bm{z} \\ &= \int q(\bm{z}) \ln q(\bm{z}) \, d\bm{z} – \int q(\bm{z}) \ln p(\bm{z}|\bm{x}) \, d\bm{z} \\ &= \int q(\bm{z}) \ln q(\bm{z}) \, d\bm{z} – \int q(\bm{z}) \ln \frac{p(\bm{x}, \bm{z})}{p(\bm{x})} d\bm{z} \\ &= \int q(\bm{z}) \ln q(\bm{z}) \, d\bm{z} – \int q(\bm{z}) \ln p(\bm{x}, \bm{z}) \, d\bm{z} + \ln p(\bm{x}) \end{align} $$
整理すると、
$$ \begin{equation} \ln p(\bm{x}) = D_{\text{KL}}(q \| p) + \underbrace{\int q(\bm{z}) \ln \frac{p(\bm{x}, \bm{z})}{q(\bm{z})} d\bm{z}}_{\text{ELBO}(q)} \end{equation} $$
ここで定義される量
$$ \begin{equation} \text{ELBO}(q) = E_{q}\left[\ln p(\bm{x}, \bm{z}) – \ln q(\bm{z})\right] \end{equation} $$
をELBO(Evidence Lower BOund)と呼びます。$\ln p(\bm{x})$ は $q$ に依存しない定数なので、KLダイバージェンスを最小化することはELBOを最大化することと等価です。
$$ \begin{equation} q^* = \arg\min_q D_{\text{KL}}(q \| p) = \arg\max_q \text{ELBO}(q) \end{equation} $$
平均場近似
平均場近似では、潜在変数を $M$ 個のグループに分割し、近似分布 $q(\bm{z})$ が各グループの独立な積で表せると仮定します。
$$ \begin{equation} q(\bm{z}) = \prod_{j=1}^{M} q_j(z_j) \end{equation} $$
イメージとしては、本来は相互に依存し合う潜在変数どうしを、独立に扱えるように簡略化するということです。
この仮定のもとでELBOを最大化すると、各 $q_j$ の最適解が次のように求まります。
$$ \begin{equation} \ln q_j^*(z_j) = E_{q_{-j}}\left[\ln p(\bm{x}, \bm{z})\right] + \text{const.} \end{equation} $$
ここで $E_{q_{-j}}$ は、$z_j$ 以外の全ての潜在変数について $q$ で期待値を取ることを意味します。
導出
ELBOに平均場近似を代入して $q_j$ について最適化します。
$$ \begin{align} \text{ELBO} &= \int \prod_{i} q_i(z_i) \ln p(\bm{x}, \bm{z}) \, d\bm{z} – \sum_{i} \int q_i(z_i) \ln q_i(z_i) \, dz_i \end{align} $$
$q_j$ に関する項だけを取り出すと(他の $q_i$, $i \neq j$ は固定)、
$$ \begin{align} \text{ELBO}_{q_j} &= \int q_j(z_j) \left[\int \prod_{i \neq j} q_i(z_i) \ln p(\bm{x}, \bm{z}) \, d\bm{z}_{-j}\right] dz_j – \int q_j(z_j) \ln q_j(z_j) \, dz_j \\ &= \int q_j(z_j) E_{q_{-j}}[\ln p(\bm{x}, \bm{z})] \, dz_j – \int q_j(z_j) \ln q_j(z_j) \, dz_j \end{align} $$
これは $q_j$ と $\exp(E_{q_{-j}}[\ln p(\bm{x}, \bm{z})])$ の間のKLダイバージェンスの負値に等しいため、最適解は
$$ \begin{equation} q_j^*(z_j) \propto \exp\left(E_{q_{-j}}[\ln p(\bm{x}, \bm{z})]\right) \end{equation} $$
となります。
座標上昇法アルゴリズム
平均場近似では、各 $q_j$ を順番に更新する座標上昇型変分推論(CAVI: Coordinate Ascent Variational Inference)を使います。
- 全ての $q_j(z_j)$ を初期化する
- 各 $j = 1, 2, \dots, M$ について、他の $q_i$ ($i \neq j$) を固定して $q_j$ を更新: $$ q_j^*(z_j) \propto \exp\left(E_{q_{-j}}[\ln p(\bm{x}, \bm{z})]\right) $$
- ELBOが収束するまで手順2を繰り返す
各更新ステップでELBOは必ず増加(または不変)するため、アルゴリズムの収束が保証されます。
Pythonでの実装
具体例として、1次元のガウス混合モデルに対する平均場変分推論を実装します。データが混合ガウス分布から生成されたとき、各データ点がどのクラスタに属するかを推論します。
import numpy as np
import matplotlib.pyplot as plt
from scipy.special import digamma, gammaln
np.random.seed(42)
# データ生成: 2成分のガウス混合
n = 200
true_means = [-3.0, 3.0]
true_pi = [0.4, 0.6]
K = 2
# 各データ点の所属クラスタを生成
z_true = np.random.choice(K, size=n, p=true_pi)
x = np.array([np.random.normal(true_means[z], 1.0) for z in z_true])
# 平均場変分推論(ガウス混合モデル)
# 近似分布: q(z) * q(mu) * q(pi)
# ハイパーパラメータ
alpha_0 = np.ones(K) # ディリクレ事前分布のパラメータ
m_0 = np.zeros(K) # 平均の事前分布の平均
beta_0 = 1.0 # 平均の事前分布の精度
a_0 = 1.0 # 精度の事前分布のパラメータ
b_0 = 1.0
# 変分パラメータの初期化
# q(z_i) = Cat(r_i): 責任度
r = np.random.dirichlet(np.ones(K), size=n)
# CAVI
n_iter = 50
elbo_history = []
for iteration in range(n_iter):
# N_k, x_bar_k の計算
N_k = r.sum(axis=0) + 1e-10
x_bar_k = (r.T @ x) / N_k
# q(pi) のパラメータ更新(ディリクレ分布)
alpha_k = alpha_0 + N_k
# q(mu_k) のパラメータ更新(ガウス分布)
beta_k = beta_0 + N_k
m_k = (beta_0 * m_0 + N_k * x_bar_k) / beta_k
# q(z_i) の更新
# E[ln pi_k]
E_ln_pi = digamma(alpha_k) - digamma(alpha_k.sum())
# E[mu_k], E[mu_k^2]
E_mu = m_k
E_mu2 = m_k**2 + 1.0 / beta_k
ln_rho = np.zeros((n, K))
for k in range(K):
ln_rho[:, k] = (E_ln_pi[k]
- 0.5 * (x**2 - 2 * x * E_mu[k] + E_mu2[k]))
# 数値安定化のためlog-sum-expトリック
ln_rho -= ln_rho.max(axis=1, keepdims=True)
r = np.exp(ln_rho)
r /= r.sum(axis=1, keepdims=True)
# ELBO計算(簡易版)
elbo = np.sum(r * ln_rho) - np.sum(r * np.log(r + 1e-10))
elbo_history.append(elbo)
# 結果の表示
print(f"推定された平均: {m_k}")
print(f"推定された混合比: {alpha_k / alpha_k.sum()}")
print(f"真の平均: {true_means}")
print(f"真の混合比: {true_pi}")
# 可視化
fig, axes = plt.subplots(1, 3, figsize=(16, 4))
# 左: データとクラスタ割り当て
cluster = r.argmax(axis=1)
for k in range(K):
mask = cluster == k
axes[0].hist(x[mask], bins=20, alpha=0.6, label=f'Cluster {k}')
axes[0].set_xlabel('x')
axes[0].set_ylabel('Count')
axes[0].set_title('Clustering Result')
axes[0].legend()
# 中央: 責任度の可視化
order = np.argsort(x)
axes[1].scatter(x[order], r[order, 0], s=5, alpha=0.7, label='q(z=0)')
axes[1].scatter(x[order], r[order, 1], s=5, alpha=0.7, label='q(z=1)')
axes[1].set_xlabel('x')
axes[1].set_ylabel('Responsibility')
axes[1].set_title('Variational Responsibilities')
axes[1].legend()
# 右: ELBOの推移
axes[2].plot(elbo_history, 'o-', markersize=3)
axes[2].set_xlabel('Iteration')
axes[2].set_ylabel('ELBO')
axes[2].set_title('ELBO Convergence')
plt.tight_layout()
plt.show()
このコードでは、2成分のガウス混合モデルに対して平均場変分推論(CAVI)を実装しています。各イテレーションで、クラスタ割り当て $q(z_i)$、混合比 $q(\pi)$、各クラスタの平均 $q(\mu_k)$ を順番に更新します。ELBOが単調に増加し収束する様子が確認できます。
まとめ
本記事では、平均場近似による変分推論について解説しました。
- 変分推論は、計算困難な事後分布を近似分布で置き換え、KLダイバージェンスの最小化(= ELBOの最大化)として定式化する
- 平均場近似は、近似分布を独立な因子の積として仮定する簡略化手法である
- 各因子の最適解は $q_j^*(z_j) \propto \exp(E_{q_{-j}}[\ln p(\bm{x}, \bm{z})])$ で与えられる
- 座標上昇法(CAVI)により各因子を順番に更新し、ELBOの収束を保証する
- ガウス混合モデルへの適用例を通じて、実装と動作を確認した
次のステップとして、以下の記事も参考にしてください。