[Paper Review] A Variational Autoencoder for Heterogeneous Temporal and Longitudinal Data

Mine Öğretir, Siddharth Ramchandran+2 more|arXiv (Cornell University)|Apr 20, 2022

Machine Learning in Healthcare4 citations

TL;DR

This paper proposes the Heterogeneous Longitudinal Variational Autoencoder (HL-VAE), a deep generative model that jointly models continuous, count, categorical, and ordinal longitudinal data using a multi-output additive Gaussian process prior and flexible likelihoods. It enables efficient, end-to-end inference for imputing missing values and predicting unseen time points, outperforming existing methods in both imputation and long-term prediction on simulated and clinical datasets.

ABSTRACT

The variational autoencoder (VAE) is a popular deep latent variable model used to analyse high-dimensional datasets by learning a low-dimensional latent representation of the data. It simultaneously learns a generative model and an inference network to perform approximate posterior inference. Recently proposed extensions to VAEs that can handle temporal and longitudinal data have applications in healthcare, behavioural modelling, and predictive maintenance. However, these extensions do not account for heterogeneous data (i.e., data comprising of continuous and discrete attributes), which is common in many real-life applications. In this work, we propose the heterogeneous longitudinal VAE (HL-VAE) that extends the existing temporal and longitudinal VAEs to heterogeneous data. HL-VAE provides efficient inference for high-dimensional datasets and includes likelihood models for continuous, count, categorical, and ordinal data while accounting for missing observations. We demonstrate our model's efficacy through simulated as well as clinical datasets, and show that our proposed model achieves competitive performance in missing value imputation and predictive accuracy.

Motivation & Objective

• To address the gap in existing variational autoencoders that cannot handle heterogeneous data types—such as continuous, count, categorical, and ordinal variables—simultaneously in longitudinal settings.
• To develop a model that accounts for temporal correlations across time points and subjects while handling missing data that are missing completely at random.
• To enable efficient inference for high-dimensional, real-world datasets with mixed data types and auxiliary covariates such as time and subject-specific factors.
• To demonstrate that assigning appropriate likelihood models to different data types significantly improves performance in imputation and prediction tasks.

Proposed method

• HL-VAE uses a multi-output additive Gaussian process prior to model the structured dynamics of latent variables across time and subjects, capturing correlations between latent embeddings.
• It employs a deep neural network-based decoder with separate likelihood models for each data type: Gaussian for continuous, Poisson for count, categorical logit for discrete, and ordinal logit for ordered categorical data.
• The inference network performs amortized variational inference, approximating the posterior over latent variables using stochastic gradient descent.
• The model integrates auxiliary covariates (e.g., time, gender, disease status) into the latent space via the GP prior, enabling conditional generation and improved representation learning.
• It supports missing data by marginalizing over unobserved values during training and inference, using likelihood-specific likelihood functions.
• The framework supports end-to-end training with backpropagation, enabling scalable learning on high-dimensional datasets.

Experimental results

Research questions

• RQ1 Can a variational autoencoder effectively model heterogeneous longitudinal data with mixed data types (continuous, count, categorical, ordinal) while handling missing values?
• RQ2 How does the inclusion of appropriate likelihood models for different data types affect imputation and prediction performance in longitudinal settings?
• RQ3 Does the use of a multi-output additive Gaussian process prior improve modeling of temporal and inter-subject correlations compared to independent latent representations?
• RQ4 To what extent does HL-VAE outperform existing VAEs and baselines in missing value imputation and long-term prediction on real-world clinical data?
• RQ5 How does the model’s performance scale with increasing missing data rates and data heterogeneity?

Key findings

• On the PPMI dataset with 50% missing data, HL-VAE achieved a normalized root mean squared error (NRMSE) of 0.091 ± 0.003 for numerical predictions, outperforming L-VAE, which had a higher error and a steeper decline in performance as missingness increased.
• For categorical data, HL-VAE achieved an accuracy error of 0.085 ± 0.001, significantly outperforming competing models, demonstrating the benefit of modeling categorical variables explicitly.
• In ordinal prediction, HL-VAE achieved a displacement error of 0.085 ± 0.001, showing robustness and improved accuracy over baselines.
• The model maintained stable performance across increasing missingness levels, with a slower decline in predictive log-likelihood compared to L-VAE, indicating better generalization under data scarcity.
• Including categorical features in the training process improved the prediction performance of numerical features, highlighting the benefit of joint modeling of heterogeneous data.
• HL-VAE achieved competitive performance in both imputation and long-term prediction tasks, with the best results across all data types on both simulated and real-world clinical datasets.