Let the Experts Speak: Improving Survival Prediction&Calibration via Mixture-of-Experts Heads

2025-11-11

AI Summary

In survival analysis, conventional group-based modeling often compromises both predictive accuracy and calibration. Method: This paper proposes a discrete-time deep Mixture-of-Experts (MoE) model that abandons the restrictive inductive bias of hard expert assignment inherent in standard MoE architectures. Instead, it introduces a learnable gating mechanism and highly expressive, patient-specific expert networks to decouple subgroup discovery from individualized risk prediction. The model is trained end-to-end to jointly optimize clustering fidelity and prediction quality. Contribution/Results: Experiments across multiple clinical survival datasets demonstrate consistent improvements over state-of-the-art baselines: lower calibration error (measured by Brier score and Expected Calibration Error) and higher predictive performance (measured by Concordance index and Integrated Brier Score). These results validate the effectiveness of co-optimizing interpretable subgroup structure learning with accurate, well-calibrated risk prediction.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Mixed Discrete/Continuous Search
• Reasoning under Uncertainty: Graphical Models

Application Category

• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Deep mixture-of-experts models have attracted a lot of attention for survival analysis problems, particularly for their ability to cluster similar patients together. In practice, grouping often comes at the expense of key metrics such calibration error and predictive accuracy. This is due to the restrictive inductive bias that mixture-of-experts imposes, that predictions for individual patients must look like predictions for the group they're assigned to. Might we be able to discover patient group structure, where it exists, while improving calibration and predictive accuracy? In this work, we introduce several discrete-time deep mixture-of-experts (MoE) based architectures for survival analysis problems, one of which achieves all desiderata: clustering, calibration, and predictive accuracy. We show that a key differentiator between this array of MoEs is how expressive their experts are. We find that more expressive experts that tailor predictions per patient outperform experts that rely on fixed group prototypes.

Problem

• Improving survival prediction accuracy and calibration in mixture-of-experts models
• Overcoming restrictive inductive bias in patient grouping strategies
• Developing expressive experts that tailor predictions per patient

Innovation

• Mixture-of-experts architecture improves survival analysis
• Expressive experts tailor predictions per patient
• Discrete-time deep models enhance calibration and accuracy