Gaussian Mean Field Regularizes by Limiting Learned Information

TL;DR

This paper explains how Gaussian mean field regularization improves neural network generalization by limiting mutual information between parameters and data, supported by theoretical analysis and experiments.

Contribution

It provides a theoretical understanding of how mean field inference regularizes models by controlling mutual information, linking it to generalization performance.

Findings

Mean field inference limits mutual information, improving generalization.

Bounding information acts as an effective regularizer in neural networks.

Theoretical bounds relate posterior variance to generalization error.

Abstract

Variational inference with a factorized Gaussian posterior estimate is a widely used approach for learning parameters and hidden variables. Empirically, a regularizing effect can be observed that is poorly understood. In this work, we show how mean field inference improves generalization by limiting mutual information between learned parameters and the data through noise. We quantify a maximum capacity when the posterior variance is either fixed or learned and connect it to generalization error, even when the KL-divergence in the objective is rescaled. Our experiments demonstrate that bounding information between parameters and data effectively regularizes neural networks on both supervised and unsupervised tasks.

Figures

b](#S2.F1.sf2). We use $p^{\prime}$ to emphasize the distinction between distributions of the modified noisy model and the original one. As in the original model, $\theta\sim\operatorname{\mathcal{N}}\left(0,\sigma_{p}^{2}I\right)$ represents the parameters (with same prior), i.e. $p(\theta)=p^{\prime}(\theta)$. We denote the noisy parameters as $\tilde{\theta}\sim\operatorname{\mathcal{N}}\left(\theta,\sigma^{2}\right)$. The likelihood remains unchanged, i.e. $p^{\prime}(D\mid\tilde{\theta})=p(D\mid\theta)$, except that it now depends on the noisy parameters instead of the ‘clean’ ones.

a](#S4.F4.sf1). Similar to the supervised case, we can observe that there is a certain model capacity range that explains the data very well while less or more capacity results in noise drowning and overfitting respectively. In the same figure, we also investigated whether the information-theoretic model capacity can predict generalization independently of the specific prior distribution. Since we merely state an upper bound on mutual information in [subsection 2.1](#S2.SS1), the prior may have an effect in practice which cannot be explained only by the capacity. [Figure 4a](#S4.F4.sf1) shows

b](#S4.F5.sf2). Similar to varying the prior distribution, we are interested in how well the total capacity predicts generalization and the role the architecture plays. It can be observed that larger networks are more resilient to larger total capacities before they start overfitting. This indicates that the total capacity is less important than the individual capacity (i.e. noise) per parameter. Nevertheless, larger networks are more prone to overfitting for very large model capacities. This makes sense as their functional form is less constrained, an aspect that is not captured by our theory

Equations

$$\begin{array}{rl}\log p(D)& \ge \log p(D)-D_{\mathrm{KL}}\left(q(\theta )||p(\theta \mid D)\right)\\ & ={\mathbb{E}}_{q(\theta )}\log p(D\mid \theta )-D_{\mathrm{KL}}\left(q(\theta )||p(\theta )\right)\end{array}$$

$${\mathbb{E}}_{\theta \sim N\left(\mu ,{\sigma }^{2}I\right)}\log p(D\mid \theta )-\frac{1}{2{\sigma }_p^2}\sum _i{\mu }_i^2+\text{const.}$$

$$\log p^{\prime }(D,\theta )$$

$$=$$

$$\ge$$

$$=$$

$$I(D,\theta )\le I(\tilde{\theta },\theta )$$

$$I(\tilde{\theta },\theta )$$

$$\begin{array}{rl}& {\mathbb{E}}_{\theta \sim N\left(\mu ,{\sigma }^2\right)}\log p(D\mid \theta )\\ & +\frac{\beta }{2}\sum _i\left(\log {\sigma }_i^2-{\sigma }_i^2-{\mu }_i^2-1\right)\end{array}$$

$$\begin{array}{rl}& \log p^{\prime }(D,\theta ,{\sigma }^2)\\ =& \log \int p^{\prime }(D\mid \tilde{\theta })p^{\prime }(\tilde{\theta }\mid \theta ,{\sigma }^2)\mathrm{d}\tilde{\theta }+\log p^{\prime }(\theta )+\log p^{\prime }({\sigma }^2)\\ \ge & {\mathbb{E}}_{\tilde{\theta }\sim N\left(\theta ,{\sigma }^2\right)}\log p^{\prime }(D\mid \tilde{\theta })+\sum _i\left(\log p^{\prime }({\theta }_i)+\log p^{\prime }({\sigma }_i^2)\right)\\ =& {\mathbb{E}}_{\tilde{\theta }\sim N\left(\mu ,{\sigma }^2\right)}\log p^{\prime }(D\mid \tilde{\theta })\\ & +\frac{\beta }{2}\sum _i\left(\log {\sigma }_i^2-{\sigma }_i^2-{\mu }_i^2\right)+\text{const.}\end{array}$$

$$I(D,(\theta ,{\sigma }^2))\le I(\tilde{\theta },(\theta ,{\sigma }^2))=\sum _{i}I({\tilde{\theta }}_i,({\theta }_i,{\sigma }_i^2))$$

$$\underset{I(X,Z)<C}{\max }I(Y,Z)$$

$$\min {\mathcal{L}}_{IB}=\min I(X,Z)-\beta I(Y,Z)$$

$$\begin{array}{rl}& I({\tilde{\theta }}_i,({\theta }_i,{\sigma }_i^2))\\ =& H({\tilde{\theta }}_i)-H({\tilde{\theta }}_i\mid {\theta }_i,{\sigma }_i^2)\\ =& -{\int }_{-\infty }^{\infty }{p}^{\prime }({\tilde{\theta }}_i)\log p^{\prime }({\tilde{\theta }}_i)\mathrm{d}{\tilde{\theta }}_i\\ & -{\int }_0^{\infty }{p}^{\prime }({\sigma }_i^2)\frac{1}{2}\log 2\pi e{\sigma }_i^2\mathrm{d}{\sigma }_i^2\end{array}$$

$$\begin{array}{rl}{p}^{\prime }({\tilde{\theta }}_i)=& {\int }_0^{\infty }\mathrm{d}{\sigma }_i^2{p}^{\prime }({\sigma }_i^2)p^{\prime }({\tilde{\theta }}_i\mid {\sigma }_i^2)\\ =& {\int }_0^{\infty }\mathrm{d}{\sigma }_i^2\frac{1}{\Gamma \left(\frac{\beta }{2}\right)}{\left(\frac{\beta }{2}{\sigma }_i^2{e}^{-{\sigma }_i^2}\right)}^{\frac{\beta }{2}}\\ & \cdot {\left(2\pi \left({\sigma }_i^2+\frac{1}{\beta }\right)\right)}^{-\frac{1}{2}}{e}^{-\frac{1}{2\left({\sigma }_i^2+\frac{1}{\beta }\right)}{\tilde{\theta }}_i^2}\end{array}$$

Full text

Machine Learning, ICML

1 Introduction

Bayesian machine learning is a popular framework for dealing with uncertainty in a principled way by integrating over model parameters rather than finding point estimates (Bishop, 2006; Barber, 2012; Ghahramani, 2015). Unfortunately, exact inference is usually not feasible due to the intractable normalization constant of the posterior. A popular alternative is variational inference (Wainwright et al., 2008), where a tractable approximate distribution is optimized to resemble the true posterior as closely as possible. Due to its amenability to stochastic gradient descent (Hoffman et al., 2013; Kingma & Welling, 2013; Titsias & Lázaro-Gredilla, 2014; Rezende & Mohamed, 2015), variational inference is scalable to large models and datasets.

The most common choice for the variational posterior is a factorized Gaussian. Outside of Bayesian inference, parameter noise has been found to be an effective regularizer (Graves et al., 2013; Plappert et al., 2018; Fortunato et al., 2018), e.g. for training neural networks. In combination with $L2$-regularization, additive Gaussian parameter noise corresponds to variational inference with a Gaussian approximate posterior with fixed variance. Interestingly, it has been observed that flexible posteriors can perform worse than simple ones (Turner & Sahani, 2011; Trippe & Turner, 2018; Braithwaite & Kleijn, 2018; Shu et al., 2018).

Variational inference follows the Minimum Description Length (MDL) principle (Rissanen, 1978, 1983; Hinton & van Camp, 1993), a formalization of Occam’s Razor. Loosely speaking, it states that of two models describing the data equally well, the ‘simpler’ one should be preferred. However, MDL is only an objective for compressing the training data and the model, and makes no formal statement about generalization to unseen data. Yet, generalization to new data is a key property of a machine learning algorithm.

Recent work (Xu & Raginsky, 2017; Bu et al., 2019; Bassily et al., 2018; Russo & Zou, 2015) has proposed upper bounds on the generalization error as a function of the mutual information between model parameters and training data. It states that the gap between train and test error can be reduced by limiting the mutual information. However, to the best of our knowledge these bounds and specific inference methods have so far not been linked.

In this work, we show that Gaussian mean field inference in models with Gaussian priors can be reinterpreted as point estimation in corresponding noisy models. This leads to an upper bound on the mutual information between model parameters and data through the data processing inequality. Our result holds for both supervised and unsupervised models. We discuss the connection to generalization bounds from Xu & Raginsky (2017) and Bu et al. (2019), suggesting that Gaussian mean field aids generalization. In our experiments, we show that limiting model capacity via mutual information is an effective measure of regularization, further supporting our theoretical framework.

2 Regularization through Mean Field

In our derivation, we denote a generic model as $p(\theta,D)=p(\theta)p(D\mid\theta)$ with unobserved variables $\theta$ and data $D$. We refer to $\theta$ as the model parameters, however in latent variable models $\theta$ can also include the per-data point latents. The model consists of a prior $p(\theta)$ and a likelihood $p(D\mid\theta)$. Ideally, one would like to find the posterior $p(\theta\mid D)=p(D\mid\theta)p(\theta)/Z$, where $Z=\int p(D\mid\theta)p(\theta)d\theta$ is the normalizer. However, calculating $Z$ is typically intractable. Variational inference finds an approximation by maximizing the evidence lower bound (ELBO)

[TABLE]

w.r.t. the approximate posterior $q(\theta)$. Our focus in this work lies on Gaussian mean field inference, so $q$ is a fully factorized normal distribution with a learnable mean $\mu$ and variance $\sigma^{2}$. The prior is also chosen to be component-wise independent $p(\theta)=\operatorname{\mathcal{N}}\left(0,\sigma_{p}^{2}I\right)$. The generative and inference models for this setting are shown in Figure 1a.

2.1 Fixed-Variance Gaussian Mean Field Inference

When the variance $\sigma^{2}$ of the approximate posterior is fixed to some constant, the ELBO can be written as

[TABLE]

which is optimized with respect to $\mu$. We use $i\in\{1,\ldots,K\}$ to denote the parameter index.

To show how Gaussian mean field implicitly limits learned information, we extend the model with a noisy version of the parameters $\tilde{\theta}\sim p(\tilde{\theta}\mid\theta)$ and let the likelihood depend on those noisy parameters. We choose the noise distribution to be the same as the inference distribution for the original model and find a lower bound on the log-joint of the noisy model. This leads to the same objective as mean-field variational inference in the original model.

Specifically, we define the noisy model $p^{\prime}(\theta,\tilde{\theta},D)=p^{\prime}(\theta)p^{\prime}(\tilde{\theta}\mid\theta)p^{\prime}(D\mid\tilde{\theta})$ as visualized in Figure 1b. We use $p^{\prime}$ to emphasize the distinction between distributions of the modified noisy model and the original one. As in the original model, $\theta\sim\operatorname{\mathcal{N}}\left(0,\sigma_{p}^{2}I\right)$ represents the parameters (with same prior), i.e. $p(\theta)=p^{\prime}(\theta)$. We denote the noisy parameters as $\tilde{\theta}\sim\operatorname{\mathcal{N}}\left(\theta,\sigma^{2}\right)$. The likelihood remains unchanged, i.e. $p^{\prime}(D\mid\tilde{\theta})=p(D\mid\theta)$, except that it now depends on the noisy parameters instead of the ‘clean’ ones.

We now show that maximizing a lower bound on the log joint probability of the noisy model results in an identical objective as for variational inference in the clean model

[TABLE]

where Equation 5 follows from Jensen’s inequality as in Equation 1. In the final equation we have replaced $\theta$ with $\mu$ (which is simply a change of names since we are maximizing the objective over this free variable) to emphasize that the objective functions are identical.

Since $D$ is independent of $\theta$ given $\tilde{\theta}$, the joint $p(\theta,\tilde{\theta},D)$ forms a Markov chain and the data processing inequality (Cover & Thomas, 2012) limits the mutual information $I(D,\theta)$ between learned parameters and data through

[TABLE]

The upper bound is given by

[TABLE]

where $K$ denotes the number of parameters. Here, we exploit that $\theta$ and $\tilde{\theta}\mid\theta$ are Gaussian with $H(\tilde{\theta})=\frac{K}{2}\log 2\pi e\left(\sigma^{2}+\sigma_{p}^{2}\right)$ and $H(\tilde{\theta}\mid\theta)=\frac{K}{2}\log 2\pi e\sigma^{2}$. This quantity is known as the capacity of channels with Gaussian noise in signal processing (Cover & Thomas, 2012). Intuitively, a high prior variance $\sigma_{p}^{2}$ corresponds to a large capacity, while a high noise variance $\sigma^{2}$ reduces it. Any desired capacity can be achieved by simply adjusting the signal-to-noise ratio $\sigma_{p}^{2}/\sigma^{2}$.

2.2 Generalization Error vs. Limited Information

Intuitively, we characterize overfitting as learning too much information about the training data, suggesting that limiting the amount of information extracted from the training data into the hypothesis should improve generalization. This idea has recently been formalized by Xu & Raginsky (2017); Bu et al. (2019); Bassily et al. (2018); Russo & Zou (2015) showing that limiting mutual information between data and learned parameters bounds the expected generalization error under certain assumptions.

Specifically, their work characterizes the following process: Assume that our training dataset is sampled from a true distribution $p_{t}(D)$. Based on this training set, a learning algorithm subsequently returns a distribution over hypotheses given by $p_{t}(\theta\mid D)$. The process defines a mutual information $I_{t}(D,\theta)$ on the joint distribution $p_{t}(D,\theta)=p_{t}(D)p_{t}(\theta\mid D)$. Under certain assumptions on the loss function, Xu & Raginsky (2017) derive a bound on the generalization error of the learning algorithm in expectation over this sampling process. Bu et al. (2019) relax the condition on the loss and prove applicability to a simple estimation algorithm involving $L2$-loss.

Exact Bayesian inference returns the true posterior $p(\theta\mid D)$ on a model $p(\theta,D)$. The theorem then states that a bound on $I(D,\theta)$ limits the expected generalization error as described in Bu et al. (2019) if the model captures the nature of the generating process in the marginal $p(D)=\int\mathop{}\!\mathrm{d}\theta p(\theta)p(D\mid\theta)$. This is a common assumption necessary to justify any (variational) Bayesian approach.

Exact inference is intractable on deep models, and instead, one typically learns variational or point estimates for the posterior. That is also true for the objective on the noisy model above, where we only used a point estimate as given by Equation 6. Therefore, the assumption of exact inference is not met. Yet, we believe that those bounds motivate the expectation that variational inference aids generalization by limiting the learned information. If we performed exact inference on the noisy model in the last section, the given mutual information would imply a bound on generalization error as implied by Xu & Raginsky (2017) and Bu et al. (2019). Therefore, we are optimistic that the gap between variational inference and those generalization bounds can be closed either by performing more accurate inference in the noisy model or by taking the dynamics of the training algorithm into account when bounding mutual information (see subsection 5.2 for further discussion).

2.3 Learned-Variance Gaussian Mean Field Inference

The variance in Gaussian mean field inference is typically learned for each parameter (Kingma et al., 2015; Rezende & Mohamed, 2015; Blundell et al., 2015). Similar to when the variance in the approximate posterior is fixed, one can obtain a capacity constraint. This is the case even for a generalization of the objective from Equation 1 where the KL-term is scaled by some factor $\beta>0$.111Higgins et al. (2017) propose using $\beta>1$ to learn ‘disentangled’ representations in variational autoencoders. Further, $\beta$ is commonly annealed from [math] to $1$ for expressive models (e.g. Bowman et al. (2015); Blundell et al. (2015); Sønderby et al. (2016)). In the following, we quantify a general capacity depending on $\beta$, where $\beta=1$ recovers the standard variational objective. For notational simplicity, we here assume a prior variance of $\sigma_{p}^{2}=1$. It is straight-forward to adapt the derivation to the general case.

In this case, the objective can be written as

[TABLE]

where now both $\mu$ and $\sigma^{2}$ represent learned vectors, and $\operatorname{\mathcal{N}}\left(\mu,\sigma^{2}\right)$ denotes a variable composed of pairwise independent Gaussian components with means and variances given by the elements of $\mu$ and $\sigma^{2}$.

Similar to the previous section, we show a lower bound on the log-joint of a new noisy model to be identical to Equation 9. Specifically, we define the noisy model $p^{\prime}(\theta,\sigma^{2},\tilde{\theta},D)=p^{\prime}(\theta)p^{\prime}(\sigma^{2})p^{\prime}(\tilde{\theta}\mid\theta,\sigma^{2})p^{\prime}(D\mid\tilde{\theta})$ (Figure 1c), with independent priors $\theta_{i}\sim\operatorname{\mathcal{N}}\left(0,\frac{1}{\beta}\right)$ and $\sigma^{2}_{i}\sim\Gamma\left(\frac{\beta}{2}+1,\frac{\beta}{2}\right)$ where $\Gamma(\cdot,\cdot)$ denotes the Gamma distribution. As previously done subsection 2.1, we define the noise-injected parameters as $\tilde{\theta}\sim\operatorname{\mathcal{N}}\left(\theta,\sigma^{2}\right)$ and likelihood as $p^{\prime}(D\mid\tilde{\theta})=p(D\mid\theta)$.

The priors are chosen so that with Jensen’s inequality, we find a lower bound on the log-joint probability of this model that recovers the objective from Equation 9

[TABLE]

In the noisy model, the data processing inequality and the independence of dimensions implies a bound

[TABLE]

where the capacity $I(\tilde{\theta}_{i},(\theta_{i},\sigma^{2}_{i}))$ per dimension is derived in Appendix A. Figure 2 shows numerical results for various values of $\beta$. Standard variational inference ($\beta=1$) results in a capacity of $0.45\text{ bits}$ per dimension. We observe that higher $\beta$ corresponds to smaller capacity, which is given by the mutual information $I(\tilde{\theta}_{i},(\theta_{i},\sigma^{2}_{i}))$ between our new latent $(\theta_{i},\sigma^{2}_{i})$ and $\tilde{\theta}_{i}$. This formalizes the intuition that a higher weight of the complexity term in our objective increases regularization by decreasing a limit on the capacity.

2.4 Supervised and Unsupervised Learning

The above derivations apply to any learning algorithm that is purely trained with Gaussian mean-field inference. This covers supervised and unsupervised tasks.

In supervised learning, the training data typically consists of pairs of inputs and labels, and a loss is assigned to each pair that depends on the trained model, e.g. neural network parameters. When all parameters are learned with one of the discussed mean-field methods, the given bounds apply.

The derivation also comprises unsupervised methods with per-data latents and even amortized inference such as VAEs (Kingma et al., 2015; Rezende & Mohamed, 2015), again as long as all learned variables are learned via Gaussian mean field inference. While this might be helpful to find generalizing representations, the focus of the experiments is on validating the generalizing behavior of mean field variational Bayes on neural network parameters for overfitting regimes, namely small dataset and complex models.

2.5 Flexible Variational Distributions

The objective function for variational inference is maximized when the approximate posterior is equal to the true one. This motivates the development of flexible families of posterior distributions (Rezende & Mohamed, 2015; Kingma et al., 2016; Salimans et al., 2015; Ranganath et al., 2016; Huszár, 2017; Chen et al., 2018; Vertes & Sahani, 2018; Burda et al., 2015; Cremer et al., 2017). In the case of exact inference, a bound on generalization as discussed in subsection 2.2 only applies if the model itself has finite mutual information between data and parameters. However, estimating mutual information is generally a hard problem, particularly in high-dimensional, non-linear models. This makes it hard to state a generic bound, which is why we focus on the case of Gaussian mean field inference.

3 Related Work

Regularization in Neural Networks

Gaussian mean field is intimately related with other popular regularization approaches in deep learning: As apparent from Equation 6, fixed-variance Gaussian mean field applied to training neural network weights is equivalent to L2-regularization (weight decay) combined with Gaussian parameter noise (Graves et al., 2013; Plappert et al., 2018; Fortunato et al., 2018) on all network weights. Molchanov et al. (2017) shows that additive parameter noise results in multiplicative noise on the unit activations. The resulting dependencies between noise components on the layer output can be ignored without significantly changing empirical results (Wang & Manning, 2013). This is in turn equivalent to scaled Gaussian dropout (Kingma et al., 2015).

Information Bottlenecks

The Information Bottleneck principle by Tishby et al. (2000); Shamir et al. (2010) aims to find a representation $Z$ of some input $X$ that is most useful to predict an output $Y$. For this purpose, the objective is to maximize the amount of information $I(Y,Z)$ the representation contains about the output under a bounded amount of information $I(X,Z)$ about the input

[TABLE]

They describe a training procedure using the softly constrained objective

[TABLE]

where $\beta>0$ controls the trade-off.

Alemi et al. (2016) suggest a variational approximation for this objective. For the task of reconstruction, where labels $Y$ are identical to inputs $X$, this results exactly in the $\beta$-VAE objective (Achille & Soatto, 2017; Alemi et al., 2018). This is in accordance with our result from subsection 2.3 that there is a maximum capacity per latent dimension that decreases for higher $\beta$. Setting $\beta>1$, as suggested by Higgins et al. (2017) for obtaining disentangled representations, corresponds to lower capacity per latent component than achieved by standard variational inference.

Both Tishby et al. (2000) and Higgins et al. (2017) introduce $\beta$ as a trade-off parameter without a quantitative interpretation. With our information-theoretic perspective, we quantify the implied capacity and provide a link to the generalization error. Further, both methods are concerned with the information in the latent representation. They do not limit the mutual information with the model parameters, leaving them vulnerable to model overfitting under our theoretical assumptions. We experimentally validate this vulnerability and explore the effect of filling this gap by applying Gaussian mean field inference to the model parameters.

Information Estimation with Neural Networks

Multiple recent techniques (Belghazi et al., 2018; van den Oord et al., 2018; Hjelm et al., 2018) propose the use of neural networks for obtaining a lower bound on the mutual information. This is useful in settings when we want to maximize mutual information, e.g. between the data and a lower-dimensional representation. In contrast, we show that Gaussian variational inference on variables with a Gaussian prior implicitly places an upper bound on the mutual information between these variables and the data, and explore its regularizing effect.

4 Experiments

In this section, we analyze the implications of applying Gaussian mean field inference of fixed scale to the model parameters in the supervised and unsupervised context. Our theoretical results suggest that varying the capacity will affect the generalization capability and we show this effect on small data regimes and how it changes with the training set size. Furthermore, we investigate whether capacity is the only predictor for generalization or whether varying priors and architectures also have an effect. Finally, we demonstrate qualitatively how the capacity bounds are reflected in fashion MNIST reconstruction.

4.1 Supervised Learning

We begin with a supervised classification task on the CIFAR10 dataset, training only on a subset of the first 5000 samples. We use 6 3x3 convolutional layers with 128 channels each followed by a ReLU activation function, every second of which implements striding 2 to reduce the input dimensionality. Finally, the last layer is a linear projection which parameterizes a categorical distribution. The capacity of each parameter in this network is set to specific values given by Equation 8.

Figure 3 shows that decreasing the model capacity per dimension (by increasing the noise) reduces the training log-likelihood and increases the test loglikelihood until both of them meet at an optimal capacity. One can observe that very small capacities lead to a signal that is too noisy and good predictions are no longer possible. In short, regimes of underfitting and overfitting are generated depending on the capacity.

4.2 Unsupervised Learning

We now evaluate the regularizing effect of fixed-scale Gaussian mean field inference in an unsupervised setting for MNIST image reconstruction. Therefore, we use a VAE (Kingma & Welling, 2013) with 2 latent dimensions and a 3-layer neural network parameterizing the conditional factorized Gaussian distribution. As usual, it is trained using the free energy objective, but different from the original work, we also use Gaussian mean field inference for the model parameters. Again, we use a small training set of 200 examples for the following experiments if not denoted otherwise.

Varying Model Capacity and Priors

In our first experiment, we analyze generalization by inspecting the test ELBO when varying the model capacity which can be seen in Figure 4a. Similar to the supervised case, we can observe that there is a certain model capacity range that explains the data very well while less or more capacity results in noise drowning and overfitting respectively. In the same figure, we also investigated whether the information-theoretic model capacity can predict generalization independently of the specific prior distribution. Since we merely state an upper bound on mutual information in subsection 2.1, the prior may have an effect in practice which cannot be explained only by the capacity. Figure 4a shows that indeed, while the general behavior remains the same for different model priors, the generalization error is not entirely independent. Furthermore, the observation that all curves descend with larger capacities, for all priors, suggests that weight decay (Krogh & Hertz, 1992) of fixed scale without parameters noise is not sufficient to regularize arbitrarily large networks. In Figure 4b we investigated the extreme case of dropping the prior entirely and switching to maximum-likelihood learning instead by using an improper uniform prior. This approach recovers Gaussian dropout (Srivastava et al., 2014; Kingma et al., 2015). Dropping the prior sets the bottleneck capacity to infinity and should lead to worse generalization. Comparing the test ELBO of this Gaussian dropout variant to the original Gaussian mean field inference in Figure 4b confirms this result for larger capacities. For larger noise scales, generalization is still working well, a result that is not explained in our information-theoretic framework, but plausible due to the deployed limited architecture.

Varying Training Set Size

Figure 5a shows how limiting the capacity affects the test ELBO for varying amounts of training data. Models with very small capacity extract less information from the data into the model, thus yielding a good test ELBO somewhat independent of the dataset size. This is visible as a graph that ascends very little with more training data (e.g. total model capacity 330 kbits). Note that we here report the capacity of the entire model, which is the sum of the capacities for each parameter. In order to improve the test ELBO, more information from the data has to be extracted into the model. But clearly, this leads to non-generalizing information being extracted when the dataset is small, leading to overfitting. Only for larger datasets the extracted information generalizes. This is visible as a strongly ascending test ELBO with larger dataset sizes and bad generalization for small datasets. We can therefore conclude that the information bottleneck needs to be chosen based on the amount of data that is available. Intuitively, when more information is available, more information should be extracted into the model.

Varying Model Size

Furthermore, we inspected how the size of the model (here in terms of number of layers) affects generalization in Figure 5b. Similar to varying the prior distribution, we are interested in how well the total capacity predicts generalization and the role the architecture plays. It can be observed that larger networks are more resilient to larger total capacities before they start overfitting. This indicates that the total capacity is less important than the individual capacity (i.e. noise) per parameter. Nevertheless, larger networks are more prone to overfitting for very large model capacities. This makes sense as their functional form is less constrained, an aspect that is not captured by our theory.

Qualitative Reconstruction

Finally, we plot test reconstruction means for the binarized fashion MNIST dataset under the same setup for various capacities in Figure 6. In accordance with previous experiments, we observe that if the capacity is chosen too small, the model is not learning anything useful, while too large capacities result in overconfidence. This can be observed in most means being close to either 0 or 1. An intermediate capacity, on the other hand, makes sensible predictions (given that it was trained only on 200 samples) with sensible uncertainty, visible through gray pixels that correspond to high entropy.

5 Discussion

In this section, we discuss how the capacity can be set, as well as the effect of model architecture and learning dynamics.

5.1 Choosing the Capacity

We have obtained a new trade-off parameter, the capacity, that has a simple quantitative interpretation: It determines how many bits to maximally extract from the training set. In contrast, for the $\beta$ parameter introduced in Tishby et al. (2000) and Higgins et al. (2017) a clear interpretation is not known. Yet, it may still be hard to set the capacity optimally. Simple mechanisms such as evaluation on a validation set to determine its value may be used. We expect that more theoretically rigorous methods could be developed.

Furthermore, in this paper, we have focused on the regularization that Gaussian mean field inference implies on the model parameters. The same concept is valid for data-dependent latent variables, for instance in VAEs, as discussed in subsection 2.4. In VAEs, Gaussian mean field inference on the latents leads to a restricted latent capacity, but leaves the capacity of the model unbounded. This leaves VAEs vulnerable to model overfitting, as demonstrated in the experiments, and setting $\beta$ as done in Higgins et al. (2017) is not sufficient to control complexity. This motivates the limitation of capacity between the data and both per-datapoint latents and model parameters. The interaction between the two is an interesting future research direction.

5.2 Role of Learning Dynamics and Architecture

As discussed in subsection 2.2, it is necessary to perform exact inference in the noisy model for the bounds on the generalization error to hold. However, this assumption is not met. In practice $p_{t}(\theta\mid D)$ encodes the complete learning algorithm, which in deep learning typically includes parameter initialization and dynamics of the stochastic gradient descent optimization.

Our experiments confirm the relevance of these other factors: $L2$-regularization works in practice, even though no noise is added to the parameters. This could be explained by the fact that noise is already implicitly added through stochastic gradient descent (Lei et al., 2018) or through the output distribution of the network. Similarly, Gaussian dropout (Graves et al., 2013; Plappert et al., 2018; Fortunato et al., 2018) without a prior on the parameters helps generalization. Again, early stopping combined with a finite reach of gradient descent steps effectively shapes a prior of finite variance in the parameter space. This could also formalize why the annealing schedule employed by Bowman et al. (2015); Blundell et al. (2015) and Sønderby et al. (2016) is effective.

This observed dependence on other factors suggests that quantifying mutual information $I_{t}(\theta,D)$ of the actual distribution created by the learning dynamics might be a promising approach to explain why neural networks often generalize well on their own. This idea is in accordance with recent work that links the learning dynamics of small neural networks to generalization behavior (Li & Liang, 2018).

On the other hand, the architecture choice also had an influence on generalization, which is expected by our theory since we only formulate a bound on mutual information that is completely agnostic to the actual model choice. Tightening this bound based on the model architecture and output distribution is usually hard, as discussed in subsection 2.5, but might be possible.

Another promising direction would be to approximately sample from the exact posterior on network parameters (i.e. as done by Marceau-Caron & Ollivier (2017)), on a capacity-limited architecture, instead of the usual approach of point estimation. In the limit of infinite training time, this would fully realize the discussed bound on the expected generalization error.

6 Conclusion

We have explained the regularizing effects observed in Gaussian mean field approaches from an information-theoretic perspective. The derivation features a capacity that can be naturally interpreted as a limit on the amount of information extracted about the given data by the inferred model. We validated its practicality for both supervised and unsupervised learning.

How this capacity should be set for parameters and latent variables depending on task and data is an interesting direction of research. We exploited a theoretical link of mutual information and generalization error. While this work is restricted to Gaussian mean field, incorporating the effect of learning dynamics on mutual information in future work might allow understanding why overparameterized neural networks still generalize well to unseen data.

Appendix A Capacity in Learned-Variance Gaussian Mean Field Inference

The capacity per dimension for the model discussed in subsection 2.3 is given by

[TABLE]

$\tilde{\theta}_{i}\mid\theta_{i},\sigma^{2}_{i}\sim\operatorname{\mathcal{N}}\left(\theta_{i},\sigma^{2}_{i}\right)$ with $\theta_{i}\sim\operatorname{\mathcal{N}}\left(0,\frac{1}{\beta}\right)$ implies $\tilde{\theta}_{i}\mid\sigma^{2}_{i}\sim\operatorname{\mathcal{N}}\left(0,\sigma^{2}_{i}+\frac{1}{\beta}\right)$. Together with $\sigma^{2}_{i}\sim\Gamma\left(\frac{\beta}{2}+1,\frac{\beta}{2}\right)$, this implies

[TABLE]

Numerical results for the capacity $I(\tilde{\theta}_{i},(\theta_{i},\sigma^{2}_{i}))$ with varying $\beta$ are given below and plotted in Figure 2.

$\beta$ $I(\tilde{\theta}_{i},(\theta_{i},\sigma^{2}_{i}))$

$0.01$ $0.68\text{ bits}$

$0.1$ $0.65\text{ bits}$

$1$ $0.45\text{ bits}$

$10$ $0.12\text{ bits}$

$100$ $0.014\text{ bits}$