[论文解读] Jigsaw-VAE: Towards Balancing Features in Variational Autoencoders

Saeid Asgari Taghanaki, Mohammad Havaei+4 more|arXiv (Cornell University)|May 12, 2020

Generative Adversarial Networks and Image Synthesis参考文献 41被引用 6

一句话总结

本文提出Jigsaw-VAE，一种变分自编码器，通过应用随机空间置换和混合先验，缓解潜在表征中的特征不平衡问题，实现对主导特征和稀有特征的均衡学习。该方法在分布偏移下的下游聚类任务中表现更优，并相较于基线VAE提升了生成样本的多样性与真实感。

ABSTRACT

The latent variables learned by VAEs have seen considerable interest as an unsupervised way of extracting features, which can then be used for downstream tasks. There is a growing interest in the question of whether features learned on one environment will generalize across different environments. We demonstrate here that VAE latent variables often focus on some factors of variation at the expense of others - in this case we refer to the features as ``imbalanced''. Feature imbalance leads to poor generalization when the latent variables are used in an environment where the presence of features changes. Similarly, latent variables trained with imbalanced features induce the VAE to generate less diverse (i.e. biased towards dominant features) samples. To address this, we propose a regularization scheme for VAEs, which we show substantially addresses the feature imbalance problem. We also introduce a simple metric to measure the balance of features in generated images.

研究动机与目标

• 为解决VAE中的特征不平衡问题，即模型过度拟合于颜色等主导特征而忽略形状等稀有特征。
• 提升VAE表征在测试数据分布偏离训练分布时的下游任务（如聚类）泛化能力。
• 开发一种方法，以促进潜在空间中多个变化因子（如颜色与结构）的均衡学习。
• 引入一种新指标——特征存在度量（FPM），以定量评估生成样本对训练数据特征的保留程度。
• 证明基于置换的正则化可有效降低VAE中的表征偏差，同时不损害重建或采样质量。

提出的方法

• 引入一种随机拼图置换层，在编码前对输入图像应用随机空间拼贴与打乱，迫使VAE学习局部结构关系。
• 在潜在空间采用混合先验，以缓解似然项与KL散度项之间的冲突，改善后验对齐。
• 对每个样本独立应用置换操作，确保每次前向传播均看到图像块的不同排列。
• 使用改进的VAE目标函数，结合置换不变的重建损失与混合先验，以促进均衡特征学习。
• 通过标准证据下界（ELBO）端到端训练模型，并引入新的正则化组件。
• 通过在VAE学习的潜在码上训练聚类头，将模型适配至下游聚类任务，并通过归一化互信息（NMI）评估特征平衡性。

实验结果

研究问题

• RQ1 在VAE中应用随机输入置换是否能降低对图像数据中主导特征（如颜色）的偏见？
• RQ2 与标准VAE及其他解耦方法相比，所提出的Jigsaw-VAE方法是否能提升潜在空间中学习特征的平衡性？
• RQ3 当测试分布向单一特征（如所有数字均被涂成黄色）偏移时，该模型在下游聚类任务中的表现如何？
• RQ4 所提出的特征存在度量（FPM）是否能有效量化生成样本中对训练数据特征的保留程度？
• RQ5 与现有VAE变体（如β-VAE和d-VAE）相比，Jigsaw-VAE是否能生成更多样化且更真实的样本？

主要发现

• 在单色MNIST测试集上，Jigsaw-VAE的NMI得分为0.9473，显著优于β-VAE（0.8766）和VAE（0.8872），表明其在颜色与形状特征之间实现了更好的平衡。
• 在单色测试集上，Jigsaw-VAE的中位NMI为0.379，而Mixup-VAE为0.114，β-VAE仅为7.8e-6，表明其对特征分布偏移具有更强鲁棒性。
• 在重建任务中，即使输入数字为均匀着色，Jigsaw-VAE仍能成功保留生成样本中的颜色与形状，而传统非Jigsaw模型则会错误地根据主导聚类分配颜色。
• 所提出的特征存在度量（FPM）显示，尽管β-VAE在标准设置下MSE和FPM表现接近，但Jigsaw-VAE在生成图像中保留了更多来自训练数据的特征。
• Jigsaw-VAE在样本间插值更加平滑，避免了特征的突变，表明其潜在空间具有更好的解耦性与连续性。
• 在CelebA数据集上，Jigsaw-VAE生成的样本比β-VAE更具真实感与多样性，尽管β-VAE在FPM和MSE上得分较高，但其生成图像视觉质量较差。