Dual Mixture-of-Experts Framework for Discrete-Time Survival Analysis

AI summary

Overview

Research area: Machine learning for healthcare, specifically discrete-time survival analysis applied to breast cancer datasets.

Technical level: Intermediate. The paper assumes familiarity with hazard functions, censoring, and the C-index, and builds on existing deep survival models (CoxPH, DeepHit, ConSurv) and the mixture-of-experts (MoE) design from Shazeer et al. (2017).

One-sentence scope: The paper introduces a "dual MoE" architecture that replaces the single feature encoder and single hazard network of a deep discrete-time survival model with two separately routed mixtures of experts, and evaluates it on the METABRIC and GBSG breast cancer datasets.

What This Paper Is About

Standard deep survival models typically use one shared feature encoder for all patients and one shared network to estimate hazards, which forces every patient and every time bin through a single shared functional form. This is a poor fit for real clinical data, where patients fall into heterogeneous subgroups with different risk profiles and where risk changes differently over time for different people. The paper's goal is to break that single-network assumption by letting separate expert networks specialize along both patient subgroups and temporal horizons, while a router decides which expert to trust for each patient and each time point.

Key Contributions

1. A feature-encoder MoE: an initial encoder extracts patient-level representations, which are then soft-routed into K expert encoders using routing probabilities produced from patient features, yielding subgroup-aware representations.
2. A hazard MoE: L hazard experts each predict hazards across all discrete time bins, with the router conditioned on the concatenation of patient features and learnable time embeddings, producing a full hazard vector over the prediction horizon.
3. Load-balancing regularizers for both stages: an NLL loss combined with a feature load-balancing loss (coefficient alpha) and a hazard load-balancing loss (coefficient beta) to prevent collapse onto a single expert.
4. Empirical validation and integration: consistent improvements in overall and time-dependent C-index on METABRIC and GBSG, plus additional gains when the dual MoE is swapped into the ConSurv framework, showing the design is portable to other deep discrete-time survival pipelines.

Main Findings

• Time-dependent C-index improves by up to 0.04: the abstract reports the method boosts the time-dependent C-index up to 0.04 on the test sets. The largest such gap in Table 1 is on METABRIC with ConSurv at the 10% horizon (0.696 with dual MoE vs. 0.656 without).
• METABRIC, ConSurv backbone: overall C-index rises from 0.657 +/- 0.020 (no dual MoE) to 0.668 +/- 0.018 (with dual MoE). Time-dependent C-index at the 10% horizon rises from 0.656 +/- 0.044 to 0.696 +/- 0.034, and the dual-MoE variant is higher at every reported horizon from 10% through 90%.
• METABRIC, naïve implementation: overall C-index rises from 0.646 +/- 0.021 to 0.654 +/- 0.015; at the 10% horizon from 0.670 +/- 0.050 to 0.669 +/- 0.032 (a slight decrease at that horizon), with gains at the remaining horizons.
• GBSG, ConSurv backbone: overall C-index rises from 0.665 +/- 0.011 to 0.668 +/- 0.011; the 10% horizon rises from 0.742 +/- 0.039 to 0.752 +/- 0.036.
• GBSG, naïve implementation: overall C-index rises from 0.662 +/- 0.012 to 0.667 +/- 0.010; the 10% horizon rises from 0.744 +/- 0.039 to 0.751 +/- 0.033.
• CoxPH baseline comparison: on METABRIC, CoxPH (Cox, 1972) reaches an overall C-index of 0.663 +/- 0.017; on GBSG, 0.659 +/- 0.012. On GBSG, both ConSurv variants at the 10% horizon (0.742 and 0.752) exceed the CoxPH 10% value of 0.739 +/- 0.046.
• Both MoE components are complementary (METABRIC ablation, naïve implementation): C-index is 0.646 +/- 0.021 with neither component, 0.649 +/- 0.023 with feature MoE only, 0.650 +/- 0.025 with hazard MoE only, and 0.654 +/- 0.015 with both.
• Hazard router needs both inputs: the ablation on router inputs (reported as Figure 4, using time-dependent C-index) shows patient features alone and time embeddings alone are both suboptimal, while their combination performs best. The hazard experts themselves always receive both features and time embeddings in all variants.
• Routing is not uniform: visualizing feature-encoder routing averaged over estrogen receptor (ER) and HER2 subgroups on METABRIC (Figure 2) shows distinct expert preferences between subgroups. Patient-level hazard routing over time for four patients (A–D) shows patterns that vary across patients but consistently shift in expert dominance between early and late time horizons.
• Results are averaged over 10 random seeds, and the paper notes results may differ from Lee et al. (2024) because dataset splits vary with different random seeds.

Methodology in Plain English

The authors work in a discrete-time survival setting. Each patient has covariates, an observed time drawn from the range 0 to T_max, and an event indicator that is 1 for an observed event and 0 for right-censoring. A conditional hazard gives the instantaneous event probability at time t given survival to t and the covariates; multiplying (1 - hazard) across time bins gives the survival function, and hazard times the previous survival value gives the event probability mass at t.

Training uses the standard negative log-likelihood, which adds the log event probability for patients with observed events and the log survival function for censored patients.

On top of that backbone, the model inserts two MoE layers:

• Feature stage: an initial encoder produces a patient representation. A router (a 1-hidden-layer MLP) takes patient features and outputs softmax mixing weights over K expert encoders. The final representation is the weighted sum of the experts' outputs. Because routing depends on patient features, the model can discover and serve hidden subgroups.
• Hazard stage: L hazard experts each output hazards for all time bins. The router here sees the concatenation of patient features and a learnable time embedding for each bin, and produces weights over the experts per time bin. The final hazard at time t is the weighted sum of expert hazards. This lets experts specialize both by patient and by how risk evolves over time.

To stop either MoE from collapsing onto one expert, each stage adds a load-balancing regularizer built from the batch-averaged routing probabilities, with coefficients alpha (feature) and beta (hazard). The overall objective is NLL plus both load-balancing terms.

Experimental setup: METABRIC has 1,981 patients and 21 variables with 55.2% censored and 44.8% uncensored; GBSG has 2,232 patients and 21 clinical and tumor-related variables, originally collected to study hormone therapy's impact on recurrence-free survival, with 43.2% censored and 56.8% uncensored. The number of feature-encoder experts (K) and hazard experts (L) is set to (4, 4) for METABRIC and (6, 3) for GBSG. Evaluation uses the overall concordance index (C-index) plus time-dependent C-index computed at the 10%–90% percentiles of observed event times. The authors compare against a "naïve implementation" trained only with the NLL loss and consisting of one encoder and one hazard network, and they also apply the dual MoE on top of ConSurv by replacing its encoder and hazard network.

Hyperparameters from Appendix A: initial encoder MLP depth 4 for METABRIC and 3 for GBSG; feature-encoder router depth 1 for both; feature-encoder expert depth 1 (METABRIC) and depth 2 (GBSG); hazard router depth 1 and hazard expert depth 1 for both; time embedding dimension 8 for both; alpha = 0.3 for both; beta = 0.5 for both.

Why This Matters

Impact on research: The paper argues that the single-encoder, single-hazard-head design shared by most deep survival models is a structural limitation, not just a capacity limitation, because it implicitly ties all patients and all time bins to one functional form. The dual MoE is presented as a drop-in module that improves existing pipelines — the ConSurv experiments support that claim — which makes it a general recipe rather than a standalone model.

Real-world applications:

• Breast cancer risk stratification using clinical and gene expression data (METABRIC) or clinical and tumor-related variables (GBSG), where subgroup-specific risk trajectories matter for prognosis.
• Identifying patient subgroups whose risk profile differs from the population average, since the feature-encoder router learns subgroup-aware representations rather than favoring dominant patterns.
• Time-varying risk monitoring, where the hazard router's shifting expert dominance between early and late horizons could support decisions about when to intensify follow-up.
• Direct integration into existing deep survival pipelines for clinical research, given the demonstrated gains on top of ConSurv.

Industry relevance: The work comes from Lunit Inc., and the authors' stated future direction is extension to mammography-based risk prediction, indicating a path toward imaging-based clinical risk tools where heterogeneous patient populations are the norm.

Future Directions

• Analyzing the role of each individual expert, which the authors explicitly list as future work.
• Extending the framework to multimodal settings, specifically mammography-based risk prediction.
• Determining how the MoE design behaves beyond the two breast cancer datasets studied and the two backbones tested (naïve implementation and ConSurv), including whether the gains hold at other censoring rates.
• Understanding why the hazard router's expert dominance shifts between early and late time horizons, and whether that structure maps onto clinically meaningful risk periods.

Target Audience

Researchers and practitioners working on deep learning for survival analysis and clinical risk prediction, particularly those already using or evaluating models such as DeepHit and ConSurv. It is also relevant to machine learning engineers interested in mixture-of-experts routing outside of large language models, and to clinical data scientists who want a modular upgrade path that can be added to an existing discrete-time survival pipeline without redesigning it.

Authors’ abstract

Survival analysis is a task to model the time until an event of interest occurs, widely used in clinical and biomedical research. A key challenge is to model patient heterogeneity while also adapting risk predictions to both individual characteristics and temporal dynamics. We propose a dual mixture-of-experts (MoE) framework for discrete-time survival analysis. Our approach combines a feature-encoder MoE for subgroup-aware representation learning with a hazard MoE that leverages patient features and time embeddings to capture temporal dynamics. This dual-MoE design flexibly integrates with existing deep learning based survival pipelines. On METABRIC and GBSG breast cancer datasets, our method consistently improves performance, boosting the time-dependent C-index up to 0.04 on the test sets, and yields further gains when incorporated into the Consurv framework.