[论文解读] Mitigating Modality Collapse in Multimodal VAEs via Impartial Optimization

Adrián Javaloy, Maryam Meghdadi+1 more|arXiv (Cornell University)|Jun 9, 2022

AI in cancer detection被引用 6

一句话总结

本文提出无偏优化以缓解多模态变分自编码器（VAEs）中的模态崩溃问题，即模型因梯度冲突而忽略某些模态。通过识别计算图中的“无偏性模块”，并应用多任务学习解决方案（如梯度重加权），该方法确保所有模态的优化保持平衡，显著提升了在多种数据集和模型架构下的重建性能、条件生成能力以及潜在空间的一致性。

ABSTRACT

A number of variational autoencoders (VAEs) have recently emerged with the aim of modeling multimodal data, e.g., to jointly model images and their corresponding captions. Still, multimodal VAEs tend to focus solely on a subset of the modalities, e.g., by fitting the image while neglecting the caption. We refer to this limitation as modality collapse. In this work, we argue that this effect is a consequence of conflicting gradients during multimodal VAE training. We show how to detect the sub-graphs in the computational graphs where gradients conflict (impartiality blocks), as well as how to leverage existing gradient-conflict solutions from multitask learning to mitigate modality collapse. That is, to ensure impartial optimization across modalities. We apply our training framework to several multimodal VAE models, losses and datasets from the literature, and empirically show that our framework significantly improves the reconstruction performance, conditional generation, and coherence of the latent space across modalities.

研究动机与目标

• 解决多模态VAE中的模态崩溃问题，即模型在训练过程中偏好某些模态而忽略其他模态。
• 将模态崩溃的根本原因归因于计算图中特定子结构内的梯度冲突，这些子结构被称为“无偏性模块”。
• 开发一种可泛化的训练框架，通过利用现有的多任务学习解决方案，实现在所有模态上的无偏优化。
• 通过实证验证该框架在多种VAE架构、损失函数和多模态数据集上的有效性。
• 提升所有模态的联合、边缘和条件分布建模能力，从而增强重建、生成与解耦性能。

提出的方法

• 识别‘无偏性模块’——即在反向传播过程中，来自不同模态的梯度发生冲突的计算图子图结构。
• 应用多任务学习中的梯度冲突缓解技术（如梯度重加权或不确定性加权），以平衡各模态间的优化。
• 将无偏优化框架无缝集成到现有的多模态VAE模型（如MVAE、MMVAE、MoPoE）中，无需修改网络架构。
• 采用不同的训练目标（ELBO、IWAE、SIWAE）以评估所提方法的鲁棒性与泛化能力。
• 通过模态维度对数似然分数进行归一化，以公平比较在异构模态上的性能提升。
• 使用多个随机种子进行训练，并报告均值与标准差，以确保统计可靠性。

实验结果

研究问题

• RQ1 多模态VAE中的模态崩溃由何原因引起？是否可追溯至计算图中特定位置的梯度冲突？
• RQ2 现有的多任务学习解决方案在梯度冲突缓解方面是否可有效适配，以提升多模态VAE训练的公平性？
• RQ3 无偏优化是否能在多种模态上带来可测量的重建、条件生成与潜在空间质量提升？
• RQ4 所提出的框架在不同VAE架构、损失函数与数据集上的表现如何？
• RQ5 无偏优化在多模态设置下在多大程度上减少了负迁移现象，并改善了联合分布建模？

主要发现

• 所提出的无偏优化框架显著提升了所有模态的重建性能，MVAE在ELBO损失下联合对数似然提升1.22（从-8.61提升至-8.07）。
• 对于MMVAE，该方法将'S'模态的对数似然从-2.18降低至-2.31（ELBO损失下），表明对以往被忽略模态的拟合能力得到改善。
• 在MoPoE模型中采用IWAE损失时，'T'模态的对数似然从-1.19降至-1.27，表明对先前欠拟合模态的建模能力得到提升。
• 该方法在条件生成与潜在空间一致性方面持续表现优异，表现为所有模态子集的联合与条件对数似然均得到改善。
• 该框架在所有评估数据集与模型上均达到最先进性能，且在5个随机种子下均表现出统计显著的性能提升。
• 在高维或优化程度较低的模态（如图文VAE中的文本模态）中，性能提升最为显著，证实了模态崩溃的有效缓解。