Title: NOISY INFORMATION BOTTLENECKS FOR GENERALIZATION

Abstract: We propose Noisy Information Bottlenecks (NIB) to limit mutual information between learned parameters and the data through noise. We show why this benefits generalization and allows mitigation of model overfitting both for supervised and unsupervised learning, even for arbitrarily complex architectures. We reinterpret methods including the Variational Autoencoder, β-VAE, network weight uncertainty and a variant of dropout combined with weight decay as special cases of our approach, explaining and quantifying regularizing properties and vulnerabilities within information theory.

1. INTRODUCTION:

Bayesian inference is at the core of many directions of machine learning research in recent years. Applications range from a principled treatment of uncertainty in neural networks (Gal, 2016) over latent variable models for high-dimensional data (Kingma & Welling, 2013; Rezende & Mohamed, 2015) to reinterpretations of popular stochastic regularizers (Kingma et al., 2015; Gal, 2016) that inspired more flexible extensions (Louizos & Welling, 2017; Molchanov et al., 2017) . In all but the simplest models, exact inference is intractable, so approximations are necessary. Among those, variational inference (Wainwright et al., 2008) has been particularly popular, as it is amenable to gradient-based stochastic optimization (Kingma & Welling, 2013; Titsias & Lázaro-Gredilla, 2014; Ranganath et al., 2014) and thus highly scalable.

Typically, researchers seek to develop more flexible approximate 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 hope of more faithfully representing the true posterior. However, recent works have suggested (Trippe & Turner, 2018; Braithwaite & Kleijn, 2018; Shu et al., 2018) that restricting the family of variational approximations can, in fact, have a positive regularizing effect, leading to better generalization.

In this work, we seek to provide an information theoretic explanation for the observed behaviour of different families variational approximations. By reinterpreting Gaussian mean-field inference as maximum a-posteriori in a noisy model, we can quantify bounds on the mutual information between the data and the parameters variables. We explore the potential of this approach for regularization. Unlike methods relying on information-based objectives (Tishby et al., 2000; Shamir et al., 2010; Agakov & Barber, 2004; Chen et al., 2016; Phuong et al., 2018) , it is compatible with standard variational inference.

2. NOISY INFORMATION BOTTLENECKS:

This section characterizes model overfitting within information theory and points out a mitigation strategy that limits the amount of information extracted by the inferred model about the data. We will show this to be implicitly deployed by a common class of inference techniques in section 3 in order to explain observed regularizing effects.

2.1. AN INFORMATION-THEORETIC CHARACTERIZATION OF MODEL OVERFITTING:

Intuitively, model overfitting can be understood as memorization or learning too much information about the training data. More formally, if D denote the random variable for all our data and θ all learned parameters and latent variables in a given probabilistic model, we can quantify the amount of information gained about the data D given a particular parameter value θ * in

I(D, θ = θ * ) := H(D) -H(D|θ = θ * ) (1)

where H(D|θ = θ * ) := -E p(D|θ=θ * ) log p(D|θ = θ * ) denotes the entropy of a variable D conditioned on a particular value θ * .

In the context of variational inference, model overfitting can now be characterized by the symptom that the learned information about the data D

E θ * ∼q(θ) I(D, θ = θ * ) (2)

is too large in expectation under the approximately inferred posterior distribution q(θ) ≈ p(θ|D).

We now motivate this objective through the extreme case of maximum-likelihood learning in the limit of unrestrictedly expressive models for unsupervised tasks. Here, q is a point-mass on the parameter values that maximize the likelihood of the data. Under the assumption of indepenent and identically distributed (iid) datapoints, the model will then store all information that can possibly be extracted from the data, as shown in Appendix A. This is all information except the datapoint identity (the information necessary to distinguish between distinct training samples, e. g. the index) which cannot be learned due to the iid assumption. This will result in perfect overfitting up to the sample identity, thereby impeding generalization.

In unsupervised learning, information that can be learned by the model about the data will not be learned by the latent, a phenomenon known as the information preference problem (Chen et al., 2016; Alemi et al., 2016; Zhao et al., 2017; Phuong et al., 2018) . Again, this can be characterized by the quantity from Equation 2 being too high.

2.2. WHY BOUNDING MUTUAL INFORMATION?:

This observation motivates placing a limit on the expected amount of extracted information given by Equation 2. We observe that the expectation of I(D, θ = θ * ) under the prior p(θ) is just the mutual information I(D, θ) as commonly defined in literature:

E θ * ∼p(θ) I(D, θ = θ * ) = H(D) -E θ * ∼p(θ) H(D|θ = θ * ) = H(D) -H(D|θ) = I(D, θ) (3)

This implies that if we could limit I(D, θ) to not be greater than some capacity C, we would obtain the following guarantee: If our dataset d was sampled from the marginal d ∼ p(D) = dθp(θ)p(D|θ), then a sample from the exact posterior θ * ∼ p(θ|D = d) will, on average, contain no more information about the data D than our capacity C due to

E d∼p(D) E θ * ∼p(θ|D=d) I(D, θ = θ * ) = E θ * ∼p(θ) I(D, θ = θ * ) = I(D, θ) ≤ C (4)

This would limit the quantity given in Equation 2 in expectation under the assumptions that the model captures the nature of the generating process in p(D) and that our posterior estimate q(θ) is close to the true posterior p(θ|D). These are common assumptions necessary to justify any variational Bayesian approach.

We derive further motivation for our approach from the fact that limiting mutual information between data and learned parameters provably bounds generalization error (Xu & Raginsky, 2017; Bassily et al., 2018) . While the approach of limiting the amount of extracted information is already used for the purpose of generalization in adaptive analysis (Feldman & Steinke, 2018; Smith, 2017; Russo & Zou, 2015) , in the following section we give an approach compatible with a variety of variational and deep learning algorithms.

How can we limit I(D, θ)? Due to nonlinearities in typical deep models, it is hard to calculate and therefore limit the mutual information between data and parameters directly. Instead, one way to achieve this is to make the data dependent only on some noisy version θ of the learned variables θ, resulting in a model p(θ, θ, D) = p(θ)p( θ|θ)p(D| θ) forming a Markov chain as shown in Figure 1a .

(b) θ θ Y (n) Ỹ (n) X (n) n ∈ {1, . . . , N } (c)

Figure 1 : A noisy information bottleneck (a) limits the mutual information I(D, θ) between the data D and the learned variables θ by making the data only depend on a corrupted version θ of the latent, with a limited I( θ, θ). It can be applied to both (b) supervised with inputs X (n) and labels Y (n) and (c) generative models with latents Y (n) and datapoints X (n) .

The prior p(θ) and corruption process p( θ|θ) are chosen to result in a certain finite I( θ, θ). The data processing inequality (Cover & Thomas, 2012) then ensures

EQUATION

for any model architecture p(D| θ). C = I( θ, θ) now acts as a capacity, bounding the expected amount of information that is extracted about the data. In the next section, we will show that such a bound is implicit in an important class of variational inference techniques. As we will see, the corruption process will be realized through injected noise, and we therefore term this architecture noisy information bottleneck (NIB). Quantifying the corresponding capacity C will help to explain observed regularizing effects in both supervised and unsupervised learning (as shown in Figure 1b and 1c) from an information-theoretic perspective.

3. REGULARIZATION THROUGH CONSTRAINED INFERENCE:

Most variational approximation schemes such as variational autoencoders (Kingma et al., 2015; Rezende & Mohamed, 2015) and network weight uncertainty (Blundell et al., 2015) use simple Gaussian mean field inference for tractable training using the reparameterization trick. In this section, we quantify information-theoretic capacity constraints implicit in Gaussian mean field inference that helps to understand observed regularizing effects. We will first illustrate the bound for a fixed-scale distribution and then discuss the more general case where the variance is learned as well.

3.1. FIXED-SCALE GAUSSIAN MEAN FIELD INFERENCE:

Let p(θ, D) denote a model with parameters θ and data D. The free energy objective now is

EQUATION

We consider the case of Gaussian mean field inference with fixed variance σ 2 on parameters with component-wise independent priors p(θ) = N 0, σ 2 p I . We can write the inferred q(θ) in its reparameterized form θ = µ * + σ , with µ * being the inferred mean and being noise sampled from N (0, I), as shown in Figure 2a . We can then write the objective as

EQUATION

)

which is optimized with respect to µ * . We use i ∈ {1, . . . , K} to denote the model parameter index.

To show how Gaussian mean field contains an implicit information bottleneck we include the inferred means µ into the generative model, without changing the resulting objective. We therefore define a model p (µ, μ, D) = p (µ)p (μ|µ)p (D|μ) (Figure 2b ), where µ ∼ N 0, σ 2 p I represents the parameters of our new model and μ = µ + σ its noise-injected version with ∼ N (0, I). We leave the rest of the model p (D|μ) = p(D|θ) unchanged. This model is an instantiation of a noisy information bottleneck as in Figure 1a , where µ and μ take the roles of θ and θ. On this new model, we do variational inference by sampling noise from the prior q ( ) = p ( ), while learning a deterministic approximate posterior q (µ) with point-mass at a particular µ * , resulting in a MAP-like objective1 

µ * θ θ D σ 2 p σ 2 Inference model q Generative model p (a) µ * µ µ μ D σ 2 p σ 2 Inference model q Generative model p (b)

EQUATION

Using above definitions, and Equation 7, this implies equivalence F (µ * ) = F (µ * ) of both objectives, resulting in exactly equivalent training algorithms.

For the new model, we get

H(μ) = K 2 log 2πe σ 2 + σ 2 p and H(μ|µ = µ * ) = K 2 log 2πeσ 2 for any µ * , inducing a noisy information bottleneck of capacity I(μ, µ) = H(μ) -E µ * ∼p(µ) H(μ|µ = µ * ) = K 2 log 1 + σ 2 p σ 2 (9)

in the sense of Equation 5, where K denotes the number of parameters. This quantity is known as the capacity of channels with Gaussian noise in signal processing (Cover & Thomas, 2012) . Intuitively, a high prior variance σ 2 p corresponds to a large capacity, while a high noise variance σ 2 reduces it. Simply adjusting the signal-to-noise ratio σ 2 p σ 2 allows to create an information bottleneck of any desired capacity. This suggests the usefulness of applying fixed-scale Gaussian mean field inference to model parameters for regularization.

There are relations to existing approaches: When applied to training neural networks, MAP training on this model can be interpreted as applying weight decay combined with a additive noise N (0, σ 2 ) on all network weights. Molchanov et al. (2017) shows that this results in multiplicative noise on the unit activations. Wang & Manning (2013) reports that empirical results do not change significantly when dependencies between the different elements of the layer output are ignored, which is then equivalent to scaled Gaussian dropout (Kingma et al., 2015) .

3.2. FLEXIBLE-SCALE 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) . We can obtain a capacity constraint for this case by regarding both the inferred mean µ and variance σ 2 as the new latent. For simplicity and different from the fixed-scale case, we assume a prior variance of σ 2 p = 1. The derivation is similar to the one for the fixed-scale case from the previous section and given in a slightly different form in Appendix C. The resulting capacity is not in closed analytic form, but we obtain a numeric result of 0.45 bits per latent component. The derivation can be generalized to approaches where the KL term from the objective is scaled by some factor β > 0 as done by β-VAE (Higgins et al., 2017) in the context of amortized inference on latent variables. Curiously, this results in priors µ i ∼ N 0, 1 β and σ 2 i ∼ Γ β 2 + 1, β 2 . We observe that higher β corresponds to smaller capacity, which is given by the mutual information I(μ i , (µ i , σ 2 i )) between our new latent (µ i , σ 2 i ) and μi . This formalizes the intuition that a higher weight of the complexity term in our objective increases regularization by limiting the capacity. Numeric results for the capacity are shown in Figure 8 in the appendix.

3.3. SHOULD VARIATIONAL INFERENCE BE FLEXIBLE OR REGULARIZING?:

As mentioned in the introduction, more flexible approximate inference is the goal of many recent approaches. However, in previous subsections we showed that training with Gaussian mean field inference can be reinterpreted as MAP learning on a noise-injected model that has a certain informationtheoretic capacity that is useful to mitigate overfitting. Therefore the goals of good inference and regularization are in conflict.

The reinterpretation necessary to retrieve the capacity also suggests a natural resolution: Instead of using MAP on the noise-injected model, we can deploy arbitrarily flexible variational inference techniques. This separates out the concern of regularization into the model and allows to combine powerful inference with good regularization. The focus of this piece of work, however, is to demonstrate that limiting the mutual information between the model parameters and the data leads to better generalization, hence we leave exploring the use of powerful inference techniques in noisy models as future research. 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 shown in subsection 2.1, and setting β as done in (Higgins et al., 2017) is not sufficient to control complexity. Relatedly, the bound given through Equation 9 was also noted by Braithwaite & Kleijn (2018) , but was again applied only to the latents and not the model parameters.

4. EXPERIMENTS:

We therefore propose to apply Gaussian mean field inference of fixed scale to the model parameters in order to mitigate model overfitting both for supervised and unsupervised tasks. This is further motivated by the regularizing effect empirically observed when using Gaussian mean field inference over point estimates reported by Blundell et al. (2015) . Fixing the variance σ 2 allows to easily control the capacity of the bottleneck, as described in subsection 3.1.

In this section we validate this idea. In order to clearly show overfitting, we train large architectures and a small number of training samples for most experiments. We explore the effect noisy information bottlenecks implicit in fixed scale Gaussian mean field inference for varying model capacities, architectures and priors as well as training set sizes.

4.1. SUPERVISED LEARNING:

We apply our approach to classification on the CIFAR10 dataset, where we train 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 parameters in this network is set to specific values given by Equation 9.

Figure 3 shows that decreasing the model capacity per channel (by increasing the noise) reduces the training log-likelihood and increases the test loglikelihood until both of them meet at an optimal capacity. As noted before, without this prior the capacity would be infinite. We will analyze the phenomenon in the unsupervised case in the next subsection. It is also observable that very small capacities lead to a signal that is too noisy and good predictions are no longer possible. All these observations are in accordance with our predictions.

4.2. UNSUPERVISED LEARNING:

We now evaluate the regularizing effect of fixed-scale Gaussian mean field inference in an unsupervised setting. 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. Again, we use a small training set of 200 examples for most experiments.

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. From Figure 4a we can conclude that while model capacity seems to influence generalization very strongly, it seems like model choice such as the prior on the parameters also affects generalization slightly. It also implies that weight decay (Krogh & Hertz, 1992) of fixed scale without parameters noise is not sufficient to regularize arbitrarily large networks. Nevertheless, the shapes of all ELBO-capacity dependencies are similar. In Figure 4b we investigated the extreme case of dropping the prior entirely and switching to ML learning instead by using an improper uniform prior. This approach is very similar to Gaussian dropout, as described in subsection 3.1. Dropping the prior sets the bottleneck capacity to infinity and should lead to worse generalization. Comparing the test ELBO to NIB 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.

Figure 5a shows how our approach 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, and generalizing information being extracted for larger datasets. 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 amounts of data that is available. This is expected as we want to extract more information into the model the more information is available.

Furthermore, we inspected how the size of the model (here in terms of number of layers) affects generalization in Figure 5b . The generalization does not decrease for the same total capacity but larger networks deteriorate the performance for larger total capacities. 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 framework.

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, which can be seen through 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. RELATED WORK:

We have shown that NIB is a practical approach to regularize supervised and unsupervised models, and that in contrast to existing approaches, it successfully regularizes models with a fixed parameter setting, largely independent of the depth of a network. Unlike existing regularization techniques, our approach 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.

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:

EQUATION

They describe a training procedure using the softly constrained objective

EQUATION

where β > 0 controls the trade-off.

Alemi et al. ( 2016) suggests a variational approximation for this objective. For the task of reconstruction, where labels Y are identical to inputs X, this results exactly in the β-VAE objective (Achille & Soatto, 2017; Alemi et al., 2018) . This is in accordance with our result from subsection 3.2 that there is a maximum capacity per latent dimension in the reinterpreted version of β-VAE that get smaller for greater β.

Reconstruction implicit in amortized inference for unsupervised learning is only one of many possible objectives for extracting representations, and it can be viewed a proxy for a not-yet-known future task. In this case, applying NIB to the latents is a way to enforce the hard constraint I(X, Z) < C of objective Equation 10 directly through the structure of the model, removing the need to augment the objective with a soft constraint, as in Equation 11.

For the supervised case where the task is already known, as in Tishby et al. (2000) , we could achieve a similar hard constraint by constructing a noisy information bottleneck layer anywhere in the network based on the channel capacity theorem (Cover & Thomas, 2012) by limiting input variance (e. g. through a tanh activation) and injecting noise as in NIB. This allows extracting representations useful for a given output task with standard variational training, while naturally constraining mutual information with the input.

The Information Bottleneck principle is concerned with the information contained in the latent representation. NIB limits mutual information with all inferred variables, namely latents and model parameters in the case of unsupervised learning. This paper focuses on characterizing and mitigating the vulnerability of existing learning algorithms due to unconstrained mutual information of the data with the model parameters.

6. CONCLUSION:

We have quantified the regularizing effects observed in Gaussian mean field approaches from an information-theoretic perspective. We have explored the usefulness of the implicit noisy information bottleneck for the purpose of generalization. Our approach 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. We have shown that the approach allows to improve generalization even for arbitrarily large networks.

While this work demonstrates the capability of NIB to mitigate model overfitting when applied to model parameters, inspecting the effect of a limited latent capacity is left for future work. We expect our approach to be compatible with powerful inference techniques while keeping up regularization guarantees. This is still to be confirmed experimentally.

A MAXIMUM-LIKELIHOOD ON UNRESTRICTED MODELS:

We consider the extreme case of maximum-likelihood learning on N identically distributed (iid) data in the limit of unrestrictedly expressive models. For simplicity, we assume discrete data to avoid an undefined empirical differential entropy. A similar argument could be made for continuous data. We can then write

EQUATION

where p(X|θ) is the modeled distribution p(X (n) |θ) of any data sample X (n) . This definition is valid because the distribution is independent of the data sample index n due to the iid assumption. Let p( X) denote the empirical distribution. The maximum-likelihood objective can then be written as

EQUATION

In the unrestricted limit (e.g. arbitrarily sized network architecture), there exists a θ * so that the distribution p(X|θ) is equal to p( X). Because D KL p( X)||p(X|θ = θ * ) is then 0 and therefore minimal, this is actually θ M L . From p(X|θ = θ M L ) = p( X) follows that the entropy conditioned on θ M L becomes the empirical entropy

EQUATION

Thus,

EQUATION

This implies that all information about the data is learned by the model up to the identity of the training sample, which remains as uncertainty in the prediction.

B MAP AND MAXIMUM-LIKELIHOOD AS VARIATIONAL INFERENCE:

Variational inference constrained to deterministic approximate posterior distributions with (continuous) probability mass function

EQUATION

(= point-mass at Z * ) can be interpreted as MAP:

EQUATION

= arg min

EQUATION

= arg min

EQUATION

=Z MAP (23) D KL (q(Z)||p(Z|X)) is not finite, but its arg max and gradient are still well-defined. The same result can be obtained more formally by regarding MAP as the limit of Gaussian mean field inference with diminishing variance σ 2 → 0. For improper uniform priors p(Z) = const. we recover Maximum-Likelihood:

X µ Y Y X σ 2 Inference model q Generative model p (a) X µ µ μ X σ 2 σ 2 β Inference model q Generative model p (b)

EQUATION

= arg max

Z * log p(Z * ) + log p(X|Z * ) (25) = arg max Z * log p(X|Z * ) (26) =Z ML (27) C LATENT CAPACITY IN β-VAE-LIKE GAUSSIAN MEAN FIELD INFERENCE

The purpose of this section is to derive a capacity for flexible-scale Gaussian mean field inference, similar to subsection 3.1. We generalize this to the case when the KL term from the objective is scaled by a factor β > 0, as done by (Higgins et al., 2017) . To illustrate the same capacity bound can be obtained for amortized inference, we derive this capacity for components Y

(n) i of the latents in a β-VAE, as this is where Gaussian mean field inference is deployed in this case. This is only done fur the purpose of familarity, the same bound would hold for when this inference approach would be applied to model parameters.

We assume Gaussian mean field inference on a latent with a unit Gaussian prior per dimension, as shown in Figure 7a . This analysis includes the standard VAE for β = 1 as a special case. The β-VAE objective is (per datapoint, to be summed over all indices (n) which are omitted for readability):

EQUATION

We can write Y in its reparameterized form Y = µ * + σ * , with µ * and σ * being the mean and variance yielded deterministically by the inference network and being noise sampled from N (0, I). We can then write the objective as

EQUATION

As in subsection 3.1, we now define a new model featuring NIB that can be trained with a standard variational objective (without need for a factor on the complexity term) so that it results in the same training algorithm (data indices (n) are again omitted). The key idea is to reinterpret the means µ and variances σ 2 from the inference distribution as the latents of our new model, as shown in Figure 7b . We keep the deterministic part of the inference network as our new deterministic approximate posterior q (µ, σ 2 |X) with point-mass at (µ * , σ * 2 ) per datapoint and dimension i, resembling a MAP approach. We define the priors µ ∼ N 0, 1 β I and ∀i : σ 2 i ∼ Γ β 2 + 1, β 2 and then define μ = µ + σ , where noise is sampled from the prior q ( ) = p ( ) = N (0, I). We leave the rest of the model and the inference network unchanged at p (X|μ, θ) = p(X|Y, θ) and q (µ, σ 2 |X) = q(µ, σ 2 |X), and as before, we perform ML learning on the model parameters θ. Table 1 .

The MAP objective of our newly defined model then is (again per datapoint) F (q , θ) = E q (µ,σ 2 |X)q ( ) log p (X|µ, σ, , θ)

EQUATION

Equation 29 now implies which is equivalent F (µ * , σ * , θ) = F (µ * , σ * , θ) + const., resulting in identical gradients. Because i is sampled from N (0, 1) in both our method and β-VAE, the training algorithms are exactly equivalent.

We can now analyze the properties of this reinterpretation of β-VAE. In particular, we are interested in

EQUATION

)

where

EQUATION

μi |µ i , σ 2 i ∼ N µ i , σ 2 i with µ i ∼ N 0, 1 β implies μi |σ 2 i ∼ N 0, σ 2 i + 1 β . Therefore,

EQUATION

Figure 8 shows numerical results of Equation 31 for various values of β. Interestingly, the setting β > 1 that is suggested by Higgins et al. (2017) for obtaining disentangled representations corresponds to lower latent capacity.

:

Due to the inferred noise q ( ), this is not pure MAP, but rather a lower bound objective similar to the free energy. We refer to it as MAP for compactness, and connections to the variational free energy are given in Appendix B.

Figures and Tables:

Figure 2: Gaussian mean field inference of fixed scale on model parameters with a Gaussian prior (a) can be reinterpreted as MAP inference on a model with injected noise (b): The mean of the original inference model correspond to the parameters of the new generative model.

Figure3: Classifying CIFAR10 with varying model capacities. Large capacities lead to overfitting while small capacities drown the signal in noise. Each configuration has been evaluated 5 times; mean and standard deviation are displayed.

The test ELBO is not invariant when varying the prior on the model parameters. Neverthless, the first increasing and then decreasing trend when changing the capacity remains. Using an improper prior, similar to just using Gaussian Dropout on the weights, leads to accelerated decreasing generalization for smaller noise scales.

Figure 4: MNIST test reconstruction with a VAE training on 200 samples for various priors and capacities.

Total model capacity 330 kbit Total model capacity 396 kbit Total model capacity 495 kbit Total model capacity 594 kbit Total model capacity 891 kbit Total model capacity 1189 kbit Total model capacity 2378 kbit Total model capacity unrestricted(a) Varying the number of samples. Depending on the size of the dataset higher capacities of the model are required to fit all the datapoints.0.0 250.0K 500.0K 750.0K 1.0M 1.2M 1.5M 1.8M 2Varying architecture. Overfitting is not getting worse for more layers if capacity is low enough. More layers do overfit only for higher capacities.

Figure 5: MNIST test reconstruction with a VAE training on varying dataset sizes, architectures, and model capacities.

Figure 6: Test reconstruction means for binarized fashion MNIST trained on 200 samples with per-parameter capacities 5, 2 and 1 bits (from top) compared to the true data (bottom).

Figure 7: Gaussian mean field inference on a latent model (a) with the (β-)VAE objective can be reinterpreted as MAP inference on mean and variance latents (b): The output mean and variance of (β-)VAE's inference model correspond to the latents of the new generative model. This is a per-datapoint-view with data indices (n) and model parameter dependencies omitted.

Figure 8: Relationship between β and capacity I(μ i , (µ i , σ 2 i )) of each latent component in flexiblescale Gaussian mean field inference with complexity term scaled by β > 0. Values are given in Table1.

Peer review (iclr2019):
 I read the paper and understand it, for the most part. The idea is to interpret some regularization technics as a from of noisy bottleneck, where the mutual information b tween learned parameters and the data is limited through the injection of noise. While, the paper is a plaisant read, I find difficult to access its importance and the applicability of the ideas presented beyond the analogy with the capacity computation. Perhaps other referee will have a clearer opinion. I*d be interested to hear if the authors see a connection between their formalism and the one of Reference prior in Bayesian inference (Bernardo et al https://arxiv.org/pdf/0904.0156) Pro: nicely written, clear interpretation of regularization as a noise injection technics, explicit link with information theoery and Shanon capacity. Con: not clear to me how strong and wide the implications are, beyond the analogies and the reinterpretation