OTSurv: A Novel Multiple Instance Learning Framework for Survival Prediction with Heterogeneity-aware Optimal Transport

2025-06-25

AI Summary

Existing multi-instance learning (MIL)-based whole-slide image (WSI) survival prediction methods struggle to model pathological heterogeneity: global morphological distributions exhibit long-tailed characteristics, while local patch-level predictions suffer from inherent uncertainty. To address this, we propose the first MIL framework integrating optimal transport (OT) theory, formulating survival prediction as a *non-equilibrium* optimal transport problem. Our method jointly enforces a global long-tail distribution constraint and a local uncertainty-aware constraint to enable heterogeneity-aware feature alignment and aggregation. We adopt a hardware-efficient matrix scaling algorithm, balancing computational efficiency with interpretability. Evaluated on six mainstream benchmarks, our approach achieves an average 3.6% improvement in concordance index (C-index) and statistically significant survival separation (log-rank test, *p* < 0.01), consistently outperforming state-of-the-art methods.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Distributed Search
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs

Abstract

Survival prediction using whole slide images (WSIs) can be formulated as a multiple instance learning (MIL) problem. However, existing MIL methods often fail to explicitly capture pathological heterogeneity within WSIs, both globally -- through long-tailed morphological distributions, and locally through -- tile-level prediction uncertainty. Optimal transport (OT) provides a principled way of modeling such heterogeneity by incorporating marginal distribution constraints. Building on this insight, we propose OTSurv, a novel MIL framework from an optimal transport perspective. Specifically, OTSurv formulates survival predictions as a heterogeneity-aware OT problem with two constraints: (1) global long-tail constraint that models prior morphological distributions to avert both mode collapse and excessive uniformity by regulating transport mass allocation, and (2) local uncertainty-aware constraint that prioritizes high-confidence patches while suppressing noise by progressively raising the total transport mass. We then recast the initial OT problem, augmented by these constraints, into an unbalanced OT formulation that can be solved with an efficient, hardware-friendly matrix scaling algorithm. Empirically, OTSurv sets new state-of-the-art results across six popular benchmarks, achieving an absolute 3.6% improvement in average C-index. In addition, OTSurv achieves statistical significance in log-rank tests and offers high interpretability, making it a powerful tool for survival prediction in digital pathology. Our codes are available at https://github.com/Y-Research-SBU/OTSurv.

Problem

• Survival prediction using whole slide images with MIL
• Capturing pathological heterogeneity in WSIs globally and locally
• Modeling heterogeneity via optimal transport with constraints

Innovation

• Heterogeneity-aware optimal transport framework
• Global long-tail constraint modeling
• Local uncertainty-aware transport constraint