Emergence of Invariance and Disentanglement in Deep Representations

TL;DR

This paper links invariance and disentanglement in deep neural networks to information theory principles, proposing regularization methods that improve generalization and explain phase transitions in training.

Contribution

It introduces a novel information-theoretic regularization framework for deep learning that promotes invariance and disentanglement, connecting these properties to the information in weights.

Findings

Regularization bounds invariance and disentanglement.

Quantifies phase transitions between underfitting and overfitting.

Experimental validation of theoretical predictions.

Abstract

Using established principles from Statistics and Information Theory, we show that invariance to nuisance factors in a deep neural network is equivalent to information minimality of the learned representation, and that stacking layers and injecting noise during training naturally bias the network towards learning invariant representations. We then decompose the cross-entropy loss used during training and highlight the presence of an inherent overfitting term. We propose regularizing the loss by bounding such a term in two equivalent ways: One with a Kullbach-Leibler term, which relates to a PAC-Bayes perspective; the other using the information in the weights as a measure of complexity of a learned model, yielding a novel Information Bottleneck for the weights. Finally, we show that invariance and independence of the components of the representation learned by the network are bounded above and below by the information in the weights, and therefore are implicitly optimized during training. The theory enables us to quantify and predict sharp phase transitions between underfitting and overfitting of random labels when using our regularized loss, which we verify in experiments, and sheds light on the relation between the geometry of the loss function, invariance properties of the learned representation, and generalization error.

Figures

Tables

Table 1. Table 1: (Left) The Small AlexNet model used in the random label experiment, adapted from Zhang et al. ( 2017 ) . All convolutions have a 5 × 5 5 5 5{\times}5 kernel. The use of batch normalization makes the training procedure more stable, but did not significantly change the results of the experiments. (Center) All Convolutional Network (Springenberg et al., 2014 ) used as a classifier in the experiments. All convolutions but the last one use a 3 × 3 3 3 3{\times}3 kernel, “s2” denotes a convolution with stride 2. The final representation we use are the activations of the last “conv 192” layer. (Right) The ResNet architecture (He et al., 2016 ) on which we test the phase transition. Each block with f filters and stride s 𝑠 s is structured as follows: BN -> ReLU -> conv f stride s -> BN -> ReLU -> conv f with a skip connection between first ReLU and the output.

Input 32x32
conv 64
ReLU
MaxPool 2x2
conv 64 + BN
ReLU
MaxPool 2x2
FC 3136x384 + BN
ReLU
FC 384x192 + BN
ReLU
FC 192x10
softmax

Equations

$$\mathrm{TC}(z)=\mathrm{KL}(p(z)\Vert {\prod }_{i}p(z_i)),$$

$$I(z;x)={\mathbb{E}}_{x\sim p(x)}\mathrm{KL}(p(z|x)\Vert p(z)).$$

$$\mathcal{L}(p(z|x))=H(y|z)+\beta I(z;x),$$

$$I(z;n)\le I(z;x)-I(x;y).$$

$$I(z;n)=I(z;x)-I(x;y)-ϵ,$$

$$\mathcal{L}(p(z|x))=H(y|z)+\beta I(z;x),$$

$$x\to z_1\to z_2,$$

$$x\to z_1\to z_2\to \dots \to z_L,$$

$$H_{p,q}(\mathbf{y}|\mathbf{x},w)={\mathbb{E}}_{\mathcal{D}=(\mathbf{x},\mathbf{y})}{\mathbb{E}}_{w\sim q(w|\mathcal{D})}\sum _{i=1}^N-\log q(y^{(i)}|x^{(i)},w),$$

$$H_{p,q}(\mathbf{y}|\mathbf{x},w)=\underset{\text{intrinsic error}}{\underbrace{H(\mathbf{y}|\mathbf{x},\theta )}}+\underset{\text{sufficiency}}{\underbrace{I(\theta ;\mathbf{y}|\mathbf{x},w)}}+{\mathbb{E}}_{\mathbf{x},w}\underset{\text{efficiency}}{\underbrace{\mathrm{KL}(p(\mathbf{y}|\mathbf{x},w)\Vert q(\mathbf{y}|\mathbf{x},w))}}-\underset{\text{overfitting}}{\underbrace{I(\mathbf{y};w|\mathbf{x},\theta )}}.$$

$$\mathcal{L}(q(w|\mathcal{D}))=H_{p,q}(\mathbf{y}|\mathbf{x},w)+\beta I(w;\mathcal{D}),$$

$${\mathbb{E}}_{\mathcal{D}}\mathrm{KL}(q(w|\mathcal{D})\Vert q(w))$$

$$={\mathbb{E}}_{\mathcal{D}}\mathrm{KL}(q(w|\mathcal{D})\Vert p(w)),$$

$$\mathcal{L}(q(w|\mathcal{D}))=H_{p,q}(\mathbf{y}|\mathbf{x},w)+\beta \mathrm{KL}(q(w|\mathcal{D})\Vert p(w))$$

$$\mathcal{L}(q(w|\mathcal{D}))=H_{p,q}(\mathbf{y}|\mathbf{x},w)+\beta \mathrm{KL}(q(w|\mathcal{D})\Vert \tilde{q}(w))$$

$$\tilde{I}(w;\mathcal{D}):=\mathrm{KL}(q(w|\mathcal{D})\Vert \tilde{q}(w))=\mathrm{KL}(q(w|\mathcal{D})\Vert {\prod }_{i}q(w_i)).$$

$$w_i|\mathcal{D}\sim {ϵ}_i{\hat{w}}_i,$$

$$I(w;\mathcal{D})\le \tilde{I}(w;\mathcal{D})=-\frac{1}{2}\sum _{i=1}^{\dim (w)}\log {\alpha }_i+C,$$

$$I(w;\mathcal{D})\le \tilde{I}(w;\mathcal{D})\le \frac{1}{2}K[\log {\Vert \hat{w}\Vert }_2^2+\log {\Vert \mathcal{H}\Vert }_{\ast }-K\log (K^2\beta /2)]$$

$$I(z;x)+\mathrm{TC}(z)=-\frac{1}{2}\sum _{i=1}^{\dim (z)}{\mathbb{E}}_x\log \frac{{\tilde{\alpha }}_i{\hat{W}}_i^2\cdot x^2}{{\hat{W}}_i\cdot \mathrm{Cov}(x){\hat{W}}_i+{\tilde{\alpha }}_i{\hat{W}}_i^2\cdot \mathbb{E}(x^2)},$$

$$g(\alpha )\le \frac{I(x;z)+TC(z)}{\dim (z)}\le g(\alpha )+c,$$

$$I(z_{L};x)\leq\min_{k<L}\left\{\dim(z_{k})\big{[}g(\alpha^{k})+1\big{]}\right\}$$

$$\mathbb{E}_{\mathcal{D}}[L^{\text{test}}(q(w|\mathcal{D}))]\leq\frac{1}{N(1-\frac{1}{2\lambda})}\Big{(}H_{p,q}(\mathbf{y}|\mathbf{x},w)+\lambda L_{\text{max}}\mathbb{E}_{\mathcal{D}}[\operatorname{KL}(\,q(w|\mathcal{D})\,\|\,p(w)\,)]\Big{)},$$

$${\mathbb{E}}_{\mathcal{D}}[\mathrm{KL}(q(w|\mathcal{D})\Vert q(w))]$$

$$\le {\mathbb{E}}_{\mathcal{D}}[\mathrm{KL}(q(w|\mathcal{D})\Vert p(w))],$$

$$I(z;n)$$

$$={\mathbb{E}}_{n\sim p(n)}{\mathbb{E}}_{z\sim p(z|n)}\log [p(z|n)/p(z)].$$

$$D^{\ast }(z;n)=\frac{p(z)}{p(z)+p(z|n)},$$

$$\log \frac{p(z|n)}{p(z)}=\log \frac{1-D^{\ast }(z;n)}{D^{\ast }(z;n)}\simeq \log \frac{1-D(z;n)}{D(z;n)}.$$

$${\mathbb{E}}_z{\mathbb{E}}_{\hat{x}\sim q(\hat{x}|z)}[\log p(\hat{x}|z)]=$$

$${\mathbb{E}}_z{\mathbb{E}}_{\hat{x}\sim q(\hat{x}|z)}[\log p(\hat{x}|z)]+H(p(\hat{x}))={\mathbb{E}}_z{\mathbb{E}}_{\hat{x}\sim q(\hat{x}|z)}[\log p(z|\hat{x})]-\mathrm{KL}(q(\hat{x})\Vert p(\hat{x}))+C,$$

$$I(z;n)\le I(z;x)-I(x;y).$$

Full text

Keywords: Representation learning; PAC-Bayes; information bottleneck; flat minima; generalization; invariance; independence;

1 Introduction

Efforts to understand the empirical success of deep learning have followed two main lines: Representation learning and optimization. In optimization, a deep network is treated as a black-box family of functions for which we want to find parameters (weights) that yield good generalization. Aside from the difficulties due to the non-convexity of the loss function, the fact that deep networks are heavily over-parametrized presents a theoretical challenge: The bias-variance trade-off suggests they may severely overfit; yet, even without explicit regularization, they perform remarkably well in practice. Recent work suggests that this is related to properties of the loss landscape and to the implicit regularization performed by stochastic gradient descent (SGD), but the overall picture is still hazy (Zhang et al., 2017).

Representation learning, on the other hand, focuses on the properties of the representation learned by the layers of the network (the activations) while remaining largely agnostic to the particular optimization process used. In fact, the effectiveness of deep learning is often ascribed to the ability of deep networks to learn representations that are insensitive (invariant) to nuisances such as translations, rotations, occlusions, and also “disentangled,” that is, separating factors in the high-dimensional space of data (Bengio, 2009). Careful engineering of the architecture plays an important role in achieving insensitivity to simple geometric nuisance transformations, like translations and small deformations; however, more complex and dataset-specific nuisances still need to be learned. This poses a riddle: If neither the architecture nor the loss function explicitly enforce invariance and disentangling, how can these properties emerge consistently in deep networks trained by simple generic optimization?

In this work, we address these questions by establishing information theoretic connections between these concepts. In particular, we show that: (a) a sufficient representation of the data is invariant if and only if it is minimal, i.e., it contains the smallest amount of information, although may not have small dimension; (b) the information in the representation, along with its total correlation (a measure of disentanglement) are tightly bounded by the information that the weights contain about the dataset; (c) the information in the weights, which is related to overfitting (Hinton and Van Camp, 1993), flat minima (Hochreiter and Schmidhuber, 1997), and a PAC-Bayes upper-bound on the test error (Section 6), can be controlled by implicit or explicit regularization. Moreover, we show that adding noise during the training is a simple and natural way of biasing the network towards invariant representations.

Finally, we perform several experiments with realistic architectures and datasets to validate the assumptions underlying our claims. In particular, we show that using the information in the weights to measure the complexity of a deep neural network (DNN), rather than the number of its parameters, leads to a sharp and theoretically predicted transition between overfitting and underfitting regimes for random labels, shedding light on the questions of Zhang et al. (2017).

1.1 Related work

The Information Bottleneck (IB) was introduced by Tishby et al. (1999) as a generalization of minimal sufficient statistics that allows trading off fidelity (sufficiency) and complexity of a representation. In particular, the IB Lagrangian reduces finding a minimal sufficient representation to a variational optimization problem. Later, Tishby and Zaslavsky (2015) and Shwartz-Ziv and Tishby (2017) advocated using the IB between the test data and the activations of a deep neural network, to study the sufficiency and minimality of the resulting representation. In parallel developments, the IB Lagrangian was used as a regularized loss function for learning representation, leading to new information theoretic regularizers (Achille and Soatto, 2018; Alemi et al., 2017a, b).

In this paper, we introduce an IB Lagrangian between the weights of a network and the training data, as opposed to the traditional one between the activations and the test datum. We show that the former can be seen both as a generalization of Variational Inference, related to Hinton and Van Camp (1993), and as a special case of the more general PAC-Bayes framework (McAllester, 2013), that can be used to compute high-probability upper-bounds on the test error of the network. One of our main contributions is then to show that, due to a particular duality induced by the architecture of deep networks, minimality of the weights (a function of the training dataset) and of the learned representation (a function of the test input) are connected: in particular we show that networks regularized either explicitly, or implicitly by SGD, are biased toward learning invariant and disentangled representations. The theory we develop could be used to explain the phenomena described in small-scale experiments in Shwartz-Ziv and Tishby (2017), whereby the initial fast convergence of SGD is related to sufficiency of the representation, while the later asymptotic phase is related to compression of the activations: While SGD is seemingly agnostic to the property of the learned representation, we show that it does minimize the information in the weights, from which the compression of the activations follows as a corollary of our bounds. Practical implementation of this theory on real large scale problems is made possible by advances in Stochastic Gradient Variational Bayes (Kingma and Welling, 2014; Kingma et al., 2015).

Representations learned by deep networks are observed to be insensitive to complex nuisance transformations of the data. To a certain extent, this can be attributed to the architecture. For instance, the use of convolutional layers and max-pooling can be shown to yield insensitivity to local group transformations (Bruna and Mallat, 2011; Anselmi et al., 2016; Soatto and Chiuso, 2016). But for more complex, dataset-specific, and in particular non-local, non-group transformations, such insensitivity must be acquired as part of the learning process, rather than being coded in the architecture. We show that a sufficient representation is maximally insensitive to nuisances if and only if it is minimal, allowing us to prove that a regularized network is naturally biased toward learning invariant representations of the data.

Efforts to develop a theoretical framework for representation learning include Tishby and Zaslavsky (2015) and Shwartz-Ziv and Tishby (2017), who consider representations as stochastic functions that approximate minimal sufficient statistics, different from Bruna and Mallat (2011) who construct representations as (deterministic) operators that are invertible in the limit, while exhibiting reduced sensitivity (“stability”) to small perturbations of the data. Some of the deterministic constructions are based on the assumption that the underlying data is spatially stationary, and therefore work best on textures and other visual data that are not subject to occlusions and scaling nuisances. Anselmi et al. (2016) develop a theory of invariance to locally compact groups, and aim to construct maximal (“distinctive”) invariants, like Sundaramoorthi et al. (2009) that, however, assume nuisances to be infinite-dimensional groups (Grenander, 1993). These efforts are limited by the assumption that nuisances have a group structure. Such assumptions were relaxed by Soatto and Chiuso (2016) who advocate seeking for sufficient invariants, rather than maximal ones. We further advance this approach, but unlike prior work on sufficient dimensionality reduction, we do not seek to minimize the dimension of the representation, but rather its information content, as prescribed by our theory. Recent advances in Deep Learning provide us with computationally viable methods to train high-dimensional models and predict and quantify observed phenomena such as convergence to flat minima and transitions from overfitting to underfitting random labels, thus bringing the theory to fruition. Other theoretical efforts focus on complexity considerations, and explain the success of deep networks by ways of statistical or computational efficiency (Lee et al., 2017; Bengio, 2009; LeCun, 2012). “Disentanglement” is an often-cited property of deep networks (Bengio, 2009), but seldom formalized and studied analytically, although Ver Steeg and Galstyan (2015) has suggested studying it using the Total Correlation of the representation, also known as multi-variate mutual information, which we also use.

We connect invariance properties of the representation to the geometry of the optimization residual, and to the phenomenon of flat minima (Dinh et al., 2017).

Following (McAllester, 2013), we have also explored relations between our theory and the PAC-Bayes framework (Dziugaite and Roy, 2017). As we show, our theory can also be derived in the PAC-Bayes framework, without resorting to information quantities and the Information Bottleneck, thus providing both an independent and alternative derivation, and a theoretically rigorous way to upper-bound the optimal loss function. The use of PAC-Bayes theory to study the generalization properties of deep networks has been championed by Dziugaite and Roy (2017), who point out that minima that are flat in the sense of having a large volume, toward which stochastic gradient descent algorithms are implicitly or explicitly biased (Chaudhari and Soatto, 2018), naturally relates to the PAC-Bayes loss for the choice of a normal prior and posterior on the weights. This has been leveraged by Dziugaite and Roy (2017) to compute non-vacuous PAC-Bayes error bounds, even for deep networks.

2 Preliminaries

A training set $\mathcal{D}=\left\{\mathbf{x},\mathbf{y}\right\}$, where $\mathbf{x}=\left\{x^{(i)}\right\}_{i=1}^{N}$ and $\mathbf{y}=\left\{y^{(i)}\right\}_{i=1}^{N})$, is a collection of $N$ randomly sampled data points $x^{(i)}_{i}$ and their associated (usually discrete) labels. The samples are assumed to come from an unknown, possibly complex, distribution $p_{\theta}(x,y)$, parametrized by a parameter $\theta$. Following a Bayesian approach, we also consider $\theta$ to be a random variable, sampled from some unknown prior distribution $p(\theta)$, but this requirement is not necessary (see Section 6). A test datum $x$ is also a random variable. Given a test sample, our goal is to infer the random variable $y$, which is therefore referred to as our task.

We will make frequent use of the following standard information theoretic quantities (Cover and Thomas, 2012): Shannon entropy $H(x)=\mathbb{E}_{p}[-\log p(x)]$, conditional entropy $H(x|y):=\mathbb{E}_{\bar{y}}[H(x|y=\bar{y})]=H(x,y)-H(y)$, (conditional) mutual information $I(x;y|z)=H(x|z)-H(x|y,z)$, Kullbach-Leibler (KL) divergence $KL(p(x)||q(x))=\mathbb{E}_{p}[\log p/q]$, cross-entropy $H_{p,q}(x)=\mathbb{E}_{p}[-\log q(x)]$, and total correlation $TC(z)$, which is also known as multi-variate mutual information and defined as

[TABLE]

where $p(z_{i})$ are the marginal distributions of the components of $z$. Recall that the KL divergence between two distributions is always non-negative and zero if and only if they are equal. In particular $\operatorname{TC}(z)$ is zero if and only if the components of $z$ are independent, in which case we say that $z$ is disentangled. We often use of the following identity:

[TABLE]

We say that $x,z,y$ form a Markov chain, indicated with $x\rightarrow z\rightarrow y$, if $p(y|x,z)=p(y|z)$. The Data Processing Inequality (DPI) for a Markov chain $x\rightarrow z\rightarrow y$ ensures that $I(x;z)\geq I(x;y)$: If $z$ is a (deterministic or stochastic) function of $x$, it cannot contain more information about $y$ than $x$ itself (we cannot create new information by simply applying a function to the data we already have).

2.1 General definitions and the Information Bottleneck Lagrangian

We say that $z$ is a representation of $x$ if $z$ is a stochastic function of $x$, or equivalently if the distribution of $z$ is fully described by the conditional $p(z|x)$. In particular we have the Markov chain $y\to x\to z$. We say that a representation $z$ of $x$ is sufficient for $y$ if y\mathbin{\rotatebox[origin={c}]{90.0}{\models}}x\ |\ z, or equivalently if $I(z;y)=I(x;y)$; it is minimal when $I(x;z)$ is smallest among sufficient representations. To study the trade-off between sufficiency and minimality, Tishby et al. (1999) introduces the Information Bottleneck Lagrangian

[TABLE]

where $\beta$ trades off sufficiency (first term) and minimality (second term); in the limit $\beta\rightarrow 0$, the IB Lagrangian is minimized when $z$ is minimal and sufficient. It does not impose any restriction on disentanglement nor invariance, which we introduce next.

2.2 Nuisances for a task

A nuisance is any random variable that affects the observed data $x$, but is not informative to the task we are trying to solve. More formally, a random variable $n$ is a nuisance for the task $y$ if y\mathbin{\rotatebox[origin={c}]{90.0}{\models}}n, or equivalently $I(y;n)=0$. Similarly, we say that the representation $z$ is invariant to the nuisance $n$ if z\mathbin{\rotatebox[origin={c}]{90.0}{\models}}n, or $I(z;n)=0$. When $z$ is not strictly invariant but it minimizes $I(z;n)$ among all sufficient representations, we say that the representation $z$ is maximally insensitive to $n$.

One typical example of nuisance is a group $G$, such as translation or rotation, acting on the data. In this case, a deterministic representation $f$ is invariant to the nuisances if and only if for all $g\in G$ we have $f(g\cdot x)=f(x)$. Our definition however is more general in that it is not restricted to deterministic functions, nor to group nuisances. An important consequence of this generality is that the observed data $x$ can always be written as a deterministic function of the task $y$ and of all nuisances $n$ affecting the data, as explained by the following proposition.

Proposition 2.1** (Task-nuisance decomposition, Appendix C.1)**

Given a joint distribution $p(x,y)$, where $y$ is a discrete random variable, we can always find a random variable $n$ independent of $y$ such that $x=f(y,n)$, for some deterministic function $f$.

3 Properties of optimal representations

To simplify the inference process, instead of working directly with the observed high dimensional data $x$, we want to use a representation $z$ that captures and exposes only the information relevant for the task $y$. Ideally, such a representation should be (a) sufficient for the task $y$, i.e. $I(y;z)=I(y;x)$, so that information about $y$ is not lost; among all sufficient representations, it should be (b) minimal, i.e. $I(z;x)$ is minimized, so that it retains as little about $x$ as possible, simplifying the role of the classifier; finally, it should be (c) invariant to the effect of nuisances $I(z;n)=0$, so that the final classifier will not overfit to spurious correlations present in the training dataset between nuisances $n$ and labels $y$. Such a representation, if it exists, would not be unique, since any bijective mapping preserves all these properties. We can use this to our advantage and further aim to make the representation (d) maximally disentangled, i.e., choose the one(s) for which $\operatorname{TC}(z)$ is minimal. This simplifies the classifier rule, since no information will be present in the higher-order correlations between the components of $z$.

Inferring a representation that satisfies all these properties may seem daunting. However, in this section we show that we only need to enforce (a) sufficiency and (b) minimality, from which invariance and disentanglement follow naturally thanks to the stacking of noisy layers of computation in deep networks. We will then show that sufficiency and minimality of the learned representation can also be promoted easily through implicit or explicit regularization during the training process.

Proposition 3.1** (Invariance and minimality, Appendix C.2)**

Let $n$ be a nuisance for the task $y$ and let $z$ be a sufficient representation of the input $x$. Suppose that $z$ depends on $n$ only through $x$ (i.e., $n\to x\to z$). Then,

[TABLE]

Moreover, there is a nuisance $n$ such that equality holds up to a (generally small) residual $\epsilon$

[TABLE]

where $\epsilon:=I(z;y|n)-I(x;y)$. In particular $0\leq\epsilon\leq H(y|x)$, and $\epsilon=0$ whenever $y$ is a deterministic function of $x$. Under these conditions, a sufficient statistic $z$ is invariant (maximally insensitive) to nuisances if and only if it is minimal.

Remark 3.2

Since $\epsilon\leq H(y|x)$, and usually $H(y|x)=0$ or at least $H(y|x)\ll I(x;z)$, we can generally ignore the extra term.**

An important consequence of this proposition is that we can construct invariants by simply reducing the amount of information $z$ contains about $x$, while retaining the minimum amount $I(z;x)$ that we need for the task $y$. This provides the network a way to automatically learn invariance to complex nuisances, which is complementary to the invariance imposed by the architecture. Specifically, one way of enforcing minimality explicitly, and hence invariance, is through the IB Lagrangian.

Corollary 3.3** (Invariants from the Information Bottleneck)**

Minimizing the IB Lagrangian

[TABLE]

in the limit $\beta\to 0$, yields a sufficient invariant representation $z$ of the test datum $x$ for the task $y$.

Remarkably, the IB Lagrangian can be seen as the standard cross-entropy loss, plus a regularizer $I(z;x)$ that promotes invariance. This fact, without proof, is implicitly used in Achille and Soatto (2018), who also provide an efficient algorithm to perform the optimization. Alemi et al. (2017a) also propose a related algorithm and empirically show improved resistance to adversarial nuisances. In addition to modifying the cost function, invariance can also be fostered by choice of architecture:

Corollary 3.4** (Bottlenecks promote invariance)**

Suppose we have the Markov chain of layers

[TABLE]

and suppose that there is a communication or computation bottleneck between $z_{1}$ and $z_{2}$ such that $I(z_{1};z_{2})<I(z_{1};x)$. Then, if $z_{2}$ is still sufficient, it is more invariant to nuisances than $z_{1}$. More precisely, for all nuisances $n$ we have $I(z_{2};n)\leq I(z_{1};z_{2})-I(x;y)$.

Such a bottleneck can happen for example because $\dim(z_{2})<\dim(z_{1})$, e.g., after a pooling layer, or because the channel between $z_{1}$ and $z_{2}$ is noisy, e.g., because of dropout.

Proposition 3.5** (Stacking increases invariance)**

Assume that we have the Markov chain of layers

[TABLE]

and that the last layer $z_{L}$ is sufficient of $x$ for $y$. Then $z_{L}$ is more insensitive to nuisances than all the preceding layers.

Notice, however, that the above corollary does not simply imply that the more layers the merrier, as it assumes that one has successfully trained the network ($z_{L}$ is sufficient), which becomes increasingly difficult as the size grows. Also note that in some architectures, such as ResNets (He et al., 2016), the layers do not necessarily form a Markov chain because of skip connections; however, their “blocks” still do.

Proposition 3.6** (Actionable Information)**

When $z=f(x)$ is a deterministic invariant, if it minimizes the IB Lagrangian it also maximizes Actionable Information (Soatto, 2013), which is ${\cal H}(x):=H(f(x))$.

Although Soatto (2013) addressed maximal invariants, we only consider sufficient invariants, as advocated by (Soatto and Chiuso, 2016).

Information in the weights

Thus far we have discussed properties of representations in generality, regardless of how they are implemented or learned. Given a source of data (for example randomly generated, or from a fixed dataset), and given a (stochastic) training algorithm, the output weight $w$ of the training process can be thought as a random variable (that depends on the stochasticity of the initialization, training steps and of the data). We can therefore talk about the information that the weights contain about the dataset $\mathcal{D}$ and the training procedure, which we denote by $I(w;{\cal D})$.

Two extreme cases consist of the trivial settings where we use the weights to memorize the dataset (the most extreme form of overfitting), or where the weights are constant, or pure noise (sampled from a process that is independent of the data). In between, the amount of information the weights contain about the training turns out to be an important quantity both in training deep networks, as well as in establishing properties of the resulting representation, as we discuss in the next section.

Note that in general we do not need to compute and optimize the quantity of information in the weights. Instead, we show that we can control it, for instance by injecting noise in the weights, drawn from a chosen distribution, in an amount that can be modulated between zero (thus in theory allowing full information about the training set to be stored in the weights) to an amount large enough that no information is left. We will leverage this property in the next sections to perform regularization.

4 Learning minimal weights

In this section, we let $p_{\theta}(x,y)$ be an (unknown) distribution from which we randomly sample a dataset $\mathcal{D}$. The parameter $\theta$ of the distribution is also assumed to be a random variable with an (unknown) prior distribution $p(\theta)$. For example $p_{\theta}$ can be a fairly general generative model for natural images, and $\theta$ can be the parameters of the model that generated our dataset. We then consider a deep neural network that implements a map $x\mapsto f_{w}(x):=q(\,{\cdot}\,|x,w)$ from an input $x$ to a class distribution $q(y|x,w)$.111 We use $p$ to denote the real (and unknown) data distribution, while $q$ denotes approximate distributions that are optimized during training.

In full generality, and following a Bayesian approach, we let the weights $w$ of the network be sampled from a parametrized distribution $q(w|\mathcal{D})$,whose parameters are optimized during training.222Note that, while the two are somewhat related, here by $q(w|\mathcal{D})$ we denote the output distribution of the weights after training with our choice algorithm on the dataset $\mathcal{D}$, and not the Bayesian posterior of the weights given the dataset, which would be denoted $p(w|\mathcal{D})$. When $q(w|\mathcal{D})$ is a Dirac delta at a point, we recover the standard loss function for a MAP estimate of the weights. The network is then trained in order to minimize the expected cross-entropy loss333 Note that for generality here we treat the dataset $\mathcal{D}$ as a random variable. In practice, when a single dataset is given, the expectation w.r.t. the dataset can be ignored.

[TABLE]

in order for $q(y|x,w)$ to approximate $p_{\theta}(y|x)$.

One of the main problems in optimizing a DNN is that the cross-entropy loss in notoriously prone to overfitting. In fact, one can easily minimize it even for completely random labels (see Zhang et al. (2017), and Figure 1). The fact that, somehow, such highly over-parametrized functions manage to generalize when trained on real labels has puzzled theoreticians and prompted some to wonder whether this may be inconsistent with the intuitive interpretation of the bias-variance trade-off theorem, whereby unregularized complex models should overfit wildly. However, as we show next, there is no inconsistency if one measures complexity by the information content, and not the dimensionality, of the weights.

To gain some insights about the possible causes of over-fitting, we can use the following decomposition of the cross-entropy loss (we refer to Section C for the proof and the precise definition of each term):

[TABLE]

The first term of the right-hand side of (8) relates to the intrinsic error that we would commit in predicting the labels even if we knew the underlying data distribution $p_{\theta}$; the second term measures how much information that the dataset has about the parameter $\theta$ is captured by the weights, the third term relates to the efficiency of the model and the class of functions $f_{w}$ with respect to which the loss is optimized. The last, and only negative, term relates to how much information about the labels, but uninformative of the underlying data distribution, is memorized in the weights. Unfortunately, without implicit or explicit regularization, the network can minimize the cross-entropy loss (LHS), by just maximizing the last term of eq. 8, i.e., by memorizing the dataset, which yields poor generalization.

To prevent the network from doing this, we can neutralize the effect of the negative term by adding it back to the loss function, leading to a regularized loss $L=H_{p,q}(\mathbf{y}|\mathbf{x},w)+I(\mathbf{y};w|\mathbf{x},\theta)$. However, computing, or even approximating, the value of $I(\mathbf{y},w|\mathbf{x},\theta)$ is at least as difficult as fitting the model itself.

We can, however, add an upper bound to $I(\mathbf{y};w|\mathbf{x},\theta)$ to obtain the desired result. In particular, we explore two alternate paths that lead to equivalent conclusions under different premises and assumptions: In one case, we use a PAC-Bayes upper-bound, which is $\operatorname{KL}(\,q(w|\mathcal{D})\,\|\,p(w)\,)$ where $p(w)$ is an arbitrary prior. In the other, we use the IB Lagrangian and upper-bound it with the information in the weights $I(w;\mathcal{D})$. We discuss this latter approach now, and look at the PAC-Bayes approach in Section 6.

Notice that to successfully learn the distribution $p_{\theta}$, we only need to memorize in $w$ the information about the latent parameters $\theta$, that is we need $I(\mathcal{D};w)=I(\mathcal{D};\theta)\leq H(\theta)$, which is bounded above by a constant. On the other hand, to overfit, the term $I(\mathbf{y};w|\mathbf{x},\theta)\leq I(\mathcal{D};w|\theta)$ needs to grow linearly with the number of training samples $N$. We can exploit this fact to prevent overfitting by adding a Lagrange multiplier $\beta$ to make the amount of information a constant with respect to $N$, leading to the regularized loss function

[TABLE]

which, remarkably, has the same general form of an IB Lagrangian, and in particular is similar to (1), but now interpreted as a function of the weights $w$ rather than the activations $z$. This use of the IB Lagrangian is, to the best of our knowledge, novel, as the role of the Information Bottleneck has thus far been confined to characterizing the activations of the network, and not as a learning criterion. Equation 3 can be seen as a generalization of other suggestions in the literature:

IB Lagrangian, Variational Learning and Dropout.

Minimizing the information stored at the weights $I(w;\mathcal{D})$ was proposed as far back as Hinton and Van Camp (1993) as a way of simplifying neural networks, but no efficient algorithm to perform the optimization was known at the time. For the particular choice $\beta=1$, the IB Lagrangian reduces to the variational lower-bound (VLBO) of the marginal log-likelihood $p(\mathbf{y}|\mathbf{x})$. Therefore, minimizing eq. 3 can also be seen as a generalization of variational learning. A particular case of this was studied by Kingma et al. (2015), who first showed that a generalization of Dropout, called Variational Dropout, could be used in conjunction with the reparametrization trick Kingma and Welling (2014) to minimize the loss efficiently.

Information in the weights as a measure of complexity.

Just as Hinton and Van Camp (1993) suggested, we also advocate using the information regularizer $I(w;\mathcal{D})$ as a measure of the effective complexity of a network, rather than the number of parameters $\dim(w)$, which is merely an upper bound on the complexity. As we show in experiments, this allows us to recover a version of the bias-variance trade-off where networks with lower information complexity underfit the data, and networks with higher complexity overfit. In contrast, there is no clear relationship between number of parameters and overfitting (Zhang et al., 2017). Moreover, for random labels the information complexity allows us to precisely predict the overfitting and underfitting behavior of the network (Section 7).

4.1 Computable upper-bound to the loss

Unfortunately, computing $I(w,\mathcal{D})=\mathbb{E}_{\mathcal{D}}\operatorname{KL}(\,q(w|\mathcal{D})\,\|\,q(w)\,)$ is still too complicated, since it requires us to know the marginal $q(w)$ over all possible datasets and trainings of the network. To avoid computing this term, we can use the more general upper-bound

[TABLE]

where $p(w)$ is any fixed distribution of the weights. Once we instantiate the training set, we have a single sample of ${\cal D}$, so the expectation over $\mathcal{D}$ becomes trivial. This gives us the following upper bound to the optimal loss function

[TABLE]

Generally, we want to pick $p(w)$ in order to give the sharpest upper-bound, and to be a fully factorized distribution, i.e., a distribution with independent components, in order to make the computation of the KL term easier. The sharpest upper-bound to $\operatorname{KL}(\,q(w|\mathcal{D})\,\|\,q(w)\,)$ that can be obtained using a factorized distribution $p$ is obtained when $p(w):=\tilde{q}(w)=\prod_{i}q(w_{i})$ where $q(w_{i})$ denotes the marginal distributions of the components of $q(w)$. Notice that. once a training procedure is fixed, this may be approximated by training multiple times and approximating each marginal weight distribution. With this choice of prior, our final loss function becomes

[TABLE]

for some fixed distribution $\tilde{q}$ that approximates the real marginal distribution $q(w)$. The IB Lagrangian for the weights in eq. 3 can be seen as a generally intractable special case of eq. 5 that gives the sharpest upper-bound to our desired loss in this family of losses.

In the following, to keep the notation uncluttered, we will denote our upper bound $\operatorname{KL}(\,q(w|\mathcal{D})\,\|\,\tilde{q}(w)\,)$ to the mutual information $I(w;\mathcal{D})$ simply by $\tilde{I}(w;\mathcal{D})$, where

[TABLE]

4.2 Bounding the information in the weights of a network

To derive precise and empirically verifiable statements about $\tilde{I}(w;\mathcal{D})$, we need a setting where this can be expressed analytically and optimized efficiently on standard architectures. To this end, following Kingma et al. (2015), we make the following modeling choices.

Modeling assumptions.

Let $w$ denote the vector containing all the parameters (weights) in the network, and let $W^{k}$ denote the weight matrix at layer $k$. We assume an improper log-uniform prior on $w$, that is $\tilde{q}(w_{i})=c/|w_{i}|$. Notice that this is the only scale-invariant prior (Kingma et al., 2015), and closely matches the real marginal distributions of the weights in a trained network (Achille and Soatto, 2018); we parametrize the weight distribution $q(w_{i}|\mathcal{D})$ during training as

[TABLE]

where $\hat{w}_{i}$ is a learned mean, and $\epsilon_{i}\sim\log\mathcal{N}(-\alpha_{i}/2,\alpha_{i})$ is i.i.d. multiplicative log-normal noise with mean 1 and variance $\exp(\alpha_{i})-1$.444For a log-normal $\log\mathcal{N}(\mu,\sigma^{2})$ mean and variance are respectively $\exp(\mu+\sigma^{2}/2)$ and $[\exp(\sigma^{2})-1]\exp(2\mu+\sigma^{2})$.

Note that while Kingma et al. (2015) uses this parametrization as a local approximation of the Bayesian posterior for a given (log-uniform) prior, we rather define the distribution of the weights $w$ after training on the dataset $\mathcal{D}$ to be $q(w|\mathcal{D})$.

Proposition 4.1** (Information in the weights, Proposition C.4)**

Under the previous modeling assumptions, the upper-bound to the information that the weights contain about the dataset is

[TABLE]

where the constant $C$ is arbitrary due to the improper prior.

Remark 4.2** (On the constant $C$)**

To simplify the exposition, since the optimization is unaffected by any additive constant, in the following we abuse the notation and, under the modeling assumptions stated above, we rather define $\tilde{I}(w;\mathcal{D}):=-\frac{1}{2}\sum_{i=1}^{\dim(w)}\log\alpha_{i}$. Neklyudov et al. (2017) also suggest a principled way of dealing with the arbitrary constant by using a proper log-uniform prior.

Note that computing and optimizing this upper-bound to the information in the weights is relatively simple and efficient using the reparametrization trick of Kingma et al. (2015).

4.3 Flat minima have low information

Thus far we have suggested that adding the explicit information regularizer $I(w;\mathcal{D})$ prevents the network from memorizing the dataset and thus avoid overfitting, which we also confirm empirically in Section 7. However, real networks are not commonly trained with this regularizer, thus seemingly undermining the theory. However, even when not explicitly present, the term $I(w;\mathcal{D})$ is implicit in the use of SGD. In particular, Chaudhari and Soatto (2018) show that, under certain conditions, SGD introduces an entropic bias of a very similar form to the information in the weights described thus far, where the amount of information can be controlled by the learning rate and the size of mini-batches.

Additional indirect empirical evidence is provided by the fact that some variants of SGD (Chaudhari et al., 2017) bias the optimization toward “flat minima”, that are local minima whose Hessian has mostly small eigenvalues. These minima can be interpreted exactly as having low information $I(w;\mathcal{D})$, as suggested early on by Hochreiter and Schmidhuber (1997): Intuitively, since the loss landscape is locally flat, the weights may be stored at lower precision without incurring in excessive inference error. As a consequence of previous claims, we can then see flat minima as having better generalization properties and, as we will see in Section 5, the associated representation of the data is more insensitive to nuisances and more disentangled. For completeness, here we derive a more precise relationship between flatness (measured by the nuclear norm of the loss Hessian), and the information content based on our model.

Proposition 4.3** (Flat minima have low information, Appendix C.5)**

Let $\hat{w}$ be a local minimum of the cross-entropy loss $H_{p,q}(\mathbf{y}|\mathbf{x},w)$, and let $\mathcal{H}$ be the Hessian at that point. Then, for the optimal choice of the posterior $w|\mathcal{D}=\epsilon\odot\hat{w}$ centered at $\hat{w}$ that optimizes the IB Lagrangian, we have

[TABLE]

where $K=\dim(w)$ and $\|\cdot\|_{*}$ denotes the nuclear norm.

Notice that a converse inequality, that is, low information implies flatness, needs not hold, so there is no contradiction with the results of Dinh et al. (2017). Also note that for $\tilde{I}(w;\mathcal{D})$ to be invariant to reparametrization one has to consider the constant $C$, which we have ignored (Remark 4.2). The connection between flatness and overfitting has also been studied by Neyshabur et al. (2017), including the effect of the number of parameters in the model.

In the next section, we prove one of our main results, that networks with low information in the weights realize invariant and disentangled representations. Therefore, invariance and disentanglement emerge naturally when training a network with implicit (SGD) or explicit (IB Lagrangian) regularization, and are related to flat minima.

5 Duality of the Bottleneck

The following proposition gives the fundamental link in our model between information in the weights, and hence flatness of the local minima, minimality of the representation, and disentanglement.

Proposition 5.1** (Appendix C.6)**

Let $z=Wx$, and assume as before $W=\epsilon\odot\hat{W}$, with $\epsilon_{i,j}\sim\log\mathcal{N}(-\alpha_{i}/2,\alpha_{i})$. Further assume that the marginals of $p(z)$ and $p(z|x)$ are both approximately Gaussian (which is reasonable for large $\dim(x)$ by the Central Limit Theorem). Then,

[TABLE]

where $W_{i}$ denotes the $i$-th row of the matrix $W$, and $\tilde{\alpha_{i}}$ is the noise variance $\tilde{\alpha}_{i}=\exp(\alpha_{i})-1$. In particular, $I(z;x)+\operatorname{TC}(z)$ is a monotone decreasing function of the weight variances $\alpha_{i}$.

The above identity is difficult to apply in practice, but with some additional hypotheses, we can derive a cleaner uniform tight bound on $I(z;x)+\operatorname{TC}(z)$.

Proposition 5.2** (Uniform bound for one layer, Appendix C.7)**

*Let $z=Wx$, where $W=\epsilon\odot\hat{W}$, where $\epsilon_{i,j}\sim\log\mathcal{N}(-\alpha/2,\alpha)$; assume that the components of $x$ are uncorrelated, and that their kurtosis is uniformly bounded.555 This is a technical hypothesis, always satisfied if the components $x_{i}$ are IID, (sub-)Gaussian, or with uniformly bounded support.

Then, there is a strictly increasing function $g(\alpha)$ s.t. we have the uniform bound*

[TABLE]

where $c=O(1/\dim(x))\leq 1$, $g(\alpha)=-\log\,(1-e^{-\alpha})/2$ and $\alpha$ is related to $\tilde{I}(w;\mathcal{D})$ by $\alpha=\exp\left\{-I(W;\mathcal{D})/\dim(W)\right\}$. In particular, $I(x;z)+TC(z)$ is tightly bounded by $\tilde{I}(W;\mathcal{D})$ and increases strictly with it.

The above theorems tells us that whenever we decrease the information in the weights, either by explicit regularization, or by implicit regularization (e.g., using SGD), we automatically improve the minimality, and hence, by Proposition 3.1, the invariance, and the disentanglement of the learner representation. In particular, we obtain as a corollary that SGD is biased toward learning invariant and disentangled representations of the data. Using the Markov property of the layers, we can easily extend this bound to multiple layers:

Corollary 5.3** (Multi-layer case, Appendix C.8)**

Let $W^{k}$ for $k=1,...,L$ be weight matrices, with $W^{k}=\epsilon^{k}\odot\hat{W}^{k}$ and $\epsilon^{k}_{i,j}=\log\mathcal{N}(-\alpha^{k}/2,\alpha^{k})$, and let $z_{i+1}=\phi(W^{k}z_{k})$, where $z_{0}=x$ and $\phi$ is any nonlinearity. Then,

[TABLE]

where $\alpha^{k}=\exp\left\{-I(W^{k};\mathcal{D})/\dim(W^{k})\right\}$.

Remark 5.4** (Tightness)**

While the bound in Proposition 5.2 is tight, the bound in the multilayer case needs not be. This is to be expected: Reducing the information in the weights creates a bottleneck, but we do not know how much information about $x$ will actually go through this bottleneck. Often, the final layers will let most of the information through, while initial layers will drop the most. **

Remark 5.5** (Training-test transfer)**

We note that we did not make any (explicit) assumption about the test set having the same distribution of the training set. Instead, we make the less restrictive assumption of sufficiency: If the test distribution is entirely different from the training one – one may not be able to achieve sufficiency. This prompts interesting questions about measuring the distance between tasks (as opposed to just distance between distributions), which will be studied in future work. **

6 Connection with PAC-Bayes bounds

In this section we show that using a PAC-Bayes bound, we arrive at the same regularized loss function eq. 5 we obtained using the Information Bottleneck, without the need of any approximation. By Theorem 2 of McAllester (2013), we have that for any fixed $\lambda>1/2$, prior $p(w)$, and any weight distribution $q(w|\mathcal{D})$, the test error $L^{\text{test}}(q(w|\mathcal{D}))$ that the network commits using the weight distribution $q(w|\mathcal{D})$ is upper-bounded in expectation by

[TABLE]

where $L_{\text{\emph{max}}}$ is the maximum per-sample loss function, which for a classification problem we can assume to be upper-bounded, for example by clipping the cross-entropy loss at chance level. Notice that right hand side coincides, modulo a multiplicative constant, with eq. 4 that we derived as an approximation of the IB Lagrangian for the weights (eq. 3).

Now, recall that since we have

[TABLE]

the sharpest PAC-Bayes upper-bound to the test error is obtained when $p(w)=q(w)$, in which case eq. 7 reduces (modulo a multiplicative constant) to the IB Lagrangian of the weights. That is, the IB Lagrangian for the weights can be considered as a special case of PAC-Bayes giving the sharpest bound.

Unfortunately, as we noticed in Section 4, the joint marginal $q(w)$ of the weights is not tractable. To circumvent the problem, we can instead consider that the sharpest PAC-Bayes upper-bound that can be obtained using a tractable factorized prior $p(w)$, which is obtained exactly when $p(w)=\tilde{q}(w)=\prod_{i}q(w_{i})$ is the product of the marginals, leading again to our practical loss eq. 5.

On a last note, recall that under our modeling assumptions the marginal $\tilde{q}(w)$ is assumed to be an improper log-uniform distribution. While this has the advantage of being a non-informative prior that closely matches the real marginal of the weights of the network, it also has the disadvantage that it is only defined modulo an additive constant, therefore making the bound on the test error vacuous under our model.

The PAC-Bayes bounds has also been used by Dziugaite and Roy (2017) to study the generalization property of deep neural networks and their connection with the optimization algorithm. They use a Gaussian prior and posterior, leading to a non-vacuous generalization bound.

7 Empirical validation

7.1 Transition from overfitting to underfitting

As pointed out by Zhang et al. (2017), when a standard convolutional neural network (CNN) is trained on CIFAR-10 to fit random labels, the network is able to (over)fit them perfectly. This is easily explained in our framework: It means that the network is complex enough to memorize all the labels but, as we show here, it has to pay a steep price in terms of information complexity of the weights (Figure 2) in order to do so. On the other hand, when the information in the weights is bounded using and information regularizer, overfitting is prevented in a theoretically predictable way.

In particular, in the case of completely random labels, we have $I(\mathbf{y};w|\mathbf{x},\theta)=I(\mathbf{y};w)\leq I(w;\mathcal{D})$, where the first equality holds since $\mathbf{y}$ is by construction random, and therefore independent of $\mathbf{x}$ and $\theta$. In this case, the inequality used to derive eq. 3 is an equality, and the IBL is an optimal regularizer, and, regardless of the dataset size $N$, for $\beta>1$ it should completely prevent memorization, while for $\beta<1$ overfitting is possible. To see this, notice that since the labels are random, to decrease the classification error by $\log|\mathcal{Y}|$, where $|\mathcal{Y}|$ is the number of possible classes, we need to memorize a new label. But to do so, we need to store more information in the weights of the network, therefore increasing the second term $I(w;\mathcal{D})$ by a corresponding quantity. This trade-off is always favorable when $\beta<1$, but it is not when $\beta>1$. Therefore, the theoretically the optimal solution to eq. 1 is to memorize all the labels in the first case, and not memorize anything in the latter.

As discussed, for real neural networks we cannot directly minimize eq. 1, and we need to use a computable upper bound to $I(w;\mathcal{D})$ instead (Section 4.2). Even so, the empirical behavior of the network, shown in Figure 1, closely follows this prediction, and for various sizes of the dataset clearly shows a phase transition between overfitting and underfitting near the critical value $\beta=1$. Notice instead that for real labels the situation is different: The model is still able to overfit when $\beta<1$, but importantly there is a large interval of $\beta>1$ where the model can fit the data without overfitting to it. Indeed, as soon as $\beta N\propto I(w;\mathcal{D})$ is larger than the constant $H(\theta)$, the model trained on real data fits real labels without excessive overfitting (Figure 1).

Notice that, based on this reasoning, we expect the presence of a phase transition between an overfitting and an underfitting regime at the critical value $\beta=1$ to be largely independent on the network architecture: To verify this, we train different architectures on a subset of 10000 samples from CIFAR-10 with random labels. As we can see on the left plot of Figure 2, even very different architectures show a phase transition at a similar value of $\beta$. We also notice that in the experiment ResNets has a sharp transition close to the critical $\beta$.

In the right plot of Figure 2 we measure the quantity information in the weights for different levels of corruption of the labels. To do this, we fix $\beta<1$ so that the network is able to overfit, and for various level of corruption we train until convergence, and then compute $I(w;\mathcal{D})$ for the trained model. As expected, increasing the randomness of the labels increases the quantity of information we need to fit the dataset. For completely random labels, $I(w;\mathcal{D})$ increases by $\sim 3$ nats/sample, which the same order of magnitude as the quantity required to memorize a 10-class labels ($2.30$ nats/sample), as shown in Figure 2.

7.2 Bias-variance trade-off

The Bias-Variance trade-off is sometimes informally stated as saying that low-complexity models tend to underfit the data, while excessively complex models may instead overfit, so that one should select an adequate intermediate complexity. This is apparently at odds with the common practice in Deep Learning, where increasing the depth or the number of weights of the network, and hence increasing the “complexity” of the model measured by the number of parameters, does not seem to induce overfitting. Consequently, a number of alternative measures of complexity have been proposed that capture the intuitive bias-variance trade-off curve, such as different norms of the weights (Neyshabur et al., 2015).

From the discussion above, we have seen that the quantity of information in the weights, or alternatively its computable upperbound $\tilde{I}(w;\mathcal{D})$, also provides a natural choice to measure model complexity in relation to overfitting. In particular, we have already seen that models need to store increasingly more information to fit increasingly random labels (Figure 2). In Figure 3 we show that by controlling $\tilde{I}(w;\mathcal{D})$, which can be done easily by modulating $\beta$, we recover the right trend for the bias-variance tradeoff, whereas models with too little information tend to underfit, while models memorizing too much information tend to overfit.

7.3 Nuisance invariance

Corollary 5.3 shows that by decreasing the information in the weights $I(w;\mathcal{D})$, which can be done for example using eq. 3, the learned representation will be increasingly minimal, and therefore insensitive to nuisance factors $n$, as measured by $I(z;n)$. Here, we adapt a technique from the GAN literature Sønderby et al. (2017) that allows us to explicitly measure $I(z;n)$ and validate this effect, provided we can sample from the nuisance distribution $p(n)$ and from $p(x|n)$; that is, if given a nuisance $n$ we can generate data $x$ affected by that nuisance. Recall that by definition we have

[TABLE]

To approximate the expectations via sampling we need a way to approximate the likelihood ratio $\log p(\mathbf{z}|\mathbf{n})/p(\mathbf{z})$. This can be done as follows: Let $D(z;n)$ be a binary discriminator that given the representation $z$ and the nuisance $n$ tries to decide whether $z$ is sampled from the posterior distribution $p(z|n)$ or from the prior $p(z)$. Since by hypothesis we can generate samples from both distributions, we can generate data to train this discriminator. Intuitively, if the discriminator is not able to classify, it means that $z$ is insensitive to changes of $n$. Precisely, since the optimal discriminator is

[TABLE]

if we assume that $D$ is close to the optimal discriminator $D^{*}$, we have

[TABLE]

therefore we can use $D$ to estimate the log-likelihood ratio, and so also the mutual information $I(z;n)$. Notice however that this comes with no guarantees on the quality of the approximation.

To test this algorithm, we add random occlusion nuisances to MNIST digits (Figure 4). In this case, the nuisance $n$ is the occlusion pattern, while the observed data $x$ is the occluded digit. For various values of $\beta$, we train a classifier on this data in order to learn a representation $z$, and, for each representation obtained this way, we train a discriminator as described above and we compute the resulting approximation of $I(z;n)$. The results in Figure 4 show that decreasing the information in the weights makes the representation increasingly more insensitive to $n$.

8 Discussion and conclusion

In this work, we have presented bounds, some of which are tight, that connect the amount of information in the weights, the amount of information in the activations, the invariance property of the network, and the geometry of the residual loss. These results leverage the structure of deep networks, in particular the multiplicative action of the weights, and the Markov property of the layers. This leads to the surprising result that reducing information stored in the weights about the past (dataset) results in desirable properties of the learned internal representation of the test datum (future).

Our notion of representation is intrinsically stochastic. This simplifies the computation as well as the derivation of information-based relations. However, note that even if we start with a deterministic representation $w$, Proposition 4.3 gives us a way of converting it to a stochastic representation whose quality depends on the flatness of the minimum. Our theory uses, but does not depend on, the Information Bottleneck Principle, which dates back to over two decades ago, and can be re-derived in a different frameworks, for instance PAC-Bayes, which yield the same results and additional bounds on the test error.

This work focuses on the inference and learning of optimal representations, that seek to get the most out of the data we have for a specific task. This does not guarantee a good outcome since, due to the Data Processing Inequality, the representation can be easier to use but ultimately no more informative than the data themselves. An orthogonal but equally interesting issue is how to get the most informative data possible, which is the subject of active learning, experiment design, and perceptual exploration. Our work does not address transfer learning, where a representation trained to be optimal for a task is instead used for a different task, which will be subject of future investigations.

Acknowledgments

Supported by ONR N00014-17-1-2072, ARO W911NF-17-1-0304, AFOSR FA9550-15-1-0229 and FA8650-11-1-7156. We wish to thank our reviewers and David McAllester, Kevin Murphy, Alessandro Chiuso for the many insightful comments and suggestions.

A Details of the experiments

A.1 Random labels

We use a similar experimental setup as Zhang et al. (2017). In particular, we train a small version of AlexNet on a 28$\times$28 central crop of CIFAR-10 with completely random labels. The dataset is normalized using the global channel-wise mean and variance, but no additional data augmentation is performed. The exact structure of the network is in Table 1. As common in practice we use batch normalization before all the ReLU nonlinearities, except for the first layer. Optimization of the IB Lagrangian loss function is performed similarly to Kingma et al. (2015) and Molchanov et al. (2017). We found that constraining the variance $\alpha_{i}$ of the weights to be the same for all weights in the same filter helps stabilizing the training process. We train with learning rates $\eta\in\{0.02,0.005\}$ and select the best performing network of the two. Generally, we found that a higher learning rate is needed to overfit when the number of training samples $N$ is small, while a lower learning rate is needed for larger $N$. We train with SGD with momentum $0.9$ for $360$ epochs reducing the learning rate by a factor of $10$ every $140$ epochs. We use a large batch-size of $500$ to minimize the noise coming from SGD. No weight decay or other regularization methods are used.

The final plot is obtained by triangulating the convex envelope of the data points, and by interpolating their value on the resulting simplexes. Outside of the convex envelope (where the accuracy is mostly constant), the value was obtained by inpainting.

To measure the information content of the weights as the percentage of corrupted labels varies, we fix $\beta=0.1$, $N=30000$ and $\eta=0.005$ and train the network on different corruption levels with the same settings as before.

To test the phase transition on multiple architectures, we train the Small AlexNet, the AllCNN network and a ResNet (see Table 1). For all architectures, we train with $N=10000$ random labels, $\eta=0.05$ and different values of $\beta$ log-uniformly spaced in $[10^{-2},10^{2}]$.

A.2 Bias-variance trade-off

For this experiment we train the AllCNN architecture (Table 1) on the CIFAR-10 dataset with ZCA whitening Krizhevsky and Hinton (2009) and without any additional data augmentation. First, we train a standard network and change the number of filters (we multiplying the number of filters of all layers by the same constant) and train with $\eta=0.05$, batch size $128$, weight decay $0.001$. Then, we use the standard number of layers and train instead with the IBL loss function with different values of $\beta$.

A.3 Nuisance invariance

The cluttered MNIST dataset is generated by adding ten $4\times 4$ squares uniformly at random on the digits of the MNIST dataset (LeCun et al., 1998). For each level of $\beta$, we train the classifier in Table 1 on this dataset. The weights of all layers, excluding the first and last one, are threated as a random variable with multiplicative Gaussian noise (Section B) and optimized using the local reparameterization trick of Kingma et al. (2015). We use the last convolutional layer before classification as representation $z$.

The discriminator network used to estimate the log-likelihood ratio is constructed as follows: the inputs are the nuisance pattern $n$, which is a $28{\times}28{\times}1$ image containing 10 random occluding squares, and the 7$\times$7$\times$192 representation $z$ obtained from the classifier. First we preprocess $n$ using the following network: conv 48 $\to$ conv 48 $\to$ conv 96 s2 $\to$ conv 96 $\to$ conv 96 $\to$ conv 96 s2, where each conv block is a 3$\times$3 convolution followed by batch normalization and ReLU. Then, we concatenate the $7{\times}7{\times}96$ result with $z$ along the feature maps, and the final discriminator output is obtained by applying the following network: conv 192 $\to$ conv 192 $\to$ conv 1$\times$1$\times$192 $\to$ conv 1$\times$1$\times$1 $\to$ AvgPooling 7$\times$7 $\to$ sigmoid.

A.4 Visualizing the representation

Even when we cannot generate data affected by nuisances like in the previous section, we can still visualize the information content of $z$ to learn what nuisances are discarded in the representation. To this end, given a representation $z$, we want to learn to sample from a distribution $q(\hat{x}|z)$ of images that are maximally likely to have $z$ as their representation. Formally, this means that we want a distribution $q(\hat{x}|z)$ that maximizes the amortized maximum a posteriori estimate of $z$:

[TABLE]

Unfortunately, the term $p(\hat{x})$ in the expression is difficult to estimate. However, Sønderby et al. (2017) notice that the modified gain function

[TABLE]

differs from the amortizes MAP only by a term $H(p(\hat{\mathbf{x}}))$, which has the positive effect of improving the exploration of the reconstruction, and contains the term $\operatorname{KL}(\,q(\hat{x})\,\|\,p(\hat{x})\,)$, which can be estimated easily using the discriminator network of a GAN Sønderby et al. (2017). To maximize this gain, we can simply train a GAN with an additional reconstruction loss $-\log p(z|\hat{x})$.

To test this algorithm, we train a representation $z$ to classify the 40 binary attributes in the CelebA face dataset (Yang et al., 2015), and then use the above loss function to train a GAN network to reconstruct an input image $\hat{x}$ from the representation $z$. The results in Figure 5 show that, as expected, increasing the value of $\beta$, and therefore reducing $I(w;\mathcal{D})$, generates samples that have increasingly more random backgrounds and hair style (nuisances), while retaining facial features. In other words, the representation $z$ is increasingly insensitive to nuisances affecting the data, while information pertaining the task is retained in the reconstruction $\hat{x}$.

More precisely, we first train a classifier on the images from the CelebA datasets resized to 32$\times$32, where the task is to recover the 40 binary attributes associated to each image. The classifier network is the same as the one in Table 1 with the following modifications: we use Exponential Linear Units (Clevert et al., 2015) for the activations, instead of ReLU, since invertible activations generally perform better when training a GAN, and we divide by two the number of output filters in all layers to reduce the training time. A sigmoid nonlinearity is applied to the final 40-way output of the network.

To generate the image $\hat{x}$ given the 8$\times$8$\times$96 representation $z$ computed by the classifier, we use a similar structure to DCGAN (Radford et al., 2016), namely $z$ $\to$ conv 256 $\to$ ConvT 256s2 $\to$ ConvT 128s2 $\to$ conv 3 $\to$ tanh, where ConvT 256s2 denotes a transpose convolution with 256 feature maps and stride 2. All convolutions have a batch normalization layer before the activations. Finally, the discriminator network is given by $\hat{x}$ $\to$ conv 64s2 $\to$ conv 128s2 $\to$ ConvT 256s2 $\to$ conv 1 $\to$ sigmoid. Here, all convolutions use batch normalization followed by Leacky ReLU activations.

In this experiment, we use Gaussian multiplicative noise which is slightly more stable during training (Section B). To stabilize the training of the GAN, we found useful to (1) scale down the “reconstruction error” term in the loss function and (2) slowly increase the weight of the reconstruction error up to the desired value during training.

B Gaussian multiplicative noise

In developing the theory, we chose to use log-normal multiplicative noise for the weights: The main benefit is that with this choice the information in the weights $I(w;\mathcal{D})$ can be expressed in closed form, up to an arbitrary constant $C$ which does not matter during the optimization process (but see also Neklyudov et al. (2017) for a principled approach to this problem that uses a proper log-uniform prior). Another possibility, suggested by Kingma et al. (2015) is to use Gaussian multiplicative noise with mean 1. Unfortunately, there is no analytical expression for $I(w;\mathcal{D})$ when using Gaussian noise, but $I(w;\mathcal{D})$ can still be approximated numerically with high precision (Molchanov et al., 2017), and it makes the training process slightly more stable. The theory holds with minimal changes also in this case, and we use this choice in some experiments.

C Proofs of theorems

Lemma C.1** (Task-nuisance decomposition)**

Given a joint distribution $p(x,y)$, where $y$ a discrete random variable, we can always find a random variable $n$ independent of $y$ such that $x=f(y,n)$, for some deterministic function $f$.

**Proof ** Fix $n\sim\mathrm{Uniform}(0,1)$ to be the uniform distribution on $[0,1]$. We claim that, for a fixed value of $y$, there is a function $\Phi_{y}(n)$ such that $x|y={\Phi_{y}}_{*}(n)$, where $({\cdot})_{*}$ denotes the push-forward map of measures. Given the claim, let $\Phi(y,n)=(y,\Phi_{y}(n))$. Since $y$ is a discrete random variable, $\Phi(y,n)$ is easily seen to be a measurable function and by construction $(x,y)\sim\Phi_{*}(y,n)$. To see the claim, notice that, since there exists a measurable isomorphism between $\mathbb{R}^{n}$ and $\mathbb{R}$ (Theorem 3.1.1 of Berberian (1988)), we can assume without loss of generality that $x\in\mathbb{R}$. In this case, by definition, we can take $\Phi_{y}(n)=F_{y}^{-1}(n)$ where $F_{y}(t)=\mathbb{P}[x<t\,|\,y]$ is the cumulative distribution function of $p(x|y)$.

Proposition C.2** (Invariance and minimality)**

Let $n$ be a nuisance for the task $y$ and let $z$ be a sufficient representation of the input $x$. Suppose that $z$ depends on $n$ only through $x$ (i.e., $n\to x\to z$). Then,

[TABLE]

Moreover, there exists a nuisance $n$ such that equality holds up to a (generally small) residual $\epsilon$

[TABLE]

where $\epsilon:=I(z;y|n)-I(x;y)$. In particular $0\leq\epsilon\leq H(y|x)$, and $\epsilon=0$ whenever $y$ is a deterministic function of $x$. Under these conditions, a sufficient statistic $z$ is invariant (maximally insensitive) to nuisances if and only if it is minimal.

**Proof ** By hypothesis, we have the Markov chain $(y,n)\to x\to z$; therefore, by the DPI, we have $I(z;y,n)\leq I(z;x)$. The first term can be rewritten using the chain rule as $I(z;y,n)=I(z;n)+I(z;y|n)$, giving us

[TABLE]

Now, since $y$ and $n$ are independent, $I(z;y|n)\geq I(z;y)$. In fact,

[TABLE]

Substituting in the inequality above, and using the fact that $z$ is sufficient, we finally obtain

[TABLE]

Moreover, let $n$ be as in Lemma 2.1. Then, since $x$ is a deterministic function of $y$ and $n$, we have

[TABLE]

and therefore

[TABLE]

with $\epsilon$ defined as above. Using the sufficiency of $z$, the previous inequality for $I(z;y|n)$, the DPI, we get the chain of inequalities

[TABLE]

from which we obtain the desired bounds for $\epsilon$.

While the proof of the following theorem is quite simple, some clarifications on the notation are in order: We assume, following a Bayesian perspective, that the data is generated by some generative model $p(\mathbf{x},\mathbf{y}|\theta)$, where the parameters $\theta$ of the model are sampled from some (unknown) prior $p(\theta)$. Given the parameters $\theta$, the training dataset $\mathcal{D}=(\mathbf{x},\mathbf{y})\sim p(x,y|\theta)$ is composed of i.i.d. samples from the unknown distribution $p(x,y|\theta)$. The output of the training algorithm on the dataset $\mathcal{D}$ is a (generally simple, e.g., normal or log-normal) distribution $q(w|\mathbf{x},\mathbf{y})$ over the weights. Putting everything together, we have a well-defined joint distribution $p(\mathbf{x},\mathbf{y},\theta,w)=p(\theta)p(\mathbf{x},\mathbf{y}|\theta)q(w|\mathbf{x},\mathbf{y})$.

Given the weights $w$, the network then defines an inference distribution $q(\mathbf{y}|\mathbf{x},w)$, which we know and can compute explicitly. Another distribution, which instead we do not know, is $p(\mathbf{y}|\mathbf{x},w)$, which is obtained from $p(\mathbf{x},\mathbf{y},\theta,w)$ and express the optimal inference we could perform on the labels $\mathbf{y}$ using the information contained in the weights. In a well trained network, we want the distribution approximated by the network to match the optimal distribution $q(\mathbf{y}|\mathbf{x},w)=p(\mathbf{y}|\mathbf{x},w)$.

Finally, recall that the conditional entropy is defined as

[TABLE]

where $z$ can be one random variable or a tuple of random variables. When not specified, it is assumed that the cross-entropy is computed with respect to unknown underlying data distribution $p(\mathbf{x},\mathbf{y},w,\theta)$. Similarly, the conditional cross-entropy is defined as

[TABLE]

Proposition C.3** (Information Decomposition)**

Let $\mathcal{D}=(\mathbf{x},\mathbf{y})$ denote the training dataset, then for any training procedure, we have

[TABLE]

**Proof ** Recall that cross-entropy can be written as

[TABLE]

so we only have to prove that

[TABLE]

which is easily done using the following identities:

[TABLE]

Proposition C.4** (Information in the weights)**

Under the previous modeling assumptions, the upper-bound to the information that the weights contain about the dataset is

[TABLE]

where the constant $C$ is arbitrary due to the improper prior.

**Proof ** Recall that we defined the upperbound $\tilde{I}(\mathbf{w};\mathcal{D})$ as

[TABLE]

where $\tilde{q}(w)$ is a factorized log-uniform prior. Since the KL divergence is reparametrization invariant, we have:

[TABLE]

where we have used the formula for the entropy of a Gaussian and the fact that the KL divergence of a distribution from the uniform prior is the entropy of the distribution modulo an arbitrary constant.

Proposition C.5** (Flat minima have low information)**

Let $\hat{w}$ be a local minimum of the cross-entropy loss $H_{p,q}(\mathbf{y}|\mathbf{x},w)$, and let $\mathcal{H}$ be the Hessian at that point. Then, for the optimal choice of the posterior $w|\mathcal{D}=\epsilon\odot\hat{w}$ centered at $\hat{w}$ that optimizes the IB Lagrangian, we have

[TABLE]

where $K=\dim(w)$ and $\|\cdot\|_{*}$ denotes the nuclear norm.

**Proof ** First, we switch to a logarithmic parametrization of the weights, and let $h:=\log|w|$ (we can ignore the sign of the weights since it is locally constant). In this parametrization, we can approximate the IB Lagrangian to second order as

[TABLE]

where $H_{0}=H(\mathbf{y}|\mathbf{x},\hat{w})$. Now, notice that since $q(w|\mathcal{D})$ is a log-normal distribution, we have $q(h|\mathcal{D})\sim\mathcal{N}(h_{0},\alpha)$.666Note that for simplicity we have ignored the offset $\alpha/2$ in the mean of the log-normal distribution. Therefore, can compute the expectation exactly as

[TABLE]

Optimizing w.r.t. $\alpha_{i}$ we get

[TABLE]

and plugging it back in the expression for $\tilde{I}(w;\mathcal{D})$ that we obtained in the previous proposition, we have

[TABLE]

Finally, by Jensen’s inequality, we have

[TABLE]

as we wanted.

Proposition C.6

Let $z=Wx$, and assume as before $W=\epsilon\odot\hat{W}$, with $\epsilon_{i,j}\sim\log\mathcal{N}(-\alpha_{i}/2,\alpha_{i})$. Further assume that the marginals of $q(z)$ and $q(z|x)$ are both approximately Gaussian (which is reasonable for large $\dim(x)$ by the Central Limit Theorem). Then,

[TABLE]

where $W_{i}$ denotes the $i$-th row of the matrix $W$, and $\tilde{\alpha_{i}}$ is the noise variance $\tilde{\alpha}_{i}=\exp(\alpha_{i})-1$. In particular, $I(z;x)+\operatorname{TC}(z)$ is a monotone decreasing function of the weight variances $\alpha_{i}$.

**Proof ** First, we consider the case in which $\dim(z)=1$, and so $w:=W$ is a single row vector. By hypothesis, $q(z)$ is approximately Gaussian, with mean and variance

[TABLE]

A similar computation gives us mean and variance of $q(z|x)$:

[TABLE]

Since we are assuming $\dim(z)=1$, we trivially have $\operatorname{TC}(z)=0$, so we are only left with $I(z;x)$ which is given by

[TABLE]

Now, for the general case of $\dim(z)\geq 1$, notice that

[TABLE]

where $q(z_{i})$ is the marginal of the $k$-th component of $z$. We can then use the previous result for each component separately, and sum everything to get the desired identity.

Proposition C.7** (Uniform bound for one layer)**

Let $z=Wx$, where $W=\epsilon\odot\hat{W}$, where $\epsilon_{i,j}\sim\log\mathcal{N}(-\alpha/2,\alpha)$; assume that the components of $x$ are uncorrelated, and that their kurtosis is uniformly bounded. Then, there is a strictly increasing function $g(\alpha)$ s.t. we have the uniform bound

[TABLE]

where $c=O(1/\dim(x))\leq 1$, $g(\alpha)=\log\,(1-e^{-\alpha})/2$ and $\alpha$ is related to $I(w;\mathcal{D})$ by $\alpha=\exp\left\{-I(W;\mathcal{D})/\dim(W)\right\}$. In particular, $I(x;z)+TC(z)$ is tightly bounded by $I(W;\mathcal{D})$ and increases strictly with it.

**Proof ** To simplify the notation we do the case $\dim z=1$, the general case being identical. Let $w:=W$ be the only row of $W$. First notice that, since $x$ is uncorrelated, we have

[TABLE]

Therefore,

[TABLE]

To conclude, we want to approximate the expectation of the logarithm using a Taylor expansion, but we first need to check that the variance of the term inside the logarithm is low, which is where we need the bound on the kurtosis. In fact, since the kurtosis is bounded, there is some constant $C$ such that for all $i$

[TABLE]

Now,

[TABLE]

Therefore, we can conclude

[TABLE]

Corollary C.8** (Multi-layer case)**

Let $W^{k}$ for $k=1,...,L$ be weight matrices, with $W^{k}=\epsilon^{k}\odot\hat{W}^{k}$ and $\epsilon^{k}_{i,j}=\log\mathcal{N}(-\alpha^{k}/2,\alpha^{k})$, and let $z_{i+1}=\phi(W^{k}z_{k})$, where $z_{0}=x$ and $\phi$ is any nonlinearity. Then,

[TABLE]

where $\alpha^{k}=\exp\left\{-I(W^{k};\mathcal{D})/\dim(W^{k})\right\}$.

**Proof ** Since we have the Markov chain $x\to z_{1}\to\ldots\to z_{L}$, by the Data Processing Inequality we have $I(z_{L};x)\leq\min\left\{I(z_{L};z_{L-1}),I(z_{L-1};x)\right\}$. Iterating this inequality, we have

[TABLE]

Now, notice that $I(z_{k+1};z_{k})\leq I(\phi(W^{k}z_{k});z_{k})\leq I(W^{k}z_{k};z_{k})$, since applying a deterministic function can only decrease the information. But $I(W^{k}z_{k};z_{k})$ is exactly the quantity we bounded in Corollary 5.2, leading us to the desired inequality.

D Q&A

It is well-known that overfitting relates to the “effective number of degrees of freedom” that can be measured in a number of ways (Friedman et al., 2001, Chapter 7). Why should we use the information in the weights?

The information in the weights is indeed one particular choice of measure of complexity. One nice aspect is that it plays a central role in many different frameworks (minimum description length, variational inference, PAC-Bayes), and correlates well with the performance of a real network.

How do you compute the nuclear norm of the Hessian? It sounds expensive!

We do not need to. What we show is that, if the optimization algorithm happens to find flat minima (nuclear norm being a proxy), then it automatically limits the information in the weights - which promotes good generalization. However, if one wanted to approximate the trace of the Hessian, it could be done in linear time (Bai et al., 1996, Prop 4.1).

A nuisance $n$ should convey no information on $y$ given $x$, so why imposing $I(y;n)=0$ rather than $I(y;n|x)=0$?

While $I(y;n|x)=0$ may seem intuitively the right condition, it is actually too weak. Suppose for example that $y$ is a deterministic function of $x$ (i.e., the labels are perfectly determined by the data, as often is the case). Then, we would have $I(y;n|x)=0$ for any $n$, which would imply that everything is a nuisance for the task, which of course is not intuitively the case.

Why is the dataset a random variable, if we only have one realization of it?

The theory is almost identical in both the case of a fixed dataset and a randomly sampled dataset. We consider the case of a randomly sampled dataset since it is simpler and at the same time more general. Some expressions simplify slightly and are easier to interpret, but in the end a fixed training set is given either way.

The use of $KL(q(w|D)||\tilde{q}(w))$ where $q(w|D)$ is a “posterior” defined by a learning algorithm that returns $w$ and $\tilde{q}(w)$ is a log-uniform prior is the basic PAC-Bayes bound, that, however, gives an vacuous generalization bound due to the improper prior. A generalization bound could be stated in terms of the length of a finite interval approximation of log-uniform prior.

Indeed, this is the case. However, in the limit of the interval length going to infinity, the KL divergence would still be infinite, and the optimization would be slightly more complex. As simpler option to have a generalization-bound would be to use Gaussian prior and posterior Dziugaite and Roy (2017). However, computing a good PAC-Bayes upper-bound is outside the scope of the paper, and the use of a non-informative, scale invariant prior matches the empirical behavior of networks and simplifies the theoretical analysis.

Of course a minimal representation should be invariant to nuisance variation since that’s what it means to be minimal for the task.

While the result may be intuitive to some, compression and invariance are not the same thing, and we are unaware of an existing proof of the claim other than for special tasks like clustering, and for small perturbations.

How do you compute the information in the weights, since you only have a sample (one set of weights that the network converged to for a given dataset)? And how do you optimize it? That looks hard!

We do not need to compute the information in the weights, since we can control it. Even if we do not do so explicitly, optimizing cross-entropy with SGD yields the right solution (the solution that minimizes the IBL for the weights). We only have one sample of the weights if we anneal the learning rate to zero, but otherwise SGD produces a posterior distribution of the weights, even if we do not impose additional stochasticity. In many cases we do, for instance using Information Dropout (Achille and Soatto (2018)) or its simpler version, Dropout. A special case of the theory can be re-derived assuming an instantiated dataset and a set of weights, with the same results.

The use of $I(w;\mathcal{D})$ does not seem to upper bound the PAC-Bayes bound. The direction of inequality is the opposite. The IB Lagrangian for the weights is a particular case that gives the sharpest PAC-Bayes bound. Choosing $q(w)$ requires some attention: Once we fix a training procedure, we can (in theory) compute/approximate the marginal distribution $q(w)$ over the stochasticity of the data (if any) as well as the stochasticity of the training procedure. This marginal can then be used in the PAC-Bayes bound: This gives the sharpest bound (McAllester, 2013) and is also equivalent to an IBL since the term $\mathbb{E}_{\mathcal{D}}\operatorname{KL}(\,q(w|\mathcal{D})\,\|\,q(w)\,)$ measures the mutual information between the training procedure and the weights. Since in general it is not possible to explicitly compute this bound, in practice we use less tight a bound based on a factorized approximation of the marginal. This opportunistic choice later turns out to play a role in disentanglement.