Re-mixing Embeddings for Patient Augmentation in Data Scarce Multiple Instance Learning

기술 요약: 데이터가 부족한 다중 인스턴스 학습(MIL)에서 환자 증강을 위한 임베딩 재혼합(Re-mixing)

문제 정의
데이터 부족은 희귀 질환, 특수하고 비용이 많이 드는 양상(modality), 또는 침습적 데이터 수집이 제한적인 시나리오를 포함한 의료 분야의 다중 인스턴스 학습(Multiple Instance Learning, MIL)에서 핵심적인 병목 현상으로 남아 있습니다. MIL에서 "백(bag)"(예: 환자)은 여러 개의 인스턴스(예: 이미지 패치, 단일 세포, 또는 유전체 특징)로 구성되며, 모델은 인스턴스 수준의 주석 없이 백 수준의 레이블을 예측합니다. 전통적인 증강 기법(회전, 뒤집기 등)은 인스턴스 수준에서 작동하며 백 수준에서의 현실적인 변이를 생성하지 못합니다. 또한, 기존의 백 수준 혼합 전략(예: PseMix, ReMix)은 다음과 같은 세 가지 주요 한계점을 가집니다:

1. 이들은 오직 온라인(online) 방식으로만 작동하여, 검증하거나 재사용할 수 없는 일시적인 합성 백을 생성합니다.
2. 이들의 국소적 혼합 범위는 **전역적 인구 통계(global population statistics)**를 포착하지 못합니다.
3. 이들은 모든 클래스로부터 실제 사례를 필요로 하므로, 특정 클래스(예: 건강한 대조군)가 데이터셋에서 누락된 경우에는 효과적으로 작동할 수 없습니다.

방법론: RECIPE
저자들은 임베딩 공간에서 새로운 환자 백을 생성하기 위한 통계적으로 근거가 있는 오프라인 프레임워크인 RECIPE(Re-mixing Embeddings via Clustering for Patient augmEntation)를 소개합니다. 이 파이프라인은 두 가지 주요 단계로 구성됩니다:

1. 클러스터 기반 비지도 샘플링 (Recipe Learning):

   • 임베딩 풀링 (Embedding Pooling): 모든 훈련 환자의 인스턴스 임베딩을 하나로 모읍니다.
   • 가우시안 혼합 모델 (Gaussian Mixture Modeling, GMM): 풀링된 임베딩에 GMM을 적합시켜 비지도 클러스터를 정의합니다.
   • 레시피 도출 (Recipe Derivation): 각 질병 클래스에 대해, 학습된 클러스터들에 걸쳐 인스턴스가 분포하는 통계적 분포인 "레시피"를 계산합니다. 이는 각 클러스터 $k$에 할당된 인스턴스의 수에 대한 평균($\mu_{c,k}$)과 표준편차($\sigma_{c,k}$)로 정의됩니다.
   • 환자 생성 (Patient Generation): 학습된 레시피를 기반으로 각 클러스터의 인스턴스 수($n_{c,k} \sim \mathcal{N}(\mu_{c,k}, \sigma^2_{c,k})$)를 샘플링한 후, 해당 클러스터에 할당된 실제 인스턴스 풀에서 그만큼의 임베딩을 추출하여 새로운 환자를 생성합니다. 이를 통해 외부 데이터셋으로부터 레시피를 도출할 수 있다면, 특정 클래스의 실제 사례가 존재하지 않더라도 환자를 생성할 수 있습니다.

2. 불확실성 인지 환자 선택 (Uncertainty-Aware Patient Selection):

   • 생성된 데이터가 모델 성능을 향상시키고 노이즈를 유발하지 않도록 하기 위해, 이 방법은 불확실성 정량화를 사용합니다.
   • 사전 훈련된 MIL 모델은 추론 시 몬테카를로(MC) 드롭아웃을 사용하여 생성된 백을 평가합니다.
   • 불확실성을 정량화하기 위해 **예측 엔트로피(Predictive Entropy)**를 계산합니다 ($U(x) = H[\bar{p}(x)]$).
   • 가장 불확실한 생성 환자들(높은 엔트로피를 가진 샘플들)만을 선정하여 훈련 세트를 증강합니다. 이는 이 샘플들이 가장 정보가 풍부한 학습 신호를 제공한다는 가설에 기반합니다.

주요 기여

• 오프라인, 통계적 근거 기반 생성: 온라인 혼합 방식과 달리, RECIPE는 전역적인 분포적 "레시피"를 학습하여 지속 가능하고 분석 가능한 합성 코호트를 생성합니다.
• 누락된 클래스 처리: RECIPE는 민감한 환자 데이터를 공유하지 않고도 외부 데이터셋으로부터 통계적 레시피를 전이함으로써, 완전히 누락된 환자 클래스(예: 건강한 대조군)를 생성할 수 있는 최초의 MIL 프레임워크입니다.
• 교차 양상 적용성 (Cross-Modality Applicability): 이 방법은 조직 병리 이미지, 단일 세포 RNA 시퀀싱(scRNA-seq), 유세포 분석(flow cytometry)을 포함한 다양한 데이터 양상에서 효과적임을 입증했습니다.
• 아키텍처 불가지론 (Architecture Agnosticism): 이 프레임워크는 ABMIL, DSMIL, Transformer와 같은 다양한 MIL 어그리게이터(aggregator)와 호환됩니다.

실험 결과
저자들은 세 가지 임상적으로 유의미한 데이터 희소성 시나리오에서 RECIPE를 평가했습니다:

1. 교차 데이터셋 전이 (누락된 클래스):

   • 시나리오: AML-Hehr 데이터셋에서 모든 건강한 대조군을 제거했습니다. 외부 cAItomorph 데이터셋으로부터 "건강한 레시피"를 도출하여 AML-Hehr 임베딩 풀에 적용했습니다.
   • 결과: 생성된 건강한 환자들은 실제 세포 유형 구성과 밀접하게 일치했습니다 (Jensen-Shannon Divergence = 0.12). 이 생성된 환자들로 훈련했을 때, 전체 데이터셋 훈련과 유사한 성능(균형 정확도 0.78 vs 0.83)을 달성했으며, 외부 데이터의 직접 전이(0.66)보다 유의미하게 높은 성능을 보였습니다.

2. 저데이터 환경 (Low-Data Regimes):

   • 시나리오: AML-Hehr의 건강한 대조군 수를 단 $n=1$개까지 줄였습니다.
   • 결과: 단 하나의 실제 건강한 대조군만 있을 때, 베이스라인은 0.26의 균형 정확도를 기록했습니다. RECIPE는 0.68을 달성하여 MixUp(0.60) 및 PseMix(0.64)를 실질적으로 앞질렀습니다.

3. 비이미지 양상 (scRNA-seq 및 유세포 분석):

   • 시나리오: 자연적으로 샘플 크기가 제한된 COVID-19 중증도 데이터셋에 대한 증강을 수행했습니다.
   • 결과: RECIPE는 scRNA-seq 데이터에서 PCA 임베딩 기준 +14%, scVI 임베딩 기준 +26%의 균형 정확도 향상을 보였으며, 유세포 분석 데이터에서는 +22%의 향상을 보이며 모든 베이스라인 증강 전략을 능가했습니다.

의의 및 주장
본 논문은 RECIPE가 풀링된 임베딩을 학습된 통계적 분포에 따라 재혼합함으로써 의료 MIL의 데이터 희소성 문제를 해결하는 새로운 솔루션을 제공한다고 주장합니다. 이 방법의 주요 의의는 다음과 같습니다:

• "누락된 클래스" 문제를 해결하여, 환자의 개인정보를 침해하지 않으면서도 외부 통계치를 사용하여 건강한 대조군이나 희귀 질환 표현형을 합성할 수 있게 합니다.
• 전통적인 증강 기법이 실패하는 극단적인 저데이터 환경에서 MIL 성능을 개선합니다.
• 단일 세포 및 유세포 데이터 외에도 적용 가능한 일반화된 프레임워크를 제공합니다.

저자들은 이들의 접근 방식이 국소적인 온라인 혼합을 넘어 전역적인 인구 통계를 모델링함으로써, 희귀 질환 연구와 개인정보 보호 중심의 데이터 증강 사이의 간극을 메울 수 있는 검증 가능하고 재사용 가능한 합성 데이터를 생성할 수 있다고 강조합니다.