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

1. Problem Statement

The paper addresses a critical gap in precision oncology: multimodal survival prediction under missing data conditions.

• Context: Accurate patient survival prediction typically requires integrating diverse data sources, specifically Whole Slide Images (WSI/pathology) and genomic profiles.
• Challenge: In clinical reality, modalities are frequently incomplete due to cost, technical limitations, or retrospective data availability (e.g., having pathology slides but no genomic sequencing, or vice versa).
• Limitations of Existing Methods:
  • Feature Alignment: Enforces output similarity but fails to explicitly model what information is missing.
  • Imputation: Generates noisy pseudo-features in high-dimensional spaces, degrading downstream performance.
  • Joint Distribution Learning: Captures interactions but fails to disentangle modality-specific information from shared information, making deterministic recovery of missing components impossible.
• Core Issue: Existing models lack explicit modeling of the unique contributions of each modality versus the information derivable from others, preventing robust reconstruction when a modality is absent.

2. Methodology: MUST Framework

The authors propose MUST (Modality-Specific representation-aware Transformer), a framework that explicitly decomposes multimodal representations into modality-specific and cross-modal contextualized components using algebraic constraints and latent diffusion.

A. Architecture Overview

1. Feature Extraction:
   • Pathology: Uses a patch encoder (UNI2) and TransMIL to extract patch-level features, aggregating them into a global representation ($g_P$) via self-attention.
   • Genomics: Organizes genes into functional groups (e.g., tumor suppression, oncogenesis) and processes them via group-specific MLPs into genomic tokens, aggregated into a global representation ($g_G$).
2. Bidirectional Cross-Attention:
   • Pathology attends to Genomics ($c_{P \leftarrow G}$) and vice versa ($c_{G \leftarrow P}$) to capture contextualized interactions.
3. Algebraic Decomposition (The Core Innovation):
   • The framework projects representations into a learned low-rank shared subspace ($P_\cap$).
   • It decomposes each modality's global representation into:
     • Shared Component ($\hat{c}$): Information common to both modalities, deterministically derivable from any available modality.
     • Modality-Specific Component ($\hat{u}$): Unique information that cannot be inferred from the other modality.
   • Constraints:
     • Shared Consistency: $\hat{c}_{P \leftarrow G} = \hat{c}_{G \leftarrow P}$.
     • Orthogonality: Ensures specific components are orthogonal to shared components and to each other ($\hat{u}_P \perp \hat{u}_G$).
   • Recovery Logic: If Genomics is missing, the shared component is recovered deterministically via $\tilde{c} = g_P - \hat{u}_P$.

B. Two-Stage Training Strategy

1. Stage 1 (Warm-up): Learns task-relevant global representations from each modality independently with Gaussian noise injection to prevent overfitting and ensure robust feature learning.
2. Stage 2 (Decomposition): Introduces the algebraic constraints (decomposition loss, shared consistency, and orthogonality) to enforce the separation of shared and specific components.

C. Missing Modality Reconstruction via Latent Diffusion

For information that is truly modality-specific and cannot be algebraically recovered (the $\hat{u}$ component):

• Conditional Latent Diffusion Models (LDM): Trained to generate the missing modality-specific component.
• Conditioning: The generation is conditioned on:
  1. The recovered shared component ($\tilde{c}$) from the available modality.
  2. A structural prior (learned class token $[CLS_{u}]$) that encodes the structural distribution of the missing modality.
• Inference: Uses DDIM sampling to generate the missing $\hat{u}$, which is then combined with the recovered shared component to predict survival risk.

3. Key Contributions

1. Explicit Algebraic Decomposition: A novel architecture that uses bidirectional cross-attention and low-rank projection to mathematically separate shared and specific information, enabling deterministic recovery of the shared component without generative uncertainty.
2. Hybrid Reconstruction Strategy: Combines deterministic algebraic projection (for shared info) with probabilistic latent diffusion (for specific info). This isolates stochastic generation to only the truly missing unique residuals, improving stability and quality.
3. State-of-the-Art Performance: Demonstrates superior performance in both complete-data scenarios and robust predictions under missing modality conditions across five diverse cancer datasets.

4. Experimental Results

Experiments were conducted on five TCGA cancer datasets (BLCA, BRCA, GBMLGG, LUAD, UCEC).

• Complete Data Performance:
  • MUST achieved an overall C-index of 0.742, outperforming the previous best (CMTA, 0.724) by 2.5%.
  • The explicit decomposition acts as a structural regularizer, preventing modality collapse and redundant feature encoding.
• Missing Modality Performance:
  • Missing Genomics: MUST achieved 0.716 (only 3.5% degradation from complete data), significantly outperforming baselines like LD-CVAE (0.697) and SMIL (0.624).
  • Missing Pathology: MUST achieved 0.739 (only 0.4% degradation), maintaining bidirectional symmetry. Notably, it outperformed baselines that could not handle this scenario (e.g., LD-CVAE failed entirely).
• Patient Stratification: Kaplan-Meier curves showed that MUST maintains statistically significant separation between high-risk and low-risk groups even with missing modalities (p-values < 0.001), whereas baselines often lost significance.

5. Significance and Impact

• Clinical Applicability: MUST addresses the "incompleteness" problem inherent in real-world clinical data, allowing for robust survival predictions even when critical data (like expensive genomic sequencing) is unavailable.
• Theoretical Advancement: The paper introduces algebraic invertibility in multimodal learning. By proving that shared components are deterministically derivable, it moves beyond implicit alignment to explicit, mathematically grounded recovery.
• Efficiency: The method maintains clinically acceptable inference latency (<1 second per patient even with diffusion sampling), making it viable for treatment planning workflows.
• Future Direction: The framework suggests a pathway for handling training-time missingness (unpaired data), as preliminary experiments showed adaptability even with 50% unpaired data.

In summary, MUST provides a rigorous, mathematically grounded solution to missing multimodal data in oncology, leveraging the strengths of algebraic constraints for shared information and diffusion models for specific information to achieve robust, state-of-the-art survival prediction.