Med-MoE: Mixture of Domain-Specific Experts for Lightweight Medical Vision-Language Models (EMNLP 2024, Jinag et al.)
Introduction
Med-MoE는 의료 의사결정에 필수적인 인간 수준의 다중 모드 이해 시스템을 구축하는 것을 목표로 한다. 현재의 다중 모드 대규모 언어 모델(MLLM)은 일반 작업에서는 뛰어난 성능을 보이지만, 웹 콘텐츠로 학습되어 의료 데이터와는 큰 차이가 있어 의료 분야에서는 효과가 떨어진다. 또한 Med-Flamingo, Med-PaLM M, LLaVA-Med와 같은 도메인 특화 모델들은 유망한 결과를 보이지만, 특정 작업(예: 개방형 또는 폐쇄형 VQA)에 특화되어 있고 매개변수 규모가 커 훈련 및 추론 비용이 높다는 한계를 가진다. 이는 “풍부한 컴퓨팅 성능을 갖추지 못한 수많은 임상 실무자에게 매력적이지 않다”는 문제로 이어진다.
Med-MoE는 이러한 문제를 해결하기 위해 경량이면서도 효과적인 의료 MLLM을 구축하는 것을 목표로 한다. 특히 “다양한 자원 제약 시나리오에서 임상적 유용성을 방해하는 방대한 매개변수와 막대한 컴퓨팅 자원” 없이도, 다양한 판별 및 생성 기반 다중 모드 의료 작업을 처리할 수 있는 프레임워크를 제안한다.
Proposed Method
Med-MoE는 "Mixture-of-Experts (MoE)" 전략을 활용하여 의료 분야의 특수성과 시너지를 고려한 경량화된 모델을 구현한다. 이는 여러 소규모 서브 모듈(전문가)을 결합하고, 각 작업에 가장 적합한 상위 k개의 전문가만 활성화함으로써 적은 컴퓨팅 비용으로도 우수한 성능을 달성할 수 있도록 한다.
Med-MoE의 학습 과정은 세 단계로 구성된다(Figure 2).
1단계: 다중 모드 의료 정렬 (Multimodal Medical Alignment)
목표: 시각적 의료 이미지와 LLM 토큰을 정렬하여, LLM이 의료 이미지를 이해할 수 있도록 한다.
방법: 시각 인코더 뒤의 MLP만 훈련하여 의료 이미지와 캡션 쌍을 통해 시각·텍스트 모드를 정렬한다. 이미지 토큰과 텍스트 토큰을 연결하고, LLM이 텍스트 토큰의 연속을 생성하도록 학습한다.
2단계: 명령어 튜닝 및 라우팅 (Instruction Tuning and Routing)
목표: 모델이 복잡한 의료 명령을 따르고 다양한 다중 모드 작업을 수행할 수 있도록 능력을 향상시키며, 동시에 전문가 선택을 위한 라우터를 훈련한다.
방법: 의료 질의응답 데이터셋으로 모델을 학습시켜 정확한 응답을 생성하도록 하고, 동시에 라우터를 훈련하여 입력 모달리티(CT, MRI, 병리학, X-ray 등)를 예측하게 한다. 이 단계에서는 시각 인코더를 고정하고 나머지 구성 요소를 훈련한다.
3단계: 도메인 특화 MoE 튜닝 (Domain-Specific MoE Tuning)
목표: 모델의 FFN을 희소하게 활성화되는 전문가(MoE) 아키텍처로 대체하여, 특정 의료 도메인에 대한 미세 조정을 수행한다.
방법: 2단계에서 학습된 라우터가 입력에 따라 전문가를 선택하고, 전역 정보를 캡처하는 메타 전문가가 항상 활성화된다. 메타 전문가는 병원에서 여러 부서가 협력하여 진단하는 워크플로우를 모방하여, 지정된 전문가의 성능을 보조한다. 추론 시에는 메타 전문가와 선택된 전문가만 활성화되어 모델이 경량성을 유지한다.
Experiments
Med-MoE는 VQA-RAD, SLAKE, PathVQA, PneumoniaMNIST, OrganCMNIST 등 다양한 의료 데이터셋에서 지속적으로 성능 향상을 입증한다.
[Main Result 1] Comparison with SOTA Methods on Med-VQA
Med-MoE 모델은 activated parameter 2.0B 또는 3.6B만으로 최고의 성능을 보이는 LLaVA-Med (7B)와 동등하거나 더 우수한 성능을 달성한다.
[Main Result 2] Zero-shot Performance on Med-VQA tasks
Zero-shot 성능에서 LLaVA-Med(LLama7B) 대비 VQA-RAD Open에서 약 1.4%, VQA-RAD Closed에서 2.6%, SLAKE Open에서 5.3%, SLAKE Closed에서 9.4% 향상된 성능을 보였다. 특히, Med-MoE (Phi2)는 SLAKE Open (85.06), SLAKE Closed (85.58), PathVQA Closed (91.98)에서 LLaVA-Med 변형을 능가하며, 나머지 작업에서도 competitive한 성능을 나타냈다.
[Main Result 3] Results on Medical Image Classification
일반적으로 LLM 기반 모델들(예: LLaVA-Med)은 Med-VQA 성능에 초점을 맞추지만, 본 연구에서는 더 종합적인 분석(comprehensive analysis)을 위해 classification task 성능도 함께 평가한다.
PneumoniaMNIST: Med-MoE (Phi-2)는 91.4% 정확도를 기록하며 BiomedGPT 및 Med-Mamba를 능가
OrganCMNIST: Med-Mamba에 이어 두 번째로 높은 성능을 달성
결국, Med-MoE는 VQA와 classification 양쪽에서 모두 기존 의료 특화 LLM들을 능가하거나 최소한 대등한 성능을 보여준다.
Ablation and Analysis
Ablation of Router
일반적인 MoE의 라우터(훈련 시점에만 학습되는 router) 대비, Med-MoE의 사전 학습된 라우터가 모든 VQA 데이터셋에서 우수한 성능을 보였으며, 특히 오픈형 질문(Open-ended VQA) 에서 개선 폭이 두드러짐. 입력 modality(CT, MRI, X-ray 등)를 미리 학습한 라우터 덕분에, 전문가 선택이 더 효율적이고 정확하게 이루어지는 것을 확인할 수 있다.
Ablation of Meta Expert
Meta Expert를 포함했을 때, 성능이 1.3~4.2% 개선되었으며, 도메인 전문가들이 각기 특정 modality(CT, MRI, Pathology 등)에 집중하는 동안, Meta Expert는 글로벌한 의료 정보를 보완적으로 제공한다. 이는 병원에서 여러 전문의가 협업할 때, 이를 종합·조율하는 총괄 주치의와 같은 역할을 하게 된다.
Ablation of Domain-Specific MoE-Tuning
MoE-Tuning을 적용했을 때, 모든 Med-VQA 벤치마크(VQA-RAD, SLAKE, PathVQA)에서 +1~4%의 성능 향상이 이루어졌으며, 순히 전체 FFN을 학습하는 SFT보다, 전문화된 전문가 집합(MoE) 을 통해 더 나은 성능을 달성할 수 있음을 입증한다.
Comparison with LoRA-based Methods
Med-MoE는 LoRA를 적용했을 때 성능 저하가 매우 적었으나 (예: SLAKE Closed에서 단 -0.49%), LLaVA-Med은 LoRA 적용 시 성능 하락이 더 컸다 (-1.97%). 이는 Med-MoE가 파라미터 효율적 학습 기법과 상호 보완적으로 작동할 수 있으며, 실제 임상 환경에서 메모리·GPU 비용 절감 가능성을 크게 높였다.
Image and Text Specialization in Experts
각 Expert는 이미지/텍스트 처리에서 명확한 전문성을 보였다 (ex: Expert 1은 텍스트, Expert 3은 이미지, Expert 4는 병리학 데이터에 특화). 이는 단순한 라우터 기반 MoE보다 해석 가능성과 효율성을 동시에 높여줌.
Domain Specialization of Images in Experts
Med-MoE의 흥미로운 특징은 전문가들이 영상 modality별로 자연스럽게 역할을 나눈다는 점입니다. CT는 Expert 1, 2, MRI는 Expert 2, 3, Pathology는 Expert 4, X-ray는 Expert 1.즉, 각 Expert가 특정 의료 영상 분야에 특화되어 활성화되며, 덕분에 모델은 더 효율적이고 해석 가능한 방식으로 작동한다.
Effect of the Number of Activated Expert
대체적으로 MoE는 보통 top-k=1~2 사용하며 sparsity 유지하면서 메모리 절약하는 경향을 보인다. 본 연구는 다양한 활성화 수 실험을 진했는데, 최종적으로 전문가 4명 중 top-2를 activate하는 것이 성능과 오버헤드의 균형을 보였다.
Effect of the Number of Experts
MoE 적용 시, 적절한 수의 expert를 지정하는 것이 성능과 computational efficiency 사이의 균형을 맞추는 데 중요하다. Top-2 activation과 함께 2개, 4개, 6개의 expert를 사용했을 때의 비교 결과, 4개의 expert를 사용하는 것이 가장 좋은 balance를 보였다. 6개의 전문가를 사용하는 것이 성능적으로는 더 우수하지만, computational time과 memory usage를 많이 요하는 것을 보였다.
Effect of Scaling Training Data
더 큰 규모의 학습 데이터셋을 사용했을 떄의 성능을 평가하기 위해, Stage3의 task로 BloodMNIST와 DermaMNIST를 사용했다 (총 36.7K). 아래의 표와 같이 분류 성능이 눈에 띄게 높아졌으며 이는 제안 모델이 다양한 modality와 더 큰 dataset에 적용 가능한 potential을 보였다.
Limitations and Future Works
Med-MoE는 경량화된 의료 MLLM의 가능성을 보여주지만, 다음과 같은 한계가 존재한다.
훈련 데이터 부족: 의료 데이터는 민감성과 개인정보 보호 문제로 인해 크게 제한된다. 합성 데이터 생성에도 한계가 있으며, 수동 주석은 비용이 크고 확장성이 낮다. 이로 인해 정확하고 전문적 지식을 요구하는 복잡한 개방형 질문에서 모델이 실패할 수 있다.
신뢰성 및 설명 가능성 부족: 의료 애플리케이션은 단순히 정확한 출력뿐 아니라 설명 가능성과 신뢰 점수를 제공해야 한다. 이는 의료 의사결정이 환자의 건강에 직접적인 영향을 미치기 때문에, 추론 과정을 명확히 설명하고 신뢰 수준을 제공할 수 있는 모델의 필요성을 강조한다.
Conclusion
Med-MoE는 판별 및 생성 기반의 다양한 다중 모드 의료 작업을 위한 경량 프레임워크를 성공적으로 제안하였으며, activated parameter 수는 줄이면서도 SOTA 성능을 유지하거나 능가했다. Resource constrain이 있는 임상 환경에서도 실용적으로 활용 가능하다.
