멀티스케일 선형시간 인코더 MARBLE로 전체 슬라이드 이미지 분석 혁신

초록

본 논문은 Mamba 기반의 상태공간 모델을 활용해 다중 배율의 병리학 슬라이드 이미지를 동시에 처리하고, 거친‑세밀한 정보를 순차적으로 결합하는 MARBLE이라는 새로운 다중인스턴스 학습 프레임워크를 제안한다. 선형 시간 복잡도로 크로스‑스케일 의존성을 포착하면서 파라미터 오버헤드를 최소화하고, 5개의 공개 데이터셋에서 AUC, 정확도, C‑index 등 주요 지표에서 기존 최첨단 방법들을 크게 능가한다.

상세 분석

MARBLE은 기존 MIL(다중인스턴스 학습) 접근법이 단일 배율에 국한되거나, 트랜스포머 기반 방법이 2차원 어텐션 비용으로 인해 확장성이 떨어지는 문제를 해결한다. 핵심 아이디어는 각 배율(level)을 독립적인 시퀀스(패치 임베딩)로 보고, Mamba‑2 블록을 이용한 선형‑시간 상태공간 모델로 인코딩한다는 점이다. 여기서 “선형‑시간”이란 토큰 수 T_k에 대해 O(T_k·D)의 연산 복잡도를 의미하며, 이는 기존 트랜스포머의 O(T_k²)와 비교해 크게 효율적이다.

다중 배율 처리는 두 단계로 이루어진다. 첫 번째 단계에서는 각 레벨별로 별도의 Mamba‑2 모듈이 병렬로 실행된다. 두 번째 단계에서는 ‘코스‑투‑파인(coarse‑to‑fine) fusion’ 메커니즘을 통해, 더 세밀한 레벨(k>0)의 토큰에 해당하는 상위 레벨(k‑1)의 부모 토큰을 연결한다. 구체적으로, 각 세밀 토큰 x^{(k)}i는 부모 토큰 y^{(k‑1)}{p_k(i)}와 concat 후 선형 프로젝션 ϕ^{(k)}를 거쳐 강화된 토큰 \tilde{x}^{(k)}_i를 만든다. 이 과정은 공간 격자 정렬에 기반해 O(T_k) 시간에 수행되며, 파라미터는 단순 선형 변환 행렬 하나뿐이다. 따라서 크로스‑스케일 의존성을 모델링하면서도 파라미터 증가가 미미하다.

인코딩된 최하위 레벨 토큰 집합 S에 대해서는 어텐션 풀링을 적용한다. 가중치 w를 학습해 각 토큰에 스코어 α(y)=exp(wᵀy)/∑exp(wᵀy′)를 부여하고, 가중합 z=∑α(y)·y 로 슬라이드 수준 임베딩을 만든다. 이 z는 분류 헤드(선형 레이어) 혹은 생존 분석을 위한 Cox 비례위험 헤드에 입력된다.

정규화 측면에서 MARBLE은 두 가지 전략을 도입한다. 첫째, ‘랜덤 코스‑브랜치 드롭’은 학습 중 일정 비율(α=0.1)로 가장 거친 레벨의 토큰과 그 하위 토큰을 무작위로 제거해 모델이 다양한 스케일 조합에 강건하도록 만든다. 둘째, 토큰 순서를 무작위로 섞는 ‘스캔‑오더 중립성’ 기법은 위치 편향을 최소화하고, 부모‑자식 매핑이 명시적이므로 순열 불변성을 보장한다.

실험에서는 PANDA(전립선암)와 TCGA‑NSCLC(폐암) 두 분류 데이터셋, 그리고 KIRP, LUAD, STAD 세 생존 분석 데이터셋을 사용했다. 모든 슬라이드에서 256×256 패치를 1024 차원 임베딩으로 변환한 뒤, 두 배율(예: 10×와 40×)을 동시에 입력했다. 비교 대상은 ABMIL, CLAM, DSMIL, TransMIL, S4‑MIL, DTFD‑MIL, MambaMIL, SRMambaMIL, 2DMambaMIL 등 최신 MIL 및 상태공간 모델이다. MARBLE은 PANDA에서 정확도 0.71(+20.3pp), AUC 0.8878(+6.94pp), TCGA‑NSCLC에서 AUC 0.9730(+0.012) 등 전반적으로 가장 높은 성능을 기록했다. 생존 분석에서도 C‑index 0.8184(KIRP), 0.6432(LUAD), 0.6510(STAD)로 최고점을 찍었다.

Ablation 연구에서는 단일 배율(코스 전용 또는 파인 전용) 대비 두 배율 결합이 모두 우수함을 확인했다. 특히 파인 배율만 사용할 경우 세밀한 특징은 잡히지만 전역적인 조직 구조를 놓치고, 코스 배율만 사용할 경우 전역 맥락은 확보하지만 미세한 병변을 놓치는 경향이 있다. 두 배율을 결합한 MARBLE은 이러한 장점을 보완한다.

전체적으로 MARBLE은 (1) 선형‑시간 상태공간 백본, (2) 경량 크로스‑스케일 결합, (3) 효율적인 정규화 메커니즘을 통해 기존 트랜스포머 기반 MIL의 계산·메모리 병목을 해소하면서도 정확도·예측력을 크게 향상시킨다. 향후 연구에서는 데이터‑드리븐 토큰 순회, 선택적 패치 라우팅, 3개 이상 배율 확장 등을 통해 더욱 효율적이고 일반화 가능한 프레임워크로 발전시킬 여지가 있다.