Challenges and Opportunities in Approximate Bayesian Deep Learning for Intelligent IoT Systems

MCMC methods provide an approximation to the intractable parameter posterior $p(\theta|\mathcal{D},\lambda)$ via a set of samples drawn for this distribution. MCMC methods simulate a Markov chain that converges to the parameter posterior as its steady state distribution. The simulated states of the Markov chain after convergence correspond to samples from the parameter posterior. Once we collect a set of parameter samples of the desired size, we can approximate expectations with respect to the parameter posterior using empirical averages over the sampled parameter values [24]. For example, the posterior predictive distribution can be approximated as shown in Equation (2.2.1). As we can see, computing the Monte Carlo approximation to the posterior predictive distribution is very similar to computing the predictive distribution of a model ensemble.

	$\displaystyle p(y|\mathbf{x},\mathcal{D},\lambda)$	$\displaystyle=\mathbb{E}_{p(\theta|\mathcal{D},\lambda)}[p(y|\mathbf{x},\theta)]$
		$\displaystyle\approx\mathbb{E}_{p_{\textrm{MC}}(\theta|\mathcal{D},\lambda)}[p(y|\mathbf{x},\theta)]$
		$\displaystyle=\frac{1}{S}\sum_{s=1}^{S}p(y|\mathbf{x},\theta_{s})$		(4)

Examples of classical MCMC methods include the Gibbs sampler [25] and the Metropolis-Hastings sampler [26]. The earliest work on MCMC samplers for neural networks traces back to the application of Hamiltonian Monte Carlo [27] methods. While a number of MCMC methods have since been developed with improved properties including slice sampling [28], elliptical slice sampling [29], and Riemann manifold sampling methods [30], these methods all require using all of the available data when computing the likelihood term needed for posterior inference. Although only linear in the number of data cases, this can be a highly expensive operation for large data sets and models and can render MCMC methods practically infeasible in scenarios where stochastic gradient descent (SGD) [31] can be usefully applied to optimize model parameters.

However, recent advances in MCMC approaches have enabled the use of SGD-like mini-batch algorithms, greatly extending the range of applicability of MCMC methods. Prominent examples of such approaches include stochastic gradient Langevin dynamics (SGLD) [32], stochastic gradient Hamiltonian Monte Carlo (SGHMC) [33] and their cyclic learning rate versions as presented in [34].

Another important property that determines the efficacy of MCMC methods is the degree of mixing. The degree of mixing refers to how efficiently the Markov chain traverses the posterior distribution after convergence [35]. Better mixing enables faster collection of a more diverse set of parameter samples. However, the mixing properties of MCMC methods depend on the dimensionality of the parameter space. Modern deep neural networks can have an extremely large number of parameters (millions or more), potentially leading to inadequate mixing.

An alternative to sampling in the original parameter space $\mathbb{R}^{K}$ is to sample in a reduced dimensional space $\mathbb{R}^{K^{\prime}}$ for $K^{\prime}\ll K$ and to project back to the full-dimensional space. This process is termed subspace inference [36]. Subspaces can be generated using singular value decomposition (SVD) applied to SGD iterates, or by generating random projection matrices, among other possible options [36]. These methods introduce bias by restricting the sampler to operate in a subspace of $\mathbb{R}^{K}$, but reduce variance by enabling the underlying Markov chain to mix faster.

While recent advances in MCMC methods are enabling the application of Bayesian inference approaches to increasingly large deep learning models, there remains a sizeable gap in terms of the practical deployment of Monte Carlo-based posterior approximations on edge hardware to support intelligent IoT systems. As noted previously, MCMC methods generate $S$ samples from the model posterior in place of the single set of parameters used by optimization-based deep learning. This means that the deployment of Monte Carlo-based approximate prediction methods requires $S$ times more storage as well as $S$ times more computation without further approximations.

However, the actual increase in prediction latency depends on a number of factors. First, in the edge setting, it is typical for an edge device to process data from on-board sensors. In this case, the edge device may only need to make predictions for a single instance at a time. For GPU or TPU accelerated edge devices, it may thus be possible to run a single input through multiple models with identical structure in parallel, effectively batching the models instead of the inputs. The feasibility of this approach depends on the size of the model to be deployed. Second, many edge platform toolchains have the ability to perform weight quantization and to use mixed precision arithmetic when deploying models. Such transformations can be applied separately to each element of a posterior ensemble to improve edge deployment efficiency. We note that the effect of such quantization has not been well studied for Bayesian posterior ensembles, but there is reason to believe that they might better tolerate more aggressive quantization than single models due to the average that is taken over the outputs of the posterior ensemble when making predictions.

Lastly, we note that the memory constraints of edge devices can also play a non-trivial role in prediction latency when using ensembles due to the time required to load models. In some instances, the time needed to load models is significantly higher than the time needed to use the loaded model to make predictions for a single input. When this is the case, it can dramatically increase prediction latency. Indeed, the number of models that can be present in memory simultaneously can currently be the effective limiting factor on the size of the MCMC posterior ensemble that can usefully be deployed.

An interesting opportunity motivated by these observations is the development of optimized model loading methods specifically for ensembles of models with identical architectures. In this case, the structure of the computation graph underlying the architecture should only need to be instantiated once, and it may be possible to more rapidly iterate through different sets of parameters for the same architecture than is possible using current libraries that do not include such optimizations [37].

An alternate approach to MCMC methods is approximating the original posterior distribution via an analytically tractable parameterized surrogate distribution. Thus, given the original posterior distribution $p(\theta|\mathcal{D}_{tr},\lambda)$, surrogate density methods aim to approximate the true posterior using an auxiliary distribution $q(\theta|\phi)$, where $\phi$ are the auxiliary parameters [38, 39, 40, 41]. A common approximation for auxiliary parameters is the use of a multivariate Gaussian distribution with a diagonal covariance matrix $\mathcal{N}(\theta;\mu,\Sigma)$, also known as mean-field variational inference. Here, the auxiliary parameters are $\phi=[\mu,\Sigma]$. The main reason why mean-field variational inference is popular is due to the simple “re-parameterization trick” that makes sampling from the auxiliary distribution straightforward [42]. With the re-parameterization trick, we can sample $\theta_{k}\sim\mathcal{N}(\mu_{k},\Sigma_{kk})$ by first drawing $\eta\sim\mathcal{N}(0,1)$, followed by the linear transformation $\theta_{k}=\mu_{k}+\eta\cdot\sqrt{\Sigma_{kk}}$. This allows us to backpropagate through the variational parameters while drawing samples of the model parameters to approximate the objective functions used for learning.

Now, for estimating the auxiliary parameters, we need an objective function that can measure the discrepancy between the surrogate distribution and the ground truth posterior. Needless to say, this objective function must also be computationally tractable. The most commonly used discrepancy function is the Kullback-Leibler (KL) divergence as shown below [43].

$$\textrm{KL}(p(\theta)||q(\theta))=\mathbb{E}_{p(\theta)}\left[\log\left(\frac{p(\theta)}{q(\theta)}\right)\right]$$

The KL divergence is a directional divergence, and thus is not symmetric in its arguments. This can result in two different measures depending on the directionality. When the surrogate posterior is used as the first argument, the result is the variational inference (VI) framework [38, 39]. When the surrogate posterior is used as the second argument, the result is the expectation propagation (EP) framework [41]. This often leads to VI methods having mode seeking behavior, as they are not forced to ensure the same support as the original posterior. EP methods on the other hand are forced to ensure exact support, but this can often lead to incorrect mode estimation as it must cover the support of the original posterior. These two extremes can also be interpolated by more generalized divergence measures, including alpha divergence [44].

These methods also suffer from scalability issues due to the need to compute the log likelihood and its gradient over the entire dataset. Similar to the stochastic gradient version of MCMC methods introduced earlier, advances in the past decade have led to more scalable methods in this family as well, including stochastic variational inference [45], that are able to accommodate large-scale datasets using mini-batch gradients.

In addition to mean field VI, other more advanced approximations are possible, such as multiplicative normalizing flows (MNF) [46], Bayesian hypernetworks [47], and Rank-1 factorization [48]. In multiplicative normalizing flows, we choose a simple density function such as the isotropic Gaussian distribution, and use a bijective function to transform the samples drawn from the simple density function to form more complex distributions. The Bayesian hypernetworks approach builds upon the MNF approach and uses a neural network model to transform the samples drawn from a simpler density function to model a more complex posterior distribution. Finally, the rank-1 factorization approach represents the model parameters as a product of two rank-1 factors thereby reducing the dimensionality of the base distribution of the approximate posterior. For example, if we have a parameter matrix $\pazocal{W}\in\mathbb{R}^{M\times N}$, this can be re-written as a matrix product of $W_{1}\in\mathbb{R}^{M\times 1}$ and $W_{2}\in\mathbb{R}^{1\times N}$. This effectively reduces the number of parameters from $\pazocal{O}(M\cdot N)$ to $\pazocal{O}(M+N)$.

MC Dropout is a particularly interesting approach, which is equivalent to approximate variational inference under specific assumptions [49]. Dropout itself was first introduced as a regularization technique where during every training iteration a pre-determined proportion of activations is randomly set to zero to reduce overfitting. At test time dropout is switched off and all units participate in making predictions [50]. In MC Dropout, by contrast, dropout is used at prediction time. This leads to a stochastic forward pass through the point-estimated model. Multiple forward passes through the model are used and the predictive distributions are averaged. This procedure is equivalent to sampling from a specific approximate variational posterior, but has the advantage that it is very easy to implement.

The major drawback of surrogate density methods is that they introduce bias into the estimation of the posterior distribution, unless the true posterior belongs to the family of auxiliary distributions. The degree of bias will depend on the functional form of the auxiliary distribution, and the divergence measure used to estimate the parameters of the auxiliary distribution.

In addition, under these methods, we still face hurdles in computing posterior expectations including the approximate posterior predictive distribution. In particular, even if $q(\theta|\phi)$ is a simple parametric distribution, the expectation $\mathbb{E}_{q(\theta|\phi)}[p(y|\mathbf{x},\theta)]$ still usually cannot be computed analytically due to the non-linearity of $p(y|\mathbf{x},\theta)$. As a result, we have to again resort to Monte Carlo approximation, and draw samples from the approximate variational posterior. Generally, these methods trade-off the potential bias in the surrogate parameter posterior for the ability to draw independent samples once the surrogate posterior parameters have been estimated.

The deployment considerations for models derived using variational inference methods are distinct from those of MCMC methods. One of the benefits of mean field variational inference over MCMC methods is that the representation of the posterior approximation only requires twice the storage of a single model. However, for this to be useful at deployment time, it must be possible to efficiently sample models from the approximate parameter posterior on the fly. The same is true of approaches like MC Dropout that require the sampling of random masks at inference time. In the case of MC Dropout, for example, current model conversion pipelines from PyTorch [51] to TensorRT [37] via ONNX [52] strip out stochastic dropout layers. It appears that for such approaches to have practically useful storage advantages over MCMC methods on current edge hardware, additional development may be needed to enable the automated deployment of computation graphs with stochastic elements. An alternative for hardware where the computational cost of on-the-fly generation of samples from the variational posterior is prohibitive is to pre-generate and deploy a fixed variational approximate posterior ensemble. However, in this case, all the considerations that apply to the generation of MCMC posterior ensembles will also apply.

With the emergence of Generative Adversarial Networks (GANs) [53], there has been work on learning implicit generative model representations of the parameter posterior that can be used at test time to draw an arbitrary but finite number of samples from this posterior approximation. The existing work in this area involves training GANs that can approximate the posterior distribution asserted by SGLD [17, 54]. To compute the approximate posterior predictive distribution, we yet again would use Monte Carlo approximation as shown in the previous subsections. An advantage of these methods is that in theory we do not need to store posterior samples for use during inference and we can also determine the number of posterior samples to use on the fly. However, the same deployment issues noted for variational inference apply here as well in terms of the feasibility of on-the-fly sampling from an auxiliary model. In addition, this approach requires allocating additional memory for the auxiliary GAN model.

Deep Ensembles [55] have also shown strong performance in terms of representing predictive distributions, and can also be considered as an approximation to Bayesian inference [56]. However, deep ensembles can be expensive during training, as we need to train each model in the ensemble from scratch. To alleviate this issue, snapshot ensembles [57] uses a cyclical learning rate schedule to generate more ensembles in less training time. However, these methods learn a mixture of posterior modes, which is different from an approximation to the full posterior. Regardless, Deep Ensembles have been shown to yield better performance relative to using small numbers of samples in a traditional MCMC approach. Deployment considerations similar to those of MCMC ensembles also apply to deep ensembles.

As described in this section, different methods for approximating Bayesian inference have different strengths and weaknesses in terms of their theoretical and practical properties. The first such property is the bias-variance tradeoff in approximate posterior estimation. MCMC methods provide an unbiased estimate of the posterior in theory, but this relies on convergence of the underlying Markov chains, which can be time-consuming. Mixing of the Markov chains is another important factor determining the diversity of the approximate posterior ensemble. On the other hand, we can learn surrogate posterior distributions much more efficiently using stochastic gradient descent, but they introduce bias that depends on the functional choice of the auxiliary distribution and the divergence measure used. However, while important, these factors largely impact training time and the quality of approximations, not their subsequent edge deployability.

In terms of edge deployability, a common aspect among all the methods discussed in this section is that they require the application of Monte Carlo-based approximations at test time. When $S$ samples are used, the computational complexity of making predictions is $S$ times higher than if a single instance of the base model is used. As noted previously, the use of mixed precision arithmetic can help to accelerate these computations, but a reduction from 32 bit to 16 bit representations will typically at most half the prediction latency of deployed models. This is likely to be far from sufficient when considering closing the prediction latency gap between single optimization-based models and model ensembles.

In addition, as noted previously, MCMC methods have $O(S\cdot K)$ storage cost where $K$ is the number of parameters in the base model and $S$ is the number of samples. By contrast, mean-field variational inference has $O(2K)$ storage cost, but there appear to be practical challenges to actually realizing the benefits of the reduced storage cost of variational methods on current edge platforms. In the next sections, we turn to possible modeling and algorithmic approaches to improving the edge scalability of approximate Bayesian inference methods in terms of both computation time and storage costs.