Derive the EM algorithm to estimate the parameters of GMM,by minimizing free energies.
1 Introduction
回顾，在生成式模型中，我们希望求出数据的一个parameterized distribution最大化下面的似然概率：
$
P(X) = \prod_{i=1}^{m} P(x_i|\theta)
$
其中 $ \theta $ 为分布的参数，总共有 $ m $ 个样本点。对上述取 $ \log $，改写成我们的目标函数：
$
L(\theta) = \sum_{i=1}^{m} \log P(x_i|\theta)
$
2 Gaussian Mixture Model
高斯混合模型假设我们的训练数据服从一系列高斯分布，每个样本点属于某个高斯分布的概率与其属于哪个高斯分布有关。具体来说，我们假设数据服从以下的联合概率分布：
\(P(x^{(i)}, z^{(i)}) = P(x^{(i)}|z^{(i)}) Q(z^{(i)})\)
其中 $ z $ 服从多项分布 $ z^{(i)} \sim \text{Multinomial}(\phi) $，且 \(\phi_j \geq 0\)，\(\sum_{j=1} \phi_j = 1\)，\(\phi_j = P(z^{(i)} = j)\)。且 $ z $ 可以取 $ k $ 个值。$ x|z $ 服从正态分布，\(x^{(i)}|z^{(i)}=j \sim \mathcal{N}(\mu_j, \Sigma_j)\)。
换句话说，我们的模型假设 $ x $ 的生成过程是，先从 $ z $ 中取一个值表示这个 $ x $ 属于哪个高斯分布，然后再从对应的高斯分布生成 $ x $。
3 Free Energies
基于高斯混合模型的假设，我们有：
$
L(\theta) = \sum_{i=1}^{m} \log P(x_i|\theta)
$
$
= \sum_{i=1}^{m} \log \int P(x_i, z| \theta) dz \quad \text{（全概率公式）}
$
$
= \sum_{i=1}^{m} \log \int Q(z) \frac{P(x_i, z|\theta)}{Q(z)} dz
$
$
\geq \sum_{i=1}^{m} \int Q(z) \log \frac{P(x_i, z|\theta)}{Q(z)} dz \quad \text{（Jensen 不等式）}
$
由上述可知，$ \sum_{i=1}^{m} \int Q(z) \log \frac{P(x_i, z|\theta)}{Q(z)} dz $ 为 $ L(\theta) $ 的一个下界，我们可以通过最大化这个下界 ELBO 来得到 $ L(\theta) $ 的一个局部最优解。
令：
$
F(Q, P) = - \sum_{i=1}^{m} \int Q(z) \log \frac{P(x_i, z|\theta)}{Q(z)} dz
$
$
= \sum_{i=1}^{m} \int Q(z) \log \frac{Q(z)}{P(x_i, z|\theta)} dz
$
我们将 $ F(Q, P) $ 称为 free energies，且最大化 ELBO 等价于最小化 free energies。
4 EM Algorithm
4.1 E Step
由 $ P(z|x) = \frac{P(x,z)}{P(x)} $，可得
$
F(Q, P) = \sum_{i=1}^{m} \int Q(z) \log \frac{Q(z)}{P(x_i, z)} dz
$
$
= \sum_{i=1}^{m} \int Q(z) \log \frac{Q(z)}{P(z|x_i)P(x_i)} dz
$
$
= \sum_{i=1}^{m} [D(Q(z)||P(z|x_i)) - \log P(x_i)]
$
在 E Step 中，我们固定 $ \theta $，优化 $ Q(z) $。因为 $ \log P(x_i) = \log \int P(x|z)P(z) dz $ 只与 $ \theta $ 有关，且 KL 散度 $ D(Q(z)||P(z|x)) \geq 0 $，仅当 $ Q(z) = P(z|x) $ 时取等，所以在 E Step 中，我们有：
$
E \text{ Step}: Q_{k+1} = \arg \min_{Q} F(Q, P(\theta)) = P(z|x, \theta_k)
$
即当我们取 $ Q(z) \equiv P(z|x) $ 时，$ F(Q, P) $ 取到最小值。
4.2 M Step
在 M Step 中，我们将 E Step 估计出的 $ Q(z) \equiv P(z|x) $ 回代到 $ F(Q, P) $ 中：
$
F(Q_{k+1}, P(\theta)) = \sum_{i=1}^{m} \int Q(z) \log \frac{Q(z)}{P(x_i, z)} dz
$
$
= \sum_{i=1}^{m} \int P(z|x_i, \theta_k) \log \frac{P(z|x_i, \theta_k)}{P(x_i, z|\theta)} dz
$
$
= \sum_{i=1}^{m} \int P(z|x_i, \theta_k) \log P(z|x_i, \theta_k) dz - \int P(z|x_i, \theta_k) \log P(x_i, z|\theta) dz
$
前一项使用的是上一步的 $ \theta $，因此是个定值，所以我们有：
$ M \text{ Step}: \theta_{k+1} = \arg \min_{\theta} F(Q_{k+1}, P(\theta)) = \arg \min_{\theta} - \sum_{i=1}^{m} \int P(z|x_i, \theta_k) \log P(x_i, z|\theta) dz
$
5. EM For GMM
现在我们将上一节推导的 EM 更新公式代入 GMM 的具体表达式，推导 GMM 的参数更新方法。
M Step 和 E Step
- M Step:
$
\theta_{k+1} = \arg\min_{\theta} - \sum_{i} \int P(z|x_i, \theta_k) \log P(x_i, z|\theta) dz
$ - E Step:
$
Q_{k+1} = P(z | x, \theta_k)
$
在 GMM 中，我们假设：
- $ P(z) \equiv \text{Multinomial}(\phi), \sum_{j=1}^c \phi_j = 1 $
- $ P(x|z=j) \equiv \mathcal{N}(\mu_j, \Sigma_j) $
5.1 E Step
E Step 直接使用贝叶斯公式即可：
$
Q(z = i) = P(z = i | x, \theta_k) = \frac{P(x|z=i) P(z=i)}{\sum_{j} P(x|z=j)P(z=j)}
$
5.2 M Step
目标：最大化：
$
\sum_{i} \int P(z|x_i, \theta_k) \log P(x_i, z|\theta) dz = \sum_{i} \int P(z|x_i, \theta_k) \log [P(x_i|z, \mu, \Sigma) P(z|\phi)] dz
$
可以拆分为：
$
= \sum_{i} \big( E[\log P(x_i|z, \mu, \Sigma)] + E[\log P(z|\phi)] \big)
$
优化 $ \phi $：
$ \phi $ 只与后项 $ E[\log P(z|\phi)] $ 有关：
$
E[\log P(z|\phi)] = \sum_{j} P(z = j | x_i) \log P(z = j) = \sum_{j} P(z = j | x_i) \log \phi_j
$
使用拉格朗日法引入约束条件：
$
L = F(Q, P) + \lambda \big(\sum_{j} \phi_j - 1\big)
$
求导，令导数等于 0：
- $ \frac{\partial L}{\partial \phi_j} = \sum_{i} P(z = j|x_i) / \phi_j + \lambda = 0 $
- $ \frac{\partial L}{\partial \lambda} = \sum_{j} \phi_j - 1 = 0 $
由此可得：
$ \phi_j = \frac{1}{m} \sum_{i} w_j^{(i)}$
其中，$ w_j^{(i)} = P(z = j | x_i) $。
优化 $ \mu $：
$ \mu $ 只与公式中的 $ E[\log P(x_i|z, \mu, \Sigma)] $ 有关：
$
E[\log P(x_i|z, \mu, \Sigma)] = \sum_{j} w_j^{(i)} \log P(x_i|z = j, \mu, \Sigma)
$
展开为：
\(= \sum_{j} w_j^{(i)} \big[\log \sigma_j + \frac{1}{2}\log 2\pi + \frac{(x_i - \mu_j)^2}{2\sigma_j^2}\big]\)
求导，令导数等于 0：
\(\frac{\partial L}{\partial \mu_j} = \sum_{i} w_j^{(i)} \frac{(x_i - \mu_j)}{\sigma_j^2} = 0\)
解得：
\(\mu_j = \frac{\sum_{i} w_j^{(i)} x_i}{\sum_{i} w_j^{(i)}}\)
优化 $ \Sigma = \sigma_1, \sigma_2, \dots, \sigma_c $：
同样只考虑公式中的 \(E[\log P(x_i|z, \mu, \Sigma)]\)：
\(E[\log P(x_i|z, \mu, \Sigma)] = \sum_{j} w_j^{(i)} \big[\log \sigma_j + \frac{1}{2}\log 2\pi + \frac{(x_i - \mu_j)^2}{2\sigma_j^2}\big]\)
求导，令导数等于 0：
\(\frac{\partial F}{\partial \sigma_j^2} = \frac{1}{2}\sum_{i} w_j^{(i)} \big[ \frac{1}{\sigma_j^2} - \frac{(x_i - \mu_j)^2}{\sigma_j^4} \big] = 0\)
解得：
\(\sigma_j^2 = \frac{\sum_{i} w_j^{(i)} (x_i - \mu_j)^2}{\sum_{i} w_j^{(i)}}\)
5.3 Summary
最终的 GMM 的 EM 更新公式如下：
$
\phi_j = \frac{1}{m} \sum_{i=1}^m w_j^{(i)}
$
$
\mu_j = \frac{\sum_{i=1}^m w_j^{(i)} x_i}{\sum_{i=1}^m w_j^{(i)}}
$
\(\sigma_j^2 = \frac{\sum_{i=1}^m w_j^{(i)} (x_i - \mu_j)^2}{\sum_{i=1}^m w_j^{(i)}}\)
