[论文解读] Disentangled Representation Learning with Wasserstein Total Correlation

Yijun Xiao, William Yang Wang|arXiv (Cornell University)|Dec 30, 2019

Generative Adversarial Networks and Image Synthesis参考文献 36被引用 8

一句话总结

本文提出Wasserstein总相关性（WTC）以通过使用更鲁棒的基于Wasserstein距离的替代方法取代基于KL散度的总相关性，来改进解耦表征学习。通过在VAE和WAE框架中训练一个判别器网络来估计WTC，该方法在最小重建误差下实现了最先进水平的解耦，优于多个基准测试中的现有方法，同时引入了一种新度量——Wasserstein依赖差距（WDG），以实现更可靠的评估。

ABSTRACT

Unsupervised learning of disentangled representations involves uncovering of different factors of variations that contribute to the data generation process. Total correlation penalization has been a key component in recent methods towards disentanglement. However, Kullback-Leibler (KL) divergence-based total correlation is metric-agnostic and sensitive to data samples. In this paper, we introduce Wasserstein total correlation in both variational autoencoder and Wasserstein autoencoder settings to learn disentangled latent representations. A critic is adversarially trained along with the main objective to estimate the Wasserstein total correlation term. We discuss the benefits of using Wasserstein distance over KL divergence to measure independence and conduct quantitative and qualitative experiments on several data sets. Moreover, we introduce a new metric to measure disentanglement. We show that the proposed approach has comparable performances on disentanglement with smaller sacrifices in reconstruction abilities.

研究动机与目标

• 解决在估计解耦表征学习中总相关性的KL散度的不稳定性与度量无关性问题。
• 提出一种基于Wasserstein距离的、更具鲁棒性且具备几何感知能力的KL散度替代方法，用于总相关性估计。
• 引入一种新的解耦度量——Wasserstein依赖差距（WDG），其原理与WTC一致，能更准确地捕捉潜在因子间的独立性。
• 在无监督表征学习中，实现解耦与重建质量之间的更好权衡。

提出的方法

• 提出Wasserstein总相关性（WTC）作为基于KL散度的总相关性的度量无关、具备几何感知能力的替代方法。
• 在VAE和Wasserstein自编码器（WAE）框架中，通过对抗训练方式训练一个判别器网络来估计WTC项。
• 将WTC整合进自编码目标函数中，并引入正则化强度超参数γ，以平衡解耦与重建性能。
• 提出基于与WTC相同原理的新型解耦度量——Wasserstein依赖差距（WDG）。
• 通过梯度惩罚训练判别器网络，以确保Lipschitz连续性，从而实现对WTC项的稳定估计。
• 采用蒙特卡洛采样与小批量训练策略，以可扩展的方式优化WTC正则化目标。

实验结果

研究问题

• RQ1 与KL散度相比，基于Wasserstein距离的总相关性公式是否能提升解耦表征学习的稳定性与鲁棒性？
• RQ2 所提出的Wasserstein总相关性（WTC）是否能在保持低重建误差的同时带来更好的解耦性能？
• RQ3 新提出的Wasserstein依赖差距（WDG）度量是否能比现有度量提供更可靠、更一致的解耦评估？
• RQ4 WTC的正则化强度如何影响解耦与重建性能之间的权衡？

主要发现

• WTC-VAE在dSprites、Shapes3D和Cars3D三个数据集上均实现了与最先进模型相当或更优的解耦得分，且在所有基于VAE的方法中重建误差最低。
• WTC-WAE在标准WAE基础上表现出一致的性能提升，表明WTC在WAE框架中的有效性。
• 所提出的Wasserstein依赖差距（WDG）度量与FactorVAE、MIG和Modularity具有高度等级相关性，表明其是一种可靠且一致的评估度量。
• 在dSprites和Shapes3D数据集上，随着WTC正则化强度γ的增加，WDG得分也相应提高，表明更强的正则化带来了更优的解耦性能。
• 在CelebA数据集上，WTC-VAE成功发现了有意义且细粒度的潜在因子，如卷发和面部宽度，展示了其定性解耦能力。
• 即使在高解耦水平下，模型仍保持了强大的重建性能，使WTC-VAE位于解耦-重建权衡图的左上区域。