[Methods 이해하기] MedCLIP-SAMv2: Towards Universal Text-Driven Medical Image Segmentation

차승우·6일 전
0

Methods


Overview

Overview

다음은 핵심 구성 요약 이미지입니다. 파이프라인은 총 3단계입니다.

    1. 새로운 DHN-NCE Loss를 사용한 Fine-Tuning
    1. 텍스트 프롬프트 기반 Zero-Shot Segmentation
    1. 잠재적인 라벨 정제를 위한 Weakly Supervised Segmentation

Pseudo-Label(잠재적인 라벨): 모델이 예측한 분할 결과를 임시 라벨처럼 활용하여, 실제 정답 라벨이 부족한 상황에서도 추가 학습에 사용되는 라벨.


1) DHN-NCE Loss

기존 CLIP은 이미지-텍스트 쌍을 같은 벡터 공간에 매핑하고 관련 쌍의 거리를 최소화하고 무관 쌍의 거리를 최대화하는 대조 학습 메커니즘을 사용합니다.

InfoNCE

Original

두 백터의 내적이 크면, 두 백터는 같은 방향을 보고, 이는 두 샘플이 의미적으로 가깝다는 걸 의미한다.

B: 배치의 크기
분자: Anchor(기존 샘플)과 Positive가 얼마나 가까운 지 계산
분모: Anchor(기존 샘플)과 Positive + Negative의 유사도를 합산
τ (Temperature): Softmax의 민감도 조절 상수

즉, InfoNCE는 Positive를 당기고, Positive + Negative는 밀어냅니다.


InfoNCE for Vision-Language Matching

Original

CLIP은 InfoNCE를 바탕으로 Vision to Text, Text to Vision에 대한 각각의 Loss Function을 결합하는 Bidirectional Contrastive Loss를 사용합니다.


🚨 Challenges in InfoNCE Loss

  1. InfoNCE는 모든 Negative 샘플에 동일한 패널티를 부여합니다.
  • 즉, 구별하기 쉬운 Negative(거리가 먼 샘플)구별하기 어려운 Hard Negative(거리가 가까운 샘플)가 똑같이 취급됩니다.
  • 의료 영상처럼 미묘한 차이를 구별해야 하는 상황에서는 학습 효율이 떨어질 수밖에 없습니다.
  1. NPC(Negative–Positive Coupling) 문제가 발생합니다.
  • Positive 쌍이 분자와 분모에 동시에 포함됩니다.
  • 결국 하나의 Positive가 학습 과정에서 모순적인 신호를 발생시킵니다.
  • 이 현상은 특히 배치 크기가 작은 경우 더욱 심각해져, 학습 효율을 크게 저하시킵니다.

DHN-NCE Loss

DHN-NCE Loss는 기존 InfoNCE의 한계를 해결하기 위해 제안되었습니다.
핵심 아이디어는 두 가지입니다.

  1. Decoupled Constrastive Learning: 분모에서 Positive 항을 제거합니다.
  2. Hard Negative Sampling: 구별하기 어려운 가까운 Negative에 더 큰 패널티를 부여합니다.

Decoupled Constrastive Learning

Original

  1. 기존 Loss Function에서 로그 항을 전개하면 분자와 분모가 따로 떨어집니다.

Original

  1. 분모의 항을 Positive 항 + Negative 항들의 합으로 분리합니다.

Original

  1. Positive 항을 분모에서 제거하여 Negative만 남깁니다.

👉 다음과 같은 방법은 더 뚜렷한 학습 신호를 만듭니다. 결과적으로 작은 배치 사이즈에서도 안정적인 학습이 가능합니다.


Hard Negative Sampling

거리가 가까운 Negative에 더 높은 가중치를 부여하기 위해 다음과 같은 가중치를 사용합니다.

Weight

분자: Negative 항Anchor(기준 샘플)얼마나 가까운 지 계산.
분모: 모든 Negative 유사도의 합
(B - 1): Negative 샘플의 수
β: 기울기 파라미터, Hard Negative를 얼마나 강조할 지 조정

즉, 다음 가중치는 진짜와 비슷하지만 사실은 다른 샘플에 더 집중하도록 유도합니다.

Weight

각 Negative에 가중치를 곱해 Hard Negative Sampling을 완성합니다.


Final

두 손실 함수를 더해, 최종 DHN-NCE 손실 함수를 완성합니다.

Result


2) Zero-Shot Medical Image Segmentation

다음 목표를 위해 앞에서 설명한 DHN-NCE Loss로 Fine-Tuning된 BiomedCLIP을 Backbone으로 사용합니다.

BiomedCLIP

BiomedCLIP은 이미지와 텍스트를 입력했을 때, 이미지 임베딩과 텍스트 임베딩을 출력합니다.

  1. LLM Prompt Generation

GPT-4 같은 LLM을 사용해 의료 영상에서 어떤 특징을 주목해야 하는지 명확히 알려주는 정교한 텍스트 프롬프트를 만듭니다.

Give a sentence caption that describes unique visual features of [TARGET] in [MODALITY]

이를 바탕으로 단순한 프롬프트부터 구체적인 프롬프트까지 다양한 프롬프트를 설계할 수 있습니다.

  1. Saliency Map Generation

이미지 임배딩과 텍스트 임배딩은 M2IB(Multi-Modal Information Bottleneck)을 통과하는데, 이 모듈은 함께 중요하게 여기는 정보는 최대화하고, 불필요한 정보는 필터링합니다.

M2IB

  • MI(Mutual Information): 상호 정보량
  • 첫 항: 이미지 임베딩과 텍스트 임베딩 사이의 공유 정보량을 최대화합니다. 즉, 이미지 임배딩이 텍스트 임배딩과 관련된 시각 정보를 더 많이 담도록 유도합니다.
  • 둘째 항: 이미지 임베딩이 원본 이미지의 모든 것(배경/노이즈/불필요 디테일)까지 과하게 보존하지 않도록 억제(압축)합니다.
  • 𝛾 : 관련성(첫 항)과 압축(둘째 항) 사이의 Trade-Off를 조절하는 하이퍼파라미터입니다.

이렇게 생성된 Saliency Map은 텍스트 프롬프트에 대한 각 픽셀의 중요도를 연속적인 확률 분포로 나타냅니다.
이는 텍스트와 이미지 사이 의미론적 간극(Semantic Gap)을 줄입니다.

  1. Post-processing for Initial Segmentation

생성된 Saliency Map은 연속 확률 분포이므로, 이를 Otsu's Thresholding을 바탕으로 이진화합니다.

Otsu's Thresholding: Foreground와 Background를 나눌 최적의 임계값 η을 자동으로 결정합니다.

Result

이진화 후에도 작은 잡음이나 불연속 부분이 존재할 수 있기 때문에 Connected Component Analysis를 적용하여 이진 마스크를 서로 연결된 성분 단위의 집합 C로 분할합니다.

C의 각 연결 성분에 대해 Saliency Map을 바탕으로 신뢰도 점수를 계산합니다.

Confidence

신뢰도 점수는 곧 전경으로 분류된 픽셀들의 Saliency 평균 확률값입니다.

Result

계산된 신뢰도 점수를 바탕으로, 신뢰도가 높은 연결 성분만 유지하여 Coarse Segmentation Map을 구성합니다.

  1. Segmentation Refinement via SAM

Coarse Segmentation Map은 관심 영역의 대략적인 윤곽을 포함합니다.

SAM

이를 SAM에서 활용할 수 있도록 Coarse Segmentation Map의 구성 요소들은 각각 Bounding Box, 혹은 Point로 변환됩니다. 이렇게 만들어진 프롬프트를 SAM에 입력하여 최종 Zero-Shot 분할 마스크를 생성합니다.


3) Weakly Supervised Segmentation

의료 영상에서는 Pixel 단위의 Ground Truth를 얻기가 매우 어렵습니다. 그래서 Zero-Shot Segmentation 결과를 임시 정답(Pseudo Label)으로 활용하고 이를 바탕으로 Weakly Supervised Segmentation을 진행합니다. 하지만, 임시 정답에는 불확실성이 존재합니다. 즉, 이 불확실성을 다루는 방법이 곧 성능 향상의 핵심입니다.

Zero_Shot

Weak Supervision: 입력 이미지 I와 최종 Zero-Shot 분할 마스크 쌍을 분할 네트워크(nnUNet)의 학습 데이터로 활용합니다.

그 다음으로, Checkpoint Ensembling을 바탕으로 불확실성을 추정합니다.

Checkpoint Ensembling: 전체 학습을 Cycle로 나눈 후, 각 Cycle마다 여러 개의 Checkpoint를 저장합니다. 최종적으로 모든 Checkpoint의 예측을 평균화합니다.

Checkpoint

  • G: Checkpoint의 개수
  • M: 가중치 상태

이런 방법은 여러 모델을 앙상블한 것과 같은 효과를 지닙니다. 픽셀의 불확실성이 높을 수록 그 픽셀의 여러 Checkpoint에서 나온 예측은 서로 크게 달라집니다.

Uncertainty Estimation: 불확실성을 추정하기 위해, 엔트로피를 활용합니다. 이를 바탕으로 불확실한 부분을 알 수 있는 Uncertainty Map을 얻을 수 있습니다.

Entropy

  • H: Uncertainty Map
  • R: 클래스 개수

Entropy

  • h(r): 픽셀이 클래스 r에 속할 최종 확률

Datasets and Experimental Setup


BiomedCLIP Fine-tuning

  • Dataset: MedPix (Siragusa et al., 2024)

    • Radiology 이미지 + 캡션 데이터셋
    • 텍스트 전처리: 특수문자 제거, 공백 정리, 캡션 길이 ≥ 20만 유지
    • 이미지: 224×224 크기로 Resize & CLIP 표준 RGB 정규화
    • Split: Train 20,292 / Val 3,515 (85% ~ 15%)
  • Fine-tuning

    • Backbone: ViT (이미지), PubMedBERT (텍스트)
    • Loss: DHN-NCE
    • Hyperparameters: LR = 1e-6, Decay = 50%, Batch = 64
  • Validation

    • Dataset: ROCO (7,042 Multi-Modal Radiology Images)
    • Metric: Top-1 / Top-2 Retrieval Accuracy (Image → Text, Text → Image)
    • 비교: InfoNCE, DCL, HN-NCE vs DHN-NCE
    • 동일 세팅으로 실험 (τ = 0.6, β1 = β2 = 0.15)
    • Baseline: Pretrained BiomedCLIP, PMC-CLIP, CLIP

Datasets for Segmentation

  • Breast Tumor Ultrasound

    • BUSI dataset (Al-Dhabyani et al., 2020): Train 600
    • UDIAT dataset (Byra et al., 2020): Val 50 / Test 113
  • Brain Tumor MRI

    • Brain Tumor dataset (Cheng, 2017): Train 1,462 / Val 1,002 / Test 600
  • Lung Chest X-ray

    • COVID-QU-Ex (Chowdhury et al., 2020; Rahman et al., 2021):
      Train 16,280 / Val 1,372 / Test 957
  • Lung CT

    • Konya (2020):
      Train 7,959 / Val 3,010 / Test 1,800
      (환자 단위 분할 -> 같은 환자의 데이터가 Train, Test에 동시에 들어가지 않는다.)

Experimental Setup and Metrics

  • Segmentation Experiments

    • 비교 대상:
      • Zero-Shot: SaLIP, SAMAug (텍스트 기반 Zero-Shot)
      • Few-Shot: UniverSeg, ProtoSAM, Self-Prompt-SAM
      • Weakly Supervised: nnUNet (Pseudo-Label, No Checkpoint Ansemble)
  • Weakly Supervised Training

    • nnUNet, 600 Epochs, 3 Cycles
    • Cyclical Learning Rate (Zhao et al., 2022)
    • 각 Cycle마다 10개 Checkpoint 저장 → 총 30개 모델 앙상블
  • Ablation Studies

    1. BiomedCLIP Fine-Tuning 효과 & XAI 기법 (Saliency Map Generation)
    2. Prompt 스타일의 영향 (단순한 프롬프트 VS 구체적인 프롬프트)
    3. 구성 요소별 기여도
    4. SAM 모델 및 Visual Prompt 전략 비교
  • Metrics

    • DSC (Dice-Sørensen Coefficient, 분할 정확도)
    • NSD (Normalized Surface Distance, 경계 정밀도)
    • 통계 검증: Paired-sample t-test
profile
Seoultech Student

0개의 댓글