Enhancing Foundation VLM Robustness to Missing Modality: Scalable Diffusion for Bi-directional Feature Restoration

TL;DR

This work addresses the vulnerability of Vision-Language Models to missing modalities by introducing a diffusion-based feature restoration framework that operates as a mid-stage training module. It combines Dynamic Modality Gating to suppress noise and Cross-Modal Mutual Learning to enforce bidirectional semantic alignment, enabling high-fidelity restoration conditioned on available modalities. Across four benchmarks and multiple missing-rate settings, the method achieves state-of-the-art zero-shot robustness, highlighting strong generalization and scalability, especially with large-scale mid-stage data. The resulting approach provides a practical, plug-and-play enhancement for foundation VLMs, improving reliability in real-world, incomplete-input scenarios.

Abstract

Vision Language Models (VLMs) typically assume complete modality input during inference. However, their effectiveness drops sharply when certain modalities are unavailable or incomplete. Current research primarily faces two dilemmas: Prompt-based methods struggle to restore missing yet indispensable features and impair generalization of VLMs. Imputation-based approaches, lacking effective guidance, are prone to generating semantically irrelevant noise. Restoring precise semantics while sustaining VLM generalization remains challenging. Therefore, we propose a general missing modality restoration strategy in this paper. We introduce an enhanced diffusion model as a pluggable mid-stage training module to effectively restore missing features. Our strategy introduces two key innovations: (I) Dynamic Modality Gating, which adaptively leverages conditional features to steer the generation of semantically consistent features; (II) Cross-Modal Mutual Learning mechanism, which bridges the semantic spaces of dual encoders to achieve bidirectional alignment. Zero-shot evaluations across benchmark datasets demonstrate that our approach outperforms existing baseline methods. Extensive experiments and ablation studies confirm our model as a robust and scalable extension for VLMs in missing modality scenarios, ensuring reliability across diverse missing rates and environments. Our code and models will be publicly available.

Figures (8)

• Figure 1: An overview of Multimodal Missing Modalities: (a) Foundation VLM working with complete modality; (b) Prompt learning methods for missing modalities; (c) (Ours) Feature restoration method for missing modalities utilizing Diffusion model.
• Figure 2: PCA visualization of 1,000 samples shows that our model aligns the restored features with ground-truth (from left to right $T=1000, 800, 500, 50$) via the denoising sampling process.
• Figure 3: The architecture of our Missing Modality Restoration Framework. (a) Feature extraction process for available modalities utilizing frozen pre-trained VLM Encoders. (b) The top section depicts the gating mechanism of the DiT backbone, and the bottom section provides a conceptual illustration of how the diffusion model reshapes the feature distribution. (c) Architectural design of the mutual learning mechanism, incorporating multi-objective loss functions within specific denoising intervals.
• Figure 4: The visualization of Dynamic Modality Gating: The top plot shows the average activation values of each channel after attention-based gating. The bottom plot shows the activation values of each channel at different DiT depths.
• Figure 5: Robustness analysis on the MM-IMDb dataset across various missing rates in terms of F1-M.