Label tree semantic losses for rich multi-class medical image segmentation

기술 요약: 풍부한 다중 클래스 의료 영상 분할을 위한 라벨 트리 의미 손실

문제 제기
풍부하고 정확한 의료 영상 분할은 수술 전 계획, 수술 중 내비게이션, 수술 후 평가에 필수적입니다. 그러나 분할 작업에 대한 표준 학습 방법은 일반적으로 모든 분류 오류를 동등하게 패널티화하여, 라벨 공간 내의 의미적 관계를 활용하지 못합니다. 라벨의 카디널리티가 세밀하게 다른 클래스들 (예: 특정 뇌 하부 영역 또는 조직 유형) 을 포함하도록 커짐에 따라 이러한 한계는 점점 더 문제가 됩니다. 또한, 수술 영상 및 대규모 주석 작업에서 데이터는 종종 "양성 전용" 라벨로 희소하게 주석되며, 라벨이 지정되지 않은 영역을 배경으로 간주할 수 없습니다. 기존 접근 방식들은 라벨 계층 구조를 활용하거나 의미 있는 방식으로 라벨이 지정되지 않은 픽셀의 모호성을 처리하는 메커니즘이 부족한 경우가 많습니다.

방법론
저자들은 두 가지 보완적인 의미 손실 함수를 통해 라벨의 계층적 조직 (라벨 트리) 을 훈련 과정에 통합하는 프레임워크를 제안합니다. 이러한 손실 함수는 밀집된 감독 (예: 전체 뇌 분할) 과 희소하며 배경이 없는 주석 (예: 신경외과 초분광 영상) 모두와 함께 작동하도록 설계되었습니다.

1. 워터스타인 거리 기반 분할 손실 ($\mathcal{L}_{wass+seg}$):
   이 손실 함수는 예측된 확률 분포를 기준 (ground truth) 분포로 변환하는 비용을 측정하기 위해 워터스타인 거리 (최적 수송) 를 활용합니다. "기준 거리" 행렬 ($M$) 은 라벨 트리에서 유도되며, 두 개의 잎 노드 간의 거리는 계층 구조 내에서의 경로 길이로 정의됩니다. 이 공식은 오류를 의미적 심각도에 따라 패널티화합니다. 즉, 클래스를 의미적으로 먼 형제 클래스로 잘못 분류하는 것은 가까운 친척으로 잘못 분류하는 것보다 더 높은 패널티를 부과합니다. 최종 손실은 워터스타인 거리와 일반 분할 손실 (예: Dice 또는 교차 엔트로피) 의 복합체입니다.

2. 트리 가중 의미 교차 엔트로피 손실 ($\mathcal{L}_{twce+seg}$):
   이 접근 방식은 표준 교차 엔트로피 (CE) 를 확장하여 잎 노드뿐만 아니라 라벨 트리의 모든 노드에서 확률을 집계합니다. 인접 행렬을 사용하여 부모 - 자식 관계를 표현함으로써, 모델은 내부 노드에 대한 집계 확률을 계산합니다. 그런 다음 손실은 모든 노드에 대한 CE 의 가중 합으로 계산되며, 가중치는 엣지 구조에 따라 할당됩니다. 이를 통해 모델은 추상화 수준 (거시적 및 미시적) 의 여러 단계에서 감독 신호를 받을 수 있습니다.

3. 희소 양성 전용 학습과의 통합:
   제안된 손실 함수는 희소하고 양성 전용 주석을 사용하여 훈련하는 프레임워크에 통합됩니다 (Wang et al., 2026 기반). 이 설정에서 모델은 주석된 양성 클래스에 대해서만 훈련됩니다. 추론 중에는 분포 외 (OOD) 감지 메커니즘이 최대 예측 확률이 낮은 픽셀을 배경으로 플래그합니다. 저자들은 선택된 계층 수준 (예: 최상위 노드) 에서 확률을 집계하여 OOD 임계값을 결정함으로써 양성 클래스의 성능을 저하시키지 않으면서 더 강력한 배경 감지를 가능하게 하는 방식으로 이를 확장했습니다.

주요 기여

• 새로운 손실 함수: 최적화 과정에서 명시적으로 계층적 라벨 관계를 인코딩하는 두 가지 트리 기반 의미 손실 (워터스타인 기반 및 트리 가중 CE) 의 도입.
• 다용도 프레임워크: 이러한 손실 함수가 밀집 감독 전체 뇌 분할 (WBP) 과 희소 감독 신경외과 초분광 영상 (HSI) 모두에 적용되며, 배경 식별을 위한 OOD 감지와의 통합이 어떻게 이루어지는지 입증.
• 성능 향상: 전문가가 정의하거나 해부학적으로 유도된 라벨 계층 구조를 통합하는 것이 표준 베이스라인에 비해 더 의미 있는 감독과 향상된 분할 성능으로 이어진다는 경험적 증거.
• 소형 구조에 대한 강건성: 표준 베이스라인이 종종 어려움을 겪는 제한된 주석을 가진 작은 해부학적 구조에서 관찰된 구체적인 개선 사항.

실험 결과
저자들은 두 가지 명확한 작업에서 자신의 접근 방식을 평가했습니다:

1. 전체 뇌 분할 (WBP): Mindboggle42, AOMIC, IXI 의 세 가지 데이터셋에서 얻은 3D T1 가중 MRI 데이터를 사용하여 완전 감독 하에 수행.
   • 결과: 워터스타인 기반 복합 손실 ($\mathcal{L}_{wass+seg}$) 은 표준 nnU-Net 베이스라인 (Dice + CE) 과 일반화된 워터스타인 Dice 손실보다 일관되게 우수한 성능을 보였습니다. 개선 사항은 작은 구조와 관련된 지표 (Dicesmall, NSDsmall) 및 표면 중첩 (NSD) 에서 가장 두드러졌습니다. 통계 분석은 거의 모든 훈련 - 테스트 설정에서 NSD 및 작은 구조 지표의 유의미한 향상을 확인했습니다.
2. 신경외과 초분광 영상 (HSI): 수술 중 수술 비디오에서 얻은 희소하고 양성 전용 주석을 사용.
   • 결과: 제안된 프레임워크는 잎 노드만으로 훈련된 베이스라인에 비해 최상위 클래스 (예: 정상 대 비정상 조직 구분) 의 성능을 향상시켰습니다. 계층 가중 최상위 감독 ($\mathcal{L}_{twce+seg}$, 특정 엣지 가중치 적용) 은 희소 설정에서 강력한 지지를 보여주었습니다. 혼동 행렬 분석은 제안된 방법이 베이스라인에 비해 의미상 비현실적인 오류 (예: 정상 및 비정상 조직 혼동) 를 줄였음을 드러냈습니다.

의의 및 주장
본 논문은 풍부한 분할 작업, 특히 미묘한 클래스 변이와 희소 데이터를 다룰 때 라벨 공간의 구조를 정보의 원천으로 활용하는 것이 중요하다고 주장합니다. 균일한 오류 패널티화를 넘어섬으로써 제안된 방법은 모델이 "더 나은 실수" (즉, 진실에 의미적으로 더 가까운 오류) 를 할 수 있게 합니다. 저자들은 자신의 접근 방식이 분할 네트워크 (예: nnU-Net 또는 U-Net) 의 아키텍처 수정 없이 표준 베이스라인을 개선하여 기존 파이프라인에 대한 실용적인 확장임을 강조합니다. 이 연구는 계층적 의미 감독이 이전에 탐구되지 않았던 수술 영상 분야의 격차를 해소하고, 대규모 희소 주석 데이터셋에 대한 훈련을 위한 확장 가능한 솔루션을 제공합니다.