[论文解读] Reweighted Expectation Maximization

Adji Bousso Dieng, John Paisley|arXiv (Cornell University)|Jun 13, 2019

Generative Adversarial Networks and Image Synthesis参考文献 46被引用 12

一句话总结

本文提出重加权期望最大化（REM），一种基于EM的深度生成模型训练方法，通过使用矩匹配提议分布，将生成模型与识别网络的学习动态解耦，从而在变分自编码器（VAE）和重要性加权自编码器（IWAE）的基础上实现改进。REM通过直接最大化边缘似然，避免后验坍缩和迁移误差，显著提升了密度估计基准上的对数似然性能。

ABSTRACT

Training deep generative models with maximum likelihood remains a challenge. The typical workaround is to use variational inference (VI) and maximize a lower bound to the log marginal likelihood of the data. Variational auto-encoders (VAEs) adopt this approach. They further amortize the cost of inference by using a recognition network to parameterize the variational family. Amortized VI scales approximate posterior inference in deep generative models to large datasets. However it introduces an amortization gap and leads to approximate posteriors of reduced expressivity due to the problem known as posterior collapse. In this paper, we consider expectation maximization (EM) as a paradigm for fitting deep generative models. Unlike VI, EM directly maximizes the log marginal likelihood of the data. We rediscover the importance weighted auto-encoder (IWAE) as an instance of EM and propose a new EM-based algorithm for fitting deep generative models called reweighted expectation maximization (REM). REM learns better generative models than the IWAE by decoupling the learning dynamics of the generative model and the recognition network using a separate expressive proposal found by moment matching. We compared REM to the VAE and the IWAE on several density estimation benchmarks and found it leads to significantly better performance as measured by log-likelihood.

研究动机与目标

• 为解决变分推断在深度生成模型中的局限性，特别是后验坍缩和迁移误差问题。
• 将深度生成模型的训练重新构型为期望最大化（EM）问题，以直接优化对数似然。
• 通过解耦生成模型与识别网络的学习动态，改进重要性加权自编码器（IWAE）。
• 证明基于EM的训练可带来更优的生成模型和更丰富的后验近似。
• 表明当识别网络与生成模型的学习动态解耦时，可作为有效的提议分布。

提出的方法

• 将深度生成模型的训练重新构型为EM算法，其中E步计算提议分布，M步通过优化模型参数以最大化边缘似然。
• 重新发现IWAE是EM的一个实例，表明重要性加权对应于使用多粒子的特定形式EM。
• 引入矩匹配过程，学习一个丰富且表达能力强的提议分布 $ s(\mathbf{z}) $，其独立于识别网络 $ r_{\eta}(\mathbf{z}|\mathbf{x}) $。
• 在E步中使用矩匹配提议 $ s(\mathbf{z}) $，将识别网络与生成模型的训练解耦，防止共适应。
• 使用通过矩匹配获得的独立且表达能力强的提议 $ s(\mathbf{z}) $，在EM更新中改进后验近似，而无需依赖识别网络。
• 在M步中采用重加权似然目标，利用多个粒子更准确地近似真实边缘似然。

实验结果

研究问题

• RQ1 期望最大化（EM）能否有效应用于深度生成模型的训练，是否优于变分推断？
• RQ2 重要性加权自编码器（IWAE）是否为EM的一个特例，该洞见能否带来更优的训练算法？
• RQ3 将识别网络与生成模型的学习动态解耦，是否能减少后验坍缩并提升模型质量？
• RQ4 矩匹配提议分布 $ s(\mathbf{z}) $ 是否可作为EM训练中比识别网络更优的后验近似？
• RQ5 REM在标准密度估计基准上是否优于VAE和IWAE实现更好的对数似然性能？

主要发现

• REM在所有测试数据集（包括MNIST、随机MNIST和Omniglot）上，对数似然性能显著优于VAE和IWAE。
• 在Omniglot数据集上，REM（v2）的测试对数似然达到114.73，超过IWAE的113.94和VAE的109.84。
• 识别网络 $ r_{\eta}(\mathbf{z}|\mathbf{x}) $ 被发现是EM更新中的有效提议，在某些情况下优于矩匹配提议 $ s(\mathbf{z}) $。
• 随着粒子数 $ K $ 增加，IWAE的后验质量下降，表现为对先验的KL散度上升，表明其后验近似存在不稳定性。
• REM学习到的提议优于VAE和IWAE，表现为对先验的KL散度更低，表明后验坍缩减少。
• 在REM中使用矩匹配提议 $ s(\mathbf{z}) $ 有效解耦了识别网络与生成模型，防止了共适应，提升了模型表达能力。