Monte Carlo DropBlock for Modelling Uncertainty in Object Detection

Bayesian neural networks (BNN) assume the functional form of the model to be same as any other neural networks $f^{W}(\mathbf{x})=W^{L}\sigma(W^{L-1}\ldots\sigma(W^{1}\mathbf{x}))$ with activation function $\sigma$ and parameters $W=[W^{1}\ldots W^{K}]$. Let us assume the dimensionality of layer $l$ to be $K^{l}$, and consequently the weight matrix $W^{l}$ has dimension $K^{l}\times K^{l-1}$. BNN assumes a distribution over the parameters and treats them as random variables. They assume a prior distribution over the weight vectors $p(W)$ which is often a Gaussian. The likelihood of generating output $y$ given $f^{W}(\mathbf{x})$ for some particular value of $W$ is modeled as a Gaussian in case of regression problems and Softmax in the case classification problems. BNNs estimate a posterior distribution over the parameters using the Bayes theorem.

The prediction is done by considering multiple samples of weight vectors from the posterior and taking an average over them which is known as Bayesian model averaging. This will allow BNNs to capture uncertainty on predictions and make them resistant to overfitting. However, posterior computation is intractable and requires using inference techniques such as Markov Chain Monte Carlo (MCMC) and Variational Inference (VI). For instance, VI assumes a particular distribution form for the posterior $q_{\theta}(W)=\prod_{i=1}^{L}q_{\theta^{l}}(W^{l})$, where $q_{\theta}(W^{l})$ is assumed to be a Gaussian with parameters $\mu^{l}$ and $\Sigma^{l}$ (variational parameters $\theta^{l}$), and learns the variational parameters by minimizing the Kullback-Leibler (KL) divergence between $q(W)$ and $p(W\mid X,\textbf{y})$. As direct minimization of the KL divergence is intractable, VI learns the variational parameters $\theta$ by maximizing evidence lower bound (ELBO) which can be shown to equivalently minimize the KL divergence. The ELBO objective function in terms of minimization can be written as follows

Dropout[32] technique was introduced to avoid over-fitting in deep neural networks. During training, dropout randomly drops some of the neurons in every layer of a neural network based on the dropout probability $p$. This will prevent co-adaptation of neurons and mitigates over-fitting and lead to good generalization performance. During testing, all the neurons are considered for prediction but each contributing according to the probability they are retained during the training phase.

[21] propose to apply dropout additionally at the testing time and showed that this will lead to a Bayesian neural network. The approach known as Monte-Carlo dropout (MC dropout) has become a very effective approach to obtain a Bayesian neural network capable of providing uncertainty estimates without the need to employ sophisticated inference techniques.

Monte Carlo dropout can be shown to be equivalent to a Bayesian neural network employing variational[33] inference with a specific variational distribution. Dropout performed on the feature space can be seen as an operation in the weight space where random columns of the weight matrix mapping one layer to the next layer are dropped. This observation leads to the following form for the variational distribution $q(W^{l})$

where $M^{l}$ is the $K^{l}\times K^{l-1}$ dimensional weight vectors, and $W^{l}$ follows a distribution defined by (3). Consequently, $\hat{W}^{l}$ sampled from $q(W^{l})$ will have some of the columns set to zero (leading to dropout) depending on $\hat{\bm{z}}^{l}_{j}$ sampled from $p(\bm{z}^{l}_{j})$ which is $Bernoulli(p)$¹¹1Please note that we are considering the parameter $p$ of the Bernoulli distribution as probability of getting zero. . We can use this variational distribution in (2) to derive the objective function as follows

We can easily observe that this objective function is the same as the one used to train the deep learning models with dropout regularization. The negative log-likelihood in (5) is cross-entropy loss for classification problems or least-squares loss for regression problems. We consider multiple $\hat{W}$ corresponding to different samples from $q(W)$ across different mini-batches during training. BNN uses the posterior $q(W)$ for making predictions by sampling $\hat{W}$ from $q(W)$ and computing an average over multiple samples of $\hat{W}$ (Monte Carlo sampling) and consequently over multiple dropout architectures.

Applying dropout at test time differentiates MC dropout from the standard dropout. The model averaging at test time allows it to model the uncertainty of the model in its predictions (epistemic uncertainty). Epistemic uncertainty estimation is especially useful when the model has to deal with out-of-distribution data.

YOLO[11] is known for its one-stage detection and state-of-the-art performance. It’s the most used network architecture in the world of object detection because of its fast speed and accuracy. In YOLO the last convolution layer splits the image into nxn cells, generally 19x19. Each of these cells is responsible for k bounding boxes (in general, k is chosen as 3). The training is faster because a single architecture performs both the tasks i.e. simultaneously predicting the multiple bounding boxes and class probabilities. In the convolutional layers of yolov3 (Figure 4), kernels of shape 1x1 are applied on feature maps of three different sizes at three different places in the network. This leads to 3 predictions at 3 different scales, which happen due to downsampling the dimensions of the image by a stride of 32, 16, 8 respectively. Downsampling is done to reduce the size of data while maintaining the same spatial resolution. Every scale uses three anchor bounding boxes per layer. The three largest boxes for the first scale, three medium ones for the second scale, and the three smallest for the last scale. This way each layer excels in detecting large, medium, or small objects. The recent version of YOLO, introduced by Alexey as YOLOv4[14] and YOLOv5[34] are significant performance improvements compared to YOLOv3[27].

Instance Segmentation tasks have to generate masks, which are spatially coherent since the nearby pixels are more likely to belong to the same instance. A fully connected layer couldn’t model the spatial features well compared to a fully convolutional layer. YOLACT (You Only Look At CoefficienTs)[35] has been experimented with because of its real-time inference capability with good accuracy for real-time instance segmentation. This method combined the qualities of single-stage methods and two-stage methods by parallelizing prototype mask generation and prediction of mask coefficients per instance. By choosing methods that are way above the typical real-time threshold (30 FPS), we have some flexibility of adding minor costs at inference without going to sub-realtime outputs.

In this section, we show that the drop block applied to the feature maps in the neural network layers at train time and test time will result in a Bayesian neural network. As already discussed in Section 3.1, approximate inference techniques for BNNs such as variational inference assumes a specific form for the posterior to be learned. We show that a particular form of the variational posterior will result in operations similar to applying drop block at training time and test time in deep learning models.

Let us assume the feature representation obtained at some particular layer $l$ to be represented as $A^{l}$. For ease of understanding, let’s assume it’s a square matrix of size $K^{l}\times K^{l}$ ²²2The discussion can be generalized to feature matrix, filter matrix, block size, and stride lengths of arbitrary size and to multiple channels.. Consider a convolution operation using a filter (kernel) $W^{l}$ of size $L\times L$. The feature representation at the next layer is obtained by performing a convolution operation over the feature matrix $A^{l}$ and filter $W^{l}$ assuming a stride length $L$. This will lead to a feature representation $A^{l+1}$ of size $\frac{K^{l}}{L}\times\frac{K^{l}}{L}$. There are ${\frac{K^{l}}{L}}^{2}$ blocks in the feature matrix $A^{l}$, each of them are convolved against the filter $W^{l}$ to obtain the $\frac{K^{l}}{L}^{2}$ feature vector. Now let us consider applying drop blocks in the layer $l$ with drop probability $\gamma$ and block size $L\times L$. Assume $i^{th}$ block of size $L\times L$ in the feature matrix $A^{l}$ got masked as a result of applying drop block. Consequently, the $i^{th}$ feature representation in layer $l+1$, $A^{l+1}(i)$ will be zero.

To apply MC dropout theory, we need to first represent drop block operations as operations using the weight vectors (filters in the convolution operation). We can observe that this can be achieved by considering a larger weight matrix $\bar{W}^{l}$ of size $K^{l}\times K^{l}$ with $\frac{K^{l}}{L}^{2}$ blocks, with each block being the $L\times L$ matrix $W^{l}$, except the $i^{th}$ block being all zeros. For brevity, we consider the feature matrix $A$ as a 3-dimensional tensor with $L\times L$ blocks stacked up in the third dimension of size $\frac{K^{l}}{L}^{2}$. Similarly, $\bar{W}^{l}$ is a 3-dimensional tensor with $W^{l}$ blocks stacked one after the other and a block with zeros as the $i^{th}$ element in the third dimension. We can now think of getting the feature vector $A^{l+1}$ as performing convolution between feature matrix $A^{l}$ and weight vector $\bar{W}^{l}$ element-wise in the third dimension (see Figure 3).

We can see that considering $\bar{W}^{l}$ with blocks of all zeros appearing as some element in the third dimension of the tensor is equivalent to considering a feature representation with some masked block. Now we see how to define a variational distribution over $\bar{W}^{l}$ which will lead to a sample of $\bar{W}^{l}$ as we considered. To define the variational distribution, we consider the tensor $\bar{W}$ to consist of blocks $W_{1},W_{2},\ldots,W_{\frac{K^{2}}{L^{2}}}$ stacked one after the other in the third dimension. We assume the variational distribution $q(\bar{W})$ can be decomposed as the product of the variational distributions over $W_{i}$’s, i.e. $q(\bar{W})=\prod_{i}q(W_{i})$. The blocks $W_{i}$ can take value $0$ with drop probability $\gamma$ or can take value $W$ with probability $1-\gamma$. Thus, we define the variational distribution $q(W_{i})$ over the block $W_{i}$ as $W_{i}=W\cdot z_{i}$ and $z_{i}$ sampled from $Bernoulli(\gamma)$³³3Dropblock paper considered the probability of getting $0$ as the parameter to the Bernoulli distribution and we follow the same convention.. Thus the variational distribution over $\bar{W}$ can be summarized as

To predict the aleatoric uncertainty in bounding boxes, Gaussian Yolo [28] proposes to model each of the four coordinates as a Gaussian (with mean and variance parameters ) and re-designing the loss function as negative log-likelihood (NLL). As a result, the number of outputs for the bounding box becomes 8($\mu$ and $\sigma$ for four coordinates) instead of 4 (x,y,w,h). This will lead to a change in detection criteria by including localization uncertainty. We consider a Generalized IoU(GIoU)[36] loss, which uses orientation and shape of the object in addition to the area covered. This gives a better result when used along with Gaussian YOLO for object detection. We combine the MC-DropBlock approach with the Gaussian Yolo approach, resulting in object detection models which are capable of modeling both epistemic and aleatoric uncertainties.

For out-of-distribution (OOD) experiments for object detection, we trained our model on the Pascal VOC dataset and then tested it on the COCO dataset. Pascal VOC [39] data has 11,530 images belonging to 20 classes. COCO dataset [40] for the object detection task consisting of approximately 164K images across 91 classes. All the 20 classes of the Pascal VOC dataset (car, person, bus, etc) are also present in the COCO dataset. COCO data set contain classes, not in Pascal VOC and hence we expect the model to exhibit high uncertainty on COCO in this experimental setup. The hyperparameters used in YOLO models are the same as described in the object detection task.

We have used entropy as the evaluation metric for measuring uncertainty. Table 1 shows the results of out-of-distribution experiments. We can see that MC-DropBlock provided the highest entropy for various YOLO models as well as ResNets. It outperformed MC-Dropout and other ways of adding DropBlock in these models for object detection and classification tasks.

We have used the same experimental setup as [14] for our experiments on YOLOv4 with COCO 2017 dataset having training/validation split as 118K/5K and the test split has 41K. The hyperparameters for YOLOv4 and Gaussian YOLOv4 experiments are: 500500 training steps, learning rate 0.01 multiplied with a factor of 0.1 at 400000 steps and 450000 steps, momentum of 0.9, weight decay of 0.0005, batch size of 64, and mini-batch size of 4. For our experiments on YOLOv5, we have used the same setup as [34]. YOLOv5 has 4 variations(s, m, l, x) for various applications. We experiment on the YOLOv5(x) model is the same as YOLOv4 and gives state-of-the-art mAP. The experimental setup for YOLOv5 and Gaussian YOLOv5 is the same. We have experimented with various drop block sizes ($3\times 3,5\times 5,7\times 7$) and drop probabilities ($0.1,0.3,0.5,0.7$). We found that $7\times 7$ drop block size with drop probability $0.1$ gave optimal mAP score (on validation data) for MC-DropBlock and training time dropBlock, While for inference time dropBlock $3\times 3$ gave the best result.

Table 2 shows the mAP score of the models trained and tested on the COCO dataset. A higher mAP value indicates better precision. From the Table, we can observe that the YOLOv4/v5 models give the best mAP value for train time DropBlock, and MC-DropBlock gives the second-best mAP value. As we see later, MC-DropBlock in addition provides a substantial improvement in the uncertainty modeling capabilities of these models. We can also observe that adding Gaussian loss to model aleatoric uncertainty has improved the mAP score of the YOLO models.

Table 3 provides the calibration ability of the models measured in terms of the Brier score. We can see from Table 3 Brier score is the lowest for both the algorithm for the Inference time DropBlock. In the calibration experiments, MC-DropBlock gives the second-best result and is better than the one without DropBlock or with training time drop block alone.

For instance segmentation, we conducted experiments on the MS COCO [40] segmentation data. It has pixel-wise mask value along with bounding box value in annotation files. For MS COCO, we train on train2017 and evaluate on val2017 and test-dev. We conducted experiments with the same setup as YOLACT-550[35] segmentation model on COCO 2017 dataset. We performed our experiments with batch size 8 on a single GPU using pre-trained weights of ImageNet [46] and with the same configuration used by the authors. For YOLACT-550 with Resnet-50-FPN, we experimented with DropBlock having block size 7x7 and drop probability of 0.1 before the prediction head. From Table 4 we can see inference time DropBlock shows an improvement of 1.1% on mask mAP value, which measures the quality of segmentation mask and 1.05% improvement on box mAP, which measures the quality of bounding boxes. MC-Drop block and training time DropBlock also give a better mAP value than the baseline model.

We also experimented with ResNet-110 [47] model for image classification on the CIFAR-10 Dataset [48]. CIFAR-10 contain images from 10 classes with 50000 images for training and 10000 images for testing. DropBlock with block size $7\times 7$ and drop probability $0.1$ is added after the final convolutional layer for MC DropBlock. For MC Dropout, drop probability is set to $0.5$. We obtained $92.08\%$ top-1 accuracy using baseline ResNet-110 model, $92.26\%$ top-1 accuracy using MC Dropout and 92.98% top-1 accuracy using MC DropBlock. We conducted out-of-the-distribution experiments for image classification by training on CIFAR-10 and testing on the Imagenet [49]. Imagenet consists of 1000 classes and the experiments are conducted on a test set consisting of 1,50,000 images. As observed in Table 1, we obtained the highest entropy score for the MC-Drop block, and consequently, it has the best uncertainty modeling capability on image classification models like Resnet-110.

We conduct qualitative analysis to study the out-of-distribution object detection capabilities of the proposed MC-DropBlock approach on YOLOv5. We train our models on Pascal VOC and tested on images from Open Images Dataset [50] and COCO. From Figure 5, we can observe that the baseline YOLOv5 model (Figure  5(a)) and MC-Dropout on YOLOv5 (Figure  5(b)) identifies aeroplane wrongly. While the MC-DropBlock (Figure  5(c)) on YOLOv5, no longer miss-classifies the object as an aeroplane. Similarly, in Figure 6, baseline YOLOv5 and MC-Dropout on YOLOv5 detected the bread loaf incorrectly as sheep, while MC-DropBlock does not make any false predictions.