[论文解读] 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 AI被引用 4

一句话总结

M-FLAG 提出了一种计算高效的医学视觉-语言预训练框架，通过冻结语言模型并引入正交性损失来优化潜在空间几何结构，在将参数量减少78%的同时，在医学图像分类、分割和目标检测任务上均优于当前最先进方法，尤其在数据极少的情况下表现突出（例如，仅使用RSNA数据集的1%）。

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.

研究动机与目标

• 解决医学影像中联合视觉-语言预训练带来的高计算成本和训练不稳定性问题。
• 探究在有效的医学视觉-语言表征学习中，微调语言模型是否为必要。
• 通过显式正则化潜在空间的几何结构，缓解视觉-语言模型中的潜在空间坍塌问题。
• 开发一种更高效、更鲁棒的预训练框架，在减少可训练参数的同时保持高性能。
• 证明冻结语言模型与几何正则化在多样化下游医学视觉任务中的有效性。

提出的方法

• 冻结预训练语言模型（如BERT），以降低训练复杂度并使参数量减少78%。
• 提出一种新颖的正交性损失，用于正则化潜在空间几何结构，促进均匀分布并防止坍塌。
• 将标准对比对齐损失与正交性损失结合，联合优化视觉-语言对齐与潜在空间结构。
• 使用视觉编码器提取图像特征，并通过对比学习将这些特征与冻结的文本特征对齐。
• 通过最小化特征空间中对正交性的偏离，使潜在空间逼近均匀超球面。
• 采用双流架构，其中语言流在预训练期间保持固定，仅更新视觉流。

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

实验结果

研究问题

• RQ1 在医学视觉-语言预训练中，是否必须微调预训练语言模型？冻结语言模型是否能在降低复杂度的同时实现更好或相当的性能？
• RQ2 潜在空间坍塌如何影响下游医学视觉任务（如分割和目标检测）？
• RQ3 与标准对齐损失相比，通过正交性损失显式正则化潜在空间几何结构，是否能提升模型的鲁棒性与性能？
• RQ4 所提出方法在保持或提升性能的同时，能在多大程度上减少参数量，并在多样化的医学视觉基准上实现性能优化？
• RQ5 与ImageNet预训练模型相比，该模型在低数据场景（如仅使用RSNA数据集1%）下的泛化能力如何？

主要发现

• 通过冻结语言模型，M-FLAG 相较于现有医学视觉-语言预训练方法，将可训练参数量减少了78%。
• 在RSNA数据集上，M-FLAG 仅使用1%的训练数据即达到13.7%的mAP，优于在100%数据上微调的ImageNet预训练模型。
• 在SIIM分割任务中，M-FLAG 仅使用1%的数据即达到52.5%的Dice分数，展现出在低数据条件下的强大泛化能力。
• 消融实验证明，同时使用对齐损失与正交性损失可获得最佳性能，其中仅使用正交性损失即可使分割性能相比基线提升4.8%。
• 解冻语言模型后，SIIM数据集上的性能最高下降4.32%，证实了冻结语言流的优势。
• 潜在空间可视化显示，M-FLAG 维持了比未冻结变体更均匀、更分散的几何结构，支持正交性正则化的有效性。

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.