A Systematic Comparison of Bayesian Deep Learning Robustness in Diabetic Retinopathy Tasks

Introduction

Deep learning is continuously transforming intelligent technologies across many fields, from advancing medical diagnostics with complex data, to enabling autonomous driving, to deciding high-stakes economic actions . However, deep learning models struggle to inform their users when they don’t know – in other words, these models fail to communicate their uncertainty in their predictions. The implications for deep models entrusted with life-or-death decisions are far-reaching: experts in medical domains cannot know whether to trust their auto-diagnostics system, and passengers in self-driving vehicles cannot be alerted to take control when the car does not know how to proceed.

Bayesian deep learning (BDL) offers a pragmatic approach to combining Bayesian probability theory with modern deep learning. BDL is concerned with the development of techniques and tools for quantifying when deep models become uncertain, a process known as inference in probabilistic modelling. BDL has already been demonstrated to play a crucial role in applications such as medical diagnostics (see Figure 1), computer vision , in the sciences , and autonomous driving .

Despite BDL’s impact on a range of real-world applications and the flourish of recent ideas and inference techniques , the development of the field itself is impeded by the lack of realistic benchmarks to guide research. Evaluating new inference techniques on real-world applications often requires expert domain knowledge, and current benchmarks used for the development of new inference tools lack consideration for the cost of development, or for scalability to real-world applications.

Advances in computer vision, natural language and reinforcement learning are usually attributed to the emergence of challenging benchmarks, e.g. ImageNet , GLUE and ALE , respectively. In contrast, many BDL papers use benchmarks such as the toy UCI datasets , which consist of only evaluating root mean square error (RMSE) and negative log-likelihood (NLL) on simple datasets with only a few hundred or thousand data points, each with low input and output dimensionality. Such evaluations are akin to toy MNIST evaluations in deep learning. Due to the lack of alternative standard benchmarks, in current BDL research it is common for researchers developing new inference techniques to evaluate their methods with such toy benchmarks alone, ignoring the demands and constraints of the real-world applications which make use of BDL tools . This means that research in BDL broadly neglects exactly the applications that neural networks have proven themselves most effective for.

In order to make significant progress in the deployment of new BDL inference tools, the tools must scale to real-world settings. And for that, researchers must be able to evaluate their inference and iterate quickly with real-world benchmark tasks without necessarily worrying about the required application-specific expertise. We require benchmarks which test for inference robustness, performance, and accuracy, in addition to cost and effort of development. These benchmarks should include a variety of tasks, assessing different properties of uncertainty while avoiding the pitfalls of overfitting quickly as with UCI. These should asses for scalability to large data and be truthful to real-world applications, capturing their constraints.

Contributions. We build on-top of previous work published at Nature Scientific Reports by Leibig et al. . We extend on their methodology and develop an open-source benchmark, building on a downstream task which makes use of BDL in a real-world application—detecting diabetic retinopathy from fundus photos and referring the most uncertain cases for further inspection by an expert (Section 2). We extend this methodology with additional tasks that assess robustness to out-of-distribution and distribution shift, using test datasets which were collected using different medical equipment and for different patient populations. Our implementation is easy to use for machine learning researchers who might lack specific domain expertise, since expert details are abstracted away and integrated into metrics which are exposed through a simple API. Improvement on this benchmark will directly be contributing to the advancement of an important real-world application. We further perform a comprehensive comparison on this new benchmark, contrasting many existing BDL techniques. We develop and tune baselines for the benchmark, including Monte Carlo dropout , mean-field variational inference and model ensembling , as well as variants of these (Section 3). We conclude by demonstrating the benchmark’s usefulness in ranking existing techniques in terms of scalability and effectiveness, and show that despite the fact that some current techniques solve benchmarks such as UCI, they either fail to scale, fail to solve our benchmark, or fail to provide good uncertainty estimates. This shows that an over-reliance on UCI has the potential to badly distort work in the field because researchers prioritize their attention on approaches to Bayesian deep learning that are not suited to large scale applications (Section 5).

It is our hope that the proposed benchmarks will make testing new inference techniques for Bayesian deep learning radically easier, leading to faster development iteration cycles, and rapid development of new tools. Progress on these benchmarks will translate to more robust and reliable tools for already-deployed decision-making systems, such as automatic medical diagnostics and self-driving car prototypes.

Diabetic Retinopathy Benchmark

We describe the dataset, the data processing, as well as the downstream task and metrics used.

The benchmark is built on the Kaggle Diabetic Retinopathy (DR) Detection Challenge data. It consists of 35,126 training images and 53,576 test images. We hold-out 20$\%$ of the training data as a validation set. Each image is graded by a specialist on the following scale: 0 – No DR, 1 – Mild DR, 2 – Moderate DR, 3 – Severe DR and 4 – Proliferative DR. We recast the 5-class classification task as binary classification which is easily applicable to any BDL classification algorithm by asking the user to classify whether each image has sight-threatening DR, which is defined as Moderate DR or greater (classes 2-4) following . Samples from both classes are provided in Figure 2. The data are unbalanced, with only $19.6\%$ of the training set and $19.2\%$ of the test set having a positive label.

Robustness to distribution shift is evaluated by training on the original Kaggle diabetic retinopathy detection challenge dataset , and testing on a completely disjoint APTOS 2019 Blindness Detection dataset collected in India with different medical equipment and on a different population.

2 Data Processing

All images are cropped and resized to $512\times 512$, while all three colour channels are used. The data is standard normalized for each colour channel separately, using the empirical statistics of the training data. Similar to Leibig et al. , we augment training dataset using affine transformations, including random zooming (by up to $\pm 10\%$), random translations (independent shifts by up to $\pm 25$ pixels) and random rotations (by up to $\pm\pi$). Finally half of the augmented data is also flipped along the vertical and/or the horizontal axis. Examples of original and their corresponding processed images are provided in Figure 3.

3 Downstream Task

Machine learning researchers often evaluate their predictions directly on the whole test set. But, in fact, in real-world settings we have additional choices available, like asking for more information when we are uncertain. Because of the importance of accurate diagnosis, it would be unreasonable not to ask for further scans of the most ambiguous cases. Moreover, in this dataset, many images feature camera artefacts that distort results. In these cases, it is critically important that a model is able to tell when the information provided to it is not sufficiently reliable to classify the patient. Just like real medical professionals, any diagnostic algorithm should be able to flag cases that require more investigation by medical experts. This task is illustrated in Figure 1, where a threshold, $\tau$, is used to flag cases as certain and uncertain, with uncertain cases referred to an expert. Alternatively, the uncertainty estimates could be used to come up with a priority list, which could be matched to the available resources of a hospital, rather than waste diagnostic resources on patients for whom the diagnosis is clear cut.

To get some insight into the dataset, Figure 4 illustrates the relation between predicted probabilities, $p(\text{disease}|\text{image})$, and our estimator for the models’ uncertainty about them, the predictive entropy $\mathcal{H}_{\text{pred}}$, for an MC dropout model. Note that the model is correct and certain about most of its predictions, as shown in sub-figure (a), while it is more uncertain when wrong, sub-figure (b).

4 Metrics

In order to simulate this process of referring the uncertain cases to experts and relying on the model’s predictions for cases it is certain of, we assess the techniques by their diagnostic accuracy and area under receiver-operating-characteristic (ROC) curve, as a function of the referral rate. We expect the models with well-calibrated uncertainty to refer their least confident predictions to experts (see Figure 5), improving their performance as the number of referrals increases.

The accuracy of the binary classifier is defined as the ratio of correctly classified data-points over the size of the population. The receiver-operating-characteristic (ROC) curve (see Figure 6) illustrates the diagnostic ability of a binary classifier system as its discrimination threshold is varied. It is created by plotting the true positive rate (a.k.a. sensitivity) against the false positive rate (a.k.a. $1-\text{sensitivity}$). The quality of such a ROC curve can be summarized by its area under the curve (AUC), which varies between $0.5$ (chance level) and $1.0$ (best possible value).

A Systematic Comparison of BDL Methods

We next present and evaluate various Bayesian deep learning techniques (i.e. baselines) on the diabetic retinopathy diagnosis benchmark. Each method is tuned separately and, in order to obtain statistically significant results, we train nine independent models for each method, using different random number generator seeds. We observe consistency and robustness for our implementations across seeds.

Class imbalance.

We compensate for the class imbalance discussed in Section 2.1 by reweighing the cross-entropy part of the cost function, attributing more weight to the minority class, given by the relative class frequencies in each mini-batch, $p(k)_{\text{mini-batch}}$ :

We also tried using a constant class weight, or artificially balancing the two classes by sub-sampling negatively labelled images, but both approaches made training slower and less stable for many baselines.

Uncertainty Estimator.

We quantify the uncertainty of our binary classification predictions by predictive entropy , which captures the average amount of information contained in the predictive distributionThe predictive uncertainty is the sum of epistemic and aleatoric uncertainty.:

summing over all possible classes $c$ that $y$ can take, in our case $c\in\{0,1\}$. This quantity is high when either the aleatoric uncertainty is high (the input is ambiguous), or when the epistemic uncertainty is high (a probabilistic model has many possible explanations for the input). In practice, we approximate the $p(y=c|\mathbf{x})$ term in (2) by $T$ Monte Carlo samples, $\frac{1}{T}\sum_{t}p_{\theta}(y=c|\mathbf{x})$, obtained by stochastic forward passes through the probabilistic networks. Note that this is a biased but consistent estimator of the predictive entropy in (2) .

We contrast several methods in BDL which we discuss in more detail next.

1 Bayesian Neural Networks

Estimating the uncertainty about a machine learning based prediction on a single observation requires a distribution over possible outcomes, for which a Bayesian perspective is principled. Bayesian approaches to uncertainty estimation have indeed been proposed to assess the reliability of clinical predictions but have only been applied to a handful of large-scale real-world problems that neural networks (NNs) have proven themselves particularly effective for.

Finite NNs with distributions placed over the weights have been studied extensively as Bayesian neural networks (BNNs) , providing robustness to over-fitting (i.e. regularization). Exact inference is analytically intractable and hence approximate inference has been applied instead .

Given a dataset $\mathcal{D}=\{(x_{n},y_{n})\}_{n=1}^{N}$, a BNN is defined in terms of a prior $p(\mathbf{w})$ on the weights, as well as the likelihood $p(\mathcal{D}|\mathbf{w})$. Variational Bayesian methods attempt to fit an approximate posterior $q(\mathbf{w})$ to maximize the evidence lower bound (ELBO):

We parameterize $q(\mathbf{w})$ with $\theta$ parameters and choose prior distribution $p(\mathbf{w})$. The (variational) inference is then recast as the optimization problem $\max_{\theta}\mathcal{L}_{q_{\theta}}$. Different methods use different prior distributions and parametric families for the approximate posterior, as well as optimization methods. We discuss these different techniques next.

Mean-field variational inference (MFVI) is an approach to learning an approximate posterior over the weights of a neural network, $q_{\theta}(\mathbf{w})$, given a prior $p(\mathbf{w})$ . In MFVI, we assume a fully-factorized Gaussian posterior (and prior). This reduces the computational complexity of estimating the evidence lower-bound (ELBO). In addition, we use a Monte Carlo estimate of the KL-divergence term of the ELBO in order to reduce the time complexity of a forward pass to $\mathcal{O}(D)$ in the number of weights. Blundell et al. applied the reparametrization trick from to perform MFVI, which they call Bayes-by-backprop. Instead, we use the Flipout Monte Carlo estimator of the KL-divergence , which reduces the variance of the estimator of the gradient. A Monte Carlo estimate of model predictions is made by taking a number of samples from the posterior distribution over the weights and averaging the predictions.

Note that the effective number of trainable parameters is doubled compared to a deterministic NN, since both the mean and scale parameters are now learnable. To allow fair comparison with the other baselines, we reduce the number of channels in the convolutional layers of the MFVI model to reach the model budget of 2.5 million parameters.

Monte Carlo Dropout.

2 Model Ensembling

Lakshminarayanan et al. proposed an alternative to BNNs, termed Deep Ensembles, that is simple to implement, readily parallelizable, requires little hyperparameter tuning, and yields high quality predictive uncertainty estimates. The method quantifies uncertainty by collecting predictions from $T$ independently trained deterministic models (ensemble components). Despite the easy parallelization of the method, the resources for training scale linearly with the required number of ensemble components $T$, making it prohibitively expensive in some cases.

We also report results on an ensemble of MC Dropout models, which performs best of all the other methods, in terms of both accuracy and AUC for all the referral rates, as illustrated in Figure 5 and Table 1. In this technique, several dropout models are separately trained in parallel. Predictions are then made by sampling repeatedly from all models in the ensemble using fresh dropout masks for each sample, before being averaged, to form a Monte Carlo estimate of the ensemble prediction.

3 Deterministic

Two naive baselines are evaluated as control, a Deterministic neural network and Random. Both are based on a deep VGG model, trained with dropout and L2-regularization, using exactly the same hyperparameters and set-up as MC Dropout. In fact, because the conditions are identical, we used the same models for the Deterministic and MC Dropout baselines—the only difference is that for MC Dropout we sample dropout mask during evaluation and average over 100 samples from the dropout posterior to estimate uncertainty. In contrast, the Deterministic baseline uses the sigmoid output $p(\text{disease}|\text{image})$ to quantify uncertainty, and uses the deterministic dropout approximation at test time . That is, a model is assumed to be more confident the closer to 1 or 0 its output is. This is the simplest way a neural network might estimate uncertainty, but it captures only the aleatoric component of uncertainty—it does not capture epistemic uncertainty about the model’s knowledge . Figure 4 (right) shows that there is a correlation between the sigmoid output $p(\text{disease}|\text{image})$ and the predictive entropy $\mathcal{H}_{\text{pred}}$, which we use to measure uncertainty. But the overall evidence in Figure 5 and Table 1 suggests that models which also capture the epistemic component of the uncertainty perform much better than the Deterministic baseline.

The Random baseline makes random referrals, without taking any kind of uncertainty (or input) into account. As expected, it has the same accuracy and AUC regardless of how much data is retained vs. referred.

Results and Analysis

Table 1 and Figures 5 and 6 summarize the quantitative performance of various methods, described in Section 3. Methods that capture meaningful uncertainty estimates show this by improving performance (i.e. AUC and accuracy) as the rate of referral increases. That is, steeper slopes in Figure 5 are making better estimates of uncertainty, all else equal, because they are able to systematically refer the datapoints where their estimates are less likely to be accurate. Note that all methods perform equally well when all data is retained, conveying that all models have converged to similar overall performance, providing a fair comparison of uncertainty.

Benchmarks are often used to compare methodology, e.g. to select which tools we should build on-top. UCI, a popular benchmark in the field, has been used to reproduce such rankings of BDL methods. Importantly, in contrast to the empirical results found in on the toy UCI benchmark and summarised in Table 2, our benchmark suggest a different ranking of methods. While in mean-field variational inference outperforms the other baselines we discuss, Table 1 and Figures 5 and 6 suggest that in the real-world application of diabetic retinopathy diagnosis both ensemble methods (Section 3.2) and Monte Carlo Dropout score consistently higher than MFVI, suggesting that some methods might be ‘overfitting’ their uncertainty to the simple dataset. That is, extensive tuning on the simple UCI tasks might have resulted in rankings which do not generalise to other tasks. Moreover, Mukhoti et al. show that UCI regression benchmarks are insufficient for drawing conclusions about the effectiveness, and surely the scalability, of the inference techniques.

Implications for the Field

Deep learning, as a whole, has had its biggest successes when handling large, high-dimensional data. It is something of a surprise, then, that the standard benchmarks for Bayesian deep learning, UCI, only has input dimensionalities between 4 and 16. Due to the lack of alternative common benchmarks with well tuned baselines, researchers find it hard to publish results in Bayesian deep learning without resorting to a comparison on UCI. As a result, there is an undue focus in Bayesian deep learning on models that perform well with very low numbers of input features and on tiny models with a single layer of only 50 hidden units. UCI plays an important role for a subset of models, but the fact that it is currently the field’s main benchmark has a distorting effect on research.

Consider, for example, the ranking of deep learning methods for uncertainty offered by Bui et al. . They compare UCI rankings from multiple papers and calculate the average rankings. They find that Hamiltonian Monte Carlo (average rank 8.80$\pm$1.38) and mean-field variational inference (average rank 7.50$\pm$1.70) using the reparametrization trick perform best of the neural network models they consider (with the best performer being Deep Gaussian Processes). However, HMC is known not to scale to datasets with large data, a property which is not captured with the benchmark. Further, MC dropout is ranked second-to-last place with average rank 12.10$\pm$0.64. Our results show that on a larger-scale dataset, MC dropout has better uncertainty estimates than mean-field variational inference and they have almost identical performance when all datapoints are retained. Moreover, HMC would not scale to this data at all.

By relying too much on UCI as a benchmark, we give a misleading impression of relative performance, which will cause researchers to prioritise the wrong approaches. A number of more computationally intensive extensions to MFVI have emerged since Bui et al. produced their analysis, while less work has gone into building on the methodology of the more computationally parsimonious Bayesian deep learning approaches like deep ensembles or MC dropout . It seems likely that this is partly shaped by the fact that UCI is the predominant benchmark.

Our new benchmark and systematic comparison of BDL tools will offer a way for new methods to demonstrate their effectiveness on large-scale problems, making it easier to publish results that engage with the sorts of problems that deep learning has proven to be effective at, and which downstream users are seeking.

References