[논문 리뷰] M-FLAG: Medical Vision-Language Pre-training with Frozen Language Models and Latent Space Geometry Optimization

Che Liu, Sibo Cheng외 6명|arXiv (Cornell University)|2023. 07. 17.

COVID-19 diagnosis using AI인용 수 4

한 줄 요약

M-FLAG는 언어 모델을 고정시키고 잠재 공간 기하학을 최적화하기 위해 직교성 손실을 도입함으로써 계산적으로 효율적인 의료 시각-언어 사전학습 프레임워크를 제안한다. 파rameter 수를 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.

연구 동기 및 목표

• 의료 영상에서 공동 시각-언어 사전학습의 높은 계산 비용과 훈련 불안정성 문제를 해결하기 위해.
• 효율적인 의료 시각-언어 표현 학습을 위해 언어 모델의 미세조정이 필수적인지 조사하기 위해.
• 잠재 공간의 붕괴를 방지하기 위해 잠재 공간 기하학을 명시적으로 정규화함으로써 문제를 완화하기 위해.
• 더 적은 훈련 가능한 파rameter를 유지하면서도 높은 성능을 유지하는 더 효율적이고 견고한 사전학습 프레임워크를 개발하기 위해.
• 고정된 언어 모델과 기하학적 정규화가 다양한 의료 시각-언어 작업에서 효과적인지 입증하기 위해.

제안 방법

• 사전 학습된 언어 모델(예: BERT)을 고정시켜 훈련 복잡도와 파arameter 수를 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 제안된 방법이 다양한 의료 영상 벤치마크에서 성능을 유지하거나 향상시키면서도 파arameter 수를 얼마나 줄이는가?
• RQ5 ImageNet 기반 사전학습 모델 대비, RSNA 데이터셋의 1%만을 사용하는 저데이터 환경에서 모델의 일반화 능력은 어떠한가?

주요 결과

• M-FLAG는 언어 모델을 고정시킴으로써 기존 의료 시각-언어 사전학습 방법 대비 훈련 가능한 파arameter 수를 78% 감소시켰다.
• RSNA 데이터셋에서 M-FLAG는 훈련 데이터의 1%만으로도 mAP 13.7%를 달성했으며, 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.