😺

たぶん大学生でもわかる~VAEの目的関数ELBO完全解説:変分推論から実装まで~中野哲平

に公開

変分オートエンコーダ(VAE)は、深層学習と確率的推論を融合した革新的なアーキテクチャです。その核心にある**ELBO(Evidence Lower BOund)**は、変分推論の理論に基づいて導出される美しい数学的構造を持っています。本記事では、ELBOの導出過程を段階的に解説し、その直感的理解と実装への橋渡しを行います。

1. VAEが解決したい問題

生成モデリングの課題

従来の生成モデルでは以下の問題がありました:

  1. 計算困難性:真の事後分布 p(z|x) が複雑で計算できない
  2. 積分困難性:周辺尤度 p(x) = ∫ p(x|z)p(z)dz の計算が困難
  3. スケーラビリティ:高次元データに対する効率的な推論が難しい

VAEのアプローチ

VAEはこれらの問題を変分推論(Variational Inference)によって解決します:

  • 計算困難な真の事後分布を、パラメータ化された近似分布で置き換える
  • 近似の質を測るKLダイバージェンスを最小化
  • 深層ネットワークによる柔軟な近似分布の表現

2. 基本的な確率的設定

生成プロセス

VAEは以下の生成プロセスを仮定します:

  1. 潜在変数の生成:z ~ p(z)(通常は標準正規分布)
  2. 観測データの生成:x ~ p(x|z)(潜在変数から条件付き生成)

解決したい問題

観測データ x が与えられたときに:

  • 推論問題:事後分布 p(z|x) を求めたい
  • 学習問題:パラメータ θ を最尤推定したい(log p(x) を最大化)

3. ELBO導出の数学的展開

ステップ1:対数尤度の分解

任意の確率分布 q(z|x) を導入して、対数尤度を分解します:

log p(x) = ∫ q(z|x) log p(x) dz    [q(z|x)で重み付け積分]
         = ∫ q(z|x) log[p(x,z)/p(z|x)] dz    [ベイズの定理]
         = ∫ q(z|x) log[p(x,z)q(z|x)/(p(z|x)q(z|x))] dz

ステップ2:項の分離

log p(x) = ∫ q(z|x) log p(x,z) dz - ∫ q(z|x) log q(z|x) dz 
         + ∫ q(z|x) log q(z|x) dz - ∫ q(z|x) log p(z|x) dz

         = ∫ q(z|x) [log p(x,z) - log q(z|x)] dz 
         + ∫ q(z|x) [log q(z|x) - log p(z|x)] dz

ステップ3:ELBOとKLダイバージェンスの出現

log p(x) = E_q(z|x)[log p(x,z) - log q(z|x)] + D_KL(q(z|x)||p(z|x))
         = ELBO(q) + D_KL(q(z|x)||p(z|x))

ここで:

  • ELBO(q) = E_q(z|x)[log p(x,z) - log q(z|x)]
  • D_KL(q(z|x)||p(z|x)) ≥ 0(KLダイバージェンスの非負性)

ステップ4:ELBOの最終形

同時分布を分解すると:p(x,z) = p(x|z)p(z)

ELBO(q) = E_q(z|x)[log p(x|z)p(z) - log q(z|x)]
        = E_q(z|x)[log p(x|z)] + E_q(z|x)[log p(z) - log q(z|x)]
        = E_q(z|x)[log p(x|z)] - D_KL(q(z|x)||p(z))

これがVAEの目的関数です!

4. ELBOの構成要素の直感的理解

再構成項:E_q(z|x)[log p(x|z)]

再構成項 = エンコーダが推論した潜在変数zから、
          デコーダが元のデータxを復元できる度合い

役割:

  • データの再構成品質を保証
  • オートエンコーダの「復元誤差」に対応
  • 大きいほど良い再構成を意味

正則化項:D_KL(q(z|x)||p(z))

正則化項 = エンコーダの出力分布と事前分布の差

役割:

  • 潜在空間の構造を制御
  • 過学習を防ぐ正則化効果
  • サンプリング可能な表現学習を保証

5. 実装における具体的形式

ガウシアンVAEの場合

**近似事後分布:**q(z|x) = N(μ(x), σ²(x)I)
**事前分布:**p(z) = N(0, I)

KLダイバージェンスの解析解:

D_KL(q(z|x)||p(z)) = ½∑_i[μᵢ² + σᵢ² - log σᵢ² - 1]

完全なELBO:

def elbo_loss(x, x_reconstructed, mu, log_var):
    # 再構成項(負の対数尤度)
    reconstruction_loss = -tf.reduce_sum(
        x * tf.log(x_reconstructed + 1e-10) + 
        (1-x) * tf.log(1-x_reconstructed + 1e-10), 
        axis=1
    )
    
    # KLダイバージェンス項
    kl_loss = -0.5 * tf.reduce_sum(
        1 + log_var - tf.square(mu) - tf.exp(log_var), 
        axis=1
    )
    
    # ELBO = 再構成項 + KL項(符号に注意)
    return tf.reduce_mean(reconstruction_loss + kl_loss)

6. Reparameterization Trick

問題:勾配計算の困難性

確率的ノード z ~ q(z|x) に対して直接的に勾配計算ができない問題。

解決策:再パラメータ化

z = μ(x) + σ(x) ⊙ ε,  where ε ~ N(0, I)

この変換により:

  • 確率的要素をεに分離
  • μ(x), σ(x)に対する勾配計算が可能
  • バックプロパゲーションの実現
def reparameterize(mu, log_var):
    """再パラメータ化トリック"""
    epsilon = tf.random.normal(tf.shape(mu))
    return mu + tf.exp(0.5 * log_var) * epsilon

7. 学習過程の解釈

最適化の観点

ELBO最大化は以下を同時に行います:

  1. 再構成品質の向上

    • E_q(z|x)[log p(x|z)] を最大化
    • デコーダの性能向上
  2. 分布の正則化

    • D_KL(q(z|x)||p(z)) を最小化
    • エンコーダ分布を事前分布に近づける

β-VAEによる調整

ELBO_β = E_q(z|x)[log p(x|z)] - β·D_KL(q(z|x)||p(z))
  • β > 1:より強い正則化、解釈しやすい潜在表現
  • β < 1:再構成重視、高品質な復元

8. ELBOと他の概念との関連

情報理論的解釈

ELBO = 相互情報量 - 条件付きエントロピー + 定数

圧縮との関連

  • 再構成項:圧縮による情報損失
  • KL項:圧縮率(コードの複雑さ)

ベイズ推論との関連

ELBOは変分ベイズ法の一般的な枠組みにおける目的関数の特殊ケース。

9. 実装例:完全なVAEクラス

class VAE(tf.keras.Model):
    def __init__(self, latent_dim=32):
        super(VAE, self).__init__()
        self.latent_dim = latent_dim
        
        # エンコーダ
        self.encoder = tf.keras.Sequential([
            tf.keras.layers.Dense(512, activation='relu'),
            tf.keras.layers.Dense(256, activation='relu'),
        ])
        
        self.mu_layer = tf.keras.layers.Dense(latent_dim)
        self.log_var_layer = tf.keras.layers.Dense(latent_dim)
        
        # デコーダ
        self.decoder = tf.keras.Sequential([
            tf.keras.layers.Dense(256, activation='relu'),
            tf.keras.layers.Dense(512, activation='relu'),
            tf.keras.layers.Dense(784, activation='sigmoid')
        ])
    
    def encode(self, x):
        h = self.encoder(x)
        mu = self.mu_layer(h)
        log_var = self.log_var_layer(h)
        return mu, log_var
    
    def reparameterize(self, mu, log_var):
        epsilon = tf.random.normal(tf.shape(mu))
        return mu + tf.exp(0.5 * log_var) * epsilon
    
    def decode(self, z):
        return self.decoder(z)
    
    def call(self, x):
        mu, log_var = self.encode(x)
        z = self.reparameterize(mu, log_var)
        x_reconstructed = self.decode(z)
        return x_reconstructed, mu, log_var
    
    def compute_loss(self, x):
        x_reconstructed, mu, log_var = self(x)
        
        # 再構成損失
        reconstruction_loss = tf.reduce_mean(
            tf.reduce_sum(
                tf.keras.losses.binary_crossentropy(x, x_reconstructed), 
                axis=1
            )
        )
        
        # KL損失
        kl_loss = tf.reduce_mean(
            -0.5 * tf.reduce_sum(
                1 + log_var - tf.square(mu) - tf.exp(log_var), 
                axis=1
            )
        )
        
        return reconstruction_loss + kl_loss

まとめ

ELBOの導出は、変分推論の美しい数学的構造を示しています:

理論的意義

  1. 厳密な下界:KLダイバージェンスの非負性による理論保証
  2. 最適化可能性:実際に計算・最適化できる目的関数
  3. 解釈可能性:再構成と正則化の明確な分離

実践的価値

  1. 効率的学習:勾配ベースの最適化が可能
  2. 柔軟性:様々な分布・アーキテクチャに対応
  3. 制御可能性:β-VAEなどによる学習制御

ELBOは単なる数学的トリックではなく、確率論・情報理論・深層学習を統合した現代機械学習の重要な理論的基盤なのです。この理解により、VAEの動作原理が明確になり、より効果的な生成モデルの設計・改良が可能になります。

Discussion