Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Deep Variational Contrastive Learning for Joint Risk Stratification and Time-to-Event Estimation

Pinar Erbil, Alberto Archetti, Eugenio Lomurno, Matteo Matteucci

TL;DR

The paper tackles the interpretability-versus-performance tension in survival analysis by introducing CONVERSE, a framework that fuses variational autoencoders with multi-view contrastive clustering and an ensemble of survival heads to yield interpretable risk subgroups and accurate time-to-event estimates. It outputs discrete-time survival probabilities $\hat{p}_{i,t}$ while enabling cluster-specific risk dynamics and subpopulation insights. Across four benchmark datasets, CONVERSE achieves competitive or superior concordance indices compared with neural baselines and improves calibration over clustering methods, with Kaplan–Meier curves and feature importances evidencing clinically meaningful stratification. The approach offers an end-to-end, self-paced training pipeline with flexible architectural choices, enabling practical deployment for data-informed clinical decision support.

Abstract

Survival analysis is essential for clinical decision-making, as it allows practitioners to estimate time-to-event outcomes, stratify patient risk profiles, and guide treatment planning. Deep learning has revolutionized this field with unprecedented predictive capabilities but faces a fundamental trade-off between performance and interpretability. While neural networks achieve high accuracy, their black-box nature limits clinical adoption. Conversely, deep clustering-based methods that stratify patients into interpretable risk groups typically sacrifice predictive power. We propose CONVERSE (CONtrastive Variational Ensemble for Risk Stratification and Estimation), a deep survival model that bridges this gap by unifying variational autoencoders with contrastive learning for interpretable risk stratification. CONVERSE combines variational embeddings with multiple intra- and inter-cluster contrastive losses. Self-paced learning progressively incorporates samples from easy to hard, improving training stability. The model supports cluster-specific survival heads, enabling accurate ensemble predictions. Comprehensive evaluation on four benchmark datasets demonstrates that CONVERSE achieves competitive or superior performance compared to existing deep survival methods, while maintaining meaningful patient stratification.

Deep Variational Contrastive Learning for Joint Risk Stratification and Time-to-Event Estimation

TL;DR

The paper tackles the interpretability-versus-performance tension in survival analysis by introducing CONVERSE, a framework that fuses variational autoencoders with multi-view contrastive clustering and an ensemble of survival heads to yield interpretable risk subgroups and accurate time-to-event estimates. It outputs discrete-time survival probabilities while enabling cluster-specific risk dynamics and subpopulation insights. Across four benchmark datasets, CONVERSE achieves competitive or superior concordance indices compared with neural baselines and improves calibration over clustering methods, with Kaplan–Meier curves and feature importances evidencing clinically meaningful stratification. The approach offers an end-to-end, self-paced training pipeline with flexible architectural choices, enabling practical deployment for data-informed clinical decision support.

Abstract

Survival analysis is essential for clinical decision-making, as it allows practitioners to estimate time-to-event outcomes, stratify patient risk profiles, and guide treatment planning. Deep learning has revolutionized this field with unprecedented predictive capabilities but faces a fundamental trade-off between performance and interpretability. While neural networks achieve high accuracy, their black-box nature limits clinical adoption. Conversely, deep clustering-based methods that stratify patients into interpretable risk groups typically sacrifice predictive power. We propose CONVERSE (CONtrastive Variational Ensemble for Risk Stratification and Estimation), a deep survival model that bridges this gap by unifying variational autoencoders with contrastive learning for interpretable risk stratification. CONVERSE combines variational embeddings with multiple intra- and inter-cluster contrastive losses. Self-paced learning progressively incorporates samples from easy to hard, improving training stability. The model supports cluster-specific survival heads, enabling accurate ensemble predictions. Comprehensive evaluation on four benchmark datasets demonstrates that CONVERSE achieves competitive or superior performance compared to existing deep survival methods, while maintaining meaningful patient stratification.
Paper Structure (14 sections, 17 equations, 2 figures, 2 tables)

This paper contains 14 sections, 17 equations, 2 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Method
  4. CONVERSE Architecture
  5. Representation Learning and Autoencoders
  6. Clustering and Contrastive Learning
  7. Ensemble Survival Prediction
  8. Training Procedure
  9. Experiments
  10. Datasets and Baseline Models
  11. Experimental Procedure and Hyperparameter Optimization
  12. Results
  13. Interpretability
  14. Conclusion

Figures (2)

  • Figure 1: Overall architecture of CONVERSE. The model consists of three components: (i) variational autoencoder(s) for learning latent representations, (ii) clustering module with contrastive learning operating in the latent space, and (iii) ensemble survival prediction heads. Dashed lines indicate optional pathways depending on hyperparameter selection.
  • Figure 2: Representation visualizations from trained CONVERSE architectures. The first column collects UMAP visualizations of the latent spaces with $K$-means cluster assignments; the second column shows Kaplan--Meier curves for the resulting clusters; the third column presents Gini-based feature importance to highlight the covariates driving each group.