# Mitigating Modality Collapse in Multimodal VAEs via Impartial Optimization

arXiv: https://arxiv.org/abs/2206.04496 · HTML: https://mlsanity.com/paper/2206.04496/

## Key facts

- Cited by 32 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 41 references that are in the map.
- First posted on arXiv 9 Jun 2022.

## Builds on (most cited first, 25 of 25)

- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Multimodal Machine Learning: A Survey and Taxonomy](https://mlsanity.com/paper/1705.09406/) (2017)
- [On the Convergence of Adam and Beyond](https://mlsanity.com/paper/1904.09237/) (2019)
- [GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks](https://mlsanity.com/paper/1711.02257/) (2017)
- [Multi-Task Learning as Multi-Objective Optimization](https://mlsanity.com/paper/1810.04650/) (2018)
- [An Overview of Multi-Task Learning in Deep Neural Networks](https://mlsanity.com/paper/1706.05098/) (2017)
- [NVAE: A Deep Hierarchical Variational Autoencoder](https://mlsanity.com/paper/2007.03898/) (2020)
- [Importance Weighted Autoencoders](https://mlsanity.com/paper/1509.00519/) (2015)
- [Just Pick a Sign: Optimizing Deep Multitask Models with Gradient Sign Dropout](https://mlsanity.com/paper/2010.06808/) (2020)
- [Multimodal Generative Models for Scalable Weakly-Supervised Learning](https://mlsanity.com/paper/1802.05335/) (2018)
- [Variational Mixture-of-Experts Autoencoders for Multi-Modal Deep Generative Models](https://mlsanity.com/paper/1911.03393/) (2019)
- [Handling Incomplete Heterogeneous Data using VAEs](https://mlsanity.com/paper/1807.03653/) (2018)
- [From Variational to Deterministic Autoencoders](https://mlsanity.com/paper/1903.12436/) (2019)
- [Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference](https://mlsanity.com/paper/1703.09194/) (2017)
- [Generalized Multimodal ELBO](https://mlsanity.com/paper/2105.02470/) (2021)
- [Tighter Variational Bounds are Not Necessarily Better](https://mlsanity.com/paper/1802.04537/) (2018)
- [VAEM: a Deep Generative Model for Heterogeneous Mixed Type Data](https://mlsanity.com/paper/2006.11941/) (2020)
- [Multimodal Generative Learning Utilizing Jensen-Shannon-Divergence](https://mlsanity.com/paper/2006.08242/) (2020)
- [MultiNet++: Multi-Stream Feature Aggregation and Geometric Loss Strategy for Multi-Task Learning](https://mlsanity.com/paper/1904.08492/) (2019)
- [Doubly Reparameterized Gradient Estimators for Monte Carlo Objectives](https://mlsanity.com/paper/1810.04152/) (2018)
- [A Variational Auto-Encoder Model for Stochastic Point Processes](https://mlsanity.com/paper/1904.03273/) (2019)
- [Relating by Contrasting: A Data-efficient Framework for Multimodal Generative Models](https://mlsanity.com/paper/2007.01179/) (2020)
- [Automatic Differentiation Variational Inference with Mixtures](https://mlsanity.com/paper/2003.01687/) (2020)
- [Medical data wrangling with sequential variational autoencoders](https://mlsanity.com/paper/2103.07206/) (2021)

## Built on by (most cited first, 30 of 32)

- [Multimodal learning with graphs](https://mlsanity.com/paper/2209.03299/) (2022)
- [A Closer Look at Multimodal Representation Collapse](https://mlsanity.com/paper/2505.22483/) (2025)
- [AlignRec: Aligning and Training in Multimodal Recommendations](https://mlsanity.com/paper/2403.12384/) (2024)
- [Multimodal Data Integration for Oncology in the Era of Deep Neural Networks: A Review](https://mlsanity.com/paper/2303.06471/) (2023)
- [M3-JEPA: Multimodal Alignment via Multi-gate MoE based on the Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2409.05929/) (2024)
- [Neurosymbolic Diffusion Models](https://mlsanity.com/paper/2505.13138/) (2025)
- [Multimodal Data Integration for Precision Oncology: Challenges and Future Directions](https://mlsanity.com/paper/2406.19611/) (2024)
- [Multi-modal Latent Diffusion](https://mlsanity.com/paper/2306.04445/) (2023)
- [Transductive Decoupled Variational Inference for Few-Shot Classification](https://mlsanity.com/paper/2208.10559/) (2022)
- [Cooperation in the Latent Space: The Benefits of Adding Mixture Components in Variational Autoencoders](https://mlsanity.com/paper/2209.15514/) (2022)
- [Multi-task View Synthesis with Neural Radiance Fields](https://mlsanity.com/paper/2309.17450/) (2023)
- [Learning multi-modal generative models with permutation-invariant encoders and tighter variational objectives](https://mlsanity.com/paper/2309.00380/) (2023)
- [Enhance the Robustness of Text-Centric Multimodal Alignments](https://mlsanity.com/paper/2407.05036/) (2024)
- [COPA: Comparing the incomparable in multi-objective model evaluation](https://mlsanity.com/paper/2503.14321/) (2025)
- [Hellinger Multimodal Variational Autoencoders](https://mlsanity.com/paper/2601.06572/) (2026)
- [Improving Multimodal Joint Variational Autoencoders through Normalizing Flows and Correlation Analysis](https://mlsanity.com/paper/2305.11832/) (2023)
- [Towards Relaxed Multimodal Inputs for Gait-based Parkinson's Disease Assessment](https://mlsanity.com/paper/2510.15748/) (2025)
- [Modality Collapse as Mismatched Decoding: Information-Theoretic Limits of Multimodal LLMs](https://mlsanity.com/paper/2602.23136/) (2026)
- [SiMO: Single-Modality-Operable Multimodal Collaborative Perception](https://mlsanity.com/paper/2603.08240/) (2026)
- [A Sobering Look at Tabular Data Generation via Probabilistic Circuits](https://mlsanity.com/paper/2603.23016/) (2026)
- [Information Router for Mitigating Modality Dominance in Vision-Language Models](https://mlsanity.com/paper/2604.16264/) (2026)
- [Robust Multi-Omics Integration from Incomplete Modalities Significantly Improves Prediction of Alzheimer's Disease](https://mlsanity.com/paper/2509.20842/) (2025)
- [When One Modality Rules Them All: Backdoor Modality Collapse in Multimodal Diffusion Models](https://mlsanity.com/paper/2603.06508/) (2026)
- [Prompts Without Evidence: How Neuroimaging Mentions Shift Clinical Vision-Language Model Predictions](https://mlsanity.com/paper/2603.28387/) (2026)
- [Token-Efficient Multimodal Reasoning via Image Prompt Packaging](https://mlsanity.com/paper/2604.02492/) (2026)
- [MiMIC: Mitigating Visual Modality Collapse in Universal Multimodal Retrieval While Avoiding Semantic Misalignment](https://mlsanity.com/paper/2604.21326/) (2026)
- [GraphPL: Leveraging GNN for Efficient and Robust Modalities Imputation in Patchwork Learning](https://mlsanity.com/paper/2604.25352/) (2026)
- [Learning Emergent Modular Representations in Multi-modality Medical Vision Foundation Models](https://mlsanity.com/paper/2605.21861/) (2026)
- [H"older++: Improving the Quality-Coherence Trade-off in Multimodal VAEs](https://mlsanity.com/paper/2606.13381/) (2026)
- [ShapKO: Shapley-Adaptive Modality Knockout for Robust Multimodal Learning](https://mlsanity.com/paper/2607.09884/) (2026)

## Abstract (arXiv)

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.