Uncertainty via 3D Bayesian Deep Learning | Towards Data Science

PAPER: https://arxiv.org/abs/1910.10793
CODE: https://github.com/sandialabs/bcnn
Key Takeaways:
· Measuring uncertainty is not possible in a regular deep neural network, but it is extremely important for interpretability and validation
· Bayesian neural networks learn probability distributions rather than point estimates, allowing them to measure uncertainty
· We designed the first-ever successful Bayesian convolutional neural network (BCNN) architecture for 3D segmentation
· Our BCNN beats the current state-of-the-art neural network in uncertainty quantification while achieving equal or better segmentation accuracy
· Our academic paper and full open-source code implementation are available online
Consider this scenario: a certain aircraft manufacturing company produces safety-critical engine parts for commercial aircraft. Since the engine parts must be guaranteed to work, the company validates each of them in a painstaking, months-long process wherein they take a 3D CT scan of the part, annotate the hundreds of millions of voxels by hand, and use the annotated scan to analyze the part for defects.
This is not time- or cost-effective, so the company hires a team of data scientists to design a deep neural network which automatically validates the parts using state-of-the-art volumetric segmentation techniques. The neural network is seemingly successful, but one day a part verified by the neural network inexplicably fails during usage, causing engine failure in the aircraft and leading to hundreds of fatalities and billions of dollars lost for the company.
So what went wrong? Well, deep neural networks are well-known for making accurate predictions, but one of their greatest weaknesses is that they lack the capacity for measuring uncertainty within their predictions. Thus, the deep learning system had no capacity to distinguish between a part that was 100% certified, and a part that just barely made the cut.
One may think that the value of the sigmoid output is usable as a measure of uncertainty, but this is not true because these values are dependent on the inferred sample being “close” to the training distribution. If one infers on a sample far from the training distribution (i.e., a part with a defect in it), the sigmoid output cannot be used as a proxy for model uncertainty3.
In other words, deep neural networks perform best when the test set is “similar” to the training set; in other words, there may be examples in the test set which the neural network is unsure about because there were no similar training examples. But because there is no “I don’t know” output available, defective parts can end up verified by the deep learning system and moved to production.
This example illustrates the critical importance of uncertainty quantification in deep neural networks, the research and development of which has skyrocketed in the last five years. Recent techniques cast neural networks, usually pointwise estimators, as probabilistic or Bayesian models. The two most common neural network architectures for this purpose are Monte Carlo dropout networks3 (MCDNs) and Bayesian convolutional neural networks1 (BCNNs).
MCDNs use dropout layers to approximate deep Gaussian processes, and while easy to implement, their statistical soundness has been called into question9. BCNNs use variational inference to learn the posterior distribution of the weights given the dataset and are much more difficult to implement but give greatly improved uncertainty results.
It has been postulated that Bayesian neural networks for large problems, including in the 3D image segmentation space, are infeasible due to prohibitive computational costs3. Here, we refute that claim and present the first-ever successful implementation of a 3D BCNN for uncertainty quantification in volumetric segmentation, detail how it works, and explain some sections of the open-source codebase.
In particular, this novel neural network architecture provides improved uncertainty quantification as compared to MCDNs while achieving equal or better segmentation accuracy. Through experimentation on CT scans of battery electrode and laser-welded metals, we show that the uncertainty maps generated by the BCNN capture continuity and visual gradients, making them interpretable as confidence intervals for segmentation.
While most neural networks learn pointwise estimates of their weights, the more rigorous view is that these estimates do not fully encapsulate the uncertainty inherent in the values of the weights. Instead, in a BCNN, the weights are each implicitly described as (multivariate) probability distributions.
Describing the weights of the neural networks as probability distributions has several consequences. First, it makes the neural network nondeterministic; every time we compute a forward pass, we must sample from each weight distribution to obtain a point estimate which can be used for inference. Repeated applications of this sampling technique, called Monte Carlo sampling, will result in different predictions which can then be analyzed for uncertainty. Second, it changes the backpropagation algorithm since we cannot backpropagate through a sampling operation (it has no gradient). In a later section, we will discuss how to fix this problem using the Bayes by Backprop algorithm. Finally, it makes the neural network much more difficult to reliably train, especially in 3D, and susceptible to vanishing/exploding gradients. We use a clever normalization technique called group normalization to solve this problem, also detailed in a later section.
How exactly do we get distributions over the weights? Well, in a perfect world, we could use Bayes’ Rule to calculate them precisely. To do this, we would begin with a prior distribution over the weights; this is our “initial guess” of what a weight distribution would look like. It is denoted p(w) and is generally a standard normal distribution. Then, we would use our data to calculate the posterior distributions of the weights given the dataset, denoted p(w|D). This is equivalent to finding a w that maximizes the likelihood of the dataset given those weights, denoted p(D|w), and we can perform this calculation via Bayes’ Rule:
However, due to the extreme overparameterization found in neural networks, the integral in the denominator is generally intractable. So, we need to learn the posterior distribution rather than calculate it exactly. Previous work by Hinton and Van Camp5 and Graves4 proposed variational learning (also called variational inference) as a method to approximate the posterior distribution. Variational learning finds the parameters θ of the distribution p(w|θ), called the variational distribution, via the minimization of the variational free energy cost function F, often called the expected lower bound (ELBO).
The variational free energy consists of the sum of the Kullback-Leibler (KL) divergence, which measures the distance between the prior and variational distributions, and the negative log-likelihood, which measures the goodness-of-fit of the model. Blundell et al.1 explains the variational free energy loss function as a tradeoff between satisfying the simplicity prior (represented by the KL term) and satisfying the complexity of the dataset (represented by the NLL term):
The longer we train the neural network, the closer we will get to minimizing this cost function, and the closer our variational distribution will get to the true posterior distribution. In practice, the KL term has a regularizing effect on the output of the neural network, preventing the learned distribution from overfitting at the cost of a lower NLL term in the training set.
This loss function is amenable to mini-batch optimization by scaling the cost for mini-batch i as follows, essentially spreading out the KL divergence penalty across the entire dataset1:
If a Bayesian neural network is implemented in TensorFlow Probability, the below is a valid Python implementation of the variational free energy loss (note that binary cross-entropy is just negated NLL):
from tensorflow.keras.losses import binary_crossentropydef vfe_loss(model, dataset_size, batch_size, alpha=1):    """Defines variational free energy loss.        Sum of KL divergence and binary cross-entropy."""     # KL Divergence should be applied once per epoch only.    kl = sum(model.losses) / (dataset_size / batch_size)    def loss(y_true, y_pred):        bce = binary_crossentropy(y_true, y_pred)        return alpha * kl + (1. / alpha) * bce    return loss
Challenges in Bayesian Learning
Variational Bayesian learning is generally considered more statistically sound than, for example, approximate Bayesian inference via dropout layers9. However, this comes at the cost of several computational challenges.
First, a random variable sampled from a distribution has no gradient, so Bayesian neural networks seem incompatible with backpropagation. However, Kingma et al.6 showed that it is possible to reparametrize the random variable as a deterministic variable for computation. As an example, let a weight w be sampled from a normal distribution with mean μ and variance σ2. Then, a valid reparameterization is w=μ+σε where ε is an auxiliary noise variable sampled from a standard normal distribution. Now, instead of a sampling operation, we have an affine combination, which is easily utilized in backpropagation. More complex calculations are needed to efficiently scale this computation; see Kingma’s paper for more. This is often called the local reparameterization trick.
Next, Bayesian learning was previously considered computationally infeasible due to the massive number of weight updates necessary if trained using an ensemble method. To solve this, Blundell et al.1 designed the Bayes by Backprop algorithm. Previous work focused on training stochastic hidden units, but there are easily two orders of magnitude more weights than hidden units, and Bayes by Backprop was the first algorithm to effectively train probabilistic weights in a neural network. Bayes by Backprop works by using the gradients calculated in backpropagation to “scale and shift” the variational parameters of the posterior, thus updating the posterior with minimal additional computation.
Since 3D training volumes can be quite large, our batch size is constrained by the amount of available GPU memory, resulting in a batch size too small for batch normalization to accurately compute batch statistics. Thus, we use a technique proposed by Wu and He12 called group normalization, which normalizes groups of channels and is shown to have accurate performance independent of batch size. Proper normalization was observed to be a critical factor in the convergence of our model since it helps avoid vanishing/exploding gradients; by tuning the number of groups used in the group normalization layers, we found that the BCNN converged most reliably when using 4 groups.
Finally, one challenge associated with probabilistic weights is that all examples in a mini-batch typically have similarly sampled weights, limiting the variance reduction effect of large mini-batches. One side effect of the local reparameterization trick6 mentioned above is that it greatly reduces the variance of stochastically sampled weights by transforming global weight uncertainty into independent local noise across examples in the mini-batch. In a similar vein, Wen et al.11 proposed the Flipout estimator, which empirically achieves ideal variance reduction by sampling weights pseudo-independently for each example. An important difference is that local reparametrization works only for fully connected networks, while Flipout can be used effectively in fully-connected, convolutional, and recurrent networks.
Our 3D BCNN architecture draws from the image segmentation literature by utilizing the common encoder-decoder setup seen in V-Net7 and 3D U-Net2 , deep neural networks originally used for 3D segmentation of human prostates and frog kidneys, respectively. In this architecture, the encoder half (left) of the network compresses the input into a latent space while the decoder half (right) decompresses the latent representation of the input into a segmentation map.
The encoder half of the BCNN uses typical 3D convolutions to maximize information transfer between the original volume and the latent space, but the decoder half of the network uses 3D Bayesian convolutional layers. Each of these is initialized with a standard normal prior and employs the aforementioned Flipout estimator11 to approximate the distribution during forward passes. Note the skip connections in yellow, which assist in feature-forwarding throughout the network. Our implementation draws from the Bayesian Layers library10 included in TensorFlow Probability, which keeps track of losses representing the KL divergence of the layer’s posterior distribution with respect to its prior and makes computing the variational free energy loss simple.
An implementation of this BCNN architecture is available at https://github.com/sandialabs/bcnn/blob/master/bayesian_vnet.py.
Using a BCNN in practice brings with it a few peculiarities that can be difficult to effectively implement. First, it is currently impossible to save the architecture and weights of a model that combines, say, Keras layers and Bayesian layers. Instead, one must save the weights only and then load them into an instantiated architecture. This gets worse when the models are designated as multi-GPU, which is virtually necessary when working with many 3D datasets. When a multi-GPU model is saved as weights only, it is difficult to re-load the model because the neural network architecture expects single-GPU weights to be imported. A clever solution is to re-save the multi-GPU weights as single-GPU by extracting the weights of the second-to-last layer in the multi-GPU model. The below code provides an example:
from tensorflow.keras.utils import multi_gpu_model# Assumes the BCNN architecture is defined in the model file.from model import bcnndef load_model(input_shape, weights_path, num_gpus):    """Loads model from .h5 file. If model is       saved as multi-gpu, re-saves it as single-gpu."""    # Loads model as multi-gpu, if possible.    try:        model = bcnn(input_shape)        model = multi_gpu_model(model, gpus=num_gpus)        # Converts .h5 file to single-gpu.        model.load_weights(weights_path)        model = model.layers[-2]        model.save_weights(weights_path)    except ValueError as e:        pass    # Loads single-gpu model.    model = bcnn(input_shape)    model.load_weights(weights_path)    # Converts to multi-gpu model if applicable.    if num_gpus > 1:        model = multi_gpu_model(model, gpus=num_gpus)    return model
In addition, 3D dataset management can become overwhelming very quickly. Many CT and MRI scans can be of dimension 1000 x 1000 x 1000 or even larger, which is infeasible to perform inference on all at once. Instead, a “chunking” technique is required, which separates the large volume into overlapping chunks for feeding into the BCNN. The neural network is then trained on the chunks and predicts on chunks of the same size, which can be reconstructed to obtain a fully-inferred complete volume by reversing the chunking process.
The chunking process involves passing a sliding rectangular prism “window” across the original volume at a certain ratio of overlap, called the “step size”. We have to be careful to avoid off-by-one errors, and we also need to save the coordinates of each chunk to use during reconstruction. The output of the algorithm is an enormous 5D numpy array which contains all of the chunks in the dataset. Furthermore, there is a Keras bug that causes an error whenever the last batch in an epoch is not divisible across all GPUs; avoiding this bug requires truncating the end of the array.
An implementation of this algorithm is in the “chunks” method here: https://github.com/sandialabs/bcnn/blob/master/dataset.py.
Recall that BCNNs are nondeterministic, so when predicting on a chunk multiple times, one will obtain many different (and possibly very wrong) predictions. In order to obtain an accurate prediction as well as uncertainty maps, we must predict many times on each chunk to obtain a distribution of sigmoid values. This is called Monte Carlo sampling. This process is heavily customizable, but here we represent the segmentation as the average of all sigmoid values (casting to 0 and 1 for binary segmentations), and the uncertainty map as the difference between the 20th and 80th percentiles.
Critically, the chunking reconstruction process can lead to heavy artifacts in the output segmentation volumes. This is because the neural network does not have enough spatial context to predict effectively on the edges of every chunk; instead, we discard a certain percentage around each chunk (about 5%) to ensure that we only keep justified predictions.
An implementation of this algorithm is in the “predict” method here: https://github.com/sandialabs/bcnn/blob/master/test.py.
The BCNN is the new state-of-the-art in uncertainty quantification for volumetric segmentation; in particular, we validated the BCNN using CT scans of graphite electrodes for lithium-ion batteries and laser-welded metal joints. The BCNN provides greatly improved uncertainty maps as compared to the previously superior MCDN while achieving an equal or better segmentation accuracy. See below for a sample image from our paper which highlights the continuity and visual gradients of the BCNN uncertainty map, while the MCDN generates an uninterpretable pointwise uncertainty map.
Furthermore, we employ the PAvPU metric8, a recent measure designed to validate uncertainty results, and find that the BCNN consistently and vastly outperforms the MCDN at encoding the relationship between uncertainty and accuracy. See our paper at https://arxiv.org/pdf/1910.10793.pdf for an in-depth analysis and validation of the BCNN and its advantages over the MCDN, and feel free to utilize and fork the open-source codebase at https://github.com/sandialabs/bcnn. While our novel contribution is in the 3D space, we also provide a 2D implementation for typical image segmentation.
[1] Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. In Proceedings of the 32nd International Conference on Machine Learning, 2015. https://arxiv.org/abs/1505.05424
[2] Ozgun Cicek, Ahmed Abdulkadir, Soeren S. Lienkamp, Thomas Brox, and Olaf Ronneberger. 3d u-net: Learning dense volumetric segmentation from sparse annotation. In Proceedings of the 19th International Conference on Medical Image Computing and Computer-Assisted Intervention, 2016. https://arxiv.org/abs/1606.06650
[3] Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of the 33rd International Conference on Machine Learning, 2016. https://arxiv.org/abs/1506.02142
[4] Alex Graves. Practical variational inference for neural networks. In Proceedings of the 24th Conference on Advances in Neural Information Processing Systems, 2011. https://papers.nips.cc/paper/4329-practical-variational-inference-for-neural-networks
[5] Geoffrey E. Hinton and Drew Van Camp. Keeping neural networks simple by minimizing the description length of the weights. In Proceedings of the 16th Conference on Learning Theory, 1993. https://www.cs.toronto.edu/~hinton/absps/colt93.html
[6] Diederik P. Kingma, Tim Salimans, and Max Welling. Variational Dropout and the Local Reparameterization Trick. In Proceedings of the 28th Conference on Advances in Neural Information Processing Systems, 2015. https://arxiv.org/abs/1506.02557
[7] Fausto Milletari, Nassir Navab, and Seyed-Ahmad Ahmadi. V-net: Fully convolutional neural networks for volumetric medical image segmentation. In Proceedings of the 4th International Conference on 3D Vision, 2016. https://arxiv.org/abs/1606.04797
[8] Jishnu Mukhoti and Yarin Gal. Evaluating bayesian deep learning methods for semantic segmentation. arXiv preprint 1811.12709, 2019. https://arxiv.org/abs/1811.12709
[9] Ian Osband. Risk versus uncertainty in deep learning: Bayes, bootstrap and the dangers of dropout. In Proceedings of the 29th Conference on Advances in Neural Information Processing Systems: Workshop on Bayesian Deep Learning, 2016. https://pdfs.semanticscholar.org/dde4/b95be20a160253a6cc9ecd75492a13d60c10.pdf
[10] Dustin Tran, Michael W. Dusenberry, Mark van der Wilk, and Danijar Hafner. Bayesian layers: A module for neural network uncertainty. 2019. https://arxiv.org/abs/1812.03973
[11] Yeming Wen, Paul Vicol, Jimmy Ba, Dustin Train, and Roger Grosse. Flipout: Efficient pseudo-independent weight perturbations on mini-batches. In Proceedings of the 6th International Conference on Learning Representations, 2018. https://arxiv.org/abs/1803.04386
[12] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the 2018 European Conference on Computer Vision, 2018. https://arxiv.org/abs/1803.08494
This work was produced during the author’s internship at Sandia National Laboratories. Sandia National Laboratories is a multimission laboratory managed and operated by National Technology & Engineering Solutions of Sandia, LLC, a wholly owned subsidiary of Honeywell International Inc., for the U.S. Department of Energy’s National Nuclear Security Administration under contract DE-NA0003525. SAND2020–1419 S.