对MMVAE中IWAE代码实现的理解

原始的IWAE

优化目标：
$${\mathcal{L}}_{\mathrm{I}\mathrm{W}\mathrm{A}\mathrm{E}}\left({\mathbf{x}}_{1:M}\right)={\mathbb{E}}_{{\mathbf{z}}^{1:K}\sim q_{\Phi }\left(\mathbf{z}\mid {\mathbf{x}}_{1:M}\right)}\left[\log \sum _{k=1}^K\frac{1}{K}\frac{p_{\Theta }\left({\mathbf{z}}^k,{\mathbf{x}}_{1:M}\right)}{q_{\Phi }\left({\mathbf{z}}^k\mid {\mathbf{x}}_{1:M}\right)}\right]（1）$$

这里 $p_{\Theta }\left(\mathbf{z},{\mathbf{x}}_{1:M}\right)=p(\mathbf{z}){\prod }_{m=1}^M{p}_{{\theta }_m}\left({\mathbf{x}}_m\mid \mathbf{z}\right)$
后验分布由推理网络近似得到 $q_{\Phi }\left({\mathbf{z}}^k\mid {\mathbf{x}}_{1:M}\right)$

MMVAE中的IWAE变体

采用MoE方法进行多模态融合的优化目标：
$${\mathcal{L}}_{\mathrm{I}\mathrm{W}\mathrm{A}\mathrm{E}}^{\mathrm{M}\mathrm{o}\mathrm{E}}\left({\mathbf{x}}_{1:M}\right)=\frac{1}{M}\sum _{m=1}^M{\mathbb{E}}_{{\mathbf{z}}_m^{1:K}\sim q_{{\varphi }_m}\left(\mathbf{z}\mid {\mathbf{x}}_m\right)}\left[\log \frac{1}{K}\sum _{k=1}^K\frac{p_{\Theta }\left({\mathbf{z}}_m^k,{\mathbf{x}}_{1:M}\right)}{q_{\Phi }\left({\mathbf{z}}_m^k\mid {\mathbf{x}}_{1:M}\right)}\right]（2）$$

根据MoE方法近似的后验分布为：$q_{\Phi }\left(\mathbf{z}\mid {\mathbf{x}}_{1:M}\right)={\sum }_m{\alpha }_m\cdot q_{{\varphi }_m}\left(\mathbf{z}\mid {\mathbf{x}}_m\right)$，这里$\alpha =\frac{1}{M}$

计算IWAE的主体代码：

• .log_prob(value)是计算value在定义的概率分布中对应的概率的对数。
• log_mean_exp(value)在后面介绍

在for循环里面一行行的分析，以r=0为例：

• lpz = $logp(z_1)$， 每个潜在变量的尺寸：[K, batch size, latent dim]，在这里用sum(-1)相当于是将潜在变量由latent dim压缩到1维
• lqz_x = $log[q(z_1|x_1)+q(z_1|x_2)]$
• lpx_z = $logp(x_1|z_1)+logp(x_2|z_1)$
• lw = lpz + lpx_z + lqz_x

最后运算：

l1 = log_mean_exp(lw, dim=0)

就可以得到:$$\frac{p(z_1)\cdot (x_1|z_1)\cdot (x_2|z_1)}{q(z_1|x_1)+q(z_1|x_2)}（3）$$

这个结果就是上述公式2中m=1时的结果，这样一行行的分析就可以很好的理解上述代码是如何实现IWAE多模态变体的。

log_mean_exp

其中log_mean_exp的代码：

def log_mean_exp(value, dim=0, keepdim=False):
    return torch.logsumexp(value, dim, keepdim=keepdim) - math.log(value.size(dim))

log_mean_exp和torch.logsumexp的区别就是字面意思，前面取平均，后者求和

• 因为MMVAE中的后验分布为$q_{\Phi }\left(\mathbf{z}\mid {\mathbf{x}}_{1:M}\right)={\sum }_m{\alpha }_m\cdot q_{{\varphi }_m}\left(\mathbf{z}\mid {\mathbf{x}}_m\right)$，这里$\alpha =\frac{1}{M}$，即需要对上述式子3中的分母取平均，所以log_mean_exp可以写成下述公式：
  $$\mathrm{logmeanexp}(x{)}_i=\log \frac{1}{j}\sum _j\exp \left(x_{ij}\right)=\log \sum _j\exp (x_{ij})-\log j$$
• torch.logsumexp的介绍截图自官网：