arXiv:2606.14608v1 [cs.LG] 12 Jun 2026
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts Farica Zhuang
Zixuan Wen
[email protected] University of Pennsylvania Philadelphia, Pennsulvania, USA
[email protected] University of Pennsylvania Philadelphia, Pennsulvania, USA
Christos Davatzikos
Li Shen∗
[email protected] University of Pennsylvania Philadelphia, Pennsulvania, USA
[email protected] University of Pennsylvania Philadelphia, Pennsulvania, USA
Abstract Survival prediction plays a central role for healthcare providers and clinical researchers. Accurate risk stratification enables early intervention and improved patient management. Most existing deep survival models learn one common feature representation for all patients, which may hide important differences between patient subgroups. In contrast, a Mixture-of-Experts (MoE) framework allows different parts of the model to focus on different patient patterns, leading to more individualized representations. Therefore, in this work, we propose a mixture-of-experts enhanced adaptive deep clustering survival framework (AdaCSM) for modeling such heterogeneous survival patterns. We introduce a routing-based expert mechanism that enables conditional specialization within a parametric survival modeling framework. The proposed architecture allocates patients to specialized risk predictors dynamically while preserving the patient survival and subtype clustering objectives. We compare our method with state-of-the-art survival and deep clustering models on multiple real-world longitudinal clinical cohorts spanning diverse disease domains. The proposed method demonstrates improved predictive performance and leads to interpretable results in survival analysis.
CCS Concepts • Computing methodologies → Machine learning; • Applied computing → Health informatics.
Keywords Mixture-of-Experts, Survival Analysis, Alzheimer’s Disease, Subtype Discovery, Interpretable Machine Learning
1
Introduction
Time-to-event prediction (survival analysis) is designed to predict when an event is likely to happen, not just whether it happens. It ∗ Correspondence. [email protected].
Accepted to the 17th ACM International Conference on Bioinformatics, Computational Biology and Health Informatics (BCB ’26), 2026. Author-created version. Final Version of Record available at https://doi.org/10.1145/3807503.3819574.
This work is licensed under a Creative Commons Attribution-NonCommercialNoDerivatives 4.0 International License.
has been widely used to predict disease or death in medicine and the public health domain. Classical statistical approaches include the accelerated failure time (AFT) models and the Cox proportional hazards (Cox PH). AFT models [4, 23] are parametric models with the assumption that covariates either accelerate or decelerate the event time via a linear model with the log-transformation. While Cox PH [6], a semi-parametric model, assumes the hazard rate for every instance is constant over time. Non-parametric models, such as [3], give a stepwise survival curve that drops only when events occur. Recently, many deep learning methods have been proposed to improve time-to-event prediction. DeepSurv [16] is a Cox proportional hazards deep neural network, which predicts the effects of a patient’s covariates on their hazard rate. Deep survival machines (DSM) [20] is a fully parametric approach that models the survival function as a weighted mixture of individual survival distributions. Besides, a non-parametric model, Random Survival Forest (RSF) [13, 14] builds many survival trees on bootstrap samples, estimates the cumulative hazard function based on an ensemble of trees. Deep multi-task Gaussian process (DMGP) [2] is used to capture complex non-linear interactions between the patients’ covariates and cause-specific survival times. In medical applications, an important goal is to identify patient subgroups with similar survival patterns and risk profiles. Traditional unsupervised clustering methods, such as KMeans [10], do not use time-to-event outcomes, so the resulting groups may not reflect clinically meaningful differences in prognosis. A simple alternative is to divide patients into groups using predefined thresholds, such as risk scores or survival quantiles, but these thresholds are often arbitrary and may miss more complex subgroup structures. Recent survival clustering methods address this problem by incorporating survival outcomes directly into subgroup discovery [5, 12, 15, 19]. For example, Survival Cluster Analysis (SCA) [5] groups patients in a learned feature space using a flexible mixture model, while Deep Clustering Survival Machines (DCSM) [11, 12, 24] combines feature learning with parametric survival modeling so that mixture components correspond to probabilistic survival subtypes. However, these methods still generally use a single encoder to transform all patient variables into one common feature representation. When patient populations are diverse, such as in multimodal clinical data or across multiple cohorts, a single common representation may hide important differences between
Zhuang et al.
subgroups. A more detailed comparison is presented in Section 2 “Related Work”. To address this limitation, we propose an adaptive clustering survival framework (AdaCSM) that captures differences across patients while preserving interpretable subtype discovery. In the AdaCSM framework (Figure 1), a Mixture-of-Experts (MoE) encoder allows different parts of the model to focus on different types of patients. At the same time, the parametric mixture survival modeling continues to learn interpretable survival subtypes. Our work makes the following key contributions: (1) MoE-enhanced survival clustering architecture. We depart from the conventional shared-encoder design by reformulating survival clustering architecture under a Mixtureof-Experts (MoE) paradigm, where representation learning is driven by adaptive expert specialization and routing. This shift enables patient-specific feature selection and individualized representations, fundamentally advancing the modeling of heterogeneous survival patterns. (2) Top-K sparse routing for patient-specific specialization. We incorporate a Top-K routing mechanism so that each patient is processed by only a small subset of experts. This promotes expert specialization, improves computational efficiency, and provides a transparent view of which experts are activated for different patient profiles. (3) Preservation of interpretable subtype discovery. By combining the MoE encoder with the original mixture parametric survival layer, as in deep clustering survival model, the proposed model retains the interpretable survival subtype assignments through clustering, while improving the flexibility of the learned patient representation. (4) Empirical validation across multiple clinical cohorts. Experiments on several real-world clinical datasets show that the proposed framework achieves stronger subgroup separation and competitive survival prediction performance compared with existing baselines.
2
Related Work
Existing survival clustering methods differ in both how they encode covariates and how they define subgroups. DSM [20] uses a deterministic MLP to learn nonlinear representations that parameterize mixture weights and survival distributions, but this encoder is primarily optimized for individualized prediction rather than explicit subtype discovery. Accordingly, DSM does not natively perform clustering and is often used with post hoc risk stratification instead. SCA [5] employs a deterministic encoder to map covariates into a latent space for Bayesian nonparametric modeling, but it does not explicitly capture uncertainty in the learned representation. Its clustering is performed through a flexible latent mixture model, which can adaptively discover subgroups but may also produce overly fragmented subgroups that are difficult to interpret. VaDeSC [19, 22] adopts a variational generative encoder with a Gaussian mixture latent prior to jointly model covariates and survival outcomes, but its learned representation can be sensitive to misspecification of the assumed feature distribution. Because its clustering is tightly coupled to that generative prior, subgroup discovery may degrade when the data deviates from the assumed model.
Mixture of Experts
Clustering Survival Machines
... Router
Mixture Top-K
Expert 1
Expert 2
Expert E
...
Weibull Distributions
Expert outputs gate weight
Figure 1: AdaCSM pipeline. The model uses a top-K routing mechanism to assign each input to a subset of experts, which are subsequently processed by a Weibull-based Deep Clustering Survival Machine. Expert outputs are combined through mixture weights produced by the router, enabling adaptive modeling of heterogeneous risk patterns.
In contrast, Neural Survival Clustering (NSC) [15] employs an MLP-based assignment network, but this encoder serves only to assign patients to clusters rather than to model individual variation within clusters. It improves flexibility and interpretability by clustering patients via assignment to population-level neural survival components, yet potentially sacrificing individual-level covariate dependence. DCSM [11, 12, 21, 24] takes a more interpretable ground by learning a compact discriminative embedding with an MLP and using it to estimate mixture weights over fixed parametric survival experts. This design yields more clinically meaningful subgroup assignments because subjects are clustered according to weights over constant expert distributions, but it still relies on a single shared encoder that may blur subgroup-specific feature patterns. These limitations suggest that, while prior methods differ in their balance between flexibility and interpretability, most still rely on a single global covariate encoder that must explain heterogeneous patient populations with shared parameters. Motivated by this gap, we introduce an MoE-based encoder enhanced clustering survival framework to enable subgroup specialized representation learning before survival clustering. Instead of forcing one shared MLP to encode all patients, the MoE encoder allows different experts to focus on distinct regions of the covariate space, while a gating network adaptively combines expert outputs for each individual. This design preserves the interpretable clustering mechanism while improving the encoder’s ability to capture heterogeneous, subgroup-specific feature patterns. In this way, our approach addresses a key limitation of existing survival clustering methods, the lack of encoder specialization for complex population heterogeneity.
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts
3
Method
In this section, we propose the AdaCSM model, a mixture-ofexperts enhanced adaptive deep clustering survival framework, which enables subgroup specialized representation learning before survival clustering. Figure 1 illustrates the overall framework of our method. For an input vector 𝑥, the latent representation 𝑥˜ is generated by a mixture-of-experts (MoE) encoder, where a gating network adaptively assigns subjects to different experts. The survival density function is modeled as a weighted mixture of Weibull distributions, where 𝛼 denotes the cluster-specific mixing weights. The MoE and Weibull parameters are learned jointly in an end-toend manner from both the covariates and the survival information.
3.1
Problem Formulation
𝑁 denote a dataset of 𝑁 individuals, where Let D = {(x𝑖 , 𝑡𝑖 , 𝛿𝑖 )}𝑖=1 x𝑖 ∈ R𝑑 represents the covariate vector describing subject 𝑖, 𝑡𝑖 denotes the observed follow-up time, and 𝛿𝑖 ∈ {0, 1} is the event indicator. Specifically, 𝛿𝑖 = 1 indicates that the event of interest (e.g., disease progression or death) was observed at time 𝑡𝑖 , while 𝛿𝑖 = 0 indicates right-censored observations where the event had not occurred by the end of follow-up. The objective of survival analysis is to estimate the conditional time-to-event distribution given the subject’s covariates. This distribution is typically characterized by the survival function
𝑆 (𝑡 | x) = P(𝑇 > 𝑡 | 𝑋 = x),
(1)
which represents the probability that the event time 𝑇 exceeds time 𝑡. An equivalent representation is the hazard function. P(𝑡 ≤ 𝑇 < 𝑡 + Δ𝑡 | 𝑇 ≥ 𝑡, 𝑋 = x) , (2) Δ𝑡 which captures the instantaneous risk of event occurrence at time 𝑡 given survival up to that time. In many clinical settings, patient populations are heterogeneous and may exhibit multiple latent disease progression patterns. To capture this heterogeneity, we assume that the conditional survival distribution can be modeled as a mixture of 𝐽 latent components (or subtypes). Each component corresponds to a specialized risk predictor representing a distinct survival trajectory. Formally, the survival distribution is modeled as 𝐽 𝐽 ∑︁ ∑︁ 𝑆 (𝑡 | x) = 𝛼 𝑗 (x)𝑆 𝑗 (𝑡) = 𝑝 ( 𝑗 | x)𝑆 𝑗 (𝑡), (3) ℎ(𝑡 | x) = lim
This representation is then passed to a clustering survival machine head, which predicts mixture weights over fixed Weibull distributions. The weighted combination of these survival components yields the final survival distribution and provides an interpretable subgroup assignment through the dominant component.
3.3
The survival model 𝑆 (𝑡 | x) of subject with covariate x is defined as Í 𝑆 (𝑡 | x) = 𝐽𝑗=1 𝛼 𝑗 (x)𝑆 𝑗 (𝑡). We predefine the number of subtypes to be 2 (i.e., 𝐽 = 2), corresponding to a straightforward partition of subjects into high-risk and low-risk groups. For each subtype, the survival function 𝑆 𝑗 (𝑡) is modeled using a Weibull distribution, which is widely adopted in survival analysis due to its flexibility and its closed-form expressions for the probability density and cumulative distribution functions. Specifically, 𝑘 𝑘 𝑡 𝑡 𝑘 𝑡 𝑘 −1 exp − , 𝐹 (𝑡) = 1 − exp − , (4) 𝑓 (𝑡) = 𝜆 𝜆 𝜆 𝜆 where 𝑘 and 𝜆 denote the shape and scale parameters, respectively. Accordingly, the survival function is given by 𝑘 𝑡 𝑆 (𝑡) = 1 − 𝐹 (𝑡) = exp − . (5) 𝜆 Thus, the survival function is explicitly modeled as a weighted mixture of subtype-specific Weibull survival functions: 𝑘 𝑗 ! 2 2 ∑︁ ∑︁ 𝑡 , (6) 𝑆 (𝑡 | x) = 𝛼 𝑗 (x)𝑆 𝑗 (𝑡) = 𝛼 𝑗 (x) exp − 𝜆 𝑗 𝑗=1 𝑗=1 where exp 𝑤 𝑇𝑗 x̃ . 𝛼 𝑗 (x) = Í 2 𝑇 𝑗=1 exp 𝑤 𝑗 x̃
Δ𝑡 →0
𝑗=1
𝑗=1
where 𝛼 𝑗 denotes the probability of assigning patient 𝑖 to subtype 𝑗, and 𝑆 𝑗 (𝑡) denotes the parametric survival function of subtype 𝑗. Í The subtype assignment probability satisfies 𝐽𝑗=1 𝛼 𝑗 = 1. The goal of the proposed framework is therefore twofold: (1) Survival prediction: accurately estimate individualized survival distributions 𝑆 (𝑡 | x). (2) Subtype discovery: identify latent patient subgroups characterized by distinct survival trajectories.
3.2
Model Overview
The model first encodes input covariates using a Mixture-of-Experts module, where a router assigns top-K weights to multiple expert networks and aggregates their outputs into the representation x̃.
Survival Modeling
(7)
Here, 𝛼 𝑗 (x) denotes the mixture weight of subtype 𝑗 for the subject, and can be interpreted as the probability of assignment that the subject belongs to the 𝑗-th survival subtype. In our framework, these subtype weights are estimated from the MoE-enhanced representation x̃, and are used to combine the subtype-specific Weibull survival functions into an individualized survival prediction.
3.4
Adaptive Gating Mechanism
To better capture heterogeneous covariate patterns across subjects, we replace the first layer of the standard multilayer perceptron encoder with a Mixture-of-Experts (MoE) layer to learn x̃. Given an input covariate vector x ∈ R𝑑 , the MoE layer consists of 𝐸 expert networks and a gating network that adaptively assigns weights to the experts. The output representation is defined as h(x) = Í𝐸 𝑒=1 𝑔𝑒 (x)𝑓𝑒 (x), where 𝑓𝑒 (x) denotes the output of the 𝑒-th expert and 𝑔𝑒 (x) is its corresponding gating weight. In our implementation, each expert is a lightweight neural transformation composed of a linear layer followed by a nonlinearity, i.e., 𝑓𝑒 (x) = 𝜎 (W𝑒 x) , where W𝑒 is the learnable weight matrix for expert 𝑒, and 𝜎 (·) denotes the activation function. In this work, we use ReLU6(·) as the activation. Thus, each expert learns a distinct transformation of the input covariates, allowing different experts to specialize in different regions of the feature space.
Zhuang et al.
The gating network maps the input x to a set of routing scores, which are then normalized by a softmax function: 𝑔𝑒 (x) = Í𝐸
exp (𝑧𝑒 (x)/𝜏) ,
𝑒 ′ =1 exp (𝑧𝑒 ′ (x)/𝜏)
(8)
(9)
where T (x) denotes the set of top-𝐾 selected experts for subject x. The final MoE representation is then computed as h(x) =
𝐸 ∑︁
𝑔˜𝑒 (x)𝑓𝑒 (x).
(10)
𝑒=1
Here, h(x) is the intermediate output of the MoE layer, whereas x̃ denotes the final encoder output used by the survival clustering head. This adaptive gating mechanism enables the encoder to move beyond a single shared transformation for all subjects. Instead, each subject is represented by a weighted combination of specialized experts, allowing the model to capture subgroup-specific covariate patterns before downstream survival modeling. The resulting MoEenhanced representation is then passed to the subsequent encoder layers and ultimately used to estimate the mixture weights over subtype-specific survival experts.
3.5
The proposed model is trained end-to-end by jointly optimizing the adaptive gating mechanism and the survival mixture model. For each subject, let 𝑡𝑖 denote the observed time and let 𝛿𝑖 ∈ {0, 1} denote the event indicator, where 𝛿𝑖 = 1 indicates an observed event and 𝛿𝑖 = 0 indicates right censoring. Given the mixture-based survival model, the corresponding probability density function is Í defined as 𝑓 (𝑡 | x) = 𝐽𝑗=1 𝛼 𝑗 (x)𝑓 𝑗 (𝑡), and the survival function Í𝐽 is 𝑆 (𝑡 | x) = 𝑗=1 𝛼 𝑗 (x)𝑆 𝑗 (𝑡), where 𝑓 𝑗 (𝑡) and 𝑆 𝑗 (𝑡) denote the subtype-specific Weibull density and survival functions, respectively. Accordingly, the likelihood contribution of subject 𝑖 is written as (11)
The overall objective is obtained by minimizing the negative loglikelihood: Lsurv = −
𝑁 ∑︁ 𝑖=1
[𝛿𝑖 log 𝑓 (𝑡𝑖 | x𝑖 ) + (1 − 𝛿𝑖 ) log 𝑆 (𝑡𝑖 | x𝑖 )] .
𝐽 𝑁 h ∑︁ ∑︁ 𝛿𝑖 log 𝛼 𝑗 (x𝑖 )𝑓 𝑗 (𝑡𝑖 )
(12)
𝑗=1
+ (1 − 𝛿𝑖 ) log
𝐽 ∑︁
(13) i
𝛼 𝑗 (x𝑖 )𝑆 𝑗 (𝑡𝑖 ) .
𝑗=1
Since both the subtype mixture weights 𝛼 𝑗 (x) and the MoE representation x̃ are differentiable functions of the input covariates, this objective can be optimized jointly with respect to all model parameters. In this way, the router, expert networks, and subtype-specific survival components are learned simultaneously in a unified framework.
3.6
Training and Optimization
We adopt a consistent evaluation pipeline across all methods. For time-to-event prediction, each dataset is split using a fixed random seed into 70% training and 30% testing sets, with the training portion further divided into training and validation subsets (60%10%-30% overall). Continuous features are normalized to [0,1] using statistics computed only from the training data, and categorical features are one-hot encoded (Appendix A.1). Hyperparameters are selected with Optuna [1] by maximizing validation C-Index, and the best configuration is evaluated on the held-out test set. The best hyperparameters are shown in Appendix (Appendix A.2, Appendix Tables 1-5). The experiment is repeated across five random seeds, and we report the mean and standard deviation of the C-Index and LogRank statistic.
4
Experiment Setup
In this section, we describe the datasets, metrics, and methods used to evaluate and benchmark our proposed framework. Our code is available at: https://github.com/PennShenLab/AdaCSM.
4.1
Joint Objective Function
L𝑖 = [𝑓 (𝑡𝑖 )] 𝛿𝑖 [𝑆 (𝑡𝑖 )] 1−𝛿𝑖 .
Lsurv = −
𝑖=1
where 𝑧𝑒 (x) is the routing logit for expert 𝑒, and 𝜏 > 0 is a temperature parameter controlling the sharpness of the routing distribution. A smaller 𝜏 produces more selective routing, while a larger 𝜏 yields smoother expert combinations. To further encourage expert specialization, we optionally apply a top-𝐾 routing strategy. Specifically, only the 𝐾 experts with the largest gating weights are retained, while the remaining weights are set to zero. The retained weights are then re-normalized: 𝑔𝑒 (x) Í 𝑔 ′ (x) , 𝑒 ∈ T (x), 𝑔˜𝑒 (x) = 𝑒 ′ ∈ T (x) 𝑒 0, otherwise,
Substituting the mixture formulation into the above objective yields
Datasets
We prepare four real-world survival datasets for experiments, and their data statistics are shown in Table 1. All datasets are analyzed under standard right-censoring, where individuals without an observed event are censored at their last follow-up time. • SUPPORT [17]: The Study to Understand Prognoses and Preferences for Outcomes and Risks of Treatments (SUPPORT) dataset sourced from a study conducted by Vanderbilt University that contains comprehensive clinical and demographic data from hospitalized adults with serious illnesses. The survival endpoint is defined as the time from hospital admission to all-cause mortality. • PBC [9]: The Primary Biliary Cholangitis (PBC) dataset contains longitudinal clinical measurements from patients enrolled in a Mayo Clinic study of liver disease progression. We use the time-dependent version of the cohort (pbc2), which includes repeated laboratory and clinical assessments collected over follow-up. The event is defined as mortality. • Framingham [7]: This dataset is derived from the wellknown Framingham Heart Study that follows participants
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts
longitudinally to model the time to the first major cardiovascular event, including incident coronary heart disease and stroke. It incorporates baseline demographic and clinical risk factors such as age, blood pressure, cholesterol levels, and smoking history. We consider death as the event, with survival time measured from the baseline exam to death or the end of follow-up. • FLCHAIN [8]: This dataset originates from a populationbased study in Olmsted County, Minnesota, designed to evaluate the prognostic value of serum free light chain (FLC) levels. The survival endpoint is time to death from any cause. Beyond standard demographics like age and sex, the dataset includes various immunological and biochemical markers for plasma cell disorders and renal function. A detailed description of the data preprocessing pipeline is provided in Appendix A.1. Table 1: Dataset statistics: including sample size (𝑛), feature dimensionality after preprocessing/encoding (𝑑), event rate, censoring rate, and median follow-up time. Data set 𝑛 𝑑 Event % Cens. % 𝑡 follow-up SUPPORT 9,105 59 68.1% 31.9% 233 days PBC 1945 25 37.28% 62.72% 4.6 years FRAMINGHAM 11,627 18 30.3% 69.7 % 17.9 years FLCHAIN 6,524 27 30.1% 69.9% 11.8 years
4.2
Evaluation Metrics
In the experiments, we denote the predicted risk scores as 𝜂ˆ𝑖 and the ground-truth time-to-event as 𝑇𝑖 . With this, we assess the patient prognostic and subtype clustering precision by adopting the following metrics for evaluation: The Concordance Index (C-Index) measures the probability that for a randomly selected pair of patients, the model predicts a higher risk for the individual who experiences the event first: Í C-Index =
𝑖,𝑗 ∈𝑃 𝐼 (𝑇𝑖 < 𝑇 𝑗 ) · 𝐼 (𝜂ˆ𝑖 > 𝜂ˆ 𝑗 )
Í
𝑖,𝑗 ∈𝑃 𝐼 (𝑇𝑖 < 𝑇 𝑗 )
(14)
where 𝑃 is the set of all comparable pairs in the dataset and 𝐼 (·) is the indicator function. The Log-Rank Statistic (𝜒 2 ) is used to quantify the degree of separation between the survival distributions of the latent subtypes discovered by the adaptive gating network. The magnitude of the Log-Rank statistic reflects the strength of the evidence against the null hypothesis that the subtypes share a common survival trajectory. The statistic is calculated as: Í𝑓 ( 𝑡 =1 𝑂 1𝑡 − 𝐸 1𝑡 ) 2 2 𝜒 = (15) Í𝑓 𝑡 =1 𝑉1𝑡 where 𝑂 1𝑡 and 𝐸 1𝑡 are the observed and expected number of events at time 𝑡, and 𝑉1𝑡 is the variance. Higher values of the 𝜒 2 statistic indicate that the model has identified subtypes with more distinct and non-overlapping clinical outcomes.
4.3
Comparison Methods
To evaluate the performance of AdaCSM, we compare it with several representative classical and state-of-the-art survival analysis methods. • Cox Proportional Hazards (Cox PH) [6]: The Cox PH model is a widely used classical survival modeling approach based on a semi-parametric regression. It models the hazard function, which represents the instantaneous risk of experiencing the event at a given time, conditional on surviving up to that time. The model assumes proportional hazards, meaning that the hazard ratio between two individuals remains constant over time. The model parameters are estimated by maximizing the partial likelihood. • DeepSurv [16]: DeepSurv, also known as Deep Cox, is a deep learning-based extension of the Cox PH model that replaces the linear functional with a deep neural network to capture non-linear feature dependencies. The model is trained by maximizing the Cox partial likelihood and retains the proportional hazards assumption, meaning that the relative risk between individuals is similarly assumed constant over time. • Deep Survival Machines (DSM) [20]: DSM is a generative, fully parametric deep survival model that represents the survival distribution as a mixture of predefined parametric components (e.g., Weibull or Log-Normal distributions). Both the mixture weights and distribution parameters are predicted by a neural network conditioned on patient covariates. Unlike the Cox PH constant proportional hazards assumption, this formulation allows flexible modeling of heterogeneous survival patterns while maintaining a likelihood-based training objective. • Deep Clustering Survival Machines (DCSM) [12]: DCSM combines discriminative feature learning with generative survival modeling to identify latent patient subgroups. This framework integrates survival modeling with representation learning to induce survival-informed clustering, where mixture components correspond to soft subtype assignments. The model is trained end-to-end using a shared monolithic neural network that simultaneously determines subtype probabilities and survival parameters for the population. Consistent with previous benchmarking protocols, we ensure that all deep learning-based baselines (DeepSurv, DSM, and DCSM) are optimized using Optuna hyperparameter tuning to ensure a fair comparison with our proposed AdaCSM. Using two subtypes ensured a fair and uniform evaluation setting across baselines. In addition, a two-group setting aligns with the common clinical interpretation of low-risk versus high-risk stratification. For DCSM and AdaCSM, subgroup assignments were obtained directly from the learned mixture components by assigning each patient to the component with the highest mixture weight. For methods that do not natively produce subgroup assignments (Cox PH, DeepSurv, and DSM), we derived two risk groups by splitting patients at the median predicted risk score. Patients above the median were assigned to the high-risk group, and those below
Zhuang et al.
Cox PH
DSM
AdaCSM (Sparse top−2)
DeepSurv
DCSM
AdaCSM
LogRank
Models
1100 1000 900 800 700 600 500 400 300 200 100 0
5.3
SUPPORT
PBC
FRAMINGHAM
FLCHAIN
Figure 2: Comparison of LogRank statistics for clinical datasets across six survival-clustering methods: Cox PH, DeepSurv, DSM, DCSM, AdaCSM (Sparse top-2), and AdaCSM. Error bars indicate 95% confidence interval (CI).
the median were assigned to the low-risk group. This post hoc stratification was used only for LogRank evaluation, allowing all methods to be compared under the same two-group setting.
5 Results 5.1 Strong Subtype Identification We first evaluate whether AdaCSM can identify clinically meaningful subtypes with distinct survival trajectories. As shown in Figure 2, AdaCSM achieves the highest LogRank statistic on all four clinical cohorts, including SUPPORT (1047.17 ± 114.04), PBC (311.80 ± 19.30), Framingham (667.83 ± 44.24), and FLCHAIN (586.56 ± 38.39). These gains are substantial relative to Cox PH, DeepSurv, DSM, and DCSM, indicating that the proposed model produces more clearly separated survival subgroups. In particular, the improvement over DCSM is consistent across all datasets, suggesting that adaptive expert specialization enhances subtype discovery beyond the original clustering survival formulation. Subtype identification is also visually supported by the KaplanMeier curves in Figure 3. The red and blue Kaplan–Meier curves are visibly separated, indicating that the discovered groups correspond to distinct survival trajectories rather than arbitrary partitions. This implies AdaCSM is able to consistently identify two survival subtypes with clearly different prognoses.
5.2
improves survival stratification while maintaining strong rankbased prognostic performance. As a result, AdaCSM remains on the Pareto frontier of the trade-off between C-index and Logrank (Appendix Figure 7).
Competitive Prognosis Performance
Although AdaCSM is designed primarily for subtype discovery, we also evaluate its prognostic accuracy using the C-Index. As shown in Figure 4, AdaCSM remains competitive across all datasets, achieving 0.8284 ± 0.0060 on SUPPORT, 0.8181 ± 0.0142 on PBC, 0.7300 ± 0.0108 on Framingham, and 0.7806 ± 0.0088 on FLCHAIN. While AdaCSM does not uniformly achieve the highest C-Index, its performance is consistently close to the strongest baseline on each cohort. This result suggests that the gains in subtype separation do not come at the expense of predictive usefulness. Rather, AdaCSM
Benefit of Sparse Expert Routing
We next examine how expert count and routing sparsity affect performance. Here, we take the optimal hyperparameters from the dense AdaCSM model trained for each dataset (Appendix Tables 15). Figure 5 shows that subtype separation, measured by LogRank, generally improves when the model uses a moderate or large number of experts together with sparse or moderately sparse routing. The best configurations are dataset-dependent, but the strongest results are consistently obtained with 16 or 32 experts and relatively small top-𝑘 values. For example, the best LogRank values are achieved with 32 experts, 𝑘 = 4 on the PBC dataset (311.51), 32 experts, 𝑘=4 on the FLCHAIN dataset (628.28). These findings indicate that increasing expert diversity is beneficial, but only when routing remains selective enough to preserve specialization. Figure 6 presents the corresponding sensitivity analysis using the C-Index. Compared with LogRank, the variation across expert configurations is smaller, indicating that routing choices affect subtype separation more strongly than rank-based risk prediction. The best C-Index values are 0.6736 on SUPPORT, 0.7944 on PBC, 0.7574 on Framingham, and 0.7927 on FLCHAIN. This suggests that the predictive accuracy of AdaCSM is relatively sensitive to architectural choices, while its clustering quality is more robust to the degree of expert specialization. Sparse routing also reduces computational cost because only a subset of experts is activated for each input. Table 2 compares the computational complexity of AdaCSM with classical and deep survival baselines using the optimal hyperparameter configurations obtained for the SUPPORT dataset, which has the largest feature dimensionality among the evaluated cohorts. We report the number of active parameters and approximate floating-point operations (FLOPs) required per sample during inference, as well as total parameters for the maximum model capacity. For dense neural models such as DeepSurv, DSM, and DCSM, all parameters are activated for every input. In contrast, while AdaCSM contains a larger number of total parameters due to the presence of multiple expert networks, the effective computation during inference is controlled by the Top-𝑘 sparsity routing mechanism, where only a small subset of experts are activated for each sample. This allows AdaCSM to have significantly fewer active parameters and lower FLOPs when sparse routing is used. For example, with 𝑘 = 1, AdaCSM activates only 4.97K effective parameters and requires 9.72K FLOPs per sample, which is lower than all dense neural survival models, while achieving a LogRank score that outperforms the best-performing baseline model (Figures 2 and 5). Increasing 𝑘 increases the computational cost but allows more experts to contribute to the representation. This conditional computation enables AdaCSM to maintain a higher representational capacity while keeping the effective inference cost comparable to or lower than conventional dense models. Overall, sparse and moderately sparse routing (𝑘 = 1, 2, 4) with a moderate number of experts (16-32) provides stable performance across datasets. The general trend indicates that balanced expert
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts
Seed 0
Seed 1
Seed 2
Seed 3
Seed 4
SUPPORT
PBC
Framingham
FLCHAIN
Figure 3: Kaplan-Meier survival curves for the discovered subtypes across 5 random seeds for each clinical survival dataset by AdaCSM. Columns correspond to different random seeds, and rows correspond to datasets. All curves are generated using the test split.
Cox PH
DSM
AdaCSM (Sparse top−2)
DeepSurv
DCSM
AdaCSM
(a) SUPPORT
(b) PBC
(c) FRAMINGHAM
(d) FLCHAIN
Models 1.0
C−Index
0.8 0.6 0.4 0.2 0.0 SUPPORT
PBC
FRAMINGHAM
FLCHAIN
Figure 4: Comparison of C-Index statistics for clinical datasets across six survival-clustering methods: Cox PH, DeepSurv, DSM, DCSM, AdaCSM (Sparse top-2), and AdaCSM. Error bars indicate 95% confidence interval (CI).
capacity and routing sparsity offer the most robust performance. We find that the behavior is broadly consistent with observations in prior MoE architectures, where sparse routing offers balanced efficiency and predictive stability.
Figure 5: Sensitivity analysis for AdaCSM measured by LogRank statistic (↑).
5.4
Ablation Studies
Zhuang et al. (a) SUPPORT
(b) PBC
(c) FRAMINGHAM
(d) FLCHAIN
Figure 6: Sensitivity analysis for AdaCSM measured by mean C-Index statistic (↑). Table 2: Complexity comparison across survival models. Model CoxPH AdaCSM AdaCSM (𝑘 = 2) AdaCSM (𝑘 = 1) DeepCoxPH DSM DCSM
Active Params
FLOPs/sample
Total Params
90 18.47K 9.47K 4.97K 7.15K 7.10K 7.10K
180 36.72K 18.72K 9.72K 14.20K 14.20K 14.20K
90 18.47K 18.47K 18.47K 7.15K 7.10K 7.10K
Table 3: Ablation study of the gating mechanism on the PBC dataset. We evaluate the impact of adaptive routing compared to non-informed weighting strategies on prognostic accuracy (C-Index) and subtype separation (LogRank). Configuration
C-Index (↑)
LogRank (↑)
Random Weights Equal Weights (Uniform) Single Expert (𝑛 = 1)
0.7366 ± 0.0039 0.7827 ± 0.0195 0.8131 ± 0.0162
75.60 ± 56.83 64.92 ± 85.23 254.65 ± 39.16
AdaCSM (Full)
0.8218 ± 0.0084
311.80 ± 19.30
To evaluate the individual contributions and the adaptive gating mechanism of the MoE component in AdaCSM framework, we conducted ablation experiments on the PBC dataset (Table 3). We compare the AdaCSM framework against three constrained baselines. First, we ran AdaCSM with random weights to the experts, where expert routing is non-informed. Second, we assigned equal weights, where all experts contribute equally to the final representation. Lastly, we employed a single expert baseline, representing a
Figure 7: Expert specialization analysis on Framingham. Top: Expert assignments vary by age, diabetes, and blood pressure. Bottom: Clinical feature distributions for Expert 5 (lowerrisk) and Expert 10 (higher-risk profile). standard deep survival model without mixture components. The results show that the adaptive gating mechanism is a critical driver of AdaCSM’s performance and stability. Replacing the learned routing with random weights leads to a worse drop in prognostic accuracy than subtyping performance. The equal weights baseline results in highly unstable subtype separation, showing low LogRank but high standard deviation. In contrast, AdaCSM achieves a superior and stable LogRank and C-index. These findings confirm that the gating network in AdaCSM is the driver of its performance. The performance gap between the single expert baseline and AdaCSM further justifies the use of the MoE framework for handling the high variance and heterogeneity present in clinical survival cohorts.
5.5
Expert Specialization Across Clinical Subgroups
To better understand how sparse AdaCSM leverages the mixture-ofexperts architecture, we examine whether the routing mechanism learns clinically meaningful expert specialization. In particular, we analyze how the gating network assigns patients to experts across clinically relevant subgroups. We hypothesize that a meaningful routing mechanism would show experts to specialize in different patient populations rather than assigning samples uniformly. Here, we use the Framingham cohort for the study of cardiovascular disease using mortality as an outcome in this case study. Figure 7 illustrates the expert assignment patterns learned by sparse AdaCSM. The routing distribution varies systematically across age groups, diabetes status, and systolic blood pressure.
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts
patients in a clinically meaningful and interpretable way, identifying specific physiological states associated with different survival trajectories. In this case study, since the outcome of our data is mortality, factors contributing to the high risk preference of Expert 5 might be apparent when compared to the low-risk preference of Expert 10. The dynamic routing of patients to specialized experts is consistently observed across all evaluated datasets. We provide additional expert assignment profiles for the PBC, SUPPORT, and FLCHAIN cohorts in Appendix A.7 (Appendix Figures 4, 5, and 6). These confirm that the adaptive gating network in AdaCSM captures known drivers for mortality in each cohort, such as liver enzymes in the liver disease cohort from the PBC dataset (Appendix Figure 4). However, as clinical outcomes become more complex and as subtypes of interest increase, for example, when predicting the outcome or stages of a disease, the granular interpretability of expert association to subtype assignments that AdaCSM offers will be essential for understanding the underlying combinations of risk factors. To further validate that these learned expert specializations are grounded in robust feature importance, we compare the intrinsic expert profiles against post-hoc SHAP [18] attributions. We observe a high correlation between our learned expert profiles and SHAP importance, hence confirming that the gating mechanism consistently routes patients based on clinically relevant features (Appendix Figure 7, Appendix A.8).
6
Figure 8: Clinical feature distributions for Expert 5 (high risk) and Expert 10 (low risk). Younger, low blood pressure, and non-diabetic subjects are routed primarily to Expert 5, whereas older, high blood pressure, and diabetic subjects are assigned more frequently to Expert 10. The distinct specializations by the experts follow clinically relevant and interpretable subgroups for risks of mortality. These findings suggest that AdaCSM does not use experts uniformly but instead allocates different clinically interpretable subgroups to experts in a structured way. This behavior supports the interpretation that the MoE encoder captures subgroup-specific feature patterns that may contribute to improved survival stratification. In a clinical setting, this model behavior is desirable since patient populations are often heterogeneous and risk factors may interact differently across demographic and comorbidity groups. While the routing patterns show that expert specialization occurs, Figure 8 further reveals the overall combination of patient conditions or the clinical logic that drives the subtype assignments. By inspecting the feature distributions, we observe that Expert 5 focuses on a protective profile, characterized by lower glucose levels and blood pressure. On the other hand, Expert 10 characterizes a high-risk profile where metabolic, cardiovascular markers, and pre-existing conditions are elevated. This deep dive into the expert profiles confirms that the gating mechanism by AdaCSM clusters
Discussion
Overall, the results suggest that AdaCSM is particularly effective at identifying survival subtypes. Across all four clinical cohorts, the model achieves stronger LogRank performance than the competing baselines, which means that the discovered groups are more clearly separated in terms of survival trajectory. At the same time, AdaCSM remains competitive in C-Index, so the improvement in subtype stratification does not come with a major loss in predictive performance. This pattern is important because our main goal is not only to rank patients by risk, but also to recover clinically meaningful subgroups with distinct prognoses. The substantial improvement comes from the MoE encoder. In the original DCSM formulation, all patients are passed through a single shared encoder, which may not be flexible enough to represent heterogeneous clinical populations. In AdaCSM, the gating network routes different patients to different experts, allowing the model to learn more specialized representations for different survival patterns. This is consistent with the design of our model, where the MoE-enhanced representation is used to estimate subtype weights and combine subtype-specific Weibull survival functions into the final prediction. Future work could further extend the intrinsic interpretability of our framework. The modular experts provide a structured interface for integration with large language models (LLMs) to provide natural language explanations of the personalized patient routing and specializations for clinical interpretability. This allows for evidence-based clinical reports that could be generated when explaining a patient’s subtype assignment. The sensitivity analysis also helps explain when the model works best. In general, increasing the number of experts improves performance up to a point, especially when routing is sparse. However,
Zhuang et al.
using too many experts together with dense routing can slightly hurt performance. We see a fairly consistent trend that sparse or moderately sparse routing, especially with 𝑘 = 1, 2, or 4, gives the most stable results across datasets. This suggests that the benefit of AdaCSM is not simply larger model capacity, but the combination of enough expert diversity with enough routing sparsity to preserve specialization. There are still several limitations. The current model fixes the number of subtypes to two, which aligns with clinical interpretability but may miss more complex disease structures. Extending the framework to automatically determine the number of clusters and evaluate larger numbers of subtypes is an important direction for future work. In addition, the Weibull survival experts impose a parametric assumption, and the current experiments do not yet fully explain the clinical meaning of individual experts. Future work will explore multimodal extensions, temporal routing, and transfer of expert modules across related disease settings.
7
Conclusion
In this work, we propose AdaCSM, a deep clustering survival model based on a mixture-of-experts framework to better capture clinical heterogeneity in survival prediction. Our approach uses a learned gating mechanism to route patients to specialized neural experts. Hence, allowing distinct disease progression patterns to be modeled within a unified framework. We evaluate and demonstrate the effectiveness of the proposed model on real-world clinical survival datasets. The experimental results show consistent improvements in survival and subtype clustering performance compared to existing deep survival models. We also evaluate the interpretability of AdaCSM with case studies. We further analyze the gating behavior to examine how expert specialization relates to clinical subgroups. An additional advantage of this design is that the computational cost of AdaCSM can be adjusted through the routing sparsity parameter 𝑘. In practice, this provides flexibility for deployment across clinical environments with different computational constraints. For example, smaller values of 𝑘 with reduced computational cost make the model suitable for resource-constrained settings, while larger values allow more experts to participate when higher computational budgets are available. Future directions include extending the model to support multimodal data fusion. Furthermore, we aim to investigate temporal routing mechanisms, such as with longitudinal recurrent layers into the gating network, to allow the model to adapt its subtype assignments as clinical biomarkers shift over time. Finally, we plan to explore the transferability of our modular experts across other neurodegenerative disorders to determine if the identified survival kernels capture generalized biological mechanisms of neurodegeneration.
Acknowledgments This work was supported in part by the NIH grants U01 AG068057, U01 AG066833, and U19 AG074879; and NSF grant 2500343.
References [1] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. 2019. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining. 2623–2631.
[2] Ahmed M Alaa and Mihaela van der Schaar. 2017. Deep multi-task gaussian processes for survival analysis with competing risks. In Proceedings of the 31st International Conference on Neural Information Processing Systems. 2326–2334. [3] J Martin Bland and Douglas G Altman. 1998. Survival probabilities (the KaplanMeier method). Bmj 317, 7172 (1998), 1572–1580. [4] Jonathan Buckley and Ian James. 1979. Linear regression with censored data. Biometrika 66, 3 (1979), 429–436. [5] Paidamoyo Chapfuwa, Chunyuan Li, Nikhil Mehta, Lawrence Carin, and Ricardo Henao. 2020. Survival cluster analysis. In Proceedings of the ACM Conference on Health, Inference, and Learning. 60–68. [6] David R Cox. 1972. Regression models and life-tables. Journal of the royal statistical society: Series B (methodological) 34, 2 (1972), 187–202. [7] Thomas R Dawber, Gilcin F Meadors, and Felix E Moore Jr. 1951. Epidemiological approaches to heart disease: the Framingham Study. American Journal of Public Health and the Nations Health 41, 3 (1951), 279–286. [8] Angela Dispenzieri, Jerry A Katzmann, Robert A Kyle, Dirk R Larson, Terry M Therneau, Colin L Colby, Raynell J Clark, Graham P Mead, Shaji Kumar, L Joseph Melton III, et al. 2012. Use of nonclonal serum immunoglobulin free light chains to predict overall survival in the general population. In Mayo Clinic Proceedings, Vol. 87. Elsevier, 517–523. [9] Thomas R Fleming and David P Harrington. 2013. Counting processes and survival analysis. John Wiley & Sons. [10] John A Hartigan and Manchek A Wong. 1979. Algorithm AS 136: A k-means clustering algorithm. Journal of the royal statistical society. series c (applied statistics) 28, 1 (1979), 100–108. [11] Bojian Hou, Hongming Li, et al. 2023. Deep clustering survival machines with interpretable expert distributions. In 2023 IEEE 20th International Symposium on Biomedical Imaging (ISBI). IEEE, 1–4. [12] Bojian Hou, Zixuan Wen, Jingxuan Bao, Richard Zhang, Boning Tong, Shu Yang, Junhao Wen, Yuhan Cui, Jason H Moore, Andrew J Saykin, et al. 2024. Interpretable deep clustering survival machines for Alzheimer’s disease subtype discovery. Medical image analysis 97 (2024), 103231. [13] Hemant Ishwaran, Thomas A Gerds, Udaya B Kogalur, Richard D Moore, Stephen J Gange, and Bryan M Lau. 2014. Random survival forests for competing risks. Biostatistics 15, 4 (2014), 757–773. [14] Hemant Ishwaran, Udaya B Kogalur, Eugene H Blackstone, and Michael S Lauer. 2008. Random survival forests. (2008). [15] Vincent Jeanselme, Brian Tom, and Jessica Barrett. 2022. Neural Survival Clustering: Non-parametric mixture of neural networks for survival clustering. In Conference on Health, Inference, and Learning. PMLR, 92–102. [16] Jared L Katzman, Uri Shaham, Alexander Cloninger, Jonathan Bates, Tingting Jiang, and Yuval Kluger. 2018. DeepSurv: personalized treatment recommender system using a Cox proportional hazards deep neural network. BMC medical research methodology 18, 1 (2018), 24. [17] William A Knaus, Frank E Harrell, Joanne Lynn, Lee Goldman, Russell S Phillips, Alfred F Connors, Neal V Dawson, William J Fulkerson, Robert M Califf, Norman Desbiens, et al. 1995. The SUPPORT prognostic model: Objective estimates of survival for seriously ill hospitalized adults. Annals of internal medicine 122, 3 (1995), 191–203. [18] Scott M Lundberg and Su-In Lee. 2017. A unified approach to interpreting model predictions. Advances in neural information processing systems 30 (2017). [19] Laura Manduchi, Ričards Marcinkevičs, Michela C Massi, Thomas Weikert, Alexander Sauter, Verena Gotta, Timothy Müller, Flavio Vasella, Marian C Neidert, Marc Pfister, et al. 2021. A deep variational approach to clustering survival data. arXiv preprint arXiv:2106.05763 (2021). [20] Chirag Nagpal, Xinyu Li, and Artur Dubrawski. 2021. Deep survival machines: Fully parametric survival regression and representation learning for censored data with competing risks. IEEE Journal of Biomedical and Health Informatics 25, 8 (2021), 3163–3175. [21] Kazi Noshin, Mary Regina Boland, Bojian Hou, Weiqing He, Victoria Lu, Li Shen, and Aidong Zhang. 2025. Integrating social determinants of health in a multi-modal deep clustering survival model for injury-risk in alzheimer’s and related dementia patients. In AAAI. [22] Jiajun Qiu and Yao andothers Hu. 2025. Deep representation learning for clustering longitudinal survival data from electronic health records. Nature Communications 16, 1 (2025), 2534. [23] Lee-Jen Wei. 1992. The accelerated failure time model: a useful alternative to the Cox regression model in survival analysis. Statistics in medicine 11, 14-15 (1992), 1871–1879. [24] Zixuan Wen, Bojian Hou, et al. 2025. Multi-Modal Deep Clustering Survival Machines for Alzheimer’s Disease Subtype Discovery. In Proc. of the IEEE/CVF Int. Conf. on Computer Vision. 2264–2272.
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts
A Appendix A.1 Data Preprocessing This section details the cohort-specific preprocessing steps used to arrive at the final statistics presented in Table 1. To ensure training stability across datasets measured in different units (e.g., days vs. years), we further apply a global max-scaling transformation to the survival times. For each dataset, every observed time 𝑡 is divided by the maximum time value present in the training split: 𝑡 𝑡𝑠𝑐𝑎𝑙𝑒𝑑 = (A.1) max(𝑡𝑡𝑟𝑎𝑖𝑛 ) A.1.1 SUPPORT. The Study to Understand Prognoses and Preferences for Outcomes and Risks of Treatments (SUPPORT) dataset consists of 9,105 patients. As each patient appears only once in the cohort, no row-level exclusions were required. Preprocessing followed these steps: • Feature Selection and Leakage Prevention: To ensure the model predicts risk from baseline patient factors rather than existing clinical scores or outcomes, we removed hospital death indicators (hospdead), existing composite severityof-illness scores (aps, sps), and subjective physician survival estimates (surv2m, surv6m). • Skewness Correction: To handle strong right-skewness and stabilize training, we applied a log-transformation to cost-related variables and the oxygenation index (PaO2 /FiO2 ). • Imputation Strategy: Missing continuous predictors were imputed using the median of the training data. For missing categorical variables, we utilized the mode of the training distribution. • Encoding and Scaling: Categorical variables—including sex, primary disease group, disease classification, income bracket, race, presence of malignancy, and functional status (SF-20)—were one-hot encoded. All continuous features were subsequently standardized to a zero mean and unit variance. • Target Labeling: The binary death column was used as the ground-truth event indicator (0 = alive, 1 = dead). A.1.2 PBC. The Primary Biliary Cholangitis (PBC) dataset contains clinical measurements with 1,945 total observations. To align this with the baseline prognostic task, we performed the following preprocessing steps: • Target Labeling: We utilized the status2 variable as the ground-truth event indicator (0 = alive, 1 = dead). The original status column, which includes liver transplantation as a separate state, was excluded to focus specifically on all-cause mortality. • Feature Encoding: Categorical clinical markers, including drug, sex, ascites, hepatomegaly, spiders, edema, and histologic stage, were transformed using one-hot encoding. Continuous features were standardized to zero mean and unit variance using the training set statistics. • Feature Dimensionality: After expanding categorical variables, the final processed dataset consists of 𝑑 = 25 input features.
• Missing Data: At baseline, we observed 9.0% missingness for serum cholesterol and 1.3% for platelets. These missing continuous predictors were imputed using the column mean of the train data prior to modeling. A.1.3 Framingham. The Framingham Heart Study dataset contains 11,627 exam records. To create a baseline prognostic model, we processed the data as follows: • Feature Selection and Leakage Prevention: To prevent data leakage, we excluded all follow-up clinical event indicators and their corresponding timestamps (e.g., Angina, Stroke, Myocardial Infarction, Hypertension, and CVD flags). The model predictors were restricted to baseline demographics and physiological measurements. • Handling Lipids: Although Total Cholesterol (TOTCHOL) was retained, the specific sub-fractions HDL and LDL were excluded from the final feature set due to 100% missingness in the baseline exam rows of the provided source. • Imputation Strategy: Missing continuous values for variables such as Glucose (8.95% missingness), BPMEDS, and BMI were imputed using the column median. For categorical variables, missing values were handled during one-hot encoding by assigning them to an all-zero vector across the dummy columns for that field. • Encoding and Scaling: Ten categorical fields—including sex, smoking status, diabetes, education level, and medical histories (e.g., previous stroke or heart disease)—were one-hot encoded. All 10 numeric features were standardized to a zero mean and unit variance. • Target Labeling: The survival outcome was defined as the time from the baseline exam to either death or the end of the study period (TIMEDTH - TIME). The DEATH column served as the binary event indicator (0 = alive, 1 = dead). A.1.4 FLCHAIN. The Medical College of Wisconsin’s Free Light Chain (FLCHAIN) dataset consists of 6,524 subjects. As the source data provides a single baseline record per individual, no row-level filtering was required. The following preprocessing steps were applied: • Feature Representation: We utilized eight predictors, including demographics (age, sex), clinical markers (kappa, lambda, creatinine), and research-specific variables (sample year, FLC group index, MGUS status). Notably, sex and FLC group were treated as numerical/ordinal inputs rather than undergoing one-hot encoding. • Missing Data and Scaling: No missing values were detected in the predictor columns for this cohort. Consistent with the other datasets, a StandardScaler was applied to all eight features to normalize the input distribution based on train data. • Time-to-Event Processing: To prevent zero-length followup intervals during model optimization, a +1 day shift was applied to the follow-up duration (futime). • Target Labeling: The survival endpoint was defined as all-cause mortality, using the death column as the binary event indicator (0 = alive, 1 = dead).
Zhuang et al.
A.2
Optuna Hyperparameter Tuning
Hyperparameters for all models were selected using Optuna. For each method and dataset, we defined a model-specific search space over key architectural and optimization parameters, such as learning rate, hidden layer size, batch size, dropout, and method-specific settings. Optuna was then used to search for the configuration that maximized the validation C-Index. The best hyperparameter settings for AdaCSM, Cox PH, DeepSurv, DSM, and DCSM are reported in Tables 1-5, respectively.
A.3
Subtype Clustering
Table A.6 shows the LogRank results of different models. AdaCSM achieves the highest LogRank statistic on all four clinical cohorts. In particular, the improvement over DCSM is consistent across all datasets, suggesting that adaptive expert specialization enhances subtype discovery beyond the original clustering survival formulation.
A.4
Prognostic Accuracy
Table A.7 presents the global C-index performance across all four cohorts. These results demonstrate that AdaCSM maintains discriminative accuracy comparable to state-of-the-art dense models like DeepSurv and DCSM. This confirms that the transition to a sparse mixture-of-experts architecture provides significant gains in subtype separation (Table A.6) and interpretability without compromising the model’s fundamental ability to rank individual patient risk.
A.5
Time-dependent metrics
To evaluate model reliability over the clinical follow-up period, we report time-dependent C-index (Table A.8) and Brier Scores (Table A.9) at the 25𝑡ℎ , 50𝑡ℎ , and 75𝑡ℎ event quantiles. AdaCSM exhibits stable longitudinal performance, particularly in the later stages of follow-up. The corresponding bar plots in Figures A.1 and A.2 provide a visual comparison of these metrics, highlighting that our sparse routing mechanism maintains predictive consistency across varying survival horizons.
A.6
Prognostic vs. Subtyping
A common challenge in survival analysis is the trade-off between prognostic precision (C-Index) and group-level subtyping (LogRank). Traditional models often achieve superior ranking capability, while clustering models may sacrifice individual accuracy for improved group separation. To evaluate this trade-off, we plot the Pareto frontier of all models across the datasets (Figure A.3). As illustrated, AdaCSM remains on the frontier, achieving the highest LogRank statistics while maintaining competitive C-Index performance. This suggests that the Mixture of Experts (MoE) architecture successfully navigates the trade-off by using specialized experts to define distinct clinical subtypes without collapsing the global predictive ranking.
A.7
Expert Assignment
To demonstrate that AdaCSM captures biologically and clinically relevant patterns across diverse disease domains, we analyze expert assignment logic for the PBC (liver disease), SUPPORT (acute care), and FLCHAIN (monoclonal gammopathy) cohorts in Figures A.4, A.5, and A.6, respectively.
A.8
Expert Interpretability across datasets
We validate the intrinsic interpretability of AdaCSM by comparing it with SHAP (Shapley Additive exPlanations) [18]. While SHAP provides a post-hoc estimate of feature importance for a fixed model, AdaCSM’s expert profiles represent the model’s internal specialized logic. Figure A.7 shows a strong Spearman rank correlation (𝑟 = 0.816) between the features that define an expert’s profile and those identified by SHAP as most influential for the patients routed to that expert. For example, in the Framingham cohort, both methods identify total cholesterol, systolic blood pressure, and age as the primary drivers for high-risk mortality among the cardiovascular cohort. This alignment demonstrates that AdaCSM’s expert specialization reliably learns feature importance of the underlying clinical data, providing a two-layered interpretability framework that is both structurally sound and post-hoc verifiable.
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts
Table A.1: Best hyperparameters selected for AdaCSM on each clinical survival dataset using Optuna tuning. Hyperparameter
SUPPORT
PBC
Framingham
FLCHAIN
Learning rate Discount Hidden layers Number of experts Batch size Dropout Gate dropout Gate temperature Load balance 𝜆
1.50e-4 0.6655 [50] 4 16 0.0294 0.0609 3.0002 0.0294
4.68e-4 0.8316 [100] 32 16 0.0477 0.2708 0.1191 0.0952
6.21e-3 0.5659 [100] 32 100 0.0544 0.0982 0.8553 0.0725
3.38e-3 0.9413 [50] 32 100 0.1447 0.0093 0.8817 0.0836
Table A.2: Best hyperparameters selected for Cox Proportional Hazards (Cox PH) on each clinical survival dataset using Optuna tuning. Hyperparameter Penalizer L1 ratio
SUPPORT
PBC
Framingham
FLCHAIN
0.0096 0.3594
0.0894 0.6936
0.0019 0.4612
0.0097 0.3861
Table A.3: Best hyperparameters for DeepSurv on clinical survival datasets. Hyperparameter
SUPPORT
PBC
FRAMINGHAM
FLCHAIN
Learning Rate Hidden Layers Batch Size
1.31e-4 [100] 100
2.14e-4 [50,50] 100
7.06e-4 [100] 16
5.88e-3 [50] 100
Table A.4: Best hyperparameters for Deep Survival Machines (DSM) on clinical survival datasets. Hyperparameter
SUPPORT
PBC
Framingham
FLCHAIN
Learning Rate Hidden Layers Batch Size Discount Distribution
1.67e-4 [50, 50] 32 0.8990 Weibull
9.04e-3 [50, 50] 32 0.4433 Weibull
1.54e-3 [50, 50] 16 0.4655 Weibull
3.79e-5 [50] 128 0.6438 Weibull
Table A.5: Best hyperparameters for Deep Clustering Survival Machines (DCSM) on clinical survival datasets. Hyperparameter
SUPPORT
PBC
Framingham
FLCHAIN
Learning rate Discount Hidden layers Batch size
1.33e-4 0.4092 [50] 100
1.60e-4 0.7662 [50] 100
4.01e-4 0.6950 [50, 50] 100
8.07e-3 0.7760 [100] 100
Table A.6: LogRank statistic (𝜒 2 ↑, mean ± std) measuring survival separation between discovered subtypes on clinical survival cohorts. For each dataset, the best performance is in bold, while the second-best is underlined. Model Cox PH DeepSurv DSM DCSM AdaCSM AdaCSM (Sparse top-2)
SUPPORT
PBC
FRAMINGHAM
FLCHAIN
507.38 ± 52.86 584.00 ± 54.16 879.19 ± 117.15 757.65 ± 46.64 1047.17 ± 114.04 1028.72 ± 121.93
192.17 ± 17.61 85.15 ± 115.00 229.78 ± 49.58 292.54 ± 25.96 311.80 ± 19.30 316.42 ± 38.88
576.42 ± 21.97 590.78 ± 22.22 613.96 ± 24.76 568.74 ± 44.43 667.83 ± 44.24 638.72 ± 44.60
401.75 ± 30.51 331.51 ± 187.13 457.62 ± 116.75 563.08 ± 95.05 586.56 ± 38.39 569.81 ± 28.61
Zhuang et al.
Table A.7: Global C-Index (↑ mean ± std) comparison on clinical survival cohorts. For each dataset, the best performance is in bold, while the second-best is underlined. Model Cox PH DeepSurv DSM DCSM AdaCSM (Sparse top-2) AdaCSM
SUPPORT
PBC
FRAMINGHAM
FLCHAIN
0.8414 ± 0.0037 0.8276 ± 0.0028 0.8239 ± 0.0029 0.8151 ± 0.0033 0.8289 ± 0.0052 0.8284 ± 0.0060
0.8106 ± 0.0152 0.8182 ± 0.0151 0.8090 ± 0.0232 0.8153 ± 0.0138 0.8218 ± 0.0099 0.8181 ± 0.0142
0.7420 ± 0.0093 0.7415 ± 0.0083 0.7390 ± 0.0080 0.7293 ± 0.0064 0.7345 ± 0.0036 0.7300 ± 0.0108
0.7812 ± 0.0165 0.7891 ± 0.0054 0.7881 ± 0.0038 0.7780 ± 0.0079 0.7845 ± 0.0099 0.7806 ± 0.0088
Table A.8: C-Index (↑ mean ± std) across different time quantiles (25𝑡ℎ , 50𝑡ℎ , and 75𝑡ℎ percentiles) on clinical survival cohorts. Model / Quantile
SUPPORT
PBC
FRAMINGHAM
FLCHAIN
25th Percentile Cox PH DeepSurv DSM DCSM AdaCSM AdaCSM (Top-2)
0.9368 ± 0.0013 0.9609 ± 0.0028 0.9596 ± 0.0027 0.9116 ± 0.0022 0.9267 ± 0.0044 0.9229 ± 0.0027
0.8630 ± 0.0320 0.8754 ± 0.0179 0.8775 ± 0.0190 0.8777 ± 0.0211 0.8800 ± 0.0300 0.8814 ± 0.0296
0.7422 ± 0.0212 0.7435 ± 0.0205 0.7412 ± 0.0161 0.7280 ± 0.0213 0.7320 ± 0.0285 0.7365 ± 0.0183
0.7770 ± 0.0163 0.7887 ± 0.0155 0.7913 ± 0.0119 0.7818 ± 0.0104 0.7881 ± 0.0130 0.7884 ± 0.0097
50th Percentile Cox PH DeepSurv DSM DCSM AdaCSM AdaCSM (Top-2)
0.9441 ± 0.0004 0.9567 ± 0.0020 0.9494 ± 0.0019 0.9260 ± 0.0009 0.9383 ± 0.0019 0.9365 ± 0.0017
0.8445 ± 0.0195 0.8525 ± 0.0182 0.8447 ± 0.0246 0.8499 ± 0.0156 0.8555 ± 0.0197 0.8539 ± 0.0256
0.7395 ± 0.0148 0.7408 ± 0.0126 0.7405 ± 0.0096 0.7294 ± 0.0132 0.7296 ± 0.0171 0.7350 ± 0.0088
0.7804 ± 0.0129 0.7896 ± 0.0057 0.7897 ± 0.0044 0.7798 ± 0.0058 0.7840 ± 0.0093 0.7863 ± 0.0082
75th Percentile Cox PH DeepSurv DSM DCSM AdaCSM AdaCSM (Top-2)
0.8382 ± 0.0026 0.8517 ± 0.0027 0.8690 ± 0.0036 0.8492 ± 0.0028 0.8622 ± 0.0045 0.8638 ± 0.0058
0.8217 ± 0.0198 0.8306 ± 0.0139 0.8254 ± 0.0234 0.8248 ± 0.0120 0.8346 ± 0.0113 0.8334 ± 0.0137
0.7412 ± 0.0131 0.7416 ± 0.0111 0.7384 ± 0.0100 0.7300 ± 0.0104 0.7306 ± 0.0124 0.7350 ± 0.0077
0.7807 ± 0.0157 0.7877 ± 0.0067 0.7872 ± 0.0061 0.7783 ± 0.0089 0.7804 ± 0.0097 0.7842 ± 0.0105
SUPPORT
PBC Models
Cox PH
FRAMINGHAM DeepSurv
DSM
DCSM
FLCHAIN
AdaCSM (Sparse top−2)
AdaCSM 0.8
0.8 0.6
0.4
0.6 0.4
0.0
0.0 25%
50%
Event Quantiles
75%
0.4
0.2
0.2
0.2
0.6
C−Index
0.6
C−Index
C−Index
C−Index
0.8
0.2
0.0 25%
50%
Event Quantiles
75%
0.4
0.0 25%
50%
Event Quantiles
75%
25%
50%
75%
Event Quantiles
Figure A.1: Bar plots of C-Index at the 25%, 50%, and 75% event quantiles for six survival models on the SUPPORT, PBC, FRAMINGHAM, and FLCHAIN datasets. Error bars represent 95% confidence intervals.
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts
Table A.9: Brier Score (↓ mean ± std) across different time quantiles (25𝑡ℎ , 50𝑡ℎ , and 75𝑡ℎ percentiles) on clinical survival cohorts. Model / Quantile
SUPPORT
PBC
FRAMINGHAM
FLCHAIN
25th Percentile Cox PH DeepSurv DSM DCSM AdaCSM (Top-2) AdaCSM
0.0621 ± 0.0008 0.0388 ± 0.0014 0.0392 ± 0.0029 0.1336 ± 0.0012 0.1308 ± 0.0018 0.1338 ± 0.0018
0.0802 ± 0.0100 0.0760 ± 0.0067 0.0763 ± 0.0098 0.1040 ± 0.0007 0.1019 ± 0.0024 0.1018 ± 0.0024
0.0713 ± 0.0022 0.0714 ± 0.0023 0.0744 ± 0.0021 0.0811 ± 0.0007 0.0773 ± 0.0011 0.0782 ± 0.0015
0.0593 ± 0.0026 0.0583 ± 0.0015 0.0587 ± 0.0020 0.0686 ± 0.0027 0.0686 ± 0.0024 0.0695 ± 0.0023
50th Percentile Cox PH DeepSurv DSM DCSM AdaCSM (Top-2) AdaCSM
0.0042 ± 0.0015 0.0059 ± 0.0016 0.0134 ± 0.0007 0.2007 ± 0.0034 0.1925 ± 0.0071 0.2016 ± 0.0071
0.1221 ± 0.0113 0.1202 ± 0.0168 0.1130 ± 0.0181 0.1767 ± 0.0020 0.1703 ± 0.0094 0.1700 ± 0.0097
0.1215 ± 0.0022 0.1219 ± 0.0025 0.1280 ± 0.0034 0.1472 ± 0.0014 0.1341 ± 0.0032 0.1363 ± 0.0040
0.1044 ± 0.0107 0.0987 ± 0.0010 0.1024 ± 0.0022 0.1226 ± 0.0081 0.1230 ± 0.0079 0.1264 ± 0.0076
75th Percentile Cox PH DeepSurv DSM DCSM AdaCSM (Top-2) AdaCSM
0.1163 ± 0.0026 0.1087 ± 0.0024 0.1102 ± 0.0012 0.2366 ± 0.0035 0.2197 ± 0.0075 0.2283 ± 0.0074
0.1487 ± 0.0240 0.1431 ± 0.0127 0.1560 ± 0.0310 0.1265 ± 0.0102 0.2135 ± 0.0262 0.2129 ± 0.0265
0.1539 ± 0.0038 0.1545 ± 0.0034 0.1662 ± 0.0044 0.1988 ± 0.0020 0.1732 ± 0.0048 0.1746 ± 0.0058
0.1392 ± 0.0264 0.1262 ± 0.0033 0.1312 ± 0.0035 0.1638 ± 0.0147 0.1651 ± 0.0147 0.1718 ± 0.0141
SUPPORT
PBC Models
Cox PH
FRAMINGHAM DeepSurv
DSM
DCSM
FLCHAIN
AdaCSM (Sparse top−2)
AdaCSM 0.2
0.2 0.2
0.0
0.0 25%
50%
Event Quantiles
75%
Brier Score
Brier Score
Brier Score
Brier Score
0.2
0.0 25%
50%
Event Quantiles
75%
0.0 25%
50%
Event Quantiles
75%
25%
50%
75%
Event Quantiles
Figure A.2: Bar plots of Brier Score at the 25%, 50%, and 75% event quantiles for six survival models on the SUPPORT, PBC, FRAMINGHAM, and FLCHAIN datasets. Error bars represent 95% confidence intervals.
Zhuang et al.
Figure A.3: Pareto frontier of Subtype Separation (LogRank) vs. Prognostic Accuracy (C-Index) on the datasets.
Figure A.4: Expert assignment logic in the PBC dataset. The gating network dynamically routes patients to specific experts based on critical liver function markers. For instance, the left figure shows a clear transition in expert preference as Serum Bilirubin levels increase, indicating that specialized experts are dedicated to high-risk cholestatic profiles.
Expert-Driven Survival Machines: Improving Stratification and Interpretability in Multiple Clinical Cohorts
Figure A.5: Expert assignment logic in the SUPPORT dataset. The gating network isolates patients with acute neurological distress (low GCS scores), routing them to specialized experts who prioritize physiological instability markers.
Figure A.6: Expert assignment logic in the FLCHAIN dataset. The model identifies distinct routing patterns for patients with abnormal immunoglobulin light chain ratios (kappa/lambda), which are critical biomarkers for plasma cell dyscrasias.
Figure A.7: Alignment between intrinsic expert importance and post-hoc SHAP values. The high correlation confirms that the gating network’s routing logic is consistent with established feature attribution methods.