[论文解读] Balancing Reconstruction Quality and Regularisation in ELBO for VAEs

Shu‐Yu Lin, Stephen Roberts+2 more|arXiv (Cornell University)|Sep 9, 2019

Advancements in Photolithography Techniques参考文献 29被引用 9

一句话总结

本文提出在变分自编码器（VAE）似然函数中学习噪声方差，作为自动平衡ELBO目标中重构质量与先验正则化的一种机制。通过将方差条件化于输入数据并联合优化，该方法提升了生成样本质量，实现了不确定性估计，并减少了对启发式超参数调优的依赖。

ABSTRACT

A trade-off exists between reconstruction quality and the prior regularisation in the Evidence Lower Bound (ELBO) loss that Variational Autoencoder (VAE) models use for learning. There are few satisfactory approaches to deal with a balance between the prior and reconstruction objective, with most methods dealing with this problem through heuristics. In this paper, we show that the noise variance (often set as a fixed value) in the Gaussian likelihood p(x|z) for real-valued data can naturally act to provide such a balance. By learning this noise variance so as to maximise the ELBO loss, we automatically obtain an optimal trade-off between the reconstruction error and the prior constraint on the posteriors. This variance can be interpreted intuitively as the necessary noise level for the current model to be the best explanation of the observed dataset. Further, by allowing the variance inference to be more flexible it can conveniently be used as an uncertainty estimator for reconstructed or generated samples. We demonstrate that optimising the noise variance is a crucial component of VAE learning, and showcase the performance on MNIST, Fashion MNIST and CelebA datasets. We find our approach can significantly improve the quality of generated samples whilst maintaining a smooth latent-space manifold to represent the data. The method also offers an indication of uncertainty in the final generative model.

研究动机与目标

• 解决VAE中长期存在的重构损失与先验正则化之间平衡的挑战，该挑战通常依赖于任意的超参数选择。
• 通过动态调整重构保真度与先验约束之间的权衡，提升样本质量和潜在空间平滑性。
• 提出一种基于数据的、原则性的方法来设定似然函数中的噪声方差，避免使用固定值或启发式数值。
• 通过学习灵活的、输入条件化的方差参数，实现对生成样本和重构样本的不确定性估计。
• 建议使用近似聚合后验分布而非先验进行采样，以缓解分布错位问题并提升生成质量。

提出的方法

• 将似然函数 $ p({\mathbf{x}}|{\mathbf{z}}) $ 参数化为一个均值由VAE参数决定的高斯分布，其方差 $ \sigma^2 $ 与VAE参数联合学习。
• 通过神经网络将方差 $ \sigma^2 $ 条件化于输入数据，使其在不同数据点和维度上可变。
• 通过最大化ELBO目标来优化 $ \sigma^2 $，该目标自然地平衡了重构与正则化。
• 使用高斯混合模型近似聚合后验 $ q_{\phi}({\mathbf{z}}) $，以更好地捕捉真实后验的结构。
• 将学习到的 $ \sigma^2 $ 用作不确定性估计器：较高的值表示重构或生成过程中更高的不确定性。
• 从近似聚合后验而非先验中生成样本，以减少分布错位并提升生成结果的真实性。

实验结果

研究问题

• RQ1 在VAE似然函数中学习噪声方差是否能自动平衡重构与正则化，而无需依赖启发式超参数调优？
• RQ2 输入条件化的灵活方差估计如何影响VAE中的样本质量与不确定性量化？
• RQ3 从聚合后置分布而非先验中采样对生成质量有何影响？
• RQ4 学习到的方差在多大程度上可作为重构与生成过程中可靠的不确定性估计器？
• RQ5 高斯混合模型中组件数量对聚合后验近似及其下游性能的影响如何？

主要发现

• 在MNIST和Fashion MNIST数据集上，优化噪声方差 $ \sigma^2 $ 显著提升了测试对数似然，尤其当 $ \sigma^2 $ 条件化于输入且在不同数据维度上变化时效果更明显。
• 在消融实验中，从近似聚合后验而非先验中采样生成CelebA图像时，FID分数从80.1降至79.9，表明生成质量有所提升。
• 当使用灵活的、输入条件化的 $ \sigma^2 $ 时，与固定标量方差相比，CelebA数据集上的FID分数提升了7.5%。
• 通过学习到的 $ \sigma^2 $ 进行不确定性估计，能有效突出高重构误差区域，例如手写数字'0'的顶部笔画或面部特征多变区域。
• 当高斯混合模型的组件数超过500个后，对聚合后验近似的提升趋于平缓，且在2000个组件以上时FID性能无显著改善。
• 所提方法减少了聚合后验与先验之间的差距，从而在不损害潜在空间平滑性的情况下，生成了更真实且多样化的样本。