Enhancing MedSAM with a Lightweight Box Predictor for Medical Image Segmentation

기술 요약: 의료 영상 분할을 위한 경량 박스 예측기를 통한 MedSAM 성능 향상

문제 정의
의료 영상에서의 시맨틱 분할(Semantic segmentation)은 데이터 부족, 모달리티(CT, MRI, 초음파) 간의 높은 가변성, 그리고 대규모 주석 데이터셋 확보의 어려움으로 인해 난관에 봉착해 있다. Segment Anything Model(SAM)과 같은 파운데이션 모델이 범용적인 분할 능력을 제공하지만, 의료 영상에 직접 적용하기에는 도메인 간 격차가 존재한다. 구체적으로, MedSAM의 적응 방식은 사용자 프롬프트의 품질에 매우 민밀하게 반응한다는 한계가 있다. 자연스러운 사용자 상호작용 방식인 표준 포인트 프롬프트(Point prompt)는 대상 구조가 불규칙하거나, 작거나, 대비가 낮은 경우 신뢰할 수 있는 분할을 위한 충분한 공간적 맥락을 제공하지 못하는 경우가 많다. 이러한 모호성은 모델이 완전히 미세 조정(Fine-tuning)되지 않았거나 사용자의 클릭이 부정확할 때 상당한 성능 저하를 초래한다.

방법론
본 연구는 MedSAM 아키텍처에 경량 박스 예측기(Box Predictor) 모듈을 통합하는 강화된 분할 프레임워크를 제안한다. 핵심 혁신은 단일 사용자 클릭(포인트 프롬프트)을 근사적인 경계 상자(Bounding box)로 변환하여, 분할 과정을 가이드하는 강력한 공간적 사전 정보(Spatial prior)로 활용하는 것이다.

이 방법은 2단계 학습 파이프라인을 통해 작동한다:

1. 1단계 (박스 예측기 학습): 입력 포인트 주변에서 추출된 $5 \times 5$ 패치 이미지 임베딩으로부터 경계 상자 파라미터($W, H, \Delta x, \Delta y$)를 예측하도록 독립적으로 학습된 경량 다층 퍼셉트론(MLP)을 학습시킨다. 학습 시 입력 포인트에는 실제 사용자의 부정확성을 시뮬레이션하기 위해 무작위 공간 섭동(평균 마스크 영역의 30%로 스케일링됨)이 가해진다. 이 모듈은 $L1$ 손실과 일반화된 IoU(Generalized Intersection over Union, GIoU) 손실을 결합한 복합 손실 함수를 사용하여 최적화된다.
2. 2단계 (통합 및 미세 조정): 사전 학습된 박스 예측기를 MedSAM 파이프라인에 통합한다. 학습된 표현을 보존하기 위해 이미지 인코더와 박스 예측기는 동결(Frozen) 상태를 유지한다. 프롬프트 인코더와 마스크 디코더는 선택적으로 미세 조정될 수 있다. 예측된 경계 상자는 최종 분할 마스크를 생성하기 위해 원래의 포인트 프롬프트와 함께(또는 독립적으로) 프롬프트 인코더에 입력된다.

박스 예측기는 단 160만 개의 파라미터만을 추가하며, 추론 오버헤드는 미미하다(RTX 4090 기준 이미지당 약 61~63ms).

주요 기여

• 경량 박스 예측기: 단일 포인트 프롬프를 정보력이 풍부한 경계 상자로 변환하여, 특히 포인트 프롬프트가 모호하거나 위치가 부적절할 때 MedSAM을 안정화하는 공간적 가이드를 제공하는 모듈이다.
• 2단계 학습 파이프라인: 박스 예측기를 별도로 학습시킨 후 고정된 공간 가이드로 통합하는 전략이다. 이를 통해 메인 MedSAM 모델이 동결된 상태에서도 견고한 분할을 달 Ach성할 수 있어, 자원이 제한된 환경에 적합하다.
• 체계적 평가: CT, MRI, 초음파 모달리티를 아우르는 네 가지 이질적인 데이터셋(FLARE22, BRISC, BUSI, LungSegDB)에 대한 포괄적인 평가를 수행하였다. 본 연구는 프롬프트 유형(포인트, 박스, 하이브리드)과 모델 상태(동결 vs 미세 조정) 사이의 상호작용을 분석하며 여섯 가지 구성을 평가한다.
• 프롬프트 강건성 분석: 입력 섭동과 박스 크기 조정이 정확도에 미치는 영향을 조사하여, 본 방법이 중심에서 벗어난 클릭에도 효과적이며 다양한 모달리티에 대해 최소한의 수동 튜닝만 필요함을 입증하였다.

실험 결과
제안된 방법은 네 가지 데이터셋에 대해 기본 MedSAM 및 다른 SAM 변형 모델(SAM2, SAM2UNet)과 비교 평가되었다. 주요 결과는 다음과 같다:

• 성능 향상: 제안된 방법은 일관되게 분할 정확도를 개선했다. BUSI 데이터셋에서 Dice 점수는 MedSAM + 포인트 방식의 0.8771에서 0.8904(MedSAM + Box)로 증가했다. FLARE22에서는 Dice 점수가 0.9169에서 0.9314로 향상되었다.
• 동결 모델 회복: MedSAM 가중치를 업데이트하지 않는 "동결" 설정에서 포인트 프롬프트에만 의존할 경우 성능이 거의 완전히 실패했다(LungSegDB에서 Dice 점수 1.64%까지 하락). 그러나 예측된 경계 상자를 도입하자 성능이 크게 회복되었다(예: LungSegDB에서 Dice가 72.58%로 상승). 이는 기하학적 사전 정보가 모델 작동에 결정적임을 보여준다.
• 모달리티 적응: 방법론은 모달리티에 따라 뚜렷한 동작 차이를 보였다. 초음파 및 MRI 데이터셋(BUSI, BRISC)은 모호한 경계를 고려하여 경계 상자를 약간 확장(10~20%)했을 때 이득을 얻었으나, CT 데이터셋(FLARE22, LungSegDB)은 높은 대비 덕분에 초기 예측된 박스에서 최적의 성능을 보였다.
• 효율성: 이 접근 방식은 계산 비용을 최소한으로 유지하면서도 기본 모델과 유사한 추론 시간을 유지하며 이러한 이득을 달성했다.

의의 및 주장
본 논문은 이 연구가 근본적인 분할 패러다임을 재정의하는 것이 아니라, MedSAM에 대한 실용적인 강화를 제공한다고 주장한다. 주요 의의는 정밀한 수동 프롬프트에 대한 의존도를 낮춤으로써 실제 임상 워크플로우에서 대화형 분할의 강건성을 개선하는 데 있다. 경량 보조 모듈을 활용함으로써, 본 방법은 파운데이션 모델의 기하학적 기대치와 희소한 사용자 입력 사이의 간극을 메운다. 저자들은 이 접근 방식이 전체 모델을 미세 조정할 필요 없이 고성능 분할을 가능하게 하므로, 라벨링된 데이터나 계산 자원이 제한된 환경에서 특히 가치가 있다고 강조한다. 결론적으로, 본 방법은 다중 불연속 영역에 대한 한계가 있으나, 다양한 의료 영상 도메인 전반에서 포인트 프롬프트의 모호성을 효과적으로 해결한다.