[Paper Review] Learning Conditional Variational Autoencoders with Missing Covariates

Siddharth Ramchandran, Gleb Tikhonov+3 more|arXiv (Cornell University)|Mar 2, 2022

Machine Learning in Healthcare4 citations

TL;DR

This paper proposes a novel method to train conditional variational autoencoders (CVAEs) with missing values in auxiliary covariates by introducing an amortized variational inference framework that marginalizes uncertainty over missing covariates while maximizing the evidence lower bound (ELBO). The approach enables effective learning of CVAEs and GP-prior VAEs on non-temporal, temporal, and longitudinal datasets with state-of-the-art performance in imputation and generation despite missing covariate data.

ABSTRACT

Conditional variational autoencoders (CVAEs) are versatile deep generative models that extend the standard VAE framework by conditioning the generative model with auxiliary covariates. The original CVAE model assumes that the data samples are independent, whereas more recent conditional VAE models, such as the Gaussian process (GP) prior VAEs, can account for complex correlation structures across all data samples. While several methods have been proposed to learn standard VAEs from partially observed datasets, these methods fall short for conditional VAEs. In this work, we propose a method to learn conditional VAEs from datasets in which auxiliary covariates can contain missing values as well. The proposed method augments the conditional VAEs with a prior distribution for the missing covariates and estimates their posterior using amortised variational inference. At training time, our method marginalises the uncertainty associated with the missing covariates while simultaneously maximising the evidence lower bound. We develop computationally efficient methods to learn CVAEs and GP prior VAEs that are compatible with mini-batching. Our experiments on simulated datasets as well as on a clinical trial study show that the proposed method outperforms previous methods in learning conditional VAEs from non-temporal, temporal, and longitudinal datasets.

Motivation & Objective

• To address the gap in conditional VAE learning when auxiliary covariates contain missing values, which existing methods fail to handle.
• To develop a probabilistic framework that jointly models missing covariates and observed data within a CVAE structure.
• To enable efficient, mini-batch compatible training of CVAEs and GP-prior VAEs with missing covariates using amortized inference.
• To demonstrate superior performance in representation learning and data imputation across diverse data types, including longitudinal and clinical datasets.

Proposed method

• The method introduces a joint variational posterior over latent variables and missing covariates, enabling end-to-end training with missing data.
• It uses amortized variational inference to approximate the posterior over missing covariates, avoiding per-sample optimization.
• The evidence lower bound (ELBO) is maximized while marginalizing over the uncertainty in missing covariates via Monte Carlo approximation.
• For GP-prior VAEs, the method derives a computationally efficient, mini-batch compatible approximation of the KL divergence term involving Gaussian process priors.
• The approach integrates missing covariate estimation into the encoder-decoder architecture, allowing joint optimization of representation and imputation.
• A novel unbiased estimator for the KL divergence term is derived, enabling stable training with mini-batches through approximate splitting of patient data.

Experimental results

Research questions

• RQ1 Can conditional VAEs be effectively trained when auxiliary covariates contain missing values, without assuming full observability?
• RQ2 How can uncertainty in missing covariates be properly marginalized during training to improve model generalization?
• RQ3 Does the proposed method outperform existing VAE-based approaches in learning from partially observed data with missing covariates?
• RQ4 Can the method scale efficiently to large datasets using mini-batching, especially for GP-prior VAEs?
• RQ5 How well does the model perform in generating realistic samples and imputing missing values across non-temporal, temporal, and longitudinal data?

Key findings

• The proposed method achieves state-of-the-art performance in learning CVAEs and GP-prior VAEs on datasets with missing auxiliary covariates, outperforming prior methods in both generation and imputation tasks.
• On a simulated temporal digit dataset, the model successfully learned the underlying generative process despite missing covariates, with accurate reconstruction and generation of rotated, shifted, and contrast-adjusted digits.
• In a clinical trial study with 25.01% observed measurements in the output data and artificially introduced missingness (5% to 40%) in auxiliary covariates, the model demonstrated robust performance across all missingness levels.
• The method achieved stable training using mini-batching by deriving an unbiased estimator for the KL divergence term, enabling efficient optimization on large-scale data.
• The model significantly outperformed baseline VAEs that substitute missing values with zeros or use heuristic imputation, especially in longitudinal and complex correlation-structured data.
• The ablation study confirmed that marginalizing uncertainty over missing covariates leads to better generalization and more accurate posterior estimates than naive imputation strategies.