Variational Inference for BNNs (e.g., Bayes by Backprop)

While MCMC methods provide asymptotically exact samples from the posterior distribution $p(w | \mathcal{D})$ of a Bayesian Neural Network's weights $w$, they often struggle with the high dimensionality and large datasets typical in deep learning. The computational cost per iteration and the number of iterations required for convergence can become prohibitive. This is where Variational Inference (VI) offers a compelling alternative. Instead of sampling, VI reframes Bayesian inference as an optimization problem, seeking an approximate distribution $q_\phi(w)$ from a tractable family that is closest to the true posterior $p(w | \mathcal{D})$, typically measured by minimizing the Kullback-Leibler (KL) divergence $KL[q_\phi(w) || p(w | \mathcal{D})]$.

As we saw in Chapter 3, minimizing this KL divergence is equivalent to maximizing the Evidence Lower Bound (ELBO):

$$\mathcal{L}(\phi) = \mathbb{E}_{q_\phi(w)}[\log p(\mathcal{D}|w)] - KL[q_\phi(w) || p(w)]$$

Here, $p(\mathcal{D}|w)$ is the likelihood of the data given the weights, $p(w)$ is the prior distribution over the weights, and $q_\phi(w)$ is the variational approximation parameterized by $\phi$. The first term encourages the approximate posterior to explain the data well, while the second term acts as a regularizer, keeping the approximation close to the prior.

However, applying traditional VI methods like Coordinate Ascent Variational Inference (CAVI) directly to BNNs is difficult due to the complex, non-conjugate relationships between the weights and the data introduced by non-linear activation functions and the sheer number of parameters.

Bayes by Backprop: Optimizing the ELBO with Gradients

A breakthrough technique for applying VI to BNNs is Weight Uncertainty in Neural Networks, often referred to as Bayes by Backprop (BBB), proposed by Blundell et al. (2015). The core idea is to optimize the ELBO with respect to the variational parameters $\phi$ using stochastic gradient ascent, much like standard deep learning models are trained using backpropagation.

The main challenge lies in calculating the gradient of the ELBO, specifically the expectation term $\mathbb{E}_{q_\phi(w)}[\log p(\mathcal{D}|w)]$, because the expectation is taken with respect to $q_\phi(w)$, and the parameters $\phi$ we want to differentiate are inside this distribution. Taking gradients of expectations is tricky.

The Reparameterization Trick

Bayes by Backprop cleverly sidesteps this issue using the reparameterization trick. Instead of directly sampling weights $w$ from $q_\phi(w)$, we introduce an auxiliary noise variable $\epsilon$ with a fixed distribution (e.g., standard Gaussian, $\epsilon \sim \mathcal{N}(0, I)$) and define a deterministic transformation $w = g(\phi, \epsilon)$ such that the resulting $w$ has the distribution $q_\phi(w)$.

For a common choice of $q_\phi(w)$ as a diagonal Gaussian distribution, where $\phi = \{\mu, \rho\}$ represents the mean vector $\mu$ and parameters $\rho$ related to the standard deviation vector $\sigma$ (often $\sigma = \log(1 + \exp(\rho))$ for positivity), the reparameterization is simple:

$$w = \mu + \sigma \odot \epsilon = \mu + \log(1 + \exp(\rho)) \odot \epsilon, \quad \text{where } \epsilon \sim \mathcal{N}(0, I)$$

Now, the expectation in the ELBO can be rewritten with respect to the fixed distribution of $\epsilon$:

$$\mathbb{E}_{q_\phi(w)}[\log p(\mathcal{D}|w)] = \mathbb{E}_{\epsilon \sim p(\epsilon)}[\log p(\mathcal{D}|g(\phi, \epsilon))]$$

The gradient $\nabla_\phi$ can now be pushed inside the expectation using the chain rule, as $g(\phi, \epsilon)$ is deterministic given $\phi$ and $\epsilon$:

$$\nabla_\phi \mathbb{E}_{q_\phi(w)}[\log p(\mathcal{D}|w)] = \mathbb{E}_{\epsilon \sim p(\epsilon)}[\nabla_\phi \log p(\mathcal{D}|g(\phi, \epsilon))]$$

This expectation is typically approximated using Monte Carlo sampling. For a mini-batch of data $\mathcal{D}_i$, we sample one (or a few) $\epsilon$ values, compute the corresponding $w = g(\phi, \epsilon)$, calculate the log-likelihood $\log p(\mathcal{D}_i | w)$, and then compute the gradients $\nabla_\phi \log p(\mathcal{D}_i | w)$ using standard backpropagation through the deterministic transformation $g$ and the network itself.

The Full Objective and Gradient

Combining the likelihood term and the KL divergence term, the objective function to maximize (the ELBO) for a single data point $(x, y)$ becomes:

$$\mathcal{L}(\phi) \approx \log p(y | x, w) - KL[q_\phi(w) || p(w)]$$

where $w = g(\phi, \epsilon)$ is sampled using the reparameterization trick. The KL divergence term $KL[q_\phi(w) || p(w)]$ can often be computed analytically if $q_\phi(w)$ and $p(w)$ are chosen conveniently (e.g., both Gaussians).

The loss function used in training is typically the negative ELBO, scaled for the full dataset or mini-batches:

$$\text{Loss} = - \sum_{i=1}^N \left( \mathbb{E}_{q_\phi(w)}[\log p(y_i|x_i, w)] - \frac{1}{N} KL[q_\phi(w) || p(w)] \right)$$

In practice with mini-batches of size $M$, the loss is approximated as:

$$\text{Loss}_{\text{mini-batch}} \approx -\sum_{j=1}^M \log p(y_j | x_j, w_j) + \frac{1}{B} KL[q_\phi(w) || p(w)]$$

Where $w_j = g(\phi, \epsilon_j)$ is sampled independently for each data point in the mini-batch (or a single sample $w$ is used for the whole batch), and $B$ is the number of mini-batches in the dataset (scaling the KL term appropriately). The gradients of this loss with respect to $\phi$ (the means $\mu$ and variance parameters $\rho$) are computed via backpropagation and used to update $\phi$ using optimizers like Adam or RMSprop.

The Bayes by Backprop training process. Variational parameters $\phi$ (means $\mu$ and variance parameters $\rho$) define the approximate posterior $q_\phi(w)$. In each step, noise $\epsilon$ is sampled to generate network weights $w$ via reparameterization. A forward pass computes the predicted output for the mini-batch data. The loss (negative ELBO) combines the log-likelihood of the data and the KL divergence between the approximate posterior $q_\phi(w)$ and the prior $p(w)$. Gradients of the loss with respect to $\phi$ are computed using backpropagation, and $\phi$ is updated using an optimizer.

Practical Considerations

• Choice of Prior $p(w)$: A standard Gaussian prior ($\mathcal{N}(0, \sigma_p^2 I)$) is common, simplifying the KL divergence calculation when $q_\phi(w)$ is also Gaussian. More complex priors can be used but might require numerical estimation of the KL term.
• Initialization: Initializing the means $\mu$ similarly to standard network weights and keeping initial variances (controlled by $\rho$) small often helps stability.
• Variance Parameterization: Using $\rho$ such that $\sigma = \log(1 + \exp(\rho))$ ensures $\sigma$ is always positive.
• Gradient Variance: The stochastic nature of the gradient estimates (due to sampling $\epsilon$) can lead to high variance during training. Techniques like using more Monte Carlo samples per step or variance reduction methods (like the local reparameterization trick, which applies reparameterization to activations instead of weights) can help, though potentially at increased computational cost.
• Computational Cost: While more scalable than MCMC, VI for BNNs is still more computationally expensive than training standard deterministic NNs. Each forward/backward pass involves sampling weights and calculating the KL term, roughly doubling the number of parameters ($\mu$ and $\rho$ for each weight) and adding complexity to the computation graph.

Advantages and Limitations of VI for BNNs

Advantages:

• Scalability: Leverages mini-batch stochastic gradient optimization, making it applicable to large datasets and complex models where MCMC is infeasible.
• Compatibility: Integrates relatively well with existing deep learning frameworks and hardware acceleration (GPUs/TPUs).
• Single Optimization: Provides a point estimate of the variational parameters $\phi$ after a single optimization run, unlike MCMC which requires multiple chains and convergence checks.

Limitations:

• Approximation Quality: The accuracy of the posterior approximation is limited by the chosen variational family $q_\phi(w)$. Simple families like mean-field Gaussians might not capture complex dependencies in the true posterior.
• Optimization Difficulties: Optimizing the ELBO can be challenging due to noisy gradients and potential local optima.
• Underestimation of Variance: VI is known to sometimes underestimate the variance of the posterior distribution, potentially leading to overconfident predictions.

Compared to Monte Carlo Dropout (discussed in a later section), Bayes by Backprop represents a more principled and flexible VI approach, explicitly defining and optimizing parameters for the approximate posterior distribution over weights. While often more computationally intensive than MC Dropout, it allows for more control over the prior and the variational approximation, potentially leading to more accurate uncertainty estimates if the optimization succeeds and the variational family is adequate.