DiffuSAM: Diffusion-Based Prompt-Free SAM2 for Few-Shot and Source-Free Medical Image Segmentation

"DIFFUSAM: DIFFUSION-BASED PROMPT-FREE SAM2 FOR FEW-SHOT AND SOURCE-FREE MEDICAL IMAGE SEGMENTATION" 논문의 상세한 기술적 요약은 다음과 같습니다.

1. 문제 제기

의료 영상 분할은 임상 진단에 필수적이지만, Segment Anything Model(SAM) 및 그 후속 모델인 SAM2 와 같은 현재 최첨단 기반 모델들과 관련하여 상당한 장애물에 직면해 있습니다.

• 도메인 격차: SAM 과 SAM2 는 자연 영상으로 훈련되었으며, 질감, 대비, 해부학적 구조의 차이로 인해 의료 데이터에 대한 일반화에 어려움을 겪습니다.
• 프롬프트 의존성: 이러한 모델은 작동하기 위해 사용자가 제공한 프롬프트 (점, 상자 또는 텍스트) 를 필요로 하므로, 자동화된 임상 워크플로우에는 비실용적입니다.
• 데이터 부족: 지도 학습 미세 조정 (fine-tuning) 은 대규모의 전문가 주석 데이터셋을 필요로 하는데, 이는 생성 비용이 높고 시간이 많이 소요됩니다.
• 기존 적응법의 한계: SAM/SAM2 를 의료 데이터에 적응시키려는 이전 시도들은 종종 계산적으로 무거운 미세 조정을 요구하거나 여전히 프롬프트에 의존합니다. 또한, 기존 프롬프트 없는 변형 모델들은 SAM2 의 볼륨 메모리 능력을 활용하지 못합니다.

2. 방법론: DiffuSAM

저자들은 **조건부 잠재 확산 모델 (conditional latent diffusion model)**을 사용하여 프롬프트 없는 의료 분할을 위해 SAM2 를 적응시키는 DiffuSAM 프레임워크를 제안합니다. 핵심 아이디어는 SAM2 가 마스크를 생성하는 데 사용할 수 있는 "메모리 임베딩"을 합성하여 사용자 프롬프트의 필요성을 효과적으로 대체하는 것입니다.

핵심 아키텍처

1. 동결된 백본 (Frozen Backbone): 훈련 중 SAM2 이미지 인코더와 메모리 인코더는 동결된 상태로 유지됩니다. 모델은 경량 확산 사전 (diffusion prior) 만 학습합니다.
2. 확산 사전 훈련 (Diffusion Prior Training):
   • 입력: 모델은 동결된 SAM2 이미지 인코더로부터의 이미지 임베딩 ($z_{img}$) 과 클래스 레이블 ($l$) 을 입력받습니다.
   • 목표: 모델은 일반적으로 프롬프트가 주어졌을 때 SAM2 가 생성할 해당 메모리 임베딩 ($z_{mem}$) 을 생성하도록 학습합니다.
   • 과정: UNet 기반 확산 모델이 Denoising Diffusion Probabilistic Model (DDPM) 스케줄 하에 훈련됩니다. 이는 이미지 임베딩과 레이블에 조건부인 노이즈가 포함된 메모리 임베딩을 제거하도록 학습합니다.
   • 손실 함수: 훈련 목표는 메모리 임베딩을 재구성하는 **사전 손실 (prior loss)**과 생성된 임베딩이 SAM2 의 마스크 디코더를 통과할 때 정확한 마스크를 생성하도록 보장하는 **분할 손실 (segmentation loss)**을 결합합니다.
3. 추론 (프롬프트 없는 생성):
   • 추론 시, 확산 모델은 순수 노이즈와 이미지 임베딩으로 시작합니다.
   • 이는 메모리 임베딩 ($\hat{z}_{mem}$) 을 합성하기 위해 반복적으로 노이즈를 제거합니다.
   • 이렇게 합성된 임베딩은 SAM2 의 메모리 어텐션과 마스크 디코더에 주입되어 사용자의 입력 없이 최종 분할 마스크를 생성합니다.

볼륨 일관성 (3 차원 조건부)

의료 볼륨 (CT/MRI) 의 3 차원 특성을 해결하기 위해 DiffuSAM 은 슬라이스 간 공간적 일관성을 강제합니다:

• 특정 슬라이스 $i$에 대한 메모리 임베딩을 생성할 때, 모델은 **인접 슬라이스의 분할 예측 ($adj$)**에 조건부입니다.
• 인접 슬라이스 예측은 메모리 인코더를 통해 인코딩되어 현재 슬라이스의 데이터와 융합됩니다.
• 이는 해부학적 연속성을 보장하고 3 차원 볼륨 내 슬라이스 간의 불일치를 방지합니다.

소스 프리 비지도 도메인 적응 (SF-UDA)

이 프레임워크는 소스 도메인 (예: CT) 에서 훈련되고 타겟 도메인 (예: MRI) 에서 테스트되되, 타겟 적응 중 소스 이미지에 접근할 수 없는 SF-UDA 를 위해 설계되었습니다.

• 이 방법은 특성 친화성 (feature affinity) 속성을 활용합니다: 동일한 클래스에 대한 SAM2 이미지 임베딩은 서로 다른 도메인 간에도 밀집하게 군집화되는 경향이 있습니다.
• 소스에서 훈련된 확산 사전은 추가 미세 조정 없이 타겟 도메인 슬라이스에 직접 적용되어 타겟 도메인에 대한 가짜 메모리 임베딩을 생성합니다.

3. 주요 기여

1. 새로운 프레임워크: 프롬프트 없는 의료 영상 분할을 위한 SAM2 의 첫 번째 확산 기반 적응인 DiffuSAM 의 도입.
2. 잠재 공간 작동: 모델은 SAM2 의 잠재 임베딩 공간에서 작동하여 무거운 백본을 수정하지 않고도 고수준 특성 분포를 포착할 수 있습니다.
3. 볼륨 통합: 슬라이스 간 조건부 처리를 통해 3 차원 슬라이스 전반에 걸쳐 해부학적 연속성을 보장하는 SAM2 의 메모리 어텐션 메커니즘과의 통합.
4. 효율성: 전체 기반 모델을 미세 조정하는 대신 경량 UNet 을 훈련하므로, 낮은 메모리 사용량 (<5GB GPU) 과 빠른 추론 (확산 단계 2 회만 필요) 으로 이어집니다.

4. 실험 결과

이 방법은 BTCV(CT 스캔) 와 CHAOS(MRI 스캔) 라는 두 가지 주요 데이터셋에서 비장, 신장, 간 분할에 초점을 맞춰 평가되었습니다.

퓨샷 설정 (BTCV)

• 설정: 데이터의 10%(3 개 볼륨) 로 훈련, 90%(27 개 볼륨) 로 테스트.
• 성능: DiffuSAM 은 평균 **Dice 점수 87.24%**를 달성했습니다.
• 비교: 프롬프트에 의존하는 네이티브 SAM/SAM2 와 동일한 작은 데이터셋으로 훈련된 지도 학습 UNet/UNETR 모델보다 우수한 성능을 보였습니다. 테스트 시 정답 슬라이스가 필요한 FATE-SAM 과는 경쟁력 있는 결과를 보였으나, 수동 프롬프트나 테스트 시 지도 학습 없이 이를 달성했습니다.
• 애블레이션: 볼륨 슬라이스 간 조건부 처리를 추가하면 평균 Dice 가 82.72%(DiffuSAM-no 3D) 에서 **87.24%**로 향상되었습니다.

소스 프리 UDA 설정 (CT $\to$ MRI)

• 설정: BTCV(CT) 로 훈련, CHAOS(MRI) 로 테스트하며 소스 이미지에 접근하지 않음.
• 성능: DiffuSAM 은 평균 Dice **85.2%**를 달성했습니다.
• 비교: 다른 소스 프리 베이스라인 (DPL: 56.0%) 보다 현저히 우수한 성능을 보였으며, 강력한 특성 친화 기반 방법인 DFG(85.2%) 와 동등한 성능을 발휘했습니다.
• 애블레이션: 볼륨 조건부 처리는 성능을 **80.4%**에서 **85.2%**로 향상시켰습니다.

5. 중요성 및 결론

• 임상 적용 가능성: 사용자 프롬프트 필요성을 제거함으로써 DiffuSAM 은 기반 모델과 자동화된 임상 워크플로우 간의 격차를 해소합니다.
• 데이터 효율성: 이 방법은 최소한의 훈련 데이터 (퓨샷) 로 고품질 분할을 달성할 수 있으며, 타겟 적응 중 소스 데이터에 접근하지 않아도 (SF-UDA) 됨을 입증합니다.
• 계산 효율성: 대규모 모델의 전체 미세 조정과 달리, DiffuSAM 은 경량 확산 사전만 훈련하므로 리소스가 제한된 환경에서도 접근 가능합니다.
• 향후 작업: 저자들은 MedSAM(프롬프트와 함께 엔드 - 투 - 엔드로 훈련됨) 이 더 나은 성능을 보이지만, DiffuSAM 은 프롬프트와 전체 미세 조정을 피함으로써 고유한 장점을 제공한다고 인정합니다. 향후 작업은 추가적인 SF-UDA 기술 통합을 탐구할 것입니다.

요약하자면, DiffuSAM 은 확산 모델의 생성적 힘을 활용하여 SAM2 에 필요한 메모리 임베딩을 "환각"시켜 강력하고 프롬프트 없으며 도메인 적응형 의료 영상 분할을 가능하게 합니다.