Learning from Noisy Prompts: Saliency-Guided Prompt Distillation for Robust Segmentation with SAM

[기술 요약] Saliency-Guided Prompt Distillation (SPD)

1. 문제 정의 (Problem Statement)

최근 Segment Anything Model (SAM)과 같은 파운데이션 모델이 의료 영상 분할(Segmentation) 분야에서 주목받고 있으나, 실제 임상 환경에서는 다음과 같은 결정적인 한계가 존재합니다.

• 프롬프트의 노이즈 (Noisy Prompts): 임상 현장에서 제공되는 프롬프트(예: 중심선 포인트, 대략적인 바운딩 박스)는 정밀한 타겟 구조를 지정하기보다는 광범위한 해부학적 영역을 나타내는 경우가 많습니다. 이는 타겟이 아닌 인접 조직을 포함하거나, 타겟의 일부를 누락하는 등 '노이즈'를 포함합니다.
• SAM의 민감도 (Sensitivity): SAM은 입력된 프롬프트에 매우 엄격하게 반응합니다. 따라서 부정확하거나 모호한 프롬프트가 입력되면 모델이 잘못된 영역을 분할하거나 불완전한 마스크를 생성하는 성능 저하가 발생합니다.
• 데이터 부족: 정밀한 픽셀 단위 주석(Annotation)을 얻는 것은 비용이 많이 들고 전문가의 숙련도가 필요하여 대규모 학습 데이터 확보가 어렵습니다.

2. 제안 방법론 (Methodology: SPD Framework)

본 논문은 전문가(방사선 전문의)가 인접한 슬라이스를 참고하여 해부학적 구조를 파악하는 추론 과정을 모방한 SPD(Saliency-Guided Prompt Distillation) 프레임워크를 제안합니다. 이는 두 단계로 구성됩니다.

[Stage 1] 해부학적 사전 지식 학습 (Anatomical Prior Learning)

• Saliency Head 도입: LoRA(Low-Rank Adaptation)로 적응된 SAM 인코더의 특징을 활용하는 경량화된 Saliency Head를 학습시킵니다.
• 목적: Ground-truth 마스크를 사용하여 타겟 구조가 위치할 가능성이 높은 영역을 나타내는 **Saliency Map(돌출도 지도)**을 생성합니다. 이는 이후 단계에서 프롬프트의 신뢰도를 검증하는 '해부학적 기준점' 역할을 합니다.

[Stage 2] 프롬프트 유도 분할 (Prompt-Guided Segmentation)
학습된 Saliency Map을 기반으로 Contextual Prompt Distillation (CPD) 과정을 거쳐 노이즈가 제거된 '합의된 프롬프트 세트(Consensus Prompt Set)'를 구성합니다.

1. Local Prompt Validation (지역 프롬프트 검증): 현재 슬라이스의 노이즈 섞인 프롬프트 중, Saliency Map에서 높은 값을 가진(즉, 해부학적으로 타당한 위치에 있는) 프롬프트만 남기고 나머지는 제거합니다.
2. Context-Aware Prompt Enrichment (문맥 인식 프롬프트 보강): 현재 슬라이스의 프롬프트가 부족할 경우를 대비해, 인접한 슬라이스(Neighboring slices)의 프롬프트 중 유효한 것들을 가져와 현재 슬라이스의 Saliency Map으로 재검증한 뒤 추가합니다.
3. Consensus Set Formulation (합의 세트 형성): 검증된 지역 프롬프트와 보강된 문맥 프롬프트를 결합하여 최적의 프롬프트 세트를 만듭니다.
4. Pairwise Slice Consistency (PSC, 슬라이스 간 일관성): 인접한 슬라이스 간의 예측 결과가 해부학적으로 연속성을 갖도록 하는 손실 함수(Loss)를 적용하여, 경계가 모호한 영역에서도 안정적인 분할을 유도합니다.

3. 주요 기여 (Key Contributions)

• 새로운 문제 제기: 의료 영상 적응 시 '정밀한 프롬프트'가 아닌 '노이즈 섞인 프롬프트'가 주어지는 실제 임상 시나리오를 정의하고 해결책을 제시했습니다.
• 전문가 추론 모방: Saliency Map을 활용한 프롬프트 필터링 및 인접 슬라이스 정보를 활용한 프롬프트 보강 메커니즘을 통해 전문가의 진단 방식을 모델링했습니다.
• 강건한 프레임워크: 단순한 모델 튜닝을 넘어, 프롬프트 자체를 정제(Distillation)하여 파운데이션 모델의 성능을 극대화하는 구조를 제안했습니다.

4. 실험 결과 (Results)

• 다양한 데이터셋 검증: MRI(TI, Scar, FUMPE) 및 CT(KiTS) 등 4개의 도전적인 의료 영상 데이터셋에서 실험을 수행했습니다.
• SOTA 성능 달성: 기존의 CNN 기반 모델(UNet, nnUNet 등)과 최신 SAM 기반 모델(MedSAM, MSA, SAM-Tuning 등)을 모두 압도하는 성능을 보였습니다.
  • 특히 TI 데이터셋에서 DSC(Dice Similarity Coefficient)가 기존 최고 모델 대비 11.08% 향상되었으며, HD95(경계 오차)는 6.28 감소하는 탁월한 결과를 보였습니다.
• 3D 모델 대비 우위: 최신 3D 기반 모델인 SAM2 및 MedSAM2와 비교했을 때도 더 높은 정확도를 기록했습니다.
• 효과 입증: Ablation Study를 통해 Saliency Map을 통한 검증, 문맥 보강(CPD), 슬라이스 간 일관성(PSC)이 각각 성능 향상에 핵심적인 역할을 함을 증명했습니다.

5. 의의 (Significance)

본 연구는 **"완벽하지 않은 데이터(Noisy Prompts)로부터 어떻게 신뢰할 수 있는 결과를 도출할 것인가?"**라는 실무적인 질문에 대한 해답을 제시합니다. 이는 고가의 정밀 주석 작업 없이도 기존의 거친 임상 데이터를 활용하여 강력한 파운데이션 모델을 실제 의료 현장에 안정적으로 배포할 수 있는 실질적인 경로를 마련했다는 점에서 큰 의의가 있습니다.