[Paper Review] M-FLAG: Medical Vision-Language Pre-training with Frozen Language Models and Latent Space Geometry Optimization

Che Liu, Sibo Cheng+6 more|arXiv (Cornell University)|Jul 17, 2023

COVID-19 diagnosis using AI4 citations

TL;DR

M-FLAG proposes a computationally efficient medical vision-language pre-training framework that freezes the language model and introduces an orthogonality loss to optimize latent space geometry, reducing parameter count by 78% while outperforming SOTA methods on medical image classification, segmentation, and object detection—especially excelling with minimal data (e.g., 1% of RSNA dataset).

ABSTRACT

Medical vision-language models enable co-learning and integrating features from medical imaging and clinical text. However, these models are not easy to train and the latent representation space can be complex. Here we propose a novel way for pre-training and regularising medical vision-language models. The proposed method, named Medical vision-language pre-training with Frozen language models and Latent spAce Geometry optimization (M-FLAG), leverages a frozen language model for training stability and efficiency and introduces a novel orthogonality loss to harmonize the latent space geometry. We demonstrate the potential of the pre-trained model on three downstream tasks: medical image classification, segmentation, and object detection. Extensive experiments across five public datasets demonstrate that M-FLAG significantly outperforms existing medical vision-language pre-training approaches and reduces the number of parameters by 78\%. Notably, M-FLAG achieves outstanding performance on the segmentation task while using only 1\% of the RSNA dataset, even outperforming ImageNet pre-trained models that have been fine-tuned using 100\% of the data.

Motivation & Objective

• To address the high computational cost and training instability of joint vision-language pre-training in medical imaging.
• To investigate whether fine-tuning the language model is necessary for effective medical vision-language representation learning.
• To mitigate latent space collapse in vision-language models by explicitly regularizing the geometry of the latent space.
• To develop a more efficient and robust pre-training framework that maintains high performance with fewer trainable parameters.
• To demonstrate the effectiveness of frozen language models and geometric regularization on diverse downstream medical vision tasks.

Proposed method

• Freezes the pre-trained language model (e.g., BERT) to reduce training complexity and parameter count by 78%.
• Introduces a novel orthogonality loss to regularize the latent space geometry, promoting uniformity and preventing collapse.
• Combines the standard contrastive alignment loss with the orthogonality loss to jointly optimize visual-language alignment and latent space structure.
• Uses a vision encoder to extract image features and aligns them with frozen text features via contrastive learning.
• Optimizes the latent space to approach a uniform hypersphere by minimizing deviation from orthogonality in the feature space.
• Employs a two-stream architecture where the language stream remains fixed during pre-training, while only the vision stream is updated.

Figure 1: M-FLAG overview. M-FLAG consists of a vision encoder $E_{V}$ for learning vision latent $z_{v}$ , a frozen language model $E_{T}$ for extracting medical text latent $z_{t}$ , and a projector $p(\cdot)$ to map $z_{v}$ to $z_{a}$ for alignment with $z_{t}$ . M-FLAG employs an alignment loss

Experimental results

Research questions

• RQ1 Is it necessary to fine-tune pre-trained language models in medical vision-language pre-training, or can freezing them yield better or comparable performance with reduced complexity?
• RQ2 How does latent space collapse affect downstream medical vision tasks such as segmentation and object detection?
• RQ3 Can explicit geometric regularization of the latent space—via an orthogonality loss—improve model robustness and performance compared to standard alignment losses?
• RQ4 To what extent does the proposed method reduce parameter count while maintaining or improving performance across diverse medical vision benchmarks?
• RQ5 How does the model generalize to low-data regimes, such as using only 1% of the RSNA dataset, compared to ImageNet-pretrained models?

Key findings

• M-FLAG reduces the number of trainable parameters by 78% compared to existing medical vision-language pre-training methods by freezing the language model.
• On the RSNA dataset, M-FLAG achieves 13.7% mAP with only 1% of the training data, outperforming ImageNet-pretrained models fine-tuned on 100% of the data.
• The model achieves 52.5% Dice score on the SIIM segmentation task using only 1% of the data, demonstrating strong generalization under low-data conditions.
• Abllation studies show that combining both the alignment loss and orthogonality loss yields the best performance, with the orthogonality loss alone improving segmentation performance by 4.8% over the baseline.
• Unfreezing the language model leads to a performance drop of up to 4.32% on the SIIM dataset, confirming the benefits of freezing the language stream.
• Visualization of the latent space shows that M-FLAG maintains a more uniform and well-distributed geometry compared to unfrozen variants, supporting the effectiveness of the orthogonality regularization.

Figure 2: Visualization of the first 3 dominant PCA dimensions of latent space on NIH dataset. M-FLAG (green) is compared to its variants (red) when the last $n$ layers of the language model are not frozen.