Improving tumor segmentation performance of MedCLIP-SAM framework using DPO


Related posts


Koleilat et al. (2025). MedCLIP-SAMv2: Towards Universal Text-Driven Medical Image Segmentation. MedIA 2025.

본 프로젝트는 범용 의료 세그멘테이션 모델 MedCLIP-SAMv2의 zero-shot segmentation 성능을 높이는, 즉 의료특화 파운데이션 모델인 BiomedCLIP의 종양/배경(혹은 특정 장기/배경) 구분 능력을 향상시키는 것입니다.

프로젝트의 핵심 아이디어는 Direct Preference Optimization, 일명 DPO를 BiomedCLIP을 미세조정하는 손실함수로 쓰자는 것입니다. DPO의 본래 인간의 선호도를 최대한 반영하기 위해 LLM을 최적화하는 수단 중 하나입니다. 우리 팀은 선호도와 같이 ‘정도’를 대상으로 하던 DPO를 양성/음성이라는 ‘이진 분류’ 문제에 적용하였고, ‘세그멘테이션’ 태스크를 최적화하는데 사용했습니다.

VQ-VAE Diagram
Koleilat, T. et al., “MedCLIP-SAMv2: Towards Universal Text-Driven Medical Image Segmentation”, MedIA, 2025, Figure 3.

본론에 앞서 MedCLIP-SAM를 간략히 소개하겠습니다. MedCLIP-SAM은 종양 양성 이미지와 의사의 진단(텍스트)를 입력으로 받으면 가장 먼저 의료특화 파운데이션 모델인 BiomedCLIP이 활성화 점수 분포를 내놓게 됩니다. 활성화 점수는 각 이미지 패치로부터 추출된 그라디언트를 정규화한 값인데, ‘종양’이라는 단어에 해당하는 패치에 높은 점수를 주게 됩니다. 이 점수 분포는 Information Bottleneck 모듈로 전달되어 입력 텍스와의 상호정보량은 최대화하고, 원본 이미지와는 상호정보량은 최소화하는, 최적의 saliency map(마스킹 이미지 초안)을 추출합니다. saliency map은 프롬프트 형식으로 전처리된 후 최첨단 이미지 분할 모델 SAM에 바운딩 박스 프롬프트로 들어가게 됩니다. 그리고 SAM은 원본 이미지와 보조 프롬프트를 활용해 정교화된 마스킹 이미지, 즉 종양의 위치를 표현한 이미지를 출력합니다. 여기까지가 zero-shot segmentation 단계에 해당하며 MedCLIP-SAM은 사전단계인 DHN-NCE fine-tuning과 최종단계인 weakly supervised segmentation도 포함하고 있습니다.

VQ-VAE Diagram
Koleilat, T. et al., “MedCLIP-SAMv2: Towards Universal Text-Driven Medical Image Segmentation”, MedIA, 2025, Figure 3.

언급한 세 단계 중에서도 저희 팀의 핵심 목표는 바로 DHN-NCE fine-tuning 단계에서 DHN-NCE loss 대신 DPO loss를 사용해 성능을 도모하는 것입니다.

alt text

다만 DPO 손실함수 특성을 고려했을 때 애로사항이 하나 있었습니다. 바로 한 종양 이미지에 대해 ‘동일한’ 환자의 정상 이미지를 구해야 한다는 것이었습니다. DPO 수식은 양성/음성 데이터를 모두 요구하였고 두 이미지 간 종양 유무만 부각해야 했기 때문에 동일한 환자의 정상 데이터가 필요했습니다. 직면한 문제는 다음과 같이 정리됩니다.

  • DPO와 더불어 automatic data augmentation도 저희 프레임워크의 핵심 요소였기 때문에 소규모 데이터만 활용해야 했습니다. (의료현장의 데이터 부족 현상을 반영한 설정)
  • 사용할 MRI 데이터셋 Figshare는 참고할 수 있는 정상 이미지가 없었습니다.
  • 특정 장기 세그멘테이션의 경우, 음성 상태의 정의가 모호했습니다.

제 핵심 기여를 요약하자면

  • 양성으로부터 음성 샘플을 생성하는 방법론을 직접 고안하고 실험을 수행
    • Biomedclip directional loss를 결합한 SD-inpaint 모델 제안
    • image to prompt를 통해 정상 이미지를 그대로 조건부 프롬프트로 활용하는 모델 제안
    • 세포 발생학에서 영감을 얻은 Neural Cellular Diffusion을 CT,MRI,Ultrasound 모달리티에 적용 제안
    • self-supervised learning 기반 디퓨전 모델 활용
  • DPO 손실함수 개선
  • 데이터 증강 기법 모색
  • 기존 MedCLIP-SAM 오류 분석 및 추가 실험 데이터 후보 선정 로 정리할 수 있습니다.

alt text alt text

위 그림에서 미세조정을 위한 데이터 구성 및 사용할 DPO 수식을 표현했습니다. DPO 보상함수 항은 양성/정상에 더불어 할루시네이션(첨자에 -표기), 비할루시네이션(첨자에 +표기) 여부도 고려하기 위해 총 4개 항으로 구성됩니다. 이렇게 구성된 수식은 같은 모델(ours,reference) 내에서 할루시네이션, 비할루시네이션 데이터 간 점수 차를 벌리고 reference 모델 대비 ours 모델의 점수를 올리는 방향으로 학습하게 합니다.

데이터 증강 측면에서는 단일 종양 이미지에 대해 음성 샘플을 생성하고, 할루시네이션/비할루시네이션 텍스트에 한번씩 매칭해주니 데이터가 4배로 증가하는 효과가 있습니다.

Self-supervised learning based diffusion

VQ-VAE Diagram
Durrer et al., “Denoising Diffusion Models for Inpainting of Healthy Brain Tissue”, arXiv, 2024, Figure 2.

Durrer et al. (2024). Denoising Diffusion Models for Inpainting of Healthy Brain Tissue. arXiv:2402.17307.

정상 데이터가 없는 Figshare를 고려해 자기 지도 학습 기반의 디퓨전 모델을 고려하게 되었고 BraTS challenge(뇌MRI 정상 조직 복원 챌린지) 공개 솔루션 중 관련 논문을 참고했습니다. 이 논문은 종양을 마스킹 이미지로 가린 뒤 나머지 영역에 대해서만 노이징, 디노이징을 수행합니다. 또한 배경 유지를 위해 3채널 구조를 사용했고, 해상도를 중요시하는 의료도메인 특성을 고려해 pixel space에서 loss를 계산했습니다. 그러나 뇌종양의 경우 종양 부위 자체와 더불어 주변 조직들에도 종양 흔적(예. 부푼 종양에 의해 찌그러짐)이 남기 마련인데 해당 논문은 그 점까지 고려하지 못한 한계가 있었습니다. 이에 저는 종양으로 의심되는 부분까지 마스킹 이미지에 포함했습니다.

alt text

종양 의심 부위는 파운데이션 모델 BiomedCLIP와 gscore-cam을 결합한 모듈을 활용하여 추출했습니다. SSL을 위한 최종 마스킹 이미지는 노이징, 디노이징 과정 그리고 reconstruction loss 계산에서도 제외했습니다. 한편 복원할 부위를 나타내는 마스킹 이미지는 마스킹되지 않은 영역 중에서 샘플링하였고, 실제 종양 모양을 반영하기 위해 종양 형태를 그대로 본땄습니다.

또한 측면/정면/윗면 총 세 구도를 골고루 학습시키기 위해 전체 데이터셋에서 세 구도가 차지하는 비율을 반영해 statified sampling하였습니다. 최종적으로 90장의 stratified samples와 10장의 소수 샘플로 훈련 데이터셋을 구성했습니다.

alt text

Image as conditional prompt

두번째로 고안한 방법은 텍스트로 된 진단 내용 대신 소수의 정상 이미지들만 SD의 조건부 프롬프트로 넣는 방법이었습니다. 의학 용어를 SD의 조건부 임베딩로 활용하기 위해선 크로스 어텐션 접합부와 SD를 full-finetuning할 필요가 있었습니다. FT의 부담을 줄이기 위해, 저는 이미지 자체를 프롬프트로 활용하기로 하고 관련 논문을 참고했습니다.

Ding et al. (2024) The CLIP Model is Secretly an Image-to-Prompt Converter. NeurIPS 2024.

해당 논문은 CLIP의 이미지 임베딩을 텍스트 임베딩으로 변환하는 closed-form 행렬이 존재함을 보입니다.

다른 한편 텍스트를 살리는 방법도 생각해 보았습니다. 저는 알아듣기 힘든 진단문을 순화하기 위해 VQA를 통과한 정상 텍스트들을 블렌딩하는 방법을 떠올렸습니다.

Biomedclip directional Loss

alt text

세번째로 고안한 방법은 CLIP 공간에서 의미론적 벡터의 선형성이 성립한다는 점을 살려 일반적인 ‘종양’ 벡터를 활용하는 것이었습니다.

  • NLP에서 사용된 intrinsic evaluation을 통해 의미론적 벡터의 선형성이 보장되는지 확인

alt text

초음파의 경우 기존의 음성->양성 생성모델을 통해 생성된 이미지를 역으로 음성으로 변환하는 태스크를 수행함으로써 원본 정상이미지와 생성물 간 FID 평가지표를 사용하게 되었습니다.

유사도 점수가 CLIP 텍스트 인코더 입력에 민감하다는 문제

  • pesudo label에 덧붙여 complementary label도 포함
  • BLIP-VQA 모델로 이미지로부터 k=5개의 프롬프트를 추출함(we utilize k = 5 prompts that are semantically equivalent yet lexically diverse) -> Training-Free Label Space Alignment for Universal Domain Adaptation, CLIP 원본 논문에 언급된 프롬프팅 방법으로, 프롬프트 민감도를 낮추는데 사용됨

DPO 손실함수 개선

alt text

저는 DPO가 상대적인 거리를 벌리는데 집중한다는 문제에도 관심을 기울이고 있었습니다. 특히 저희는 ‘정도’ 회귀보다 ‘분류’에 가까운 문제를 다루기 때문에, 저는 양성/음성을 가르는 threshold가 존재한다는 가정과 DPO가 그 threshold를 간과할 수 있겠다는 가정을 세웠습니다.

Hard negative prompt generation

✅ 원본 마스킹 이미지를 5x5 grid로 보고(텍스트로 위치 설명 가능한 범위 제한) 종양 부위(하얀 부위)가 속하지 않은 grid를 할루시네이션 위치로 선정

alt text

  • MedTrinity-25M: A Large-scale Multimodal Dataset with Multigranular Annotations for Medicine
  • TripletCLIP: Improving Compositional Reasoning of CLIP via Synthetic Vision-Language Negatives

alt text

BENIGN case

POSITIVE PROMPT:

A mammogram showing a sharply marginated, round mass suggestive of a benign breast tumor. The tumor is located in the upper middle part of the image, horizontally oriented, with a clear boundary.

✅ NEGATIVE PROMPT created by Mistral-7B-Instruct-v0.2 based on TripletCLIP

A mammogram displaying a round, irregularly marginated mass indicative of a malignant breast tumor. The tumor is situated in the lower left part of the image, vertically oriented, with an indistinct boundary.

MALIGNANT case

POSITIVE PROMPT:

A mammogram displaying an irregular, invasive mass suggestive of a malignant breast tumor. The tumor is located in the upper middle part of the image, horizontally oriented, with a blurry boundary.

✅ NEGATIVE PROMPT created by Mistral-7B-Instruct-v0.2 based on TripletCLIP

A mammogram displaying a well-defined, round mass indicative of a benign breast tumor. The tumor is situated in the lower left part of the image, vertically oriented, with a clear boundary.

Experiements

Datasets

  • csv파일, images, images_local_noisy
  • csv파일 형식
    • [종양 이미지, 종양 캡션, 종양 이미지, N종양 캡션]
    • [N 종양 이미지, N 종양 캡션, N 종양 이미지, 종양 캡션]

Experimental settings

  • 어떤 데이터셋이나 모델(Baseline/DHN/DPO)이 들어와도 동일하게 작동하는 표준 스크립트를 제작
  • 각 테스트 이미지 파일명에 맞춰 알맞은 텍스트 프롬프트(Benign/Malignant)를 실시간 생성하여, 기존에 발생하던 nan 에러(KeyError)를 방지
  • 추론 시 GPU 병렬처리를 활용해 단일 실험 시간 단축
  • epoch 10 기본, 두 번째 표는 epoch 10 훈련 후 loss가 가장 낮은 가중치는 best_params 폴더로 따로 저장

Results

Breast Ultrasound

alt text

  • test set 프롬프트는 benign, malignant 각각 한 문장 선정 후 통일
  • train set 프롬프트는 text_prompt.py 참고하여 랜덤 배정

UDIAT set only case

  • 7:3 임의 분할
  • 기존 json에 프롬프트가 있는 이미지는 그대로 쓰고, 프롬프트가 없는 이미지는 saliency_maps/text_prompts.py에서 임의로 가져옴 (baseline,ref,dpo 세 케이스에 대해 동일한 로직을 사용하도록 함)
  • 훈련셋의 경우 사전에 saliency_maps/text_prompts.py를 통해 pair를 다 만들어 놓고 세 모델이 동일한 pair를 공유하게 함

UDAIT (train), BUSI (test) case

  • 테스트셋으로 BUSI benign, malignant 데이터(총 647개) 사용

Brain MRI

alt text

  • Figshare brain tumor(train data 50/100/150/200/500/800, test 200 고정, text prompt 병변별)
  • 데이터생성모델 세팅 : train 데이터 1462개, diffusion_training_step 170k, diffusion_inference_step 1k

alt text

  • Figshare brain tumor_v2(train data 50/100/150/200/500/800, test 200 고정, text prompt 병변별)
  • 데이터생성모델 세팅 : train 데이터 100개, diffusion_training_step 122k, diffusion_inference_step 1k

Lung CT / X-ray

추가로 시도한 방법들

모달리티 갭 극복을 위한 latent drifting

Latent Drifting in Diffusion Models for Counterfactual Medical Image Synthesis, CVPR 2025

Neural Cellular Diffusion을 통한 음성 샘플 생성

발생생물학에 영감을 받은 디퓨전 모델인 Neural Cellular Diffusion의 성능을 위해 실험을 진행 중입니다.

후보 reference model 고안

Biomedclip was fine-tuned by adopting the EgoBridge architecture. When co-training paired images, Optimal Transport was applied between tumor regions, and contrast learning was applied between tumor and background.

BlobGAN

TACE: Tumor-Aware Counterfactual Explanations (2024)

종양 의심 부위를 Blob이라는 클러스터를 이용해 세그멘테이션 및 정상 조직 복원 태스크를 수행한 논문입니다. 외부 분류기를 학습 파이프라인에 포함시키기 때문에 BiomedCLIP의 지식을 주입하기 좋은 방법론이라고 생각했습니다.

앞서 여러 생성 방법론들이 언급되었지만 사실 논문 투고를 위해 통일된 음성 샘플 생성모델을 선정할 필요가 있다고 판단하였다. 이에 추가적인 실험을 진행하여 CT,MRI,Ultrasound 세 모달리티에 동시에 적용가능한 방법론을 선정하려고 합니다.

생성모델을 DPO로 미세조정하기

Further research