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

TL;DR

CONVERSE is a novel deep survival analysis model that combines variational autoencoders and contrastive learning to achieve interpretable patient risk stratification without sacrificing predictive accuracy.

Contribution

It introduces a unified framework that integrates variational autoencoders with contrastive learning for improved interpretability and performance in survival analysis.

Findings

Achieves competitive or superior performance on benchmark datasets.

Provides meaningful patient risk stratification.

Supports cluster-specific survival predictions.

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.