MAGE: Color-Invariant and Spatial Knowledge Distillation for Gastric Neoplasm Classification

기술 요약: MAGE – 위 종양 분류를 위한 색상 불변 및 공간 지식 증류

문제 정의
내시경 검사 중 위 선종(adenoma)과 암(carcinoma)을 정확하게 구별하는 것은 임상적 의사결정에 있어 매우 중요하지만 여전히 매우 어려운 과제이다. 이 작업은 클래스 간 유사성이 높고, 경계가 모호하며, 딥 뉴럴 네트워크가 내재적인 형태학적 특징보다는 허위 상관관계(예: 출혈, 조명 아티팩트 또는 배경 맥락)에 의존하는 경향 때문에 어려움을 겪는다. 기존 방식들은 다음과 같은 구체적인 한계에 직면해 있다:

• ROI 기반 방식: 불완전한 탐지 또는 세그멘테이션으로 인한 오류 전파 문제를 겪을 수 있으며, 주변의 전역적 맥락(global context)을 놓치기 쉽다.
• 전체 이미지 분류: 필요한 공간적 집중력이 부족하며, 병변 자체가 아닌 배경 영역을 바탕으로 예측하는 경우가 많다.
• 색상 의존성: 모델이 도메인 특유의 텍스처 편향(texture bias)으로 기울어지는 경향이 있으며, 여기서 색상은 주요 진단 지표가 아닌 취약한 혼란 변수(confounder)로 작용한다.

방법론: MAGE 프레획워크
저자들은 색상 불변 형태학적 학습을 강제하면서도 추론 시에는 마스크 없이 전체 이미지에 배포 가능한 MAGE(Masked Achromatic Guidance Expert)라는 이중 브랜치 훈련 프레임워크를 제안한다.

1. 마스크된 무채색 전문가 (보조 브랜치):

   • 훈련 중에 보조 "전문가" 브랜치는 입력의 **마스크된 무채색 뷰(masked achromatic view)**를 처리한다.
   • 입력은 백색광 이미지(WLI)를 그레이스케일로 변환하고 정답 세그멘테이션 마스크($\tilde{x} = m \odot \text{Gray}(x)$)를 적용하여 생성되며, 관심 영역(ROI) 외부의 픽셀은 0으로 설정된다.
   • 이 브랜치는 오직 구조적, 색상 불변 특징만을 학습하도록 훈련되어, 색상 단서나 배경 아티팩트가 아닌 형태학적 패턴을 학습하도록 강제된다.
   • 이 브랜치는 특권 정보(privileged information, 마스크)를 활용하며, 증류 과정 동안 **동결(frozen)**된다.

2. 메인 모델 (배포 가능 브랜치):

   • 메인 브랜치는 **전체 원본 WLI (RGB)**를 입력으로 받으며, 추론 시 사용되는 유일한 구성 요소이다.
   • 이 모델은 이중 목적 증류 전략을 통해 전문가를 모방하도록 학습된다:
     • 분류 로짓 증류 (Classification Logit Distillation): 메인 모델의 출력 로짓은 전문가의 로짓에 대한 KL 발산을 통해 정규화되며, 이를 통해 전문가의 분류 확신도를 전달받는다.
     • 공간 주의력 전이 (Spatial Attention Transfer): 메인 모델의 특징 맵 주의력이 전문가의 주의력과 일치하도록 정규화된다. 채널 차원을 축소하고 메인 모델의 주의력 맵과 (전문가에 의해 제한된) 전문가의 주의력 맵 사이의 제곱 프로베니우스 거리(squared Frobenius distance)를 계산함으로써, 메인 모델은 전체 이미지 내에서 종양 구조에 집중하도록 암묵적으로 유도된다.

3. 아키텍처 선택:

   • 본 프레임워크는 Vision Transformer(ViT) 대신 합성곱 신경망(CNN) 백본을 활용한다. 저자들은 CNN이 데이터셋 크기에 적합한 귀납적 편향(Inductive Bias, 이동 불변성, 국소성)을 제공하며, 주의력 손실 계산에 필요한 정밀한 공간 정렬을 보장한다고 주장한다.

4. 손실 함수:
   전체 손실 $L$은 다음을 결합한다:

   • 정답 레이블에 대한 가중치 적용 이진 교차 엔트로피 (Weighted Binary Cross-Entropy, BCE).
   • 전문가와 메인 모델의 로짓 간 온도 조절된 KL 발산 (Temperature-scaled KL divergence).
   • 주의력 정규화 항 (주의력 맵 간의 프로베니우스 거리).

주요 기여

1. 훈련 전용 보조 전문가: 배경 맥락과 색상 정보를 억제함으로써 색상 불변 형태학적 학습을 촉진하는 마스크된 무채색 전문가를 도입하였다.
2. 공간 증류 프레임워크: 추론 시 세그멘테이션 마스크 없이도 작동할 수 있도록, 전문가로부터 ROI 중심의 구조적 단서와 공간적 주의력을 배포 가능한 전체 이미지 분류기로 전이하는 메커니즘을 제공한다.
3. 성능 및 강건성: 강력한 탐지 기반(예: YOLO) 및 분류 기반(예: Swin-Transformer, GastroNet) 베이스라인 모델들을 상회하는 일관된 개선을 입증하였으며, 설명 가능성과 색상 변화에 대한 강건성을 향상시켰다.

실험 결과
실험은 개인 위 내시경 데이터셋(서울 성모 병원)과 외부 대장 데이터셋(PICCOLO)을 대상으로 수행되었다.

• 분류 성능: MAGE는 위 테스트 세트에서 AUC 0.786, F1 0.733을 달성하여 GastroNet-5M (AUC 0.744), Swin-Tiny (AUC 0.719) 및 YOLO 기반 파이프라인을 포함한 베이스라인들을 능가하였다.
• 설명 가능성: MAGE는 Grad-CAM 시각화에서 정답 마스크와 우수한 정렬을 보여주었다.
  • Attention-in-Mask (AiM): MAGE는 0.5303을 기록하였다 (베이스라인 EffNet의 0.3176 대비).
  • Pointing Game Accuracy (PGA): MAGE는 0.7324를 기록하였다 (베이스라인의 0.4355 대비).
  • 이는 모델이 허위 전역 패턴보다 종양의 형태학에 더 많이 의존함을 나타낸다.
• 색상 강건성: 합성된 색조 변화(hue shifts) 하에서, MAGE는 EfficientNet 및 GastroNet-사전 학습된 ResNet50 베이스라인들에 비해 성능 저하가 현저히 적게 나타났으며, 이는 색상 불변 훈련의 효능을 확인시켜 준다.
• 교차 도메인 일반화: 외부 PICCOLO 데이터셋에서 MAGE는 가장 높은 **AUC (0.9434)**를 달성하였으며, 이는 학습된 형태학적 특징이 서로 다른 해부학적 도메인 간에 전이 가능하다는 것을 시사한다.
• 어블레이션 연구 (Ablation Studies): 비대칭 설계(Grayscale Expert + RGB Main)가 핵심적임을 확인하였다. 전문가는 그레이스케일 입력에서 더 우수한 성능을 보인 반면, 메인 모델은 최적의 배포를 위해 전체 RGB 입력이 필요했다.

의의 및 주장
본 논문은 MAGE가 특권 훈련 신호를 통해 모델이 구조적, 색상 불변 특징을 명시적으로 학습하도록 강제함으로써 위 종양 분류의 "지름길 학습(shortcut learning)" 문제를 해결한다고 주장한다. 마스크에 대한 훈련 요구사항을 추론 시의 요구사항으로부터 분리함으로써, MAGE는 진단 정확도와 임상적 해석 가능성을 모두 개선하는 실용적인 솔루션을 제공한다. 저자들은 MAGE가 우수한 분류 지표뿐만 아니라 임상적 신뢰성에 필수적인 해석 가능한 주의력 맵을 제공한다는 점을 강조한다.

한계점
저자들은 본 연구가 전용 탐지/세그멘테이션 모듈을 훈련시키지 않고 분류와 강건성에 집중했음을 언급하며, 이는 특정 임상 워크플로우에서 직접적인 로컬라이제이션(localization) 출력을 제한할 수 있다고 밝혔다. 향후 연구에서는 학습된 분류기를 경량 로컬라이제이션(예: CAM 기반)에 활용하고, 더 광범위한 다기관 검증을 수행할 것을 제안하였다.