Search

[논문리뷰]Few-Shot Adaptation of Training-Free Foundation Model for 3D Medical Image Segmentation

카테고리
VCMI논문리뷰
Index
날짜
2025/03/22

1. Introduction

1.1 SAM2 기반의 Zero-shot Segmentation의 한계

•
SAM(Segment Anythig Model)와 후속 모델인 SAM2의 등장으로 user prompt를 활용해 Zero-shot segmentation이 수행할 수 있게 됨
•
Medical Image Segmentation 분야에서도 SAM기반 모델을 활용하여 Zero-shot training을 적용하려는 시도가 있었지만, 일반 이미지에서 의료 이미지 도메인으로 transfer 하는 과정에서의 한계가 존재
◦
Medical image 특성상 불분명한 경계, 복잡한 질감, 해부학적 디테일들이 존재하기 때문이 일반 이미지와 비교 했을 때 분할해야하는 객체들이 불분명하게 표현됨
→ 수작업으로 정확한 prompt를 제공하더라도 Zero-shot 방식으로 의료 영상 분할에 어려움이 존재
•
대규모 데이터셋에서 SAM 모델을 fine-tunining 한 뒤에 의료 도메인 정보를 반영하려는 연구
◦
계산 비용이 많이 들고, 대규모의 레이블 데이터셋을 구하기가 어렵다는 한계가 존재

1.2 FATE-SAM의 주요 목표

•
FATE-SAM(Few-shot Adaptation of Training-frEe SAM)은 Few-shot example들을 SAM2 모델에 적용해 3D 의료 영상 분할
•
FATE-SAM의 주요 목표는 Fine-Tuning 과정과 user-prompt 없이 자동으로 분할을 수행하는 것
◦
SAM2의 pretrained memory 메커니즘을 사용해서 소량의 support example만 사용해서 해부학적 정보를 충분히 학습
◦
3D 의료 영상 내에서도 공간적 일관성을 유지할 수 있도록 함(3D Volumetric consistency)
•
CT, MRI를 포함한 다양한 3D medical image dataset에서 평가한 결과 FATE-SAM은 지도학습 모델, zero-shot SAM 방식, Fine-Tuning 된 의료 SAM 모델과 비교하여 우수한 성능을 달성

2. Method

2.1 Model Architecture

FATE-SAM은 사전학습된 SAM2을 기반으로 구현되었다.

FATE-SAM 파이프라인

1.
Image Encoder(E)\Epsilon)를 사용해 test slice와 support slice를 이미지 임베딩으로 변환
•
데이터 셋에서 3D volume 슬라이스와 대응되는 mask 조합인 support set을 입력해야함
→ support set들은 서로 다른 support volume으로 구성되었기때문에 분할 과정에서 해부학적 정보를 제공하는 역할
◦
test slices x={ x1,x2,…xn }x = \set{x^1, x^2, … x^n}
◦
support set {Xs,Ys{X_s, Y_s}}
◦
support volumes Xs={ xs1,….,xsl }X_s = \set {x_{s1}, …. ,x_{sl}}, support masks Ys={ ys1,…,ysl }Y_s = \set {y_{s1}, … ,y_{sl}}
•
Test slice의 이미지 임베딩과 support set에 있는 모든 슬라이스 imbedding 와 유사도 점수 비교
1.
support set에 있는 모든 슬라이스 XsX_s 를 support image embedding FsF_s 으로 변경
•
Fs=E(Xs)F_s = E(X_s)
2.
test slice xix_i 도 동일한 인코더를 사용하여 test image embedding fi=E(xi) f^i = E(x_i) 로 변환
3.
코사인 유사도를 사용해서 test slice와 모든 support slice간 유사도 계산
•
Sim(Fs,fi)=fi×Fs∣∣f∣∣×∣∣Fs∣∣Sim(F_s, f^i) = \frac{f^i \times F_s}{||f^||\times||F_s||}
4.
유사도가 가장 높은 상위 j 개의 support example(slice + mask 조합)을 선택해서 해부학적 정보를 제공
•
fsij,ysij=argmaxfsij∈Fs,ysij∈YsSim(Fs,fi){f^{ij}_s, y^{ij}_s} = arg max_{f^{ij}_s \in F_s, y^{ij}_s \in Y_s} Sim(F_s, f^i)
2.
메모리 인코딩 및 3D 볼륨 일관성 유지(Volumetric Consistency)
•
선택된 support example들을 Memory Encoder MEME 에 입력해서 anatomical memory embedding 으로 변환
◦
anatomical memory embedding f^sij=ME(ysij)+fsij\hat f^{ij}_s = ME(y^{ij}_s) + f^{ij}_s
•
이전 슬라이스의 prediction mask를 Memory Encoder에 입력해서 Volume memory embedding 생성
◦
volume memory embedding f~i−1=ME(yi−1)+fi−1\tilde f^{i-1} = ME(y^{i-1}) + f^{i-1}
◦
첫번째 슬라이스의 경우, yi−1 y^{i-1}이 존재하지 않기 때문에 볼륨 메모리를 0으로 설정
•
두 가지 메모리 임베딩을 결합해서 최종 Memory embedding 을 생성
◦
Memory Embedding f˙si=[f^sij,f~i−1]\dot f^i_s = [\hat f^{ij}_s, \tilde f^{i-1}]
FATE-SAM과 기존 SAM의 Memory 구조 차이
•
SAM2의 Memory bank 구조
◦
SAM에서는 memory bank를 사용해서 사전학습된 Feature를 저장해 재사용하는 구조
•
FATE-SAM은 별도의 memory bank를 사용하지 않고, memory encoder를 사용해서 동적으로 memory embedding 생성
•
또한, 이전 슬라이스의 prediction mask를 활용해서 Volumetric 정보를 고려하고 3D 일관성을 유지
◦
이전 슬라이스에서 예측된 mask는 모델이 이전 슬라이스에서 어떤 구조를 분할했는지 알려주게 되고, 현재 슬라이스에서 어떤 구조가 어느 위치에서 존재할 수 있는지 예측할 수 있는 단서 제공
◦
현재 슬라이스의 feature와 비교하여 volume 내에서 어떤 변화가 발생하는지 추적 가능
3.
Memory Attention 및 Mask Decoding
•
생성된 Memory embedding과 테스트 이미지 임베딩을 결합해서 최종 분할 마스크 생성
•
SAM2의 Memory Attention 모듈 MAMA 를 사용해 test image embedding과 메모리 임베딩 간의 cross attention 어텐션 연산을 수행
◦
f˙i=MA(f˙si,fi)=CrossAttention(f˙si,SelfAttention(fi))\dot f^i = MA(\dot f^i_s, f^i) = Cross Attention(\dot f^i_s, SelfAttention(f^i))
▪
SelfAttention(v)=softmax(QKTd)V+VSelf Attention (v) = softmax(\frac{QK^T}{\sqrt{d}})V+ V
▪
CrossAttention(v1,v2)=softmax(Q1K2Td)V2+V2Cross Attention (v_1, v_2) = softmax(\frac{Q_1K^T_2}{\sqrt{d}})V_2 + V_2
→ Query는 test slice fif^i, Key와 Value 는 memory embedding f˙si\dot f^i_s
◦
기존 메모리 임베딩의 공간 정보를 유지하기 위해서 Attention 과정에서 V,V2V, V_2 를 더해 줌 (Residual Connection)
•
Mask Decoder DD를 사용해서 최종 분할 마스크 생성
◦
segmentation mask yi=D(f˙i)y^i = D(\dot f^i)

3. Experiment Results

Dataset

Comparison studies

Ablation Studies Settings

Evaluation Metric and Implementation

3.1 Comparison Study Results

Overall Evaluation Results

•
FATE-SAM이 가장 높은 Dice Score Median을 기록했고, IQR 범위도 작은 범위를 기록
→ 8개의 comparison model과 비교했을 때 높은 정확도와 일관성을 가짐
•
UNet, UNETER 와 같은 Supervised 방법들은 데이터가 제한된 경우, 성능 변동성이 크고 불안정함을 보임
•
Zero-shot SAM은 Dice 점수가 현저히 낮게 나와서 의료 영상 분할에 대한 적응력이 부족함을 보임
•
Fine-Tuned Medical SAM은 특정 데이터셋에서는 Dice 점수가 높게 나왔지만 다양한 task에서 성능 편차가 크게 나옴

SKl10 Dataset Results (무릎 MRI)

•
FATE-SAM average Dice score : 79.72%
→ 전반적으로 가낭 좋은 성능을 보임
•
Femur Bone(대퇴골), Tibia Bone(경골연골)에서 95% 이상의 높은 성능을 보임
•
Femoral Cartilage, Tibia Bone 분할에서 SAM 기반 모델이 매우 낮은 성능을 보임

ACDC Dataset Results (심장 MRI)

•
FATE-SAM average Dice score : 83.30%
◦
Supervised Methods (UNet: 73.95%) 및 Fine-Tuned SAM (MedSAM2: 67.59%)보다 우수한 성능을 보임
•
Left Ventricle(좌심실), Right Ventricle(우심실)은 80% 이상의 높은 성능을 보였으나, Mycrocardium(심근)은 얇고 복잡한 구조로 인해 분할 성능이 약간 떨어짐
•
심장 분할에 있어서도 FATE-SAM이 Zero-shot SAM 및 Supervised Methods 대비 높은 일관성과 정확도를 보임

BTCV Dataset Results (복부 CT)

•
FATE-SAM average Dice score : 67.60%
◦
Fine-Tuned SAM(SAM-Med3D : 78.22%)보다 약간 낮은 성능을 보임
•
Liver(82.17%)과 Aorta(87.38%)에서는 Fine-Tuned SAM과 유사한 성능을 보임
•
복부 혈관(Portal Vein & Splenic Vein)과 같은 작은 구조에서는 Fine-Tuned SAM보다 더 나은 성능을 기록
•
FATE-SAM이 Zero-Shot SAM과 지도학습 모델보다는 좋은 성능을 보였지만 Fine-Tuned SAM SOTA 모델이 전반적으로 가장 높은 성능을 보임

BraTS 2017 Dataset Results (뇌종양 MRI)

•
FATE-SAM average Dice score : 41.60%
◦
Supervised Methods (UNet: 20.65%) 및 일부 Fine-Tuned SAM (MedSAM: 46.96%)보다 높은 성능 보임
•
Non-Enhancing Tumor(41.65%), Enhancing Tumor(44.08%) 에서 Fine-Tuned SAM과 유사한 성능을 보임
•
FATE-SAM이 Supervised Method와 Zero-shot SAM에 비해 높은 성능과 일관성을 보임

MSD Dataset

•
해마(Hippocampus, 72.75%) 및 전립선(Prostate, 72.06%)에서 Supervised Methods보다 우수한 성능을 기록.
•
췌장(Pancreas) 및 비장(Spleen)에서는 해부학적으로 일관된 마스크를 생성
•
대장암(Colon Cancer)에서는 변동성이 있었지만, 전반적으로 부드러운 경계를 유지하며 정확한 분할을 수행

3.2 Ablation Study Results

1. Support Set 크기 및 Support Example 개수

•
Support set 크기, Support sample이 많을 수록 분할 성능 향상 됨
•
해부학적으로 복잡한 구조(대퇴연골, 경골연골)에서 더 큰 개선 효과 발생
•
가장 좋은 성능은 support set 10개(전체 데이터셋 10%), support example 5개일때 달성됨

2. Similarity Metrics Selection

•
Feature level 유사도 방법이 Image-Level 유사도 보다 일관되게 더 높은 성능을 기록
•
코사인 유사도(Cosine Similarity, CS)가 가장 높은 Average Dice Score 기록(79.72%)하여 최종 모델에서 사용됨

3. Prompt Comparision

•
Zero-shot SAM에서 사용하는 Random points 가 가장 낮은 성능을 기록하고 FATE-SAM의 Few-shot 방법이 가장 좋은 성능을 보임
→ Support Example를 사용하는 Few-shot 방법이 기존 SAM에서 사용하는 수동 prompt보다 해부학적 적정보를 더 효과적으로 제공

4. 3D Volumetric Consistency

•
3D Volume Consistency를 적용한 경우(79.72%) 가 적용하지 않은 경우보다 더 좋은 성능을 보임(70.06%)
•
특히 연골 구조에서 큰 성능 향상이 관찰됨
◦
Fermoral Cartilage: 46.37% → 66.81%, Tibial Cartilage: 46.12% → 61.48%
•
3D Volume Consistency가 유지될 때, 해부학적 구조를 더 정확하게 분할할 수 있음

5. Initial Slice Selection

•
초기 슬라이스 선택이 성능에 미치는 영향은 미미했으며, 평균 Dice 점수는 79.21%~80.03% 로 일정함.
•
Q3(75% 위치)에서 시작할 때 가장 높은 성능(80.03%)을 기록
•
모델에서는 실제 활용 가능성과 직관성을 고려하여 중앙 슬라이스(Center)를 기본 설정으로 사용.

6. Pre-trained Weights Selection

•
SAM2의 다양한 모델 크기(Tiny, Small, Base+, Large)에서 성능 비교.
•
모델 크기에 관계없이 FATE-SAM의 성능은 일관되게 높았으며, SAM2 Large 모델(79.72%)에서 가장 높은 성능을 기록.
•
그러나 모델 크기와 연산 자원 사이의 균형을 고려하여, 필요에 따라 더 작은 모델도 선택 가능 함.
Made with