CPS4: Class Prompt driven Semi-Supervised Spine Segmentation with Class-specific Consistency Constraint

기술 요약: CPS4 – 클래스별 일관성 제약을 이용한 클래스 프롬프트 기반 반지도 학습 척추 분할

문제 정의

MRI에서 척추체(Vertebral Bodies, VBs)와 추간판(Intervertebral Discs, IVDs)의 다중 클래스 분할을 포함하는 자동 척추 분할은 정형외과적 진단 및 수술 계획 수립에 매우 중요하다. 딥러닝 방식이 크게 발전했음에도 불구하고, 이들은 여에 많은 노동력을 필요로 하며 전문가의 지식이 요구되는 픽셀 수준의 주석(annotation)에 크게 의존한다. 반지도 학습(Semi-supervised learning, SSL)은 제한된 라벨 데이터와 풍부한 라벨 미지 데이터를 함께 활용함으로써 해결책을 제시한다. 그러나 일관성 정규화(consistency regularization)와 의사 라벨링(pseudo-labeling)에 주로 의존하는 기존의 SSL 방식은 종종 품질이 낮은 의사 라벨을 생성하는 문제를 겪는다. 이러한 노이즈가 섞인 라벨은 잘못된 감독(supervision)을 유도하여 모델이 무관한 특징에 집중하게 만들고, 라벨 미지 데이터를 효과적으로 활용하는 것을 방해한다.

시각-언어 모델(Vision-Language Models, VLMs)은 텍스트 클래스 프롬프트를 활용하여 세그멘테이션 맵의 품질을 높일 수 있는 잠재력을 보여주지만, 현재의 VLM은 클래스 프롬프트와 특정 이미지 객체 간의 일관성을 보장하기 위한 명시적인 제약이 부족하다. 이러한 결함은 약한 클래스별 의미론적 결합(semantic binding)과 클래스 간 간섭을 초래하여, 다중 클래스 세그멘테이션 맵 생성 성능을 저하시킨다.

방법론

저자들은 클래스 프롬프트를 반지도 척추 분할에 통합하는 새로운 프레임-워크인 CPS4를 제안한다. 이 모델은 2단계 훈련 전략을 사용한다.

1단계: 클래스별 일관성 제약이 적용된 VLM 사전 학습

기존의 VLM(예: CLIP, MedCLIP)은 이미지 수준에서 교차 모달 정렬(cross-modal alignment)을 수행하므로, 서로 다른 클래스 프롬프트가 각각의 대응하는 객체에 집중하도록 보장하지 못한다. 이를 해결하기 위해 CPS4는 시각(ViT 사용) 및 텍스트(BioClinicalBERT 사용) 인코더를 사용하는 VLM을 활용하여, 클래스 프롬프트와 척추 단위 간의 일관성을 강제하는 두 가지 새로운 손실 함수를 사용하여 사전 학습 단계를 도입한다.

1. 토큰 수준 어텐션 손실 ($L_{token}$): 이 손실은 클래스 프롬프트 토큰($F_{T,i}$)과 그에 대응하는 이미지 패치 토큰($F_I$) 사이의 정렬을 최적화한다. 이는 유사도 점수를 계산하여 거친(coarse) 어텐션 맵($A_i$)을 생성하고, 예측된 척추 단위 영역($B_i$)을 향해 활성화(activation)를 집계한다. 목표는 텍스트 프롬프트가 의미 공간에서 대상 척추 단위와 밀접하게 결합되도록 강제하는 것이다.
2. 픽셀 수준 어텐션 손실 ($L_{pixel}$): 모델이 대상 단위의 특정 하위 영역에만 어텐션을 과도하게 집계하는 것을 방지하기 위해 픽셀 수준의 교차 엔트로피 목적 함수가 추가된다. 이는 어텐션 맵이 전체 대상 척추 단위를 균일하게 덮도록 보장한다.

전체 사전 학습 손실은 $L_{token}$과 $L_{pixel}$의 평균이다.

2단계: 클래스 프롬프트 기반 반지도 척추 분할

이 단계에서는 사전 학습된 시각-텍스트 인코더를 사용하여 라벨 미지 데이터에 대한 의사 라벨링 과정을 강화한다.

• 클래스 프롬프트 가이드 세그멘테이션 맵 생성: 각 라벨 미지 척추 이미지에 대해, 이미지 패치 토큰과 클래스 프롬프트 토큰 사이의 내적(dot product)을 계산하여 각 클래스(예: "요추 1번")에 대한 이진 세그멘테이션 맵을 생성한다. 이 맵들은 업샘플링된 후 OTSU 자동 임계값 방법을 사용하여 이진 마스크로 변환된다. 모든 클래스별 맵은 하나의 통합된 다중 클래스 세그멘테이션 맵($y^{text}$)을 형성하기 위해 쌓인다.
• 의사 라벨 생성 및 최적화: 프레임워크는 학생(student) 네트워크와 교사(teacher) 네트워크를 갖춘 평균 교사(Mean Teacher, MT) 구조를 채택한다. 교사 네트워크는 지수 이동 평균(EMA)을 통해 라벨 미지 이미지에 대한 의사 라벨($y^t_u$)을 생성한다.
  • 지도 학습 손실 ($L_{sup}$): 라벨 데이터에 대한 학생의 예측과 실제 정답(ground truth) 사이에서 계산된다.
  • 반지도 학습 손실 ($L_{semi}$): 두 부분으로 구성된다.
    1. $L^{text}_{semi}$: 라벨 미지 데이터에 대한 학생의 예측과 클래스 프롬프트 생성 맵($y^{text}$) 사이의 일관성 손실.
    2. $L^{enhance}_{semi}$: 교사의 의사 라벨($y^t_u$)과 클래스 프롬프트 맵($y^{text}$)을 통합하여 생성된 강화된 의사 라벨과 학생의 예측 사이의 일관성 손실.

주요 기여

1. 최초의 텍스트 가이드 기반 척추 분할 SSL: 본 연구는 의사 라벨의 품질을 높이기 위해 텍스트 클래스 프롬프트를 반지도 척추 분할에 도입한 최초의 사례이다.
2. 새로운 클래스별 일관성 제약: 저자들은 각 클래스 프롬프트와 대상 척추 단위 간의 명시적인 제약을 강제하여, 긴밀한 의미론적 결합을 보장하는 토큰 및 픽셀 수준 어텐션 손실을 사용하는 특정 사전 학습 메커니즘을 제안한다.
3. 최첨단 성능(SOTA): 이 방법은 다양한 라벨 데이터 비율, 특히 데이터가 적은 환경에서 우수한 성능을 입증했다.

실험 결과

본 방법은 215명의 피험자로부터 얻은 20개 카테고리(19개 척추 구조 및 배경)의 T2 강조 MR 볼륨 이미지를 포함하는 MRSpineSeg 데이터셋을 통해 평가되었다.

• 성능: 단 5%의 라벨 데이터만을 사용하여, CPS4는 80.44%의 Dice 점수와 70.63%의 mIoU를 달ach했다.
• 비교: 이러한 결과는 인기 있는 반지도 학습 방법(예: MT, BCP, SS-Net) 및 기존 VLM 기반 접근 방식(예: CLIP, MedCLIP, MGCA, EGMA)을 능가한다.
  • 5% 라벨 데이터 조건에서 텍스트 프롬프트를 사용한 하위 수준의 GraphCL 방법과 비교했을 때, CPS4는 mDice를 2.23%, mIoU를 2.42% 개선했다.
• 강건성: 어블레이션 연구(Ablation studies)를 통해 토큰 수준 및 픽셀 수준 손실 모두가 성능에 크게 기여함을 확인했다. 이 방법은 라벨 데이터 비율이 5%에서 50%로 증가함에 따라 MT 베이스라인보다 확연히 높은 차이로 안정적인 성능 향상을 보여주었다.
• 시각화: CPS4에 의해 생성된 어텐션 맵은 대상 척추 단위와 강력한 정렬을 보여준 반면, 베이스라인 VLM(예: MedCLIP)은 산재된 활성화와 약한 의미론적 결합을 나타냈다.

의의

본 논문은 CPS4가 텍스트의 의미론적 힘을 활용하여 고품질의 의사 라벨 생성을 유도함으로써 척추 분할의 라벨 부족 문제를 효과적으로 완화한다고 주장한다. 사전 학습 중에 클래스 프롬프트와 척추 단위 간의 일관성을 명시적으로 제약함으로써, 이 프레임워크는 의료 영상 작업에서의 현재 VLM 한계를 극복한다. 결과는 이 접근 방식이 더 신뢰할 수 있는 클래스 프롬프트 기반 의사 라벨 생성을 가능하게 하여, 라벨링된 데이터가 제한적인 시나리오에서 반지도 학습의 효용성을 크게 향상시킨다는 것을 시사한다. 저자들은 추가 연구를 촉진하기 위해 코드를 공개할 예정이라고 밝혔다.