Quantile Regression for Uncertainty Estimation in VAEs with Applications to Brain Lesion Detection

The Variational AutoEncoder (VAE) has become one of the most popular models for anomaly detection in applications such as lesion detection in medical images. The VAE is a generative graphical model that is used to learn the data distribution from samples and then generate new samples from this distribution. By training on normal samples, the VAE can be used to detect inputs that deviate from this learned distribution. The VAE models the output as a conditionally independent Gaussian characterized by means and variances for each output dimension. VAEs can therefore use reconstruction probability instead of reconstruction error for anomaly detection. Unfortunately, joint optimization of both mean and variance in the VAE leads to the well-known problem of shrinkage or underestimation of variance. We describe an alternative VAE model, Quantile-Regression VAE (QR-VAE), that avoids this variance shrinkage problem by estimating conditional quantiles for the given input image. Using the estimated quantiles, we compute the conditional mean and variance for input images under the Gaussian model. We then compute reconstruction probability using this model as a principled approach to outlier or anomaly detection. We also show how our approach can be used for heterogeneous thresholding of images for detecting lesions in brain images.