KL,MLE,GMM,EM,VAE

很多统计估计和回归（regression）过程，都可以理解为用参数方程拟合真实概率分布。

基本假设

假设：对于给定的数据 $\{X_1, \cdots, X_n\}$（统计估计）或 $\{(X_1, Y_1), \cdots, (X_n, Y_n)\}$（回归），都存在一个真实概率分布使得 $X_i \sim p(x)$ 或 $(X_i, Y_i) \sim p(x, y)$。

我们的目标是通过一个参数化的概率分布 $p_\theta$ 来拟合 $p$。我们使用 KL 散度作为衡量两个概率分布之间差距的指标，因此，我们希望找到一个能够最小化两个分布之间 KL 散度的参数 $\theta$。将 KL 散度展开后，发现只有 $-\mathbb{E}_{X \sim p(x)}[\log p_\theta(X)]$ 这一项与 $\theta$ 有关，很自然地，我们用观测数据对这一项进行蒙特卡洛估计：

$$ \begin{align} D_{KL}[p(\cdot) || p_\theta(\cdot)] &= \mathbb{E}_{X \sim p(x)} \left[\log \frac{p(X)}{p_\theta(X)}\right] \\ &= -\mathbb{E}_{X \sim p(x)}[\log p_\theta(X)] + \mathbb{E}_{X \sim p(x)}[\log p(X)] \\ &\approx -\frac{1}{N} \sum_{i=1}^N \log p_\theta(X_i) + c, \quad X_i \sim^{iid} p(x) \end{align} $$

发现最终的损失其实就是最大似然估计（MLE）。对于回归问题，参数模型是 $Y$ 对 $X$ 的条件概率分布 $p_\theta(Y|X)$，损失则是 KL 散度在观测值上的期望：

$$ \begin{align} \mathbb{E}_{X \sim p(x)} \left\{ D_{KL}[p(\cdot|X) || p_\theta(\cdot|X)] \right\} &= \mathbb{E}_{X \sim p(x)} \mathbb{E}_{Y \sim p(y|X)} \left[ \log \frac{p(Y|X)}{p_\theta(Y|X)} \right] \\ &\approx -\frac{1}{N} \sum_{i=1}^N \log p_\theta(Y_i|X_i) + c, \quad (X_i, Y_i) \sim^{iid} p(x, y) \end{align} $$

在最大似然估计中，对真实分布的假设还加上了“真实分布与 $p_\theta$ 处于同一参数空间”。

GMM（高斯混合模型）

GMM 其实就是 $p_\theta$ 的一个特例，具体来说：

$$ \begin{align} p_\theta(x) &= \sum_{i=1}^n p_i \mathcal{N}(x, \mu_i, \sigma_i), \\ &\sum_i^n p_i = 1, \quad p_i \geq 0, \\ &\theta = (\mu_1, \cdots, \mu_n, \sigma_1, \cdots, \sigma_n, p_1, \cdots, p_n) \end{align} $$

由于 $p$ 关于 $\theta$ 是非凸的（non-convex），人们试图通过另一种方式解决优化 $\log p_\theta$ 的问题。以上模型还可以等价为：

$$ \begin{align} p_\theta(X|Z) &\sim \mathcal{N}(\mu_Z, \sigma^2_Z), \\ Z &\sim \text{cat}(p_1, \cdots, p_n) \end{align} $$

因为 $X$ 的边缘分布（marginal distribution）与第一种定义完全一致。接下来，ELBO（Evidence Lower Bound）为我们提供了一个 $\log p_\theta$ 关于 $\theta$ 的下界，优化这个下界便可以优化 $\log p_\theta$：

$$ \begin{align} p_\theta(X) &= \mathbb{E}_{Z \sim p_\theta(z)} \left[ p_\theta(X|Z) \right] \\ &= \mathbb{E}_{Z \sim q(z)} \left[ \frac{p_\theta(X, Z)}{q(Z)} \right], \\ \log p_\theta(X) &= \log \mathbb{E}_{Z \sim q(z)} \left[ \frac{p_\theta(X, Z)}{q(Z)} \right] \\ &\geq \mathbb{E}_{Z \sim q(z)} \left[ \log \frac{p_\theta(X, Z)}{q(Z)} \right] \end{align} $$

其中，$q(z)$ 可以是任何概率分布，并且这个 ELBO 下界关于 $\theta$ 是凸的（convex），相对好优化。接下来可以延展到两个话题：

EM 算法
VAE（变分自编码器）

EM 算法

如果我们把 $q(z)$ 定义为 $p_\phi(z|X)$，其中 $p_\phi(z|X) = \frac{p_\phi(X|z) p_\phi(z)}{p_\phi(X)}$，且 $\phi, \theta$ 属于同一个参数空间：

$$ \begin{align} \log p_\theta(X) &= \log \mathbb{E}_{Z \sim p_\phi(z|X)} \left[ \frac{p_\theta(X, Z)}{p_\phi(Z|X)} \right] \\ &\geq \mathbb{E}_{Z \sim p_\phi(z|X)} \left[ \log \frac{p_\theta(X, Z)}{p_\phi(Z|X)} \right] \end{align} $$

并规定：

$$ \begin{align} \theta_{t+1} &= \arg\max_\theta \mathbb{E}_{Z \sim p_{\theta_t}(z|X)} \left[ \log \frac{p_\theta(X, Z)}{p_{\theta_t}(Z|X)} \right] \\ &= \arg\max_\theta \mathbb{E}_{Z \sim p_{\theta_t}(z|X)} \left[ \log p_\theta(X, Z) \right] \\ &\approx \arg\max_\theta \frac{1}{N} \sum_{i=1}^N \log p_\theta(X, Z_i), \quad Z_i \sim p_{\theta_t}(z|X) \end{align} $$

就能一步步优化 ELBO。这个迭代过程要求所有的 $\theta_t$ 属于同一参数空间，并且 $p_\theta(X, Z)$ 和 $p_\theta(Z|X)$ 要源于同一个参数方程。

具体来说，EM 的 E-step 就是计算 $\mathbb{E}_{Z \sim p_{\theta_t}(z|X)} \left[ \log p_\theta(X, Z) \right] \approx \frac{1}{N} \sum_{i=1}^N \log p_\theta(X, Z_i)$，即蒙特卡洛估计；M-step 则是从 $\theta_t$ 求得 $\theta_{t+1}$ 的过程。因此，在 EM 的每一步，我们先使用最新的 $\theta_t$ 对应的 $p_{\theta_t}(z|X)$ 估计 ELBO，然后求其关于 $\theta$ 的最优解。

VAE（变分自编码器）

在 GMM 中，我们的目标是找到能够最准确用 $p_\theta$ 拟合数据 $\{X_1, \cdots, X_N\}$ 的 $\theta$。于是先定义了 $p_\theta(X)$ 的形式，然后由此定义了 $p_\theta(X|Z), p_\theta(Z)$。在这里，$Z$ 能够辅助生成 $X$（可以通过生成 $Z$ 然后从 $X|Z$ 中生成 $X$）。注意到，我们也可以通过先定义一组生成模型 $p_\theta(X|Z), p_\theta(Z)$，然后由此定义 $p_\theta(X)$：

$$ \begin{align} p_\theta(X|Z) &\sim \mathcal{N}(\mu_Z, \sigma^2_Z), \\ Z &\sim \text{cat}(p_1, \cdots, p_n), \\ &\downarrow \\ p_\theta(X|Z) &\sim D_\theta^{(1)}, \\ Z &\sim D_\theta^{(2)}, \end{align} $$

此时 $X$ 的分布是唯一确定的。我们可以考虑将 $D_\theta^{(1)}$ 定义为一个从 $\mathcal{R}^{|Z|} \to \mathcal{R}^{|X|}$ 的神经网络，并规定如何将该网络解释为一个概率密度函数 $p_\theta(X|Z)$。对于 $D_\theta^{(2)}$，可以简单假设为 $\mathcal{N}(0, I)$ 或其他与 $\theta$ 有关（或部分有关）的分布。我们的目标仍是找到 $\arg\max_\theta \log p_\theta(X)$，根据 ELBO，只需要根据问题选择 ELBO 中 $q(z)$ 的具体形式：

$$ \begin{align} ELBO &= \mathbb{E}_{Z \sim q(z)} \left[ \log \frac{p_\theta(X, Z)}{q(Z)} \right] \end{align} $$

我们的目标是得到一个 $Z \to X$ 的生成模型，可以将 $Z$ 理解为 $X$ 在另一（更小）空间上的表示。

在 GMM 的例子中，当 ELBO 被充分优化，$\theta$ 稳定下来后，我们往往可以观察到一个能够刻画数据中主要聚类的混合正态分布，稳定下的 $\theta$ 也恰好指导了 $Z \to X$ 的过程。首先，我们随机从 $n$ 个类别中生成 $z$，根据 $p_\theta(z)$。然后，生成 $X \sim \mathcal{N}(\mu_z, \sigma^2_z)$。$z$ 其实就是 $X$ 在 $\{1, \cdots, n\}$ 中的表示！但这无法回答一个问题：为什么可以通过 ELBO 的最大化来得到 $p_\theta(X|Z)$ 的良好估计？我们可以考虑 ELBO 的等价形式：

$$ ELBO = \log p_\theta(X) - D_{KL}[q(\cdot) || p_\theta(\cdot|X)] $$

其中，$\log p_\theta(X)$ 为原始目标，$-D_{KL}[q(\cdot) || p_\theta(\cdot|X)]$ 为 $q$ 和 $p_\theta(\cdot|X)$ 之间的负 KL 散度。当 ELBO 足够大，即 KL 散度足够小时，$ELBO \approx \log p_\theta(X)$，我们的最优解也更接近最大似然估计。而 $X$ 的生成机理（$p_\theta(X)$）由 $p_\theta(X|Z), p_\theta(z)$ 决定，因此 $X$ 分布的良好估计也是生成机理的良好估计。

因此，我们有理由相信，在神经网络的参数表达下，我们也可以通过最大化 ELBO 来得到 $p_\theta(X|Z)$ 的良好估计。根据 EM 的条件，我们也可以与 GMM 做一样的选择，在 ELBO 中使得 $q(z) = p_\phi(z|X)$（与 $p_\theta(z|X)$ 属于同族但参数不同），然后使用 EM 优化。但问题是，我们只定义了 $p_\phi(Z)$ 和 $p_\phi(X|Z)$，在 GMM 的例子中，我们通过贝叶斯公式得到了 $p_\phi(z|X) = \frac{p_\phi(X|z) p_\phi(z)}{p_\phi(X)}$，因此保证了 $\theta, \phi$ 属于同一参数空间，而当我们使用更复杂的神经网络定义 $p_\phi(X|Z)$ 时，$p_\phi(X)$ 是写不出的（在 GMM 中，我们先定义了 $p_\phi(X)$）。为了解决 $p_\phi(Z|X)$ 求不出的问题，在 VAE 中，我们使用另一个神经网络 $q_\phi(Z|X): \mathcal{R}^{|X|} \to \mathcal{R}^{|Z|}$ 来代替（估计）$p_\phi(Z|X)$。带入 ELBO 中的 $q(z)$，我们最终得到：

$$ \begin{align} ELBO &= \mathbb{E}_{Z \sim q_\phi(z|X)} \left[ \log \frac{p_\theta(X, Z)}{q_\phi(Z|X)} \right] \\ &= \log p_\theta(X) - D_{KL}[q_\phi(\cdot|X) || p_\theta(\cdot|X)] \end{align} $$

此时，由于 $\theta, \phi$ 属于不同的空间，我们就不能使用 EM 了。不过在 $p_\theta(X|Z), q_\phi(Z|X), p_\theta(Z)$ 的正态假设下，我们是可以直接写出 ELBO 的解析形式的，然后就可以用梯度算法求解了。最终的损失函数（-ELBO）还可以理解为在对 $-\log p_\theta(X)$ 进行梯度下降的基础上，加入 $D_{KL}[q_\phi(\cdot|X) || p_\theta(\cdot|X)]$ 作为一个正则项，使得每次更新参数后的分布变化不那么大。

本文回答了以下几个问题：