[Paper Review] Longitudinal Variational Autoencoder

Siddharth Ramchandran, Gleb Tikhonov+3 more|arXiv (Cornell University)|Jun 17, 2020

Gaussian Processes and Bayesian Inference13 references4 citations

TL;DR

This paper proposes the Longitudinal Variational Autoencoder (L-VAE), a deep generative model that uses a multi-output additive Gaussian process prior to learn structured, low-dimensional representations from longitudinal data with missing values. By deriving a tighter KL divergence upper bound and enabling mini-batch SGD training, L-VAE captures both shared and subject-specific temporal dynamics, outperforming state-of-the-art methods in data imputation, reconstruction, and long-term prediction on synthetic and clinical datasets with up to 80% missing data.

ABSTRACT

Longitudinal datasets measured repeatedly over time from individual subjects, arise in many biomedical, psychological, social, and other studies. A common approach to analyse high-dimensional data that contains missing values is to learn a low-dimensional representation using variational autoencoders (VAEs). However, standard VAEs assume that the learnt representations are i.i.d., and fail to capture the correlations between the data samples. We propose the Longitudinal VAE (L-VAE), that uses a multi-output additive Gaussian process (GP) prior to extend the VAE's capability to learn structured low-dimensional representations imposed by auxiliary covariate information, and derive a new KL divergence upper bound for such GPs. Our approach can simultaneously accommodate both time-varying shared and random effects, produce structured low-dimensional representations, disentangle effects of individual covariates or their interactions, and achieve highly accurate predictive performance. We compare our model against previous methods on synthetic as well as clinical datasets, and demonstrate the state-of-the-art performance in data imputation, reconstruction, and long-term prediction tasks.

Motivation & Objective

• To address the limitation of standard VAEs in capturing temporal correlations and subject-specific dynamics in longitudinal datasets with missing values.
• To develop a scalable, efficient inference method for structured latent spaces using additive Gaussian process priors.
• To enable disentanglement of individual covariate effects and their interactions in the latent space.
• To improve predictive performance in data imputation, reconstruction, and long-term forecasting on high-dimensional longitudinal data.
• To provide a generative model that incorporates auxiliary covariates beyond time, such as age, sex, and clinical indicators.

Proposed method

• L-VAE employs a multi-output additive Gaussian process (GP) prior to model structured dependencies in the latent space, capturing both shared and individual-specific temporal dynamics.
• It derives a novel, tighter upper bound for the KL divergence between the approximate posterior and the GP prior, improving variational inference stability and accuracy.
• The method enables mini-batch stochastic gradient descent (SGD) training by exploiting the structure of additive GP covariance functions, enhancing scalability to large datasets.
• The model jointly learns low-dimensional latent representations that evolve smoothly over time, conditioned on auxiliary covariates such as age, sex, and time to event.
• It supports end-to-end training via backpropagation, with the decoder reconstructing observed data from the latent codes and the encoder inferring posterior distributions.
• A Bayes classifier is used for downstream tasks like mortality prediction, leveraging marginal log-likelihoods computed from the generative model.

Experimental results

Research questions

• RQ1 Can a variational autoencoder effectively model temporal correlations in longitudinal data with missing values by incorporating auxiliary covariate information?
• RQ2 Does the use of an additive GP prior lead to a tighter and more computationally efficient variational inference bound compared to standard VAEs?
• RQ3 Can the model disentangle the effects of individual covariates and their interactions in the latent space while maintaining high reconstruction and prediction accuracy?
• RQ4 How does L-VAE compare to state-of-the-art methods in data imputation, reconstruction, and long-term prediction on real-world clinical datasets with high missingness?
• RQ5 Can the model scale efficiently to large longitudinal datasets using mini-batch SGD, unlike previous GP-based VAEs that require full-batch training?

Key findings

• On the PhysioNet Challenge 2012 dataset with 80% missing data, L-VAE achieved an AUROC of 0.888 for in-hospital mortality prediction, outperforming GP-VAE, BRITS, and GRUI-GAN.
• With a latent dimension of 32, L-VAE achieved the lowest mean squared error (MSE) of 0.025 ± 0.0015 in reconstruction tasks, surpassing GPPVAE (0.057 ± 0.003) and VRNN (0.049 ± 0.004).
• The model demonstrated superior performance in long-term prediction and data imputation, particularly in capturing complex temporal dynamics and subject-specific patterns.
• The derived KL divergence upper bound is theoretically tighter than the free-form variational bound of Titsias (2009), improving variational inference quality.
• Mini-batch SGD training was successfully enabled through the structured additive GP prior, allowing efficient scaling to large datasets unlike full-epoch training in GPPVAE and GP-VAE.
• L-VAE effectively disentangled the effects of auxiliary covariates such as age, sex, and time to mortality, enabling interpretable and accurate downstream classification.