[Paper Review] Robust Variational Autoencoder

Haleh Akrami, Anand A. Joshi+3 more|arXiv (Cornell University)|May 23, 2019

Generative Adversarial Networks and Image Synthesis37 references15 citations

TL;DR

This paper proposes a Robust Variational Autoencoder (RVAE) that uses β-divergence instead of standard Kullback-Leibler divergence to improve robustness to outliers in training data. By replacing the KL divergence with β-divergence in the variational inference objective, the model achieves comparable computational complexity while effectively suppressing outlier encoding, enabling reliable representation learning and improved anomaly detection in image data, including medical brain MRI.

ABSTRACT

Machine learning methods often need a large amount of labeled training data. Since the training data is assumed to be the ground truth, outliers can severely degrade learned representations and performance of trained models. Here we apply concepts from robust statistics to derive a novel variational autoencoder that is robust to outliers in the training data. Variational autoencoders (VAEs) extract a lower-dimensional encoded feature representation from which we can generate new data samples. Robustness of autoencoders to outliers is critical for generating a reliable representation of particular data types in the encoded space when using corrupted training data. Our robust VAE is based on beta-divergence rather than the standard Kullback-Leibler (KL) divergence. Our proposed lower bound lead to a RVAE model that has the same computational complexity as the VAE and contains a single tuning parameter to control the degree of robustness. We demonstrate the performance of our $β$-divergence based autoencoder for a range of image datasets, showing improved robustness to outliers both qualitatively and quantitatively. We also illustrate the use of our robust VAE for outlier detection.

Motivation & Objective

• To address the problem of performance degradation in variational autoencoders due to outliers in training data.
• To develop a probabilistic generative model that remains robust to corrupted or mislabeled data without requiring clean data or complex architectural modifications.
• To enable effective outlier detection by ensuring that outliers are not reconstructed or encoded in the latent space.
• To provide a single, tunable hyperparameter (β) that controls the degree of robustness in the model.
• To demonstrate the effectiveness of the approach on benchmark image datasets and real-world medical imaging data with simulated and actual lesions.

Proposed method

• Replaces the standard Kullback-Leibler (KL) divergence in the VAE objective with β-divergence, a robust statistical divergence from robust statistics.
• Derives a lower bound for the evidence lower bound (ELBO) using β-divergence, resulting in a new objective function that downweights the influence of outliers.
• Applies the β-divergence framework to both Bernoulli and Gaussian likelihood models, enabling use with binary and continuous image data.
• Uses stochastic gradient descent to optimize the resulting objective, maintaining the same computational complexity as standard VAEs.
• Introduces a visual-based method to select the optimal β value by inspecting generated and reconstructed samples for desired variability and realism.
• Extends the framework to medical imaging by training on nominally normal brain MRI data and evaluating performance on simulated and real lesions.

Experimental results

Research questions

• RQ1 Can replacing the KL divergence with β-divergence in the VAE objective improve robustness to outliers in the training data?
• RQ2 How does the choice of β parameter affect the model’s ability to reconstruct normal samples while avoiding encoding of outliers?
• RQ3 Can the proposed RVAE detect anomalies in test images by identifying large reconstruction errors, even when outliers are present in the training set?
• RQ4 How does the performance of the RVAE compare to standard VAEs in terms of reconstruction quality and anomaly detection AUC on real medical imaging data?
• RQ5 Can the β-divergence-based approach be effectively tuned using visual inspection of generated and reconstructed samples to achieve optimal trade-off between robustness and generative diversity?

Key findings

• The RVAE achieved an area under the ROC curve (AUC) of 0.85 for lesion detection in brain MRI data with simulated outliers, compared to only 0.20 for standard VAEs, demonstrating strong quantitative superiority in anomaly detection.
• At β = 0.0007, the reconstructed images of outliers and generated samples from the latent space became visually indistinguishable, indicating that outliers were projected onto the normal data manifold.
• For the Fashion-MNIST dataset with 10% outliers, the model achieved optimal performance in terms of reconstruction quality and generative diversity at β = 0.0007, as confirmed by heatmap analysis of reconstruction similarity.
• The RVAE successfully suppressed reconstruction of both simulated and real lesions in brain MRI, enabling clear detection via reconstruction error, whereas the standard VAE reconstructed lesions as plausible normal samples.
• The model maintains the same computational complexity as standard VAEs while introducing only one tunable hyperparameter (β), making it practical for real-world deployment.
• Visual inspection of generated and reconstructed samples was effective in identifying the optimal β value, suggesting a practical, non-automated method for hyperparameter selection that could be extended with adversarial training in future work.