Good Initializations of Variational Bayes for Deep Models

Variational training starts from a distribution over network weights. Its initialization must specify both a plausible prediction and an appropriate amount of uncertainty. We introduce Iterative Bayesian Linear Model (I-BLM) initialization, which uses the training data to choose these distributions before optimizing the full deep model.

Initializing a distribution over weights

A factorized Gaussian approximation assigns a mean and variance to each weight. The evidence lower bound combines expected data fit with a Kullback-Leibler penalty relative to the prior. Initializing the variances therefore affects both the distribution of predictions and the balance between these two objective terms.

Methods developed for deterministic networks mainly control signal propagation through their initial weight values. Extending them to variational inference requires an additional choice for the weight variances. I-BLM derives both quantities from Bayesian linear regression, connecting the initial uncertainty to the prediction task.

Fitting one layer at a time

Before applying its activation function, a neural-network layer is linear in its weights. I-BLM uses this structure to fit a Bayesian linear regression model for each output feature, using a minibatch of inputs and training targets.

After initializing the first layer, we sample its weights and propagate the inputs through its activation function. These transformed inputs become the design matrix for initializing the next layer. Repeating the procedure gives each layer an initial mapping informed by the prediction task.

Matching the variational family

Bayesian linear regression produces a correlated Gaussian posterior. We approximate it within the chosen variational family by minimizing the KL divergence from the approximation to that posterior. For a factorized Gaussian, the optimal initial mean equals the linear-model posterior mean. If its covariance is $\Sigma$, the initial variance of coordinate $i$ is

$$s_i^2=\frac{1}{(\Sigma^{-1})_{ii}}.$$

This gives an analytical initialization for the factorized distribution. Other Gaussian covariance structures can use the corresponding posterior approximation.

Regression, classification, and convolutions

For classification, we transform class labels into regression targets through a Gaussian approximation related to the Dirichlet distribution. For convolutional layers, image patches supply the linear-model inputs.

We evaluate regression and classification with fully connected and convolutional networks, comparing convergence and predictive performance against several initialization rules. I-BLM reaches better solutions faster in these experiments. The additional layerwise regression fits introduce an initialization cost in exchange for a starting distribution informed by the training data.