MUST: Modality-Specific Representation-Aware Transformer for Diffusion-Enhanced Survival Prediction with Missing Modality

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 정밀 암 치료 (Precision Oncology) 를 위해서는 조직 병리 (WSI, Whole Slide Images) 와 유전체 데이터 (Genomics) 등 다양한 모달리티 (Multimodal) 데이터를 통합한 생존 예측이 필수적입니다.
• 문제점: 임상 현장에서는 비용, 기술적 한계, 또는 후향적 데이터의 부재로 인해 모달리티가 불완전한 (Missing Modality) 경우가 빈번합니다.
• 기존 방법의 한계:
  • 특징 정렬 (Feature Alignment): 서로 다른 모달리티 간 특징을 유사하게 맞추려 하지만, 어떤 정보가 결여되었는지 명시적으로 모델링하지 못합니다.
  • 보간 (Imputation): 생성 모델을 통해 결손된 특징을 채우려 하지만, 고차원 공간에서 노이즈가 많은 가짜 특징 (Pseudo-features) 을 생성하여 성능을 저하시킵니다.
  • 공동 분포 학습: 모달리티 간 상호작용은 학습하나, 각 모달리티가 고유하게 기여하는 정보와 다른 모달리티로부터 유추 가능한 정보를 명확히 분리 (Disentanglement) 하지 못합니다.
• 핵심 과제: 결손된 모달리티가 있을 때, 어떤 정보가 결여되었는지 정확히 식별하고, 이를 확정적 (Deterministic) 으로 복구하거나 확률적 (Probabilistic) 으로 생성하여 생존 예측의 강건성을 유지하는 것.

2. 제안 방법론: MUST (Methodology)

저자들은 MUST (Modality-Specific representation-aware Transformer) 라는 새로운 프레임워크를 제안합니다. 이는 대수적 제약 (Algebraic Constraints) 과 잠재 확산 모델 (Latent Diffusion Model, LDM) 을 결합한 2 단계 전략을 사용합니다.

A. 대수적 분해 아키텍처 (Algebraic Decomposition)

각 모달리티의 전역 표현 (Global Representation) 을 모달리티 고유 성분 (Modality-specific) 과 교차 컨텍스트 공유 성분 (Cross-modal Shared) 으로 명시적으로 분해합니다.

1. 이중 방향 교차 어텐션 (Bidirectional Cross-Attention): 병리 (P) 와 유전체 (G) 데이터가 서로의 정보를 참조하여 컨텍스트화된 성분 ($c_{P \leftarrow G}, c_{G \leftarrow P}$) 을 추출합니다.
2. 저랭크 공유 부분 공간 (Low-Rank Shared Subspace): 학습된 저랭크 투영 행렬 ($P_\cap$) 을 통해 공유 정보를 낮은 차원의 부분 공간으로 매핑합니다.
3. 대수적 제약 (Algebraic Constraints):
   • 공유 일관성: $c_{P \leftarrow G} \approx c_{G \leftarrow P}$ (공유 성분은 모달리티에 관계없이 동일해야 함).
   • 직교성: 모달리티 고유 성분 ($u_P, u_G$) 간, 그리고 공유 성분과 고유 성분 간 직교성을 enforced 하여 정보의 중복을 제거합니다.
   • 분해 식: $g_P = \hat{u}_P + \hat{c}_{G \leftarrow P}$, $g_G = \hat{u}_G + \hat{c}_{P \leftarrow G}$.
   • 핵심 이점: 이 구조를 통해 결손된 모달리티의 공유 성분은 결여된 모달리티가 아닌, 이용 가능한 모달리티로부터 대수적으로 확정적으로 유도 (Deterministic Recovery) 할 수 있습니다.

B. 점진적 학습 전략 (Progressive Training)

1. Stage 1 (Warm-up): 노이즈 주입 (Gaussian Noise Injection) 을 통해 각 모달리티가 생존 예측에 필요한 핵심 신호를 먼저 학습하도록 합니다.
2. Stage 2 (Decomposition): 대수적 분해 손실 (Decomposition Loss), 공유 일관성 손실, 직교성 손실을 추가하여 명시적인 분해 구조를 학습합니다.

C. 결손 모달리티 복원을 위한 잠재 확산 모델 (LDM for Missing Modality)

• 문제: 공유 성분은 대수적으로 복구되지만, 모달리티 고유 성분 ($\hat{u}$) 은 다른 모달리티로부터 유추할 수 없으므로 생성이 필요합니다.
• 해결: 조건부 잠재 확산 모델 (Conditional Latent Diffusion Model) 을 사용합니다.
  • 조건 (Condition): 복구된 공유 성분 ($\hat{c}$) 과 학습된 구조적 사전 지식 (Structural Prior, [CLS] 토큰) 을 입력으로 받습니다.
  • 생성: 결손된 모달리티의 고유 성분 ($\hat{u}_{missing}$) 을 확률적으로 생성합니다.
• 추론 (Inference): 결손된 모달리티가 있을 경우, 대수적 식으로 공유 성분을 추출한 후, LDM 을 통해 고유 성분을 생성하여 최종 생존 점수를 계산합니다.

3. 주요 기여 (Key Contributions)

1. 명시적 분해 아키텍처: 저랭크 공유 부분 공간 내의 대수적 제약과 양방향 교차 어텐션을 통해, 결손 시 어떤 정보가 손실되었는지 정확히 식별하고 공유 성분을 확정적으로 복구합니다.
2. 2 단계 재구성 전략: 공유 정보는 대수적 투영으로 확정적으로, 모달리티 고유 정보는 LDM 을 통해 확률적으로 생성하는 하이브리드 방식을 도입하여 생성의 불확실성을 고유 성분에만 국한시켰습니다.
3. SOTA 성능: 5 가지 TCGA 암 데이터셋에서 완전한 데이터와 결손 데이터 모두에서 기존 최첨단 방법 (State-of-the-Art) 을 능가하는 성능을 입증했습니다.

4. 실험 결과 (Results)

• 데이터셋: TCGA 의 BLCA, BRCA, GBMLGG, LUAD, UCEC 총 5 개 암종 데이터셋.
• 완전 데이터 성능: 전체 C-index 가 0.742로, 기존 최강 모델인 CMTA (0.724) 보다 2.5% 향상되었습니다.
• 결손 데이터 성능:
  • 유전체 결손 시: C-index 0.716 (기존 LD-CVAE 0.697 대비 우월).
  • 병리 결손 시: C-index 0.739 (기존 방법들이 큰 성능 저하를 보인 반면, MUST 는 0.4% 미만의 경미한 감소만 기록).
  • 특히, 병리 이미지가 결손된 경우 오히려 일부 데이터셋에서 성능이 향상되기도 했으며, 이는 LDM 의 노이즈 제거 (Denoising) 효과가 병리 데이터의 고주파 노이즈를 필터링하여 정규화된 특징을 제공했기 때문으로 분석됩니다.
• 환자 층화 (Stratification): Kaplan-Meier 생존 곡선 분석에서 결손 상황에서도 고위험군과 저위험군의 통계적 유의성 (Log-rank test, p < 0.05) 을 유지하여 임상적 신뢰성을 입증했습니다.

5. 의의 및 결론 (Significance)

• 임상 적용 가능성: 실제 임상 환경에서 데이터 불완전성이 흔한 문제를 해결하여, 결손된 데이터가 있더라도 신뢰할 수 있는 생존 예측을 제공할 수 있습니다.
• 방법론적 혁신: 단순한 특징 정렬이나 보간을 넘어, 대수적 가역성 (Algebraic Invertibility) 을 통해 공유 정보를 확정적으로 복구하고, 불확실성이 필요한 부분에만 생성 모델 (Diffusion) 을 적용하는 새로운 패러다임을 제시했습니다.
• 효율성: 추론 시 결손 모달리티가 있을 경우에도 1.8GB 메모리 내에서 약 879ms (50 단계 DDIM 샘플링) 내에 처리가 가능하여 임상적 수용 가능한 지연 시간을 보입니다.

이 논문은 다중 모달리티 생존 예측 분야에서 데이터 결손 문제를 해결하기 위한 이론적으로 엄밀하고 실용적으로 효과적인 프레임워크를 제시했다는 점에서 중요한 의의를 가집니다.