Together, Then Apart: Balancing Alignment and Distinctiveness for Multimodal Survival Analysis¶
Conference: ECCV 2026
Paper: ECCV 2026 Poster
Code: https://github.com/Y-Research-SBU/TTA
Area: Computational Biology
Keywords: Multimodal Survival Analysis, Computational Pathology, Unbalanced Optimal Transport, Prototype Alignment, Contrastive Learning
TL;DR¶
Addressing the issue of over-alignment collapse where aggressive cross-modal fusion suppresses modality-specific evidence in multimodal survival analysis, this paper proposes the Together-Then-Apart (TTA) framework: it first discovers shared prognostic patterns via shared prototypes and semi-relaxed unbalanced optimal transport with a curriculum mass schedule, and then preserves modality-specific features using learnable anchors and contrastive regularization, consistently outperforming recent SOTA models across five TCGA cohorts.
Background & Motivation¶
Multimodal survival analysis aims to improve cancer prognosis by integrating heterogeneous biomedical data sources, primarily whole-slide histopathology images (WSIs) and genomic profiles such as transcriptomic expressions. WSIs capture high-resolution tissue morphology, cellular heterogeneity, and tumor microenvironment spatial architectures, conventionally aggregated through Multiple Instance Learning (MIL). Concurrently, transcriptomic measurements grouped into biological pathways characterize molecular programs and mutation-driven alterations invisible in tissue histology. Consequently, fusing these two complementary patient views has emerged as a cornerstone in computational oncology.
However, existing multimodal models largely rely on attention-based fusion or optimal transport to enforce strong alignment across modalities in a shared latent space. This common practice introduces a fundamental failure mode: over-alignment collapse. When heterogeneous modalities with fundamentally disparate biological properties are prematurely forced to agree everywhere, fine-grained morphological structures in WSIs and subtle pathway variations in genomics are drastically diluted. Empirically, this leads to a counter-intuitive phenomenon where multimodal methods underperform strong unimodal WSI baselines (notably in colorectal cancer where morphological markers dominate prognosis). While existing disentanglement and dynamic reweighting techniques attempt to alleviate redundancy, they treat alignment and distinctiveness independently rather than unifying them under a coherent representation principle.
The guiding insight of this work is intuitive yet powerful: an effective multimodal model should first uncover what modalities share, and subsequently protect what makes them distinct. Core idea: introduce the Together-Then-Apart (TTA) framework, which aligns multimodal tokens onto shared semantic prototypes via semi-relaxed unbalanced optimal transport (UOT) with a curriculum mass schedule (Together), and then explicitly enforces modality separation via learnable anchors and an InfoNCE-style contrastive regularizer (Apart), preserving modality-specific prognostic evidence while harvesting cross-modal synergies.
Method¶
Overall Architecture¶
The input to TTA consists of two modality-specific token sets for each patient: pathology patch tokens \(X_n^p \in \mathbb{R}^{N_n^p \times D}\) extracted by a pretrained pathology foundation model (UNI) and linear projection, and pathway tokens \(X_n^g \in \mathbb{R}^{N_n^g \times D}\) derived from gene expression vectors filtered by 50 Hallmark biological pathways. The pipeline progresses through four sequential stages: Tokenization, Together (joint prototype alignment under UOT), Apart (modality anchor refinement and contrastive separation), and Prediction (prototype-level co-attention and hazard regression).
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["Input Data<br/>WSI patches & Gene expression"] --> B["Modality Tokenization<br/>Patch embeddings & Pathway pooling"]
B --> C["Shared Prototype Semi-Relaxed UOT<br/>Concatenated cost matrix & Dummy sink"]
C --> D["Curriculum Mass Transport<br/>Smooth dynamic mass schedule rho(t)"]
D --> E["Modality Anchor Refinement & Contrastive Regularization<br/>Modality anchors & InfoNCE separation"]
E --> F["Prototype-Level Co-Attention & Risk Prediction<br/>Cross-modal fusion to log-risk"]
Key Designs¶
1. Shared Prototype Semi-Relaxed UOT: Cross-Modal Alignment under Unified Geometry
Conventional multimodal approaches either maintain distinct prototype banks for each modality or rely on unconstrained softmax similarities that are easily destabilized by long-tailed WSI bag sizes and extreme genomic sparsity. TTA establishes \(K\) learnable prototypes \(P \in \mathbb{R}^{K \times D'}\) in a unified latent coordinate system. Pathology and genomic tokens are concatenated into a joint set of \(N_{\mathrm{tot}} = N_n^p + N_n^g\) instances. The pairwise transport cost matrix \(C_n = [C_n^p; C_n^g] \in \mathbb{R}^{N_{\mathrm{tot}} \times K}\) is formed from negative log-softmax similarities. To permit patient-specific variations while preventing mass collapse onto dominant clusters, assignments are solved via semi-relaxed unbalanced optimal transport:
where the source marginal \(a = \frac{1}{N_{\mathrm{tot}}}\mathbf{1}_{N_{\mathrm{tot}}}\) guarantees uniform token contribution, whereas the prototype-side marginal penalizes deviations from a uniform prior \(b = \frac{1}{K}\mathbf{1}_K\) governed by relaxation coefficient \(\gamma > 0\).
2. Curriculum Mass Transport: Taming Early Spurious Correspondences
In the early epochs of representation learning, noisy or weakly informative tokens can trigger misleading cross-modal associations. TTA regulates the total mass delivered to true prototypes by introducing a fraction parameter \(\rho(t) \in [0, 1]\), routing the residual \(1-\rho(t)\) mass to an added zero-cost dummy sink:
with \(\rho_{\mathrm{base}}=0.1\) and \(\rho_{\mathrm{upper}}=1.0\) over horizon \(T\). Early in training, a conservative \(\rho\) diverts ambiguous evidence into the sink; as representations stabilize, \(\rho\) smoothly approaches 1.0 to enforce decisive matching. The resulting row-scaled plan \(Q^\star \in \mathbb{R}^{N_{\mathrm{tot}} \times K}\) is blended with softmax weights via coefficient \(\beta\), aggregating tokens into modality prototype representations \(H_n^p, H_n^g \in \mathbb{R}^{K \times D}\).
3. Modality Anchor Refinement & Contrastive Regularization: Preventing Over-Alignment Collapse
Mapping both modalities into the shared prototype bank risks blurring modality identities. In the Apart stage, TTA explicitly introduces learnable modality anchors \(a^p, a^g \in \mathbb{R}^{D'}\). Prototype tokens are projected, concatenated with their respective anchor, and refined through a lightweight self-attention layer \(\mathcal{R}_\theta\) to produce \(\hat{H}_n^p, \hat{H}_n^g \in \mathbb{R}^{K \times D'}\). The mean tokens \(\bar{h}_n^m = \frac{1}{K}\sum_{k=1}^K \hat{h}_{n,k}^m\) are then constrained via an InfoNCE contrastive objective:
where \(s_+^m = \langle \phi(\bar{h}_n^m), \phi(a^m) \rangle / \tau_r\) attracts a modality representation toward its own anchor, while \(s_-^m = \langle \phi(\bar{h}_n^m), \phi(a^{\bar{m}}) \rangle / \tau_r\) repels it from the opposite anchor. This pulls identical-modality representations into a coherent subspace while actively repelling cross-modal collapse, safeguarding vital modality-specific morphological or genomic traits.
Loss & Training¶
The overall architecture is trained end-to-end under a multi-task objective: 1. Cox Partial Likelihood \(\mathcal{L}_{\mathrm{surv}}\): Standard ranking objective comparing relative patient risk over uncensored survival times; 2. Instance-Level UOT Supervision \(\mathcal{L}_{\mathrm{instance}}\): Cross-entropy aligning token prediction distributions with normalized UOT assignment pseudo-labels; 3. Contrastive Regularization \(\mathcal{L}_{\mathrm{contrast}}\): Preserving modality distinctiveness.
The total optimization loss is:
Default hyperparameter settings: \(K=32\) shared prototypes, \(\gamma=0.1\), \(\lambda_{\mathrm{contrast}}=\lambda_{\mathrm{inst}}=0.5\), \(\lambda_{\mathrm{wsi}}=\lambda_{\mathrm{gen}}=1.0\). Slide feature representations are pre-extracted using DINOv2-based UNI (ViT-Large), genomic features span 50 Hallmark pathways (4,241 genes), and performance is evaluated via 5-fold site-stratified cross-validation.
Key Experimental Results¶
Main Results¶
On five TCGA cancer cohorts (BRCA, BLCA, STAD, CRC, KIRC), TTA is evaluated against two genomic baselines, seven histopathology MIL baselines, and eight recent multimodal survival models using C-Index (mean ยฑ std):
| Dataset | Modality | TTA (Ours) | Runner-up SOTA (Model) | Best WSI-Only (Model) | Gain vs. Runner-up |
|---|---|---|---|---|---|
| BRCA (n=868) | g.+h. | 0.726 ยฑ 0.039 | 0.724 ยฑ 0.067 (MMP) | 0.643 ยฑ 0.128 (PANTHER) | +0.2% |
| BLCA (n=359) | g.+h. | 0.662 ยฑ 0.079 | 0.676 ยฑ 0.035 (LD-CVAE) | 0.637 ยฑ 0.065 (OTSurv) | -1.4% (Rank 2) |
| STAD (n=318) | g.+h. | 0.613 ยฑ 0.079 | 0.594 ยฑ 0.066 (MMP) | 0.566 ยฑ 0.033 (ABMIL) | +1.9% |
| CRC (n=296) | g.+h. | 0.685 ยฑ 0.131 | 0.639 ยฑ 0.101 (MRePath) | 0.723 ยฑ 0.111 (AttnMISL) | +4.6% (Multimodal Rank 1) |
| KIRC (n=340) | g.+h. | 0.778 ยฑ 0.117 | 0.751 ยฑ 0.139 (LD-CVAE) | 0.739 ยฑ 0.149 (OTSurv) | +2.7% |
| Overall (Avg) | g.+h. | 0.693 | 0.667 (MMP) | 0.644 (OTSurv) | +2.6% |
Ablation Study¶
The ablation investigations across the Together and Apart stages (Tables 2, 3, and 4) demonstrate the necessity of each design component:
| Config | Stage Configuration | Overall C-Index | BRCA | CRC | Note |
|---|---|---|---|---|---|
| Baseline (None) | Neither Together nor Apart | 0.643 | 0.682 | 0.561 | Modality-specific prototypes without UOT or anchors |
| Together Only | Enable Together stage only | 0.662 | 0.675 | 0.587 | +1.9% over baseline, but CRC suffers from over-alignment |
| Apart Only | Enable Apart stage only | 0.666 | 0.683 | 0.625 | +2.3% over baseline, protects modality individuality |
| Full TTA | Together + Apart (Default) | 0.693 | 0.726 | 0.685 | +5.0% over baseline; validates synergy of both stages |
| Separate Prototypes | Separate prototypes + separate UOT | 0.681 | 0.755 | 0.609 | Lacks shared geometric alignment space |
| Shared + Sep. UOT | Shared prototypes + separate UOT | 0.684 | 0.713 | 0.672 | +0.3% over separate prototypes |
| Shared + Joint UOT | Concatenated tokens + joint UOT | 0.693 | 0.726 | 0.685 | Full design; achieves highest cross-cohort average |
Key Findings¶
- Resolving the CRC Pathology Anomaly: In colorectal cancer (CRC), where morphological signals dominate prognosis, prior multimodal methods collapsed drastically (e.g., MCAT 0.579 vs. AttnMISL 0.723). Ablation reveals that dropping the Apart stage causes a -9.8% drop on CRC (0.685 to 0.587). The contrastive anchors in TTA effectively insulate histology features, securing 0.685 and outperforming all prior multimodal baselines.
- Impact of Curriculum Mass Scheduling: Comparing token assignment approaches shows that UOT without curriculum mass yields 0.686, linear ramp-up achieves 0.687, and the smooth sigmoid-like ramp-up schedule attains 0.693, confirming that routing ambiguous evidence into a sink during early optimization is critical.
- Statistically Significant Risk Stratification: Kaplan-Meier survival curves demonstrate sharp separation between predicted high-risk and low-risk groups across all five cohorts, with log-rank test \(p\)-values far below the 0.05 significance threshold (\(p = 2.61 \times 10^{-3}\) on BRCA, \(2.51 \times 10^{-4}\) on BLCA, \(8.93 \times 10^{-3}\) on STAD, \(3.02 \times 10^{-2}\) on CRC, and \(1.86 \times 10^{-10}\) on KIRC).
Highlights & Insights¶
- Revisiting Cross-Modal Over-Alignment: Challenges the common heuristic that deeper latent alignment is always superior. Formalizing over-alignment collapse and introducing the "Together, Then Apart" philosophy provides a sound conceptual paradigm for multimodal biomedical learning.
- Unifying Assignment Geometry via Dummy Sink UOT: Concatenating tokens from both modalities into a single transport problem with a shared prototype bank and a dynamic curriculum sink elegantly reconciles data heterogeneity, scale imbalance, and early training instability.
- Modular and Transferable Modality Anchors: The anchor refinement and InfoNCE contrastive separation mechanism is task-agnostic and computationally lightweight, making it easily adaptable to other asymmetrical multimodal tasks (e.g., vision-audio or sensor-fusion robotics).
Limitations & Future Work¶
- Residual Gap on Morphologically Dominant Cancers: While TTA marks a significant step forward on CRC among multimodal models (0.685), it still slightly trails the unimodal AttnMISL baseline (0.723), pointing to the need for adaptive modality weighting when one data source is overwhelming.
- Dependence on Fully Paired Inputs: TTA assumes complete availability of both WSI and genomic modalities during training and inference. In real-world clinics, missing modalities are frequent; integrating missing-modality imputation or masked modeling is an essential next step.
- Static Biological Pathway Prior: Genomic tokenization relies on predefined Hallmark gene sets, which might filter out unannotated non-coding variations or rare mutations outside canonical pathways.
Related Work & Insights¶
- vs. MMP [Song et al., ICML 2024]: MMP maintains separate prototype dictionaries for histology and genomics before calculating transport plans. TTA establishes a shared prototype bank across concatenated tokens, ensuring a unified assignment geometry, and subsequently applies Apart contrastive separation to retain distinctiveness.
- vs. MOTCat [Xu & Chen, ICCV 2023]: MOTCat solves dense instance-to-gene transport plans that scale quadratically and are prone to overfitting noise. TTA leverages prototype-level semi-relaxed UOT with an absorptive dummy sink, offering vastly superior noise tolerance and computational efficiency.
- vs. PIBD [Zhang et al., ICLR 2024]: PIBD attempts information-theoretic disentanglement via mutual information bounds, which often suffers from optimization volatility. TTA provides a more transparent geometric solution via optimal transport and anchor contrastive learning.
Rating¶
- Novelty: โญโญโญโญ [The Together-Then-Apart concept is conceptually lucid and effectively tackles over-alignment collapse via curriculum UOT and anchor contrast]
- Experimental Thoroughness: โญโญโญโญโญ [Evaluated on 5 diverse TCGA cohorts against 17 baselines, supported by detailed ablations, parameter sensitivity analyses, and KM curves]
- Writing Quality: โญโญโญโญโญ [Clear motivation, rigorous mathematical formulation, and excellent contextualization within computational pathology]
- Value: โญโญโญโญโญ [Provides a solid, open-sourced blueprint for robust multimodal prognosis modeling with strong clinical implications]