Source-linked AI summary
Multimodal Optimal Transport-based Co-Attention Transformer with Global Structure Consistency for Survival Prediction
Yingxue Xu, Hao Chen
TL;DR
Survival prediction must integrate histology and genomics while representing gigapixel WSIs and capturing global consistency between TME interactions and genomic co-expression. MOTCat addresses this with optimal-transport co-attention and an efficient micro-batch approximation, and experiments on five benchmark datasets show superiority over state-of-the-art methods.
Problem
Existing multimodal survival methods struggle to represent gigapixel WSIs and capture global structural consistency between histological TME interactions and genomic co-expression.
Method
MOTCat uses optimal-transport co-attention to match WSI patches with gene embeddings globally, with unbalanced mini-batch transport approximating full transport efficiently.
Results
The method shows significant improvement over state-of-the-art methods on five benchmark datasets.
Takeaways & Limitations
Global optimal matching provides a way to select informative WSI patches while modeling histological interactions and genomic co-expression for survival prediction.
Abstract
from arXiv · showhide
Survival prediction is a complicated ordinal regression task that aims to predict the ranking risk of death, which generally benefits from the integration of histology and genomic data. Despite the progress in joint learning from pathology and genomics, existing methods still suffer from challenging issues: 1) Due to the large size of pathological images, it is difficult to effectively represent the gigapixel whole slide images (WSIs). 2) Interactions within tumor microenvironment (TME) in histology are essential for survival analysis. Although current approaches attempt to model these interactions via co-attention between histology and genomic data, they focus on only dense local similarity across modalities, which fails to capture global consistency between potential structures, i.e. TME-related interactions of histology and co-expression of genomic data. To address these challenges, we propose a Multimodal Optimal Transport-based Co-Attention Transformer framework with global structure consistency, in which optimal transport (OT) is applied to match patches of a WSI and genes embeddings for selecting informative patches to represent the gigapixel WSI. More importantly, OT-based co-attention provides a global awareness to effectively capture structural interactions within TME for survival prediction. To overcome high computational complexity of OT, we propose a robust and efficient implementation over micro-batch of WSI patches by approximating the original OT with unbalanced mini-batch OT. Extensive experiments show the superiority of our method on five benchmark datasets compared to the state-of-the-art methods. The code is released.
1. Introduction
Survival prediction integrates histology and genomics but must represent gigapixel WSIs and capture global structural relationships in tumor microenvironment and genomic data. MOTCat uses optimal-transport co-attention and micro-batch approximation to address these challenges, with experiments on five benchmark datasets showing improvement over state-of-the-art methods.
- Survival prediction estimates relative death risk from heterogeneous pathology and genomic data.
- Gigapixel WSIs contain roughly 500,000 × 500,000 pixels, while TME-related patches occupy only a tiny proportion of each slide.
- TME interactions and genomic co-expression form potential structures, but existing co-attention emphasizes dense local similarity and neglects global coherence.
- MOTCat matches histology instances with genomic instances through optimal transport, selecting patches with high global structure consistency to represent WSIs.
- OT-based co-attention models within-modality structure, provides label-free optimal matching flow, and reduces cross-modal heterogeneity compared with conventional co-attention.
- Micro-batch unbalanced optimal transport approximates the full WSI transport solution with a robust and efficient procedure.
- Extensive experiments on five benchmark datasets report significant improvement over state-of-the-art methods.
2. Related Work
Related work represents large WSIs with multiple-instance learning and combines multimodal data through early, late, or intermediate fusion. Histology and genomics are established inputs for survival prediction, while MOTCat instead uses globally consistent optimal matching to identify informative pathology instances.
- Multimodal learning combines complementary clinical modalities but must overcome heterogeneity for effective feature fusion.
- Existing multimodal methods use early, late, or intermediate fusion, with early fusion potentially neglecting intra-modality dynamics.
- Multiple-instance learning treats a WSI as a bag of patches and either selects instances or learns a bag-level representation from embedded instances.
- MOTCat uses optimal match flow between pathological and genomic instances to identify WSI instances with globally consistent potential structure.
- Survival prediction estimates the probability of an event such as death before a specified time under right-censored and uncensored data.
3. Method
MOTCat represents pathology and genomic data as instance bags, uses OT-based co-attention to select globally structure-consistent informative instances, and fuses their aggregated representations for ordinal survival prediction. Micro-batch optimization makes this matching practical for gigapixel WSIs.
- Bag Formulation: Pathology WSIs are modeled as bags of patch-level features, while genomic data are organized as bags of biologically functional gene categories.CNN and category-specific genomic encoders produce instance-level embeddings for the two modalities.
- OT-based Co-Attention: OT-based co-attention matches WSI and genomic instances to identify informative pathology instances associated with tumor-microenvironment structure and genomic co-expression.The matching uses global structure consistency rather than only independent pairwise similarity.
- Micro-Batch Optimization: MOTCat samples a micro-batch from the WSI bag and applies OT-based co-attention to that subset before selecting instances and updating model parameters.The algorithm extracts features, computes transport using Sinkhorn-Knopp scaling, selects instances, aggregates representations, and optimizes the loss.
- Multimodal Survival Prediction: The selected instances are aggregated by Transformer encoders into modality-level representations, which are concatenated to predict the hazard function and ordinal survival risk.The pathology and genomic bag-level embeddings are fused before survival prediction.
- Micro-Batch Optimization: Unbalanced mini-batch OT approximates the original full-WSI transport problem, reducing the computational burden caused by the massive number of WSI patches.The method averages results over sampled subsets as a proxy for transport over all WSI instances.
4. Experiment
Experiments on five TCGA cancer datasets evaluate MOTCat against unimodal and multimodal baselines, then examine its components, micro-batch robustness, computational speed, and patient stratification.
- Datasets: Experiments use five paired TCGA cancer datasets with diagnostic WSIs, genomic data, and ground-truth survival outcomes.The datasets are BLCA, BRCA, UCEC, GBMLGG, and LUAD.
- Ablation Study: Micro-batch fusion improves performance, while OT-based co-attention further improves the averaged results and yields the best overall performance.The ablation compares variants without micro-batching, with micro-batching, and with the full OT-based design.
- Micro-Batch Robustness: MOTCat obtains the best averaged performance across micro-batch sizes, with particularly robust results on UCEC and LUAD.On BLCA and BRCA, the co-attention variant is slightly more robust, but MOTCat achieves substantially higher performance.
- Computational Speed: Training processes 6540 p/s and inference processes 11885 p/s, making OT-based histology-genomics application practicable.The speed is measured on 10 WSIs containing about 150k patches using one NVIDIA GeForce RTX 3090 GPU.
- Statistical Analysis: Kaplan-Meier analysis separates patients into low-risk and high-risk groups, with Logrank tests assessing statistical differences between their survival outcomes.Lower P-values indicate better patient stratification performance.
5. Conclusion
The conclusion presents MOTCat as a multimodal framework for survival prediction that combines OT-based matching with global structural modeling. Its micro-batch approximation is intended to make the approach practical while leaving end-to-end patch-extractor updates for future work.
- Contributions: MOTCat uses OT-based co-attention to match histology and genomic instances and select informative instances related to the tumor microenvironment.This addresses representation of gigapixel WSIs through informative-instance selection.
- Global Structure: Optimal transport provides global awareness for modeling pathological interactions and genomic co-expression.The framework targets potential structure within each modality while matching information across modalities.
- Efficiency: A robust micro-batch implementation approximates OT and supports practical application to WSI bags.The authors identify end-to-end updating of patch-image extractor parameters as future work.
A. Outline
The supplementary materials provide qualitative visualizations of micro-batch-size effects and co-attention behavior.
- Supplementary Materials: The supplementary materials are organized to document these additional visual analyses.
- Supplementary Materials: The supplement visualizes how changing the Micro-Batch Strategy affects results.
- Supplementary Materials: The supplement provides visualization cases of Co-Attention.
B. Effect on Size of Micro-Batch
Qualitative comparisons examine co-attention across micro-batch sizes 128, 256, and 512. UMBOT shows more consistent activation across sizes than the original OT and MCAT-based variants.
- Setup: The qualitative analysis compares micro-batch sizes 128, 256, and 512 for several method variants.It visualizes co-attention between the first 300 WSI instances and six genomic instances from the same patient.
- Qualitative Results: UMBOT-based co-attention shows the best activation consistency across micro-batch sizes, whereas the original OT variant shows considerably poorer consistency.This supports greater robustness of UMBOT to micro-batch size.
- Qualitative Results: Replacing UMBOT with MCAT co-attention produces a substantially different activation pattern at micro-batch size 512.
C. Visualization of Co-Attention
The visualizations compare how OT-based co-attention and dense co-attention distribute attention across histology regions for different genomic functional instances. The proposed method attends to different WSI areas across genomic instances, whereas MCAT concentrates on similar regions.
- Visualization setup: The analysis visualizes co-attention values for all WSI instances in high- and low-risk cases, comparing the proposed method with MCAT.The Tumor Suppression genomic instance serves as a reference for showing differences in co-attention values.
- Visualization findings: The proposed OT-based co-attention focuses on different WSI areas for different genomic functional instances.The visualization includes the corresponding top-4 highest-attention patches for each genomic instance of a unique functional category.
- Visualization findings: MCAT’s dense co-attention focuses on similar histology regions across different genomic functional instances.This comparison contrasts the region selectivity of OT-based co-attention with dense co-attention.
- Interpretation: The authors suggest that these different attention patterns may contribute to the proposed method’s better performance.