[논문 리뷰] Customized Segment Anything Model for Medical Image Segmentation

Kaidong Zhang, Dong Liu|arXiv (Cornell University)|2023. 04. 26.

Advanced Neural Network Applications인용 수 121

한 줄 요약

SAMed는 LoRA 미세조정을 사용하여 Segment Anything Model (SAM)을 의료 영상 분할에 맞춰 커스터마이즈하고 최소한의 배포/저장 오버헤드로 경쟁력 있는 결과를 달성합니다.

ABSTRACT

We propose SAMed, a general solution for medical image segmentation. Different from the previous methods, SAMed is built upon the large-scale image segmentation model, Segment Anything Model (SAM), to explore the new research paradigm of customizing large-scale models for medical image segmentation. SAMed applies the low-rank-based (LoRA) finetuning strategy to the SAM image encoder and finetunes it together with the prompt encoder and the mask decoder on labeled medical image segmentation datasets. We also observe the warmup finetuning strategy and the AdamW optimizer lead SAMed to successful convergence and lower loss. Different from SAM, SAMed could perform semantic segmentation on medical images. Our trained SAMed model achieves 81.88 DSC and 20.64 HD on the Synapse multi-organ segmentation dataset, which is on par with the state-of-the-art methods. We conduct extensive experiments to validate the effectiveness of our design. Since SAMed only updates a small fraction of the SAM parameters, its deployment cost and storage cost are quite marginal in practical usage. The code of SAMed is available at https://github.com/hitachinsk/SAMed.

연구 동기 및 목표

• 의미 있는 조직 라벨을 가진 의학 영상 세그먼테이션을 위한 SAM의 확장.
• SAM 매개변수의 소수 분율만 업데이트하여 효율적인 미세조정을 가능하게 한다.
• 훈련 안정화와 수렴 개선을 위해 warmup과 AdamW를 활용한다.
• Synapse 다기관 데이터셋에서 경쟁력 있는 성능을 시연한다.
• SAM-compatible한 상태를 유지하면서 배포/저장 오버헤드가 미미함을 보여준다.

제안 방법

• SAM 이미지 인코더를 동결하고 변환기 블록에 LoRA를 적용하여 의학 특징을 학습한다.
• 의미론적 분할을 위해 프롬프트 인코더와 마스크 디코더를 미세조정한다(선택적으로 LoRA 적용).
• SAM의 출력을 k개의 클래스(배경 포함)에 해당하는 k개의 의미적 마스크를 예측하도록 적응하고, 클래스 차원에서 Softmax와 ArgMax를 통해 최종 S를 계산한다.
• 학습 감독을 위해 축소된 그라운드 트루스와 함께 교차 엔트로피 및 Dice 손실을 사용한다.
• 훈련 안정화와 수렴 개선을 위해 warmup 단계와 AdamW 옵티마이저를 채택한다.
• SAM과의 호환성을 입증하고 업데이트된 매개변수 수가 감소했음을 보여준다(예: 이미지 인코더에만 LoRA를 적용했을 때 18.81M).

실험 결과

연구 질문

• RQ1 의료 데이터에서 LoRA로 미세조정된 SAM이 의료 영상에 대한 의미론적(segmentation) 분할을 수행할 수 있는가?
• RQ2 SAM 매개변수의 일부분만 업데이트하여 더 낮은 배포/저장 오버헤드로 경쟁력 있는 분할 정확도를 얻을 수 있는가?
• RQ3 의료 데이터에 대한 SAM 적응에서 어떤 교육 전략(warmup, AdamW)이 수렴과 성능을 향상시키는가?
• RQ4 Synapse에서 DSC와 HD 기준으로 SAMed가 최신 의료 분할 모델과 어떻게 비교되는가?
• RQ5 SAM의 호환성을 유지하면서 서로 다른 조직에 대한 의미 있는 의미론적 라벨링을 SAMed가 제공할 수 있는가?

주요 결과

방법	DSC ↑	HD ↓	대동맥	담낭	신장(L)	신장(R)	간	췌장	비장	위
U-Net	76.85	39.70	89.07	69.72	77.77	68.60	93.43	53.98	86.67	75.58
Att-UNet	77.77	36.02	89.55	68.88	77.98	71.11	93.57	58.04	87.30	75.75
TransUnet	77.48	31.69	87.23	63.13	81.87	77.02	94.08	55.86	85.08	75.62
SwinUnet	79.13	21.55	85.47	66.53	83.28	79.61	94.29	56.58	90.66	76.60
MissFormer	81.96	18.20	86.99	68.65	85.21	82.00	94.41	65.67	91.92	80.81
TransDeepLab	80.16	21.25	86.04	69.16	84.08	79.88	93.53	61.19	89.00	78.40
HiFormer	80.39	14.70	86.21	65.69	85.23	79.77	94.61	59.52	90.99	81.08
DAE-Former	82.43	17.46	88.96	72.30	86.08	80.88	94.98	65.12	91.94	79.19
SAMed	81.88	20.64	87.77	69.11	80.45	79.95	94.80	72.17	88.72	82.06

• SAMed는 Synapse 다기관 데이터셋에서 81.88 DSC와 20.64 HD를 달성하여 최첨단 기준선과 대등하다.
• SAMed는 Synapse 결과 중 췌장 및 위 분할에서 최첨단 성능을 달성한다.
• 작은 부분의 SAM 매개변수만 업데이트(예: 원래 358M 중 18.81M; 원래 크기의 5.25%)로 배포/저장 오버헤드를 미미하게 유지.
• LoRA를 이미지 인코더에 적용하고(선택적으로 마스크 디코더도) 마스크 디코더만 업데이트하는 것보다 더 나은 성능을 얻는다.
• Warmup 및 AdamW 옵티마이저가 훈련 안정화와 수렴 및 최종 손실 개선에 크게 기여한다.
• SAMed는 여전히 SAM과 완전히 호환되며 의학 영상 분할을 위한 SAM의 플러그인으로 사용할 수 있다.