MMSF: Multitask and Multimodal Supervised Framework for WSI Classification and Survival Analysis

기술 요약: MMSF – WSI 분류 및 생존 분석을 위한 멀티태스크 및 멀티모달 지도 학습 프레임워크

문제 정의
계산 병리학은 기가픽셀 규모의 이미지로서 풍부한 형태학적 정보를 포함하고 있지만 막대한 계산 부담을 주는 전체 슬라이드 이미지(WSI)를 분석하는 데 있어 상당한 어려움에 직면해 있습니다. 전통적인 파이프라인은 WSI를 수십만 개의 패치로 분할하곤 하는데, 이는 임상 배포를 저해하는 방대한 계산 요구량을 초과하게 만듭니다. 또한, 기존 방법론들은 다음과 같은 세 가지 주요 한계점을 가집니다:

1. 계산 복잡도: 트랜스포머 기반의 다중 인스턴스 학습(MIL) 모델은 시퀀스 길이에 대해 이차 복잡도(quadratic complexity)를 보이기 때문에, 대규모 패치 처리에 있어 비효요율적입니다.
2. 공간적 맥락 무시: 많은 접근 방식이 패치를 독립적인 인스턴스로 취급하여, 조직 병리학에 내재된 결정적인 공간적 관계와 조직 구조를 포착하지 못합니다.
3. 단일 모달리티의 한계: 대부분의 방법은 병리학적 이미지 데이터에만 의존하며, 보완적인 예후 맥락을 제공하는 일상적으로 수집되는 임상 기술자(예: 연령, 병기)를 간과합니다. 멀티오믹스 통합이 정확도를 높이기는 하지만, 이는 비용과 데이터 가용성 문제로 인해 제한적입니다.

방법론
저자들은 WSI 분류와 생존 분석을 동시에 수행하도록 설계된 선형 복잡도 프레임워크인 MMSF(Multitask and Multimodal Supervised Framework)를 제안합니다. 이 아키텍처는 구조화된 파이프라인을 통해 멀티모달 데이터를 통합합니다:

1. 특징 추출 및 그래프 구축:

   • WSI를 패치로 크롭한 후, 병리학적 파운데이션 모델인 UNI2를 사용하여 1536차원의 특징을 추출합니다.
   • 패치 수준 그래프 구축 모듈은 정점(vertex)이 패치를 나타내는 공간 그래프 $G=(V, E)$를 생성합니다. 에지(edge)는 공간적 근접성과 조직 유사성(특징 벡터의 코사인 유사도)을 기반으로 설정됩니다.
   • 그래프 신경망(특히 절제 연구에서 최적임이 확인된 GAT)이 이러한 그래프를 처리하여 공간 맥락 임베딩을 추출합니다.

2. 임상 데이터 임베딩 (CDE):

   • 전용 임상 데이터 임베딩(Clinical Data Embedding) 모듈은 이질적인 환자 수준 데이터(연령과 같은 연속형 변수 및 T/N 병기와 같은 범주형 변수)를 처리합니다.
   • 각 데이터 유형에 대해 별도의 인코더-디코더 쌍을 사용하여 자기 지도 재구성 목적 함수를 통해 통일된 임상 임베딩 벡터를 생성합니다.

3. 계층적 융합 메커니즘:

   • 조기 융합 (패치 수준): 추출된 패치 특징과 그래프 임베딩은 **특징 융합 모듈(FFM)**을 사용하여 융합됩니다. 이 모듈은 연결된 특징에 대한 Squeeze-and-Excitation (SE) 어텐션 메커니즘을 활용하여 파이프라인 초기 단계에서 공간적 맥락을 통합합니다.
   • 후기 융합 (인스턴스 수준): MIL 인코더로부터 집계된 백(bag) 수준의 표현은 두 번째 FFM을 통해 임상 임베딩과 융합됩니다. 이는 환자 수준의 맥락을 포함하면서 데이터의 계층적 특성을 보존합니다.

4. Mamba 기반 MIL 인코더:

   • 계산상의 병목 현상을 해결하기 위해, MMSF는 이차 자가 주의 집중(quadratic self-attention)을 Mamba 상태 공간 모델(EfficientMIL에서 유래)로 대체합니다.
   • **적응형 패치 선택기(Adaptive Patch Selector, APS)**가 가장 정보량이 많은 패치(상위 512개)를 식별하며, 이후 Mamba2 블록에 의해 처리됩니다. 이는 복잡도를 패치 수에 대해 이차에서 선형으로 감소시킵니다.

5. 멀티태스크 예측:

   • 최종 융합된 특징은 작업별 헤드로 전달됩니다: 분류를 위한 선형 레이어(Binary Cross-Entropy with Logits 사용)와 생존 분석을 위한 시그모이드 활성화 함수를 포함한 선형 레이어(Cox proportional hazards loss 사용).

주요 기여

• 최초의 선형 복잡도 멀티태스크 프레임워크: MMSF는 WSI 분석을 위해 선형 복잡도 상태 공간 모델(Mamba)을 사용하는 멀티태스크 학습(분류 및 생존)과 멀티모달 융합을 결합한 최초의 프레임워크입니다.
• 새로운 임상 임베딩: 멀티오믹스 데이터 없이도 다운스트림 작업을 위해 이질적인 임상 속성을 표준화하는 플러그 앤 플레이 방식의 임상 데이터 임베딩(CDE) 모듈을 도입했습니다.
• 계층적 융합 전략: 그래프 기반 공간 특징의 조기 융합과 임상 임베딩의 후기 융합을 수행하는 이중 단계 융합 메커니즘을 통해 로컬 공간 맥락과 글로벌 환자 속성을 효과적으로 결합합니다.
• 플러그 앤 플레이 모듈: 멀티모달 표현 학습을 향상시키는 패치 수준 그래프 특징 추출 및 임상 데이터 통합을 위한 특정 모듈을 설계했습니다.

실험 결과
프레임워크는 분류를 위한 CAMELYON16 및 TCGA-NSCLC, 그리고 생존 분석을 위한 5개의 TCGA 코호트(BLCA, COAD, LUAD, STAD, KIRC) 등 7개의 공개 데이터셋에서 평가되었습니다.

• 분류 성능: CAMELYON16 및 TCGA-NSCLC에서 MMSF는 경쟁력 있는 베이스라인(TransMIL, Snuffy, EfficientMIL 포함) 대비 정확도를 2.1–6.6%, AUC를 2.2–6.9% 개선하며 최첨단(SOTA) 결과를 달ет했습니다.
• 생존 분석 성능: 5개의 TCGA 생존 코호트에서 MMSF는 전체 C-index 0.6911을 달성했습니다. 이는 단일 모달리티 방법 대비 7.1–9.8%, 기존 멀티모달 대안(예: HSFSurv, MCAT) 대비 5.6–7.1% 향상된 수치입니다.
• 절제 연구: 실험을 통해 그래프 브랜치, 임상 임베딩 모듈, 그리고 SE 기반 융합 전략을 추가하는 것이 각각 성능 향상에 유의미하게 기여함을 확인했습니다. SE 기반 융합 전략은 선형 융합 및 융합 없음 베이스라인보다 우수한 성능을 보였습니다.
• 효율성: Mamba 백본과 적응형 패치 선택을 활용함으로써, 모델은 선형 복잡도를 유지하여 트랜스포머 기반 MIL 접근 방식보다 계산 효율성이 높습니다.

의의 및 주장
저자들은 MMSF가 계산 병리학의 임상 적용을 가로막는 핵심 장벽들을 해결한다고 주장합니다. 값비싼 멀티오믹스 데이터 없이도 최첨단 성능을 달 Die achieve 함으로써, 이 프레임워크는 일상적으로 사용 가능한 WSI와 임상 기록을 활용합니다. 선형 복잡도 설계는 대규모 임상 배포를 위한 확장성을 보장하며, 계층적 융합 메커니즘은 이질적인 데이터 유형의 효과적인 통합을 가능하게 합니다. 패치 수준 어텐션 시각화를 통한 해석 가능성과 결합된 이 모델의 동시 분류 및 생존 분석 능력은 멀티모달 계산 병리학을 위한 강력하고 일반화 가능한 솔루션으로서의 입지를 확고히 합니다.