AAAI2024

Structured Probabilistic Coding

Dou Hu, Lingwei Wei, Yaxin Liu, Wei Zhou, Songlin Hu

Abstract

Survival analysis is critical in many real-world do-mains where predicting time-to-event outcomes under censoring is essential. While deep learning has advanced survival modeling, existing methods struggle to jointly optimize discrimination, cali-bration, and uncertainty quantification. Probabilistic approaches like variational autoencoders offer uncertainty estimates but suffer from information loss due to encoder-decoder architectures and lack principled treatment of censored data. We propose Survival Structured Probabilistic Coding (Survival-SPC), a novel framework with two key mechanisms. First, an encoder-only architecture directly encodes features to probabilistic hazard representations, reducing information loss while preserving uncer-tainty. Second, censoring-aware structured regularization lever-ages partial information from censored observations to encourage diversity in latent representations. Unlike previous approaches, Survival-SPC enables efficient end-to-end training of calibrated survival distributions. Comprehensive experiments on six real-world datasets demonstrate superior performance, improving concordance index by up to 4.2% over strongest baselines while providing well-calibrated estimates. The method shows particular advantages under limited data and high censoring, establishing a robust solution for clinical applications. Code is available at Suvival-SPC Repository.