OOD-SEG: Exploiting out-of-distribution detection techniques for learning image segmentation from sparse multi-class positive-only annotations

논문 제목: OOD-SEG: 희소 다중 클래스 양의 (Positive-only) 주석만으로 이미지 분할을 학습하기 위한 분포 외 (OOD) 검출 기법 활용

1. 문제 정의 (Problem)

의료 및 수술 영상 분할 분야에서 딥러닝 기반 방법은 두 가지 주요한 도전에 직면해 있습니다.

1. 주석 획득의 어려움: 의료 이미지에 대한 픽셀 수준의 완전한 분할 마스크를 수동으로 작성하는 것은 시간이 많이 소요되고 전문 지식이 필요하여, 임상 환경에서 레이블이 지정된 데이터가 부족합니다. 특히, 배경 (Background) 클래스에 대한 주석이 전혀 없는 희소 다중 클래스 양의 (Positive-only) 주석 (예: 특정 장기만 표시하고 나머지는 비표시) 이 많이 존재합니다.
2. 분포 외 (OOD) 검출 부재: 기존 분할 파이프라인은 훈련 데이터와 테스트 데이터가 동일한 분포 (In-Distribution, ID) 에 있다고 가정합니다. 따라서 훈련 중 보지 못한 클래스 (예: 새로운 수술 도구, 새로운 병변, 배경) 가 입력되면 모델은 이를 임의의 알려진 클래스로 잘못 분류하여 위험한 결과를 초래할 수 있습니다. 의료 안전을 위해 모델이 자신이 모르는 데이터를 식별하고 경고할 수 있어야 합니다.

기존의 약지도 학습 (Weakly Supervised Learning) 방법들은 이러한 배경 주석의 부재를 해결하지 못하거나, OOD 검출 기능을 통합하지 못했습니다.

2. 방법론 (Methodology)

저자들은 양의 - 레이블 없는 (Positive-Unlabelled, PU) 학습 패러다임을 픽셀 단위 분할 문제에 적용하고, 이를 분포 외 (OOD) 검출 기법과 결합한 새로운 프레임워크인 OOD-SEG를 제안합니다.

핵심 개념 및 프레임워크

• 문제 재정의: 주석이 없는 픽셀을 '배경 (Negative)'으로 간주하는 대신, 양의 클래스일 수도 있고 배경/알 수 없는 클래스일 수도 있는 '레이블 없는 (Unlabelled)' 데이터로 간주합니다.
• 학습 전략:
  • 훈련 단계: 오직 양의 클래스 (ID) 로 주석된 픽셀만을 사용하여 표준 다중 클래스 분할 네트워크 (U-Net 기반) 를 훈련시킵니다. 배경 클래스에 대한 주석은 사용하지 않습니다.
  • 추론 단계: 네트워크의 출력 (Logits) 을 기반으로 각 픽셀이 ID 클래스에 속할 확률 점수 ($S_c$) 를 계산합니다.
  • OOD 검출: 계산된 점수가 미리 정의된 임계값 ($\tau$) 보다 낮으면 해당 픽셀을 OOD (배경 또는 알려지지 않은 클래스) 로 간주하고 0 으로 분류합니다.
  • 수식: $\hat{y}_{ij} = \arg\max_c S_{ij}^c$ if $\max_c S_{ij}^c > \tau$, otherwise $0$.

적용된 OOD 검출 기법

프레임워크는 분류 작업을 위해 설계된 다양한 OOD 검출 기법을 픽셀 단위로 통합할 수 있습니다. 본 논문에서는 다음 네 가지 방법을 비교 평가했습니다.

1. Baseline: Softmax 출력 확률 직접 사용.
2. ODIN: 온도 스케일링 (Temperature Scaling) 과 적대적 교란 (Adversarial Perturbation) 을 적용하여 신뢰도 점수를 보정. (본 연구에서는 교란은 생략하고 온도 스케일링만 적용).
3. Mahalanobis Distance: 특징 공간에서 클래스별 평균과 공분산 행렬을 기반으로 마할라노비스 거리를 계산하여 점수 산출.
4. GODIN (Generalized ODIN): 학습 가능한 온도 스케일링 함수와 유사도 측정 계층을 도입. 본 논문에서는 픽셀 단위 연산을 3x3 컨볼루션 레이어로 확장하여 공간적 맥락을 포착하도록 개선했습니다.

평가 전략 (Two-level Cross-Validation)

OOD 검출을 평가하기 위해 별도의 OOD 테스트셋이 필요하다는 기존 한계를 극복하기 위해 이중 레벨 교차 검증 (Two-level Cross-Validation) 전략을 제안했습니다.

• 주체 분할 (Subject Partitions, SP): 환자/주체 간 중복을 방지.
• 클래스 분할 (Class Partitions, CP): 훈련 시 일부 양의 클래스를 제외하고, 이를 검증 단계에서 'OOD 클래스'로 간주하여 평가합니다.
• 이를 통해 단일 데이터셋 내에서 ID 성능과 OOD 검출 능력을 동시에 평가할 수 있습니다.

평가 지표

• AUROC & AUPR: ID 클래스와 OOD 클래스를 이진 분류 문제로 간주하여 (One-vs-Rest가 아닌, 양의 클래스 vs 나머지) 계산된 ROC 곡선 아래 면적과 정밀도 - 재현율 곡선 아래 면적을 사용합니다.
• 임계값 선택: ID 분류 정확도 (TPR) 와 OOD 검출 성능 (TNR) 을 균형 있게 최대화하는 최적 임계값 $\tau_m$을 자동으로 선택합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 프레임워크: 희소 다중 클래스 양의 주석만으로 의료 이미지 분할을 학습하는 새로운 PU 학습 기반 프레임워크를 제안했습니다. 이 방법은 배경 주석 없이도 OOD 픽셀을 효과적으로 식별하며, 다중 클래스 ID 성능을 저하시키지 않습니다.
2. 평가 프로토콜: 별도의 OOD 데이터셋 없이도 단일 데이터셋 내에서 시스템적 평가를 가능하게 하는 이중 레벨 교차 검증 전략과 관련 지표를 제안했습니다. 이는 의료 영상 분야에서 OOD 평가의 표준을 마련하는 데 기여합니다.
3. GODIN의 확장: 분류용 GODIN 방법을 분할 작업에 맞게 **컨볼루션 기반 (Convolutional adaptation)**으로 확장하여 공간적 맥락을 고려한 OOD 검출을 가능하게 했습니다.

4. 실험 결과 (Results)

• 데이터셋:
  • Heiporspectral: 돼지 장기 20 종에 대한 16 밴드 초분광 (HSI) 이미지 (5,758 장).
  • ODSI-DB: 치과 조직 35 종에 대한 초분광 이미지 (316 장, 9 종 클래스 선택).
  • DSAD: 복강경 수술 영상 (RGB) 으로 11 종 복부 장기 (13,195 장).
• 성능:
  • OOD 검출 성능: ODIN 기반 점수 함수가 모든 데이터셋에서 가장 높은 AUROC 및 AUPR 점수를 기록했습니다. 이는 신뢰도 보정 (Calibration) 이 OOD 검출에 중요함을 시사합니다.
  • ID 분류 성능: OOD 검출 임계값 ($\tau_m$) 을 적용하더라도 알려진 클래스 (ID) 에 대한 분류 정확도는 Baseline ($\tau=0$) 과 비교해 유의미한 저하가 없었습니다.
  • 비교: Mahalanobis 방법은 고차원 특징 공간에서의 공분산 추정 불안정성으로 인해 상대적으로 낮은 성능을 보였습니다.
• 시각화: OOD-SEG 는 배경이나 알려지지 않은 구조물을 정확히 '배경 (0)'으로 분류하는 반면, Baseline 모델은 이를 잘못된 클래스로 과신하여 분류하는 경향이 있음을 확인했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 안전성 강화: 의료 자동화 시스템이 훈련 데이터에 없는 상황 (새로운 장비, 예기치 않은 조직 등) 에서 실패를 감지하고 경고할 수 있게 하여 임상 적용의 안전성을 높입니다.
• 주석 효율성: 배경 주석이 불필요하므로, 전문가의 주석 부담을 크게 줄이면서도 고품질의 분할 모델을 학습할 수 있습니다.
• 연구 방향 제시: 의료 영상 분할 분야에서 OOD 검출을 위한 표준적인 평가 방법론과 벤치마크를 제시하여, 향후 관련 연구의 기반을 마련했습니다.

한계점 및 향후 과제:

• 현재는 레이블 없는 데이터를 학습 과정에 활용하지 않음 (향후 반지도 학습 등으로 확장 가능).
• OOD 성능 향상을 위한 적대적 교란 (Adversarial Perturbation) 을 포함하지 않음 (계산 비용 및 복잡성 고려).
• ID 분류 정확도와 OOD 검출 성능 간의 트레이드오프를 더 잘 균형 잡는 방법 모색 필요.

이 논문은 의료 영상 분석의 실용적인 제약 (희소 주석) 과 안전 요구사항 (OOD 검출) 을 동시에 해결하는 혁신적인 접근법을 제시했다는 점에서 큰 의의가 있습니다.