UNCERTAINTY QUANTIFICATION USING VARIATIONAL INFERENCE FOR BIOMEDICAL IMAGE SEGMENTATION

Deep learning motivated by convolutional neural networks has been highly successful in a range of medical imaging problems like image classification , image segmentation , image synthesis etc . However for validation and interpretability , not only do we need the predictions made by the model but also how confident it is while making those predictions . This is important in safety critical applications for the people to accept it . In this work , we used an encoder decoder architecture based on variational inference techniques for segmenting brain tumour images . We evaluate our work on the publicly available BRATS dataset using Dice Similarity Coefficient ( DSC ) and Intersection Over Union ( IOU ) as the evaluation metrics . Our model is able to segment brain tumours while taking into account both aleatoric uncertainty and epistemic uncertainty in a principled bayesian manner . 1 INTRODUCTION . Medical image segmentation is a challenging task for medical practitioners . It is costly , takes time and is prone to error . Hence there is a need to automate the manually done segmentation . Lately Neural Networks have shown great potential on a variety of medical image segmentation problems . The challenge with the approaches used in literature is that the model doesn ’ t predict the uncertainty associated . This is where Bayesian methods come into play as it gives a principled way of measuring uncertainty from the model predictions . Measuring uncertainty in the output predictions made by neural networks is important for interpretation and validation . Rather than learning the point estimates , Bayesian Neural Networks ( BNN ) learns the distribution over the weights . The training process of BNN involves first initializing the parameters of the neural network . Next the weights are sampled from some distribution ( like gaussian with zero mean and unit variance ) and both the forward pass and backward pass is done to update the weights using the conventional backpropagation algorithm . Monte Carlo dropout networks ( Kingma et al. , 2015 ) use dropout layers to approximate deep Gaussian processes which still lack theoretical understanding . Bayesian Convolutional Neural Network ( Gal and Ghahramani , 2015 ) use variational inference to learn the posterior distribution over the weights given the dataset . The problem with this approach is that it requires a lot of computation involving a lot of parameters , making this technique not scalable in practice . Variational Autoencoder ( Kingma et al. , 2015 ) which is based on generative models solves the above problems and has been successful in a number of tasks like generating images , texts , recommender systems etc . This approach comes with several challenges in its own right which have been successfully tackled in the literature . A random variable sampled from posterior distribution has no gradient so the conventional backpropagation techniques can ’ t be applied to it . Local Reparameterization Trick ( Kingma et al. , 2015 ) was proposed to tackle this by converting the random variable to a deterministic one for computation . The second challenge was the huge computational requirement since it required weight updates in every iteration . Bayes by Backprop algorithm ( Blundell et al. , 2015 ) tackled this by calculating gradients in backpropagation using a scale and shift approach by updating the posterior distribution in the backward pass . 2 RELATED WORK . 2.1 MEDICAL IMAGE SEGMENTATION . The problem of segmenting medical images have been successfully tackled in literature using mainly two techniques , first using a Fully Convolutional Network ( FCN ) ( Long et al. , 2015 ) and second those which are based on U-Net ( Ronneberger et al. , 2015 ) . The main characteristic of FCN architectures is that it doesn ’ t use fully connected layers at the end which have been used successfully for image classification problems . U-Net methods on the other hand uses an encoder-decoder architecture with pooling layers in the encoder and upsampling layers in the decoder . Skip connections connect the encoder layers to the decoder layer to create an additional path for the flow of gradients back in the backpropagation step . This helps in reducing overfitting due to many parameters involved while training the network . 2.2 BAYESIAN NEURAL NETWORK . Lately , there has been a revival of interest in bayesian methods as some of the inherent problems with deep learning could be solved using it . It is a scalable approach of avoiding overfitting in neural networks and at the same time gives us a measure of uncertainty . This is very important in critical applications where not only do we require the predictions made from the model , but also how confident it is while making its predictions . BNN can be considered as an ensemble of neural networks ( Gal , 2016 ) . It has two advantages over the standard neural networks , first it avoids overfitting and second it gives a measure of uncertainty involved . Instead of point estimates , the neural network learns posterior distribution over the weights given the dataset as defined in Equation 1. p ( ω|D ) = p ( D|ω ) p ( ω ) p ( D ) = ∏N i=1 p ( yi|xi , ω ) p ( ω ) p ( D ) ( 1 ) The predictive distribution can be calculated by approximating the integral as defined in Equation 2. p ( y∗|x∗ , D ) = ∫ Ω p ( y∗|x∗ , ω ) p ( ω|D ) dω ( 2 ) The challenge is that the posterior is often intractable in nature . To combat this , ( Neal , 1993 ) used Markov Chain Monte Carlo ( MCMC ) for learning the weights over the bayesian neural networks . Also ( Graves , 2011 ) , ( Blundell et al. , 2015 ) and ( Louizos and Welling , 2016 ) proposed independently a technique using variational inference techniques for approximating the posterior distribution . KL Divergence between the posterior and the true distribution can be calculated using Equation 3 . KL { qθ ( ω ) ‖p ( ω|D ) } : = ∫ Ω qθ ( ω ) log qθ ( ω ) p ( ω|D ) dω ( 3 ) Alternatively minimizing the KL divergence can be written in another form by maximizing the Evidence Lower Bound ( ELBO ) which is tractable . This is shown in Equation 4 . − ∫ Ω qθ ( ω ) log p ( y|x , ω ) dω +KL { qθ ( ω ) ‖p ( ω ) } ( 4 ) 2.3 VARIATIONAL INFERENCE . Variational inference finds the parameters of the distribution by maximizing the Evidence Lower Bound . ELBO consists of sum of two terms Kullback-Leibler ( KL ) divergence between two distributions and the negative log-likelihood ( NLL ) as defined in Equation 5. minKL ( qθ ( w ) ) ‖p ( w|D ) ) ( 5 ) The KL divergence is defined in equation 6 . KL ( q ( x ) ) ‖p ( x ) ) = − ∫ q ( x ) log ( p ( x ) q ( x ) ) ( 6 ) The posterior in the above equation contains an integral which is intractable in nature . The equation can be re written in Equation 7 . KL ( qθ ( w ) ) ‖p ( w|D ) ) = Eqθ ( w ) log qθ ( w ) p ( D ) p ( D|w ) p ( w ) = = log p ( D ) + Eqθ ( w ) log qθ ( w ) p ( w ) − Eqθ ( w ) log p ( D|w ) = log p ( D ) − L ( θ ) ( 7 ) The above equation can be decomposed into two parts one of which is the KL divergence between the exact posterior and its variational approximation which needs to be minimized and the second is ELBO term which needs to be maximized . This is shown in Equation 8. max θ log p ( D ) = max θ [ KL ( qθ ( w ) ) ‖p ( w|D ) ) + L ( θ ) ] ( 8 ) KL divergence is zero if exact posterior is equal to variational approximation . Since the KL divergence is always greater than zero hence the equation can be approximated by maximizing only the ELBO ( Kingma et al. , 2015 ) as defined in equation 9 . L ( θ ) = Eqθ ( w ) log p ( D|w ) − Eqθ ( w ) log qθ ( w ) p ( w ) = LD −KL ( qθ ( w ) ‖p ( w ) ) ( 9 ) 2.4 ALEATORIC UNCERTAINTY AND EPISTEMIC UNCERTAINTY . There are two types of uncertainty - aleatory and epistemic uncertainty where variance is the sum of both these . Bayesian Neural Networks can be considered an ensemble of neural networks initialized randomly which averages the test results in parallel ( Gal , 2016 ) . For final predictions , single mean and variance can be estimated as shown in Equation 10 and Equation 11 respectively . µc ( x ) = 1 M M∑ i=1 µ̂i ( x ) ( 10 ) σ̂2c ( x ) = 1 M M∑ i=1 σ̃2i ( x ) + [ 1 M M∑ i=1 µ̂2i ( x ) − µ̂2 ( x ) ] ( 11 ) The first term in variance denotes aleatoric uncertainty while the second denotes epistemic uncertainty . Bayesian Neural Network model for uncertainty estimation was done by ( Kendall and Gal , 2017 ) with the last layer representing the mean and variance of logits . The predictive distribution approximating the posterior distribution which gives a measure of uncertainty is defined in Equation 12. qθ̂ ( y ∗|x∗ ) = ∫ Ω p ( y∗|x∗ , ω ) qθ̂ ( ω ) dω ( 12 ) Aleatoric uncertainty is a measure of the variability of the predictions from the dataset hence it is inherent in the data present . The aleatoric uncertainty is the uncertainty that exists inside the dataset . It captures noise inherent in the observations , which arises from the distribution of data . It is highly dependent on bias and distribution inside the input data but not the size of training samples . Epistemic uncertainty on the other hand is a measure of the variability of predictions from the model which is tied to various metrics used for evaluation like accuracy , loss etc . It represents the uncertainty inside the model . The epistemic uncertainty will decrease with enough training samples . It is resulting from the limitation of knowledge and data of the system . 3 METHOD . 3.1 DATASET . To validate the performance of our proposed approach to generalization , publicly available datasets were used for brain tumour segmentation BRATS18 ( Menze et al. , 2015 ) and ( Bakas et al. , 2018 ) . It contains MRI scans of 175 patients with glioblastoma and lower grade glioblastoma . The images were of resolution 240×240×155 pixels . The ground truth labels were created by expert neuroradiologists . The sample from the dataset is shown in Figure 2 . 3.2 DATA AUGMENTATION . The following data augmentation methods were used to increase the size of dataset : 1 . Rescaling : We rescale the pixels values by rescaling factor 1/255 . 2 . Rotation : Random rotations with the setup degree range between [ 0 , 360 ] were used . 3 . Height and Width Shift : Shifting the input to the left or right and up or down was performed . 4 . Shearing Intensity : It refers to shear angle ( unit in degrees ) in counter-clockwise direction . 5 . Brightness : It uses a brightness shift value from the setup range .

Summary: The paper proposes a method for uncertainty quantification for biomedical image segmentation. The proposed method takes mean and std of segmentation from a backbone segmentation model and trains a VAE on top of it. The method is evaluated on only BRATS dataset.