Let $x_{i}\in\mathbb{R}^{D}$ be an observed sample of random variable ${X}$ where $i\in\{1,\cdots,N\}$, $D$ is the number of features and $N$ is the number of samples; and let ${z_{i}}$ be an observed sample for latent variable ${Z}$. Given a sample $x_{i}$ representing the input data, the VAE is a probabilistic graphical model that estimates the posterior distribution $p_{\theta}({Z}|{X})$ as well as the model evidence $p_{\theta}({X})$, where $\theta$ are the generative model parameters (Kingma and Welling, 2013). The VAE approximates the posterior distribution of ${Z}$ given ${X}$ by a tractable parametric distribution and minimizes the ELBO (evidence lower bound) loss (An and Cho, 2015). It consists of the encoder network that computes $q_{\phi}(Z|X)$, and the decoder network that computes $p_{\theta}(X|Z)$ (Wingate and Weber, 2013), where ${\phi}$ and ${\theta}$ are model parameters. Since we use the neural network for learning the distributions $q_{\phi}(Z|X)$ and $p_{\theta}(X|Z)$, the parameters $\theta$ and $\phi$ are modeled by the weights of the encoder and decoder networks and will be learned from the data during training. The marginal likelihood of an individual data point can be rewritten as follows:

where

The first term (log-likelihood) in equation 2 can be interpreted as the reconstruction loss and the second term (KL divergence) as the regularizer. The total loss over all samples can be written as:

where $L_{REC}\coloneqq\mathbb{E}_{q_{\phi}({Z|X})}[\log(p_{\theta}({X}|{Z}))]$ and $L_{KL}\coloneqq D_{KL}(q_{\phi}({Z|X})||p_{\theta}({Z}))$.

Assuming the posterior distribution is Gaussian and using a 1-sample approximation (Skafte et al., 2019), the likelihood term simplifies to:

where $Z\sim p(Z)={N}(0,\,I)$ ($I$ is identity matrix), $X|Z\sim p_{\theta}(X|Z)={N}(X|\mu_{\theta}(Z),\,\sigma_{\theta}(Z))$, and $Z|X\sim q_{\phi}(Z|X)={N}(Z|\mu_{\phi}(X),\,\sigma_{\phi}(X))$. $\mu_{\theta}(Z),\sigma_{\theta}(Z)$ are posterior mean and variance; $\mu_{\phi}(X)$, and $\sigma_{\phi}(X)$ are encoder mean and variance. Optimizing VAEs over mean and variance with a Gaussian posterior is challenging (Skafte et al., 2019). If the model has sufficient capacity that there exists $(\phi,\theta)$ for which $\mu_{\theta}(z)$ provides a sufficiently good reconstruction, then the second term pushes the variance to zero before the term $\frac{-1}{2}\log(\sigma_{\theta}^{2}(x_{i})))$ catches up (Blaauw and Bonada, 2016; Skafte et al., 2019).

One practical example of this behavior is in speech processing applications (Blaauw and Bonada, 2016). The input is a spectral envelope which is a relatively smooth 1D curve. Representing this as a 2D image produces highly structured and simple training images. As a result, the model quickly learns how to accurately reconstruct the input. Consequently, reconstruction errors are small and the estimated variance becomes vanishingly small. Another example is 2D reconstruction of MRI images where the images from neighbouring 2D slices are highly correlated leading again to variance shrinkage (Volokitin et al., 2020). To overcome this problem, variance estimation networks can be avoided using a Bernoulli distribution or by simply setting variance to a constant value (Skafte et al., 2019).

In contrast to classical parameter estimation where the goal is to estimate the conditional mean of the response variable given the feature variable, the goal of quantile regression is to estimate conditional quantiles based on the data (Yu and Moyeed, 2001). The most common application of quantile regression models is in cases in which a parametric likelihood cannot be specified (Yu and Moyeed, 2001). Another motivation for quantile regression is that quantiles are robust to outliers (John, 2015).

Quantile regression can be used to estimate the conditional median (0.5 quantile) or other quantiles of the response variable conditioned on the input data. The $\alpha$-th conditional quantile function is defined as $q_{\alpha}(x)\coloneqq\inf\{y\in\mathbb{R}:F(y|X=x)\geq\alpha\}$ where $F=P(Y\leq y)$ is a strictly monotonic cumulative distribution function. Similar to classical regression analysis which estimates the conditional mean, the $\alpha$-th quantile regression $(0<\alpha<1)$ seeks a solution to the following minimization problem for input $x$ and output $y$ (Koenker and Bassett Jr, 1978; Yu and Moyeed, 2001):

where $x_{i}$ are the inputs, $y_{i}$ are the responses, $\rho_{\alpha}$ is the check function or pinball loss (Koenker and Bassett Jr, 1978) and $f$ is the model paramaterized by $\theta$. The goal is to estimate the parameter $\theta$ of the model $f$. The pinball loss is defined as:

Due to its simplicity and generality, quantile regression is widely applicable in classical regression and machine learning to obtain a conditional prediction interval (Rodrigues and Pereira, 2020). It can be shown that minimization of the loss function in equation 5 is equivalent to maximization of the likelihood function formed by combining independently distributed asymmetric Laplace densities (Yu and Moyeed, 2001):

where $\sigma$ is the scale parameter. Individual quantiles can be shown to be maximum likelihood estimates of Laplacian density. In this paper we estimate two quantiles jointly and therefore our loss function can be seen as a summation of two Laplacian likelihoods.

Instead of estimating the conditional mean and conditional variance directly at each pixel (or feature), the outputs of our QR-VAE are multiple quantiles of the output distributions at each pixel. This is achieved by replacing the Gaussian likelihood term in the VAE loss function with the quantile loss (check or pinball loss). For the QR-VAE, if we assume a Gaussian output, then only two quantiles are needed to fully characterize the Gaussian distribution. Specifically, we estimate the median and $0.15$-th quantile, which corresponds to approximately one standard deviation (more precisely 1.036 std dev.) from the mean under the Gaussian model. Our QR-VAE ouputs, $Q_{L}$ (low quantile) and $Q_{H}$ (high quantile), are then used to calculate the mean and the variance. To find these conditional quantiles, fitting is achieved by minimization of the pinball loss for each quantile. The resulting reconstruction loss for the proposed model can be calculated as:

where $\theta_{L}$ and $\theta_{H}$ are the parameters of the models corresponding to the quantiles $Q_{L}$ and $Q_{H}$, respectively. The minimization of this loss results in the desired quantile estimates for each output pixel.

We reduce the chance of quantile crossing (consistency in the quantiles defined by: $Q_{\tau_{1}}\subset Q_{\tau_{2}}$ when $\tau_{1}>\tau_{2}$) by limiting the flexibility of independent quantile regression. This is done by simultaneous estimation of both quantiles with one neural network, rather than training separate networks for each quantile (Rodrigues and Pereira, 2020). Note that the estimated quantiles share network parameters except for the last layer.

While quantile regression guarantees coverage of data (based on the quantiles chosen) in the training set, performance on a held-out validation data is not guaranteed. In order to have a coverage guarantee for finite samples on unseen data, we deployed conformalized quantile regression using a calibration set as explained in Romano et al. (2019). Conformal predictions provide a non-asymptotic, distribution-free coverage guarantee (Shafer and Vovk, 2008). The main idea of conformal prediction is to fit a model on the training data, then use the residuals on held-out calibration data to quantify the uncertainty in future predictions. This offers finite sample distribution-free performance guarantees. The conformalized quantile regression approach combines conformal prediction with quantile regression (Sesia and Candès, 2020). For this, we use the approach presented in Romano et al. (2019) that combines conformal prediction with quantile regression. This approach inherits both the finite sample, distribution-free validity of conformal prediction and the statistical efficiency of quantile regression.

For calculating quantiles of a binary response, consider the following model:

where $Y^{*}$ is a hidden variable, $h(x)$ is the true model, and $\epsilon$ is the noise. No distribution $\epsilon$ is assumed for the model (Kordas, 2006). Since the indicator function is monotone increasing, and since quantiles are invariant under monotone transform, we have:

$Q_{\tau}(Y|X)$ is the $\tau th$ conditional quantile of $Y$ given $X$. By modeling $Q_{\tau}(Y^{*}|X)=f_{\tau}(X,\beta)$ with $\beta$ as the parameter, the $\beta$ parameter can be estimated by:

In this paper we us the BQR loss to solve the lesion detection and segmentation task. We use a U-Net architecture with multiple heads (output branches), where each head estimates a specific quantile for the labels at the pixel level. We observed that joint estimation of multiple quantiles is computationally faster than solving for each separately, and also avoids the quantile crossing problem.

A standard U-Net would use a cross-entropy loss for this segmentation task. Here we replace this with the BQR loss. To estimate the $n$-th quantile, the BQR loss is given by:

where each $f$ correspond to a head of U-Net, $\tau_{1}...\tau_{n}$ shows different quantiles and $\beta_{1}...\beta_{n}$ are estimated parameter for each quantile respectively. A single network is used to estimate all quantiles. We chose to train a single network with output branches for each quantile since there is a common network across quantiles except for the last layer. This makes training of the quantiles consistent avoiding crossing of the estimated quantiles. We choose $K(t)$ to be the sigmoid function.

Following Skafte et al. (2019), we first evaluate variance estimation using VAE, Comb-VAE, and QR-VAE on a simulated dataset. The two moon dataset inspires the data generation process for this simulation¹¹1https://scikit-learn.org/stable/modules/generated/sklearn.datasets.make_moons. First, we generate 500 points in $\mathbb{R}^{2}$ in a two-moon manner to generate a known two-dimensional latent space. These are then mapped to four dimensions ($v_{1}$, $v_{2}$, $v_{3}$, $v_{4}$) using the following equations:

where $\epsilon$ is sampled from a normal distribution. For more details about the simulation, please refer to Skafte et al. (2019)²²2https://github.com/SkafteNicki/john/blob/master/toy_vae.py. After training the models, we first sample from $z$ using a Gaussian prior, and input that sample to the decoder to generate parameters of the posterior $p_{\theta}(x|z)$, and then sample again from the posteriors using the estimated means and variances from the decoder. The distribution of these generated samples represents the learned distribution in the generative model.

In Figure 1, we plot the pairwise joint distribution for the input data as well as the generated samples using various models. We used Gaussian kernel density estimation to model the distributions from 1000 samples in each case. We observe that the standard VAE underestimates the variance resulting in insufficient learning of the data distribution. The samples from our QR-VAE model capture a data distribution more similar to the ground truth than either standard VAE or Comb-VAE. Our model also outperforms VAE and Comb-VAE in terms of KL divergence between input samples and generated samples as can be seen in Figure 1. The KL-divergence is calculated using universal-divergence, which estimates the KL divergence based on k-nearest-neighbor (k-NN) distance (Wang et al., 2009)³³3https://pypi.org/project/universal-divergence.

We evaluated our BQR supervised approach on the LIDC-IDRI dataset (Armato III et al., 2011). This data consists of 1018 3D thorax CT scans annotated by four radiologists tasked with finding multiple lung nodules in each scan. The data is ideal for capturing inherent uncertainty in the data labels that comes from disagreement between experts. The data was preprocessed as described by Kohl et al. (2018). They extracted 2D slices centred around the annotated nodules and generated 180 x 180 images when at least one expert has segmented a nodule. This process resulted in a dataset of 8882 images in the training set, 1996 images in the validation set, and 1992 images in the test set. We reshaped the data into 128 x 128 images for input to the neural network. We used a U-Net architecture (Ronneberger et al., 2015) with 2D convolutional layers. The output layer was modified to generate four quantiles with four output branches using softmax activation. We compared the performance of BQR U-Net with the deterministic U-Net (Ronneberger et al., 2015). As ground truth for the comparison, first we estimated agreement maps for each test image by combining the lesion annotations of the four raters. This generates, for each image, an annotation with $P(Y=1|X)$ values greater than or equal to 0, 0.25, 0.5, 0.75,1. $X$ here represents the input image. Given an input image, the BQR U-Net generates output regions where probabilities of the label $Y=1$ are at or above the given quantile thresholds. The BQR U-Net was trained with the loss function in eq. 8. For comparison, the deterministic U-Net was trained using the binary cross-entropy loss. The BQR U-Net was trained to output 0.125, 0.375, 0.625, 0.875 quantiles. These quantile values represent the centroids of the intervals between the test data quantiles (0, 0.25, 0.5, 0.75,1) representing 0-4 rater agreements respectively. We used these thresholds rather than the same quantiles as the test data to avoid operating at the boundary points between operators resulting from combining data from four raters only. To generate the corresponding estimated quantiles for the deterministic U-Net we thresholded the softmax probability at 0.125, 0.375, 0.625, 0.875.

Both deep BQR and deterministic U-Net diverged to the trivial solution of predicting all labels as zero due to the extreme class imbalance when initialized with a cold start. We therefore warmed up both models using a weighted cross-entropy loss for one epoch weighting samples by $1\div 166$, the ratio of the zero to one labels in the training set. We trained both models for 5 epochs. Our results show no significant improvement for deep BQR compared to the cross-entropy loss in terms of Dice coefficients for different agreement areas. The Dice coefficients between estimated probability areas for deterministic U-Net and BQR U-Net are plotted in Figure 8 and summarized in Table 2. In the figure we show the distribution of the Dice coefficients between the BQR and ground truth quantiles, DT (deterministic U-Net) and ground truth, and also between DT and BQR. In some cases, particularly for higher quantiles, there was low agreement between human raters in the ground-truth test data. For example, 64 percent of the test data showed zero pixels in common between all four human raters, leaving the 0.875 quantile empty for those images. We therefore computed the Dice coefficients only over those regions in which the test data had a non-empty data set for that quantile. The results show reasonable agreement for the 0.125, 0.375 and 0.625 quantiles, but relatively poor results for the 0.875 quantile. Surprisingly, results for BRQ and the determinsic U-Net (DT U-Net) are very similar, even though the actual degree of overlap between the two is no better than between each of them and the ground truth labels. The fact that both methods perform poorly for the 0.875 quantile reflects both that the data are of relatively poor quality in this region and also that performance is likely limited by the imbalance in the training data between non-lesional areas and lesions confidently identified by all four raters. Here we used the deterministic U-Net as a backbone since it is arguably the most commonly used network for medical image segmentation task. Other networks could also be used as a backbone. To investigate whether we would expect further improvements using a probabilistic backbone, we also implemented the Probabilistic U-Net (Prob. U-Net) (Kohl et al., 2018) to capture rater uncertainty. Our results show that for the quantile estimation task, the performance of Prob. U-Net and U-Net are comparable. As a result, it appears unlikely that replacing the U-Net with its proabilistic form in the backbone would lead to significant improvements.