变分自编码器（VAE）的损失函数推导与优化过程
题目描述
变分自编码器（VAE）是一种生成模型，结合了自编码器的结构和概率图模型的推断方法。其核心目标是通过学习数据的潜在分布，生成新的样本。VAE的难点在于如何推导其损失函数（即证据下界ELBO），并通过重参数化技巧实现梯度下降优化。本题将详细讲解VAE的损失函数推导步骤，包括概率图模型假设、ELBO的分解、重参数化技巧的应用，以及优化过程的具体实现。
解题过程
-
概率图模型与问题定义
- 假设观测数据 \(\mathbf{x}\) 由潜在变量 \(\mathbf{z}\) 生成，其联合概率为 \(p(\mathbf{x}, \mathbf{z}) = p(\mathbf{z}) p(\mathbf{x} \mid \mathbf{z})\)。
- 目标：学习生成模型 \(p(\mathbf{x} \mid \mathbf{z})\) 和推断模型 \(q(\mathbf{z} \mid \mathbf{x})\)（近似真实后验 \(p(\mathbf{z} \mid \mathbf{x})\)）。
- 挑战：直接最大化对数似然 \(\log p(\mathbf{x})\) 难解，因为边缘概率 \(p(\mathbf{x}) = \int p(\mathbf{x} \mid \mathbf{z}) p(\mathbf{z}) d\mathbf{z}\) 涉及高维积分。
-
证据下界（ELBO）的推导
- 通过詹森不等式，将对数似然转化为下界：
\[ \log p(\mathbf{x}) = \log \mathbb{E}_{q(\mathbf{z} \mid \mathbf{x})} \left[ \frac{p(\mathbf{x}, \mathbf{z})}{q(\mathbf{z} \mid \mathbf{x})} \right] \geq \mathbb{E}_{q(\mathbf{z} \mid \mathbf{x})} \left[ \log \frac{p(\mathbf{x}, \mathbf{z})}{q(\mathbf{z} \mid \mathbf{x})} \right] = \text{ELBO}. \]
- ELBO可分解为两项：
\[ \text{ELBO} = \mathbb{E}_{q(\mathbf{z} \mid \mathbf{x})} [\log p(\mathbf{x} \mid \mathbf{z})] - D_{\text{KL}}(q(\mathbf{z} \mid \mathbf{x}) \| p(\mathbf{z})). \]
第一项是重构损失，衡量生成数据与真实数据的相似度；第二项是KL散度，约束近似后验 $ q(\mathbf{z} \mid \mathbf{x}) $ 接近先验 $ p(\mathbf{z}) $（通常假设为标准正态分布）。
- 重参数化技巧
- 直接采样 \(\mathbf{z} \sim q(\mathbf{z} \mid \mathbf{x})\) 会导致梯度无法反向传播（采样操作不可导）。
- 解决方法：将采样过程分离为可导的确定性部分和随机噪声。例如，若 \(q(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}(\mu, \sigma^2)\)，则令：
\[ \mathbf{z} = \mu + \sigma \cdot \epsilon, \quad \epsilon \sim \mathcal{N}(0, 1), \]
其中 $ \mu, \sigma $ 由编码器网络输出，$ \epsilon $ 为随机变量。梯度可通过 $ \mu $ 和 $ \sigma $ 传递，而 $ \epsilon $ 不参与梯度计算。
- 损失函数的具体形式
- 假设先验 \(p(\mathbf{z}) = \mathcal{N}(0, I)\)，近似后验 \(q(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}(\mu, \sigma^2 I)\)，则KL散度有闭式解：
\[ D_{\text{KL}} = -\frac{1}{2} \sum_{j=1}^J \left( 1 + \log \sigma_j^2 - \mu_j^2 - \sigma_j^2 \right), \]
其中 $ J $ 是潜在空间的维度。
- 重构损失通常用均方误差（连续数据）或交叉熵（离散数据）表示，例如：
\[ \mathbb{E}_{q} [\log p(\mathbf{x} \mid \mathbf{z})] \approx -\|\mathbf{x} - \text{Decoder}(\mathbf{z})\|^2. \]
- 最终损失函数为：
\[ \mathcal{L} = \text{重构损失} + D_{\text{KL}}. \]
- 优化过程
- 编码器网络输入 \(\mathbf{x}\)，输出 \(\mu\) 和 \(\log \sigma^2\)（确保 \(\sigma > 0\)）。
- 通过重参数化得到 \(\mathbf{z} = \mu + \exp(\frac{1}{2} \log \sigma^2) \cdot \epsilon\)。
- 解码器网络将 \(\mathbf{z}\) 映射为重构数据 \(\mathbf{x}'\)。
- 梯度下降同时优化编码器和解码器参数，最小化总损失 \(\mathcal{L}\)。
关键点总结
- ELBO将不可解的对数似然转化为可优化的下界。
- 重参数化技巧使随机采样可导，实现端到端训练。
- KL散度项充当正则化器，确保潜在空间具有连续性（相邻点解码后内容相似）。