Temporal Supervised Contrastive Learning for Modeling Patient Risk Progression

Shahriar Noroozizadeh, Jeremy C. Weiss, George H. Chen
Proceedings of the 3rd Machine Learning for Health Symposium, PMLR 225:403-427, 2023.

Abstract

We consider the problem of predicting how the likelihood of an outcome of interest for a patient changes over time as we observe more of the patient’s data. To solve this problem, we propose a supervised contrastive learning framework that learns an embedding representation for each time step of a patient time series. Our framework learns the embedding space to have the following properties: (1){~}nearby points in the embedding space have similar predicted class probabilities, (2){~}adjacent time steps of the same time series map to nearby points in the embedding space, and (3){~}time steps with very different raw feature vectors map to far apart regions of the embedding space. To achieve property (3), we employ a nearest neighbor pairing mechanism in the raw feature space. This mechanism also serves as an alternative to “data augmentation”, a key ingredient of contrastive learning, which lacks a standard procedure that is adequately realistic for clinical tabular data, to our knowledge. We demonstrate that our approach outperforms state-of-the-art baselines in predicting mortality of septic patients (MIMIC-III dataset) and tracking progression of cognitive impairment (ADNI dataset). Our method also consistently recovers the correct synthetic dataset embedding structure across experiments, a feat not achieved by baselines. Our ablation experiments show the pivotal role of our nearest neighbor pairing.

Cite this Paper

BibTeX
@InProceedings{pmlr-v225-noroozizadeh23a, title = {Temporal Supervised Contrastive Learning for Modeling Patient Risk Progression}, author = {Noroozizadeh, Shahriar and Weiss, Jeremy C. and Chen, George H.}, booktitle = {Proceedings of the 3rd Machine Learning for Health Symposium}, pages = {403--427}, year = {2023}, editor = {Hegselmann, Stefan and Parziale, Antonio and Shanmugam, Divya and Tang, Shengpu and Asiedu, Mercy Nyamewaa and Chang, Serina and Hartvigsen, Tom and Singh, Harvineet}, volume = {225}, series = {Proceedings of Machine Learning Research}, month = {10 Dec}, publisher = {PMLR}, pdf = {https://proceedings.mlr.press/v225/noroozizadeh23a/noroozizadeh23a.pdf}, url = {https://proceedings.mlr.press/v225/noroozizadeh23a.html}, abstract = {We consider the problem of predicting how the likelihood of an outcome of interest for a patient changes over time as we observe more of the patient’s data. To solve this problem, we propose a supervised contrastive learning framework that learns an embedding representation for each time step of a patient time series. Our framework learns the embedding space to have the following properties: (1){~}nearby points in the embedding space have similar predicted class probabilities, (2){~}adjacent time steps of the same time series map to nearby points in the embedding space, and (3){~}time steps with very different raw feature vectors map to far apart regions of the embedding space. To achieve property (3), we employ a nearest neighbor pairing mechanism in the raw feature space. This mechanism also serves as an alternative to “data augmentation”, a key ingredient of contrastive learning, which lacks a standard procedure that is adequately realistic for clinical tabular data, to our knowledge. We demonstrate that our approach outperforms state-of-the-art baselines in predicting mortality of septic patients (MIMIC-III dataset) and tracking progression of cognitive impairment (ADNI dataset). Our method also consistently recovers the correct synthetic dataset embedding structure across experiments, a feat not achieved by baselines. Our ablation experiments show the pivotal role of our nearest neighbor pairing.} }
Endnote
%0 Conference Paper %T Temporal Supervised Contrastive Learning for Modeling Patient Risk Progression %A Shahriar Noroozizadeh %A Jeremy C. Weiss %A George H. Chen %B Proceedings of the 3rd Machine Learning for Health Symposium %C Proceedings of Machine Learning Research %D 2023 %E Stefan Hegselmann %E Antonio Parziale %E Divya Shanmugam %E Shengpu Tang %E Mercy Nyamewaa Asiedu %E Serina Chang %E Tom Hartvigsen %E Harvineet Singh %F pmlr-v225-noroozizadeh23a %I PMLR %P 403--427 %U https://proceedings.mlr.press/v225/noroozizadeh23a.html %V 225 %X We consider the problem of predicting how the likelihood of an outcome of interest for a patient changes over time as we observe more of the patient’s data. To solve this problem, we propose a supervised contrastive learning framework that learns an embedding representation for each time step of a patient time series. Our framework learns the embedding space to have the following properties: (1){~}nearby points in the embedding space have similar predicted class probabilities, (2){~}adjacent time steps of the same time series map to nearby points in the embedding space, and (3){~}time steps with very different raw feature vectors map to far apart regions of the embedding space. To achieve property (3), we employ a nearest neighbor pairing mechanism in the raw feature space. This mechanism also serves as an alternative to “data augmentation”, a key ingredient of contrastive learning, which lacks a standard procedure that is adequately realistic for clinical tabular data, to our knowledge. We demonstrate that our approach outperforms state-of-the-art baselines in predicting mortality of septic patients (MIMIC-III dataset) and tracking progression of cognitive impairment (ADNI dataset). Our method also consistently recovers the correct synthetic dataset embedding structure across experiments, a feat not achieved by baselines. Our ablation experiments show the pivotal role of our nearest neighbor pairing.
APA
Noroozizadeh, S., Weiss, J.C. & Chen, G.H.. (2023). Temporal Supervised Contrastive Learning for Modeling Patient Risk Progression. Proceedings of the 3rd Machine Learning for Health Symposium, in Proceedings of Machine Learning Research 225:403-427 Available from https://proceedings.mlr.press/v225/noroozizadeh23a.html.

Related Material