Scalable Bayesian Deep Learning with Kernel Seed Networks

Modern deep neural networks (DNNs) are capable of learning complex representations of high dimensional data (He et al., 2016; Simonyan and Zisserman, 2014). This has enabled DNNs to exceed human performance in a growing number of decision making tasks (He et al., 2015; Mnih et al., 2015; Silver et al., 2017). However, these algorithms are notoriously difficult to implement in real world decision making systems because conventional DNNs are prone to overfitting and thus have poorly calibrated measures of uncertainty (Blundell et al., 2015). A well-calibrated measure of uncertainty provides insight into the reliability of model predictions (Guo et al., 2017). This is crucial for applications such as autonomous vehicle navigation (Kendall and Gal, 2017), and computer assisted diagnosis (Jiang et al., 2012); where confidently incorrect decisions can potentially have fatal consequences. Hence, there has been a resurgent interest in exploring Bayesian methods for deep learning as they naturally provide well-calibrated representations of uncertainty.

Unlike conventional DNNs, which learn a fixed point estimate of model weights, Bayesian Neural Networks (BNNs) learn a weight distribution, parameterizing by the mean and standard deviation of the weights at each layer. The parameter distributions are optimized using variational sampling during training, which forces BNNs to be robust to perturbations of model weights.

Effectively, this makes BNN optimization equivalent to learning an infinite ensemble of models (Blundell et al., 2015). With an ensemble, the variance of predictions among constituent models can be used to approximate the epistemic uncertainty (Lakshminarayanan et al., 2017). Similarly, epistemic uncertainty in BNNs can be estimated using the variance of predictions among weights that are randomly sampled from the learned distributions.

While BNNs have proven to be effective at inducing uncertainty into the weights of a model, they require significantly more parameters than their DNN counterparts (Shridhar et al., 2019). This makes it impractical to scale variational Bayesian learning to the depth of modern DNN architectures. To overcome this problem, Gal and Ghahramani (2016) proposed the use of Monte-Carlo Dropout (MC-Drop); a method for approximating Bayesian inference by applying Dropout (Srivastava et al., 2014) during both model training and inference.

A major advantage of MC-Drop, is that it requires no additional parameters to approximate Bayesian inference. MC-Drop represents model uncertainty by sampling sub-networks with randomly dropped out units and connections to perform variational inference (Gal and Ghahramani, 2016). Although this approach captures the uncertainty of model activations, it is unable to represent the uncertainty of model weights since all sub-network perturbations sampled using the Dropout method will share the same parameters (Srivastava et al., 2014). This limits the scope in which MC-Drop can be applied, since many applications such as continual learning (Kirkpatrick et al., 2017) and model pruning (Graves, 2011; Blundell et al., 2015), require representations of uncertainty over the entire parameter space. To address these limitations, Khan et al. (2018) propose the Variational Online Gauss-Newton (VOGN) algorithm. VOGN parameterizes a distribution of weights by modelling the output of a DNN as the mean, and the $2^{nd}$ raw moment uncentered variance of computed gradients as the variance.

Similar to BNNs, VOGN optimization involves randomly sampling a set of parameters from a weight distribution to compute variational gradients, and perform variational inference. However, the key advantage of VOGN optimization, is that it does not require learning a separate set of mean and standard deviation parameters, since it models the variance with a vector that is already used in adaptive optimization methods such as Adam (Kingma and Ba, 2014). But it is important to note that VOGN optimization only finds an approximate solution to the BNN objective (Blundell et al., 2015). Tseran et al. (2018) demonstrate how this approximation can result in a trade-off between accuracy and ease of implementation when applied to Variational Continual Learning (VCL) methods (Nguyen et al., 2017). It was shown that using the VOGN approximation for VCL, results in worse performance in certain tasks comparing to conventional VCL, which optimizes the variance of the weight distributions directly (Nguyen et al., 2017), rather than employing the online estimate of the diagonal of the Fisher Information Matrix to update the variance of the parameters (Tseran et al., 2018; Khan et al., 2018). Hence, despite the high computational cost of BNNs, their are benefits to directly optimizing the mean and standard deviation parameters of a weight distribution during training.

In this paper, we introduce a novel class of BNNs, namely, Kernel Seed Networks (KSNs), which significantly reduce the number of parameters required for Bayesian deep learning (BDL). In contrast with conventional BNNs, which require doubling the size of the network in order to learn a separate set of mean and standard deviations weights for each of the parameter distributions, our method applies $1\times 1$ convolutional kernels to trainable “seed” vectors to decode the mean and standard deviation weight vectors. The ability of $1\times 1$ convolutional layers to map inputs to a higher-dimensional space allows for further reduction of parameters, since they can be used to upsample downsized seed kernels to the required filter size. In our experiments on MNIST (LeCun et al., 2010), FMNIST (Xiao et al., 2017) and CIFAR (Krizhevsky et al., 2009) datasets, we demonstrate how our proposed method can be used to conduct BDL with a $44$–$70\%$ reduction in the number of parameters (compared to conventional BNNs) without compromising on model performance.

As described in Section 1, the purpose of KSNs is to reduce the number of parameters required to construct the BNNs first described by Blundell et al. (2015). To this end, we propose to replace the $N$ dimensional conventional kernel filters in DNNs with self-germinating kernel filters (Figure 1) comprising three main components: (1) an $N$ dimensional kernel seed, (2) a kernel germination layer, and (3) a variational kernel layer. We describe the details of these components and our optimization protocol below.

The aim of the KSN is to directly optimize a set of $\mu$ and $\rho$ for a weight distribution without a 2-fold increase in the number of parameters, as required by conventional BNNs (Blundell et al., 2015; Shridhar et al., 2019). Thus, to validate the effectiveness of the KSN we evaluate the relationship between model size and classification accuracy. Specifically, in Figure 3 we observe how changes in the number of model parameters affects model performance on MNIST, FMNIST, CIFAR10 and CIFAR100 classification tasks. Furthermore, in Tables 1, 2, 3, and 4, we evaluate how the number of parameters (Params) and model size relative to a ResNet18 baseline (RS) affects the Accuracy (ACC), Negative Log Loss (NLL), Expected Calibration Error (ECE), Average Calibration Error (ACE) and Maximum Calibration Error (MCE) rates when performing variational inference. A well calibrated model should have low NLL, ECE, ACE and MCE rates on unseen test data (Osawa et al., 2019). Hence, they provide a means of comparing each method’s ability to produce well calibrated models with Bayesian inference.

The results in Figure 1 demonstrate how KSNs can approximate a distribution of network weights without a 2-fold increase in the number of parameters. From the results tabulated in Section 3, our KSN model achieves a higher classification accuracy than the conventional BNN in all classification tasks. Most notably, in the CIFAR$10$ (Table 3) and CIFAR$100$ (Table 4) classification tasks, a $\delta$ of $0.25$ was sufficient to outperform the standard BNN. This means that the proposed KSN could achieve higher classification accuracy than BNNs with only $15\%$ of the required parameters. Moreover, KSNs consistently produce better calibrated models overall compared to BNNs, as evident by the lower NLL, ECE, ACE and MCE in Tables 1, 2, 3, and 4.

The VOGN method could also outperform the standard BNNs with significantly fewer parameters. However, VOGN often failed to match the performance of MC-Drop and our proposed KSN method; particularly on the FMNIST (Table 2), CIFAR10 (Table 3), and CIFAR100 (Table 4) classification tasks.

Overall, MC-Drop consistently achieves high classification accuracy and can produce well calibrated models. The performance improvements with MC-Drop are most prominent on CIFAR10 (Table 3), and CIFAR100 (Table 4) classification tasks, where the other methods, including the proposed KSN method, suffer significant performance degradation compared to the baseline. However, since the MC-Drop method is unable to quantify uncertainty over the entire weight space, the benefits of MC-Drop for performing variational inference are limited to certain tasks (Section 1).

The proposed KSN method can significantly reduce the number of parameters required for BDL (compared to conventional BNNs) without compromising on performance. While conventional BNNs require a 2-fold increase in the number of parameters compared to its DNN equivalent, BDL can be applied using KSNs with up to $39\%$ fewer parameters than the equivalent DNN. This enables BDL to scale up to the depths of modern DNNs.