The Conditional Entropy Bottleneck

1. Introduction

Despite excellent progress in classical generalization (e.g., accuracy on a held-out set), the field of Machine Learning continues to struggle with the following issues:

• Vulnerability to adversarial examples. Most machine-learned systems are vulnerable to adversarial examples. Many defenses have been proposed, but few have demonstrated robustness against a powerful, general-purpose adversary. Many proposed defenses are ad-hoc and fail in the presence of a concerted attacker [1,2].
• Poor out-of-distribution detection. Most models do a poor job of signaling that they have received data that is substantially different from the data they were trained on. Even generative models can report that an entirely different dataset has higher likelihood than the dataset they were trained on [3]. Ideally, a trained model would give less confident predictions for data that was far from the training distribution (as well as for adversarial examples). Barring that, there would be a clear, principled statistic that could be extracted from the model to tell whether the model should have made a low-confidence prediction. Many different approaches to providing such a statistic have been proposed [4,5,6,7,8,9], but most seem to do poorly on what humans intuitively view as obviously different data.
• Miscalibrated predictions. Related to the issues above, classifiers tend to be overconfident in their predictions [4]. Miscalibration reduces confidence that a model’s output is fair and trustworthy.
• Overfitting to the training data. Zhang et al. [10] demonstrated that classifiers can memorize fixed random labelings of training data, which means that it is possible to learn a classifier with perfect inability to generalize. This critical observation makes it clear that a fundamental test of generalization is that the model should fail to learn when given what we call information-free datasets.

We consider these to be problems of robust generalization, which we define and discuss in Section 2.1. In this work, we hypothesize that these problems of robust generalization all have a common cause: models retain too much information about the training data. We formalize this by introducing the Minimum Necessary Information (MNI) criterion for evaluating a learned representation (Section 2.2). We then introduce an objective function that directly optimizes the MNI, the Conditional Entropy Bottleneck (CEB) (Section 2.3) and compare it with the closely-related Information Bottleneck (IB) objective [11] in Section 2.5. In Section 2.6, we describe practical ways to optimize CEB in a variety of settings.

Finally, we give empirical evidence for the following claims:

• Better classification accuracy. MNI models can achieve superior accuracy on classification tasks than models that capture either more or less information than the minimum necessary information (Section 3.1.1 and Section 3.1.6).
• Improved robustness to adversarial examples. Retaining excessive information about the training data results in vulnerability to a variety of whitebox and transfer adversarial examples. MNI models are substantially more robust to these attacks (Section 3.1.2 and Section 3.1.6).
• Strong out-of-distribution detection. The CEB objective provides a useful metric for out-of-distribution (OoD) detection, and CEB models can detect OoD examples as well or better than non-MNI models (Section 3.1.3).
• Better calibration. MNI models are better calibrated than non-MNI models (Section 3.1.4).
• No memorization of information-free datasets. MNI models fail to learn in information-free settings, which we view as a minimum bar for demonstrating robust generalization (Section 3.1.5).

2. Materials and Methods

2.1. Robust Generalization

In classical generalization, we are interested in a model’s performance on held-out data on some task of interest, such as classification accuracy. In robust generalization, we want: (RG1) to maintain the model’s performance in the classical generalization setting; (RG2) to ensure the model’s performance in the presence of an adversary (unknown at training time); and (RG3) to detect adversarial and non-adversarial data that strongly differ from the training distribution.

Adversarial training approaches considered in the literature so far [12,13,14] violate (RG1), as they typically result in substantial decreases in accuracy. Similarly, provable robustness approaches (e.g., Cohen et al. [15], Wong et al. [16]) provide guarantees for a particular adversary known at training time, also at a cost to test accuracy. To our knowledge, neither approaches provide any mechanism to satisfy (RG3). On the other hand, approaches for detecting adversarial and non-adversarial out-of-distribution (OoD) examples [4,5,6,7,8,9] are either known to be vulnerable to adversarial attack [1,2], or do not demonstrate that the approach provides robustness against unknown adveraries, both of which violate (RG2).

Training on information-free datasets [10] provides an additional way to check if a learning system is compatible with (RG1), as memorization of such datasets necessarily results in maximally poor performance on any test set. Model calibration is not obviously a necessary condition for robust generalization, but if a model is well-calibrated on a held-out set, its confidence may provide some signal for distinguishing OoD examples, so we mention it as a relevant metric for (RG3).

To our knowledge, the only works to date that have demonstrated progress on robust generalization for modern machine learning datasets are the Variational Information Bottleneck [17,18] (VIB), and Information Dropout [19]. Alemi et al. [17] presented preliminary results that VIB improves adversarial robustness on image classification tasks while maintaining high classification accuracy ((RG1) and (RG2)). Alemi et al. [18] showed that VIB models provide a useful signal, the Rate, R, for detecting OoD examples ((RG3)). Achille and Soatto [19] also showed preliminary results on adversarial robustness and demonstrated failure to train on information-free datasets.

In this work, we do not claim to “solve” robust generalization, but we do show notable improvement on all three conditions simply by changing the training objective. This evidence supports our core hypothesis that problems of robust generalization are caused in part by retaining too much information about the training data.

2.2. The Minimum Necessary Information

We define the Minimum Necessary Information (MNI) criterion for a learned representation in three parts:

• Information. We would like a representation Z that captures useful information about a dataset $(X,Y)$. Entropy is the unique measure of information [20], so the criterion prefers information-theoretic approaches. (We assume familiarity with the mutual information and its relationships to entropy and conditional entropy: $I(X;Y)=H\left(X\right)-H\left(X\right|Y)=H(Y)-H(Y\left|X\right)$ [21] (p. 20).)
• Necessity. The semantic value of information is given by a task, which is specified by the set of variables in the dataset. Here we will assume that the task of interest is to predict Y given X, as in any supervised learning dataset. The information we capture in our representation Z must be necessary to solve this task. As a variable X may have redundant information that is useful for predicting Y, a representation Z that captures the necessary information may not be minimal or unique (the MNI criterion does not require uniqueness of Z).
• Minimality. Given all representations $\mathcal{Z}$ that can solve the task, we require one that retains the smallest amount of information about the task: ${inf}_{Z\in \mathcal{Z}}I(Z;X,Y)$.

Necessity can be defined as $I(X;Y)\le I(Y;Z)$. Any less information than that would prevent Z from solving the task of predicting Y from X. Minimality can be defined as $I(X;Y)\ge I(X;Z)$. Any more information than that would result in Z capturing information from X that is either redundant or irrelevant for predicting Y. Since the information captured by Z is constrained from above and below, we have the following necessary and sufficient conditions for perfectly achieving the Minimum Necessary Information, which we call the MNI Point:

$$\begin{array}{c}\hfill I(X;Y)=I(X;Z)=I(Y;Z)\end{array}$$

(1)

The MNI point defines a unique point in the information plane. The geometry of the information plane can be seen in Figure 1. The MNI criterion does not make any Markov assumptions on the models or algorithms that learn the representations. However, the algorithms we discuss here all do rely on the standard Markov chain $Z\leftarrow X\leftrightarrow Y$. See Fischer [22] for an example of an objective that doesn’t rely on a Markov chain during training.

A closely related concept to Necessity is called sufficiency by Achille and Soatto [19] and other authors. We avoid the term due to potential confusion with minimum sufficient statistics, which maintain the mutual information between a model and the data it generates [21] (p. 35). The primary difference between necessity and sufficiency is the reliance on the Markov constraint to define sufficiency. Ref. [19] also does not identify the MNI point as an idealized target, instead defining the optimization problem: minimize $I(X;Z)\mathrm{s}.\mathrm{t}.H\left(Y\right|Z)=H(Y\left|X\right)$.

In general it may not be possible to satisfy Equation (1). As discussed in Anantharam et al. [23,24,25], for any given dataset $(X,Y)$, there is some maximum value for any possible representation Z:

$$1\ge {\eta }_{KL}=\underset{Z\leftarrow X\to Y}{sup}\frac{I(Y;Z)}{I(X;Z)}$$

(2)

with equality only when $X\to Y$ is a deterministic map. Training datasets are often deterministic in one direction or the other. e.g., common image datasets map each distinct image to a single label. Thus, in practice, we can often get very close to the MNI on the training set given a sufficiently powerful model.

MNI and Robust Generalization

To satisfy (RG1) (classical generalization), a model must have $I(X;Z)\ge I(X;Y)=I(Y;Z)$ on the test dataset. Shamir et al. [26] show that $|I(X;Z)-\widehat{I}(X;Z)|\approx O\left(\frac{2^{\widehat{I}(X;Z)}}{\sqrt{N}}\right)$, where $\widehat{I}(·)$ indicates the training set information and N is the size of the training set. More recently, Bassily et al. [27] gave a similar result in a PAC setting. Both results indicate that models that are compressed on the training data should do better at generalizing to similar test data.

Less clear is how an MNI model might improve on (RG2) (adversarial robustness). In this work, we treat it as a hypothesis that we investigate empirically rather than theoretically. The intuition behind the hypothesis can be described in terms of the idea of robust and non-robust features from Ilyas et al. [28]: non-robust features in X should be compressed as much as possible when we learn Z, whereas robust features should be retained as much as is necessary. If Equation (1) is satisfied, Z must have “scaled” the importance of the the features in X according to their importance for predicting Y. Consequently, an attacker that tries to take advantage of a non-robust feature will have to change it much more in order to confuse the model, possibly exceeding the constraints of the attack before it succeeds.

For (RG3) (detection), the MNI criterion does not directly apply, as that will be a property of specific modeling choices. However, if the model provides an accurate way to measure $I(X=x;Z=z)$ for a particular pair $(x,z)$, Alemi et al. [18] suggests that can be a valuable signal for OoD detection.

2.3. The Conditional Entropy Bottleneck

We would like to learn a representation Z of X that will be useful for predicting Y. We can represent this problem setting with the Markov chain $Z\leftarrow X\leftrightarrow Y$. We would like Z to satisfy Equation (1). Given the conditional independence Z⫫$Y|X$ in our Markov chain, $I(Y;Z)\le I(X;Y)$, by the data processing inequality. Thus, maximizing $I(Y;Z)$ is consistent with the MNI criterion.

However, $I(X;Z)$ does not clearly have a constraint that targets $I(X;Y)$, as $0\le I(X;Z)\le H\left(X\right)$. Instead, we can notice the following identities at the MNI point:

$$I(X;Y|Z)=I(X;Z\left|Y\right)=I(Y;Z|X)=0$$

(3)

The conditional mutual information is always non-negative, so learning a compressed representation Z of X is equivalent to minimizing $I(X;Z|Y)$. Using our Markov chain and the chain rule of mutual information [21]:

$$I(X;Z|Y)=I(X,Y;Z)-I(Y;Z)=I(X;Z)-I(Y;Z)$$

(4)

This leads us to the general Conditional Entropy Bottleneck:

$$\begin{array}{cc}\hfill CEB& \equiv \underset{Z}{min}I(X;Z|Y)-\gamma I(Y;Z)\hfill \\ \hfill & =\underset{Z}{min}H\left(Z\right)-H\left(Z\right|X)-H\left(Z\right)+H\left(Z\right|Y)\hfill \end{array}$$

(5)

$$\begin{array}{cc}\hfill & -\gamma \left(H\right(Y)+H(Y\left|Z\right))\hfill \end{array}$$

(6)

$$\begin{array}{cc}\hfill & \iff \underset{Z}{min}-H\left(Z\right|X)+H\left(Z\right|Y)+\gamma H\left(Y\right|Z)\hfill \end{array}$$

(7)

In line 7, we can optionally drop $H\left(Y\right)$ because it is constant with respect to Z. Here, any $\gamma >0$ is valid, but for deterministic datasets (Section 2.2), $\gamma =1$ will achieve the MNI for a sufficiently powerful model. Further, we should expect $\gamma =1$ to yield consistent models and other values of $\gamma$ not to: since $I(Y;Z)$ shows up in two forms in the objective, weighing them differently forces the optimization procedure to count bits of $I(Y;Z)$ in two different ways, potentially leading to a situation where $H\left(Z\right)-H\left(Z\right|Y)\ne H(Y)-H(Y\left|Z\right)$ at convergence. Given knowledge of those four entropies, we can define a consistency metric for Z:

$$\begin{array}{c}\hfill C_{I(Y;Z)}\left(Z\right)\equiv |H\left(Z\right)-H\left(Z\right|Y)-H\left(Y\right)+H\left(Y\right|Z)|\end{array}$$

(8)

2.4. Variational Bound on CEB

We will variationally upper bound the first term of Equation (5) and lower bound the second term using three distributions: $e\left(z\right|x)$, the encoder which defines the joint distribution we will use for sampling, $p(x,y,z)\equiv p(x,y)e\left(z\right|x)$; $b\left(z\right|y)$, the backward encoder, an approximation of $p\left(z\right|y)$; and $c\left(y\right|z)$, the classifier, an approximation of $p\left(y\right|z)$ (the name is arbitrary, as Y may not be labels). All of $e(·)$, $b(·)$, and $c(·)$ may have learned parameters, just like the encoder and decoder of a VAE [29], or the encoder, classifier, and marginal in VIB.

In the following, we write expectations $〈loge\left(z\right|x)〉$. They are always with respect to the joint distribution; here, that is $p(x,y,z)\equiv p(x,y)e\left(z\right|x)$. The first term of Equation (5):

$$\begin{array}{cc}\hfill I(X;Z|Y)=& -H\left(Z\right|X)+H\left(Z\right|Y)=〈loge\left(z\right|x)〉-〈logp\left(z\right|y)〉\hfill \end{array}$$

(9)

$$\begin{array}{cc}\hfill =& 〈loge\left(z\right|x)〉-〈logb\left(z\right|y)〉-〈KL\left[p\right(z\left|y\right)\left|\right|b\left(z\right|y\left)\right]〉\hfill \end{array}$$

(10)

$$\begin{array}{cc}\hfill \le & 〈loge\left(z\right|x)〉-〈logb\left(z\right|y)〉\hfill \end{array}$$

(11)

The second term of Equation (5):

$$\begin{array}{cc}\hfill I(Y;Z)=& H\left(Y\right)-H\left(Y\right|Z)\propto -H\left(Y\right|Z)=〈logp\left(y\right|z)〉\hfill \end{array}$$

(12)

$$\begin{array}{cc}\hfill =& 〈logc\left(y\right|z)〉+〈KL\left[p\right(y\left|z\right)\left|\right|c\left(y\right|z\left)\right]〉\hfill \end{array}$$

(13)

$$\begin{array}{cc}\hfill \ge & 〈logc\left(y\right|z)〉\hfill \end{array}$$

(14)

These variational bounds give us a tractable objective function for amortized inference, the Variational Conditional Entropy Bottleneck (VCEB):

$$\begin{array}{cc}\hfill VCEB\equiv & \underset{e,b,c}{min}〈loge\left(z\right|x)〉-〈logb\left(z\right|y)〉-\gamma 〈logc\left(y\right|z)〉\hfill \end{array}$$

(15)

There are a number of other ways to optimize Equation (5). We describe a few of them in Section 2.6 and Appendix B and Appendix C.

2.5. Comparison to the Information Bottleneck

The Information Bottleneck (IB) [11] learns a representation Z from X subject to a soft constraint:

$$IB\equiv \underset{Z}{min}I(X;Z)-\beta I(Y;Z)$$

(16)

where ${\beta }^{-1}$ controls the strength of the constraint. As $\beta \to \infty$, IB recovers the standard cross-entropy loss.

In Figure 2 we show information diagrams comparing which regions IB and CEB maximize and minimize. See Yeung [30] for a theoretical explanation of information diagrams.CEB avoids trying to both minimize and maximize the central region at the same time. In Figure 3 we show the feasible regions for CEB and IB, labeling the MNI point on both. CEB’s rectification of the information plane means that we can always measure in absolute terms how much more we could compress our representation at the same predictive performance: $I(X;Z|Y)\ge 0$. For IB, it is not possible to tell a priori how far we are from optimal compression.

From Equations (4), (5) and (16), it is clear that CEB and IB are equivalent for $\gamma =\beta -1$. To simplify comparison of the two objectives, we can parameterize them with:

$$\rho =log\gamma =log(\beta -1)$$

(17)

Under this parameterization, for deterministic datasets, sufficiently powerful models will target the MNI point at $\rho =0$. As $\rho$ increases, more information is captured by the model. $\rho <0$ may capture less than the MNI. $\rho >0$ may capture more than the MNI.

Amortized IB

As described in Tishby et al. [11], IB is a tabular method, so it is not usable for amortized inference. The tabular optimization procedure used for IB trivially applies to CEB, just by setting $\beta =\gamma +1$. Two recent works have extended IB for amortized inference. Achille and Soatto [19] presents InfoDropout, which uses IB to motivate a variation on Dropout [31]. Alemi et al. [17] presents the Variational Information Bottleneck (VIB):

$$VIB\equiv 〈loge\left(z\right|x)〉-〈logm\left(z\right)〉)-\beta 〈logc\left(y\right|z)〉$$

(18)

Instead of the backward encoder, VIB has a marginal posterior, $m\left(z\right)$, which is a variational approximation to $e\left(z\right)=\int dxp\left(x\right)e\left(z\right|x)$.

Following Alemi et al. [32], we define the Rate (R):

$$\begin{array}{c}\hfill R\equiv 〈loge\left(z\right|x)〉-〈logm\left(z\right)〉\ge I(X;Z)\end{array}$$

(19)

We similarly define the Residual Information ($R{e}_X$):

$$\begin{array}{c}\hfill R{e}_X\equiv 〈loge\left(z\right|x)〉-〈logb\left(z\right|y)〉\ge I(X;Z|Y)\end{array}$$

(20)

During optimization, observing R does not tell us how tightly we are adhering to the MNI. However, observing $R{e}_X$ tells us exactly how many bits we are from the MNI point, assuming that our current classifier is optimal.

For convenience, define $CE{B}_x\equiv CE{B}_{\rho =x}$, and likewise for VIB. We can compare variational CEB with VIB by taking their difference at $\rho =0$:

$$\begin{array}{cc}\hfill CE{B}_0-VI{B}_0=& 〈logb\left(z\right|y)〉-〈logm\left(z\right)〉\hfill \end{array}$$

(21)

$$\begin{array}{cc}\hfill & -〈logc\left(y\right|z)〉+〈logp\left(y\right)〉\hfill \end{array}$$

(22)

Solving for $m\left(z\right)$ when that difference is 0:

$$m\left(z\right)=\frac{b\left(z\right|y\left)p\right(y)}{c\left(y\right|z)}$$

(23)

Since the optimal $m^{*}\left(z\right)$ is the marginalization of $e\left(z\right|x)$, at convergence we must have:

$$m^{*}\left(z\right)=\int dxp\left(x\right)e\left(z\right|x)=\frac{p\left(z\right|y\left)p\right(y)}{p\left(y\right|z)}$$

(24)

This solution may be difficult to find, as $m\left(z\right)$ only gets information about y indirectly through $e\left(z\right|x)$. For otherwise equivalent models, we may expect $VI{B}_0$ to converge to a looser approximation of $I(X;Z)=I(Y;Z)=I(X;Y)$ than CEB. Since VIB optimizes an upper bound on $I(X;Z)$, $VI{B}_0$ will report R converging to $I(X;Y)$, but may capture less than the MNI. In contrast, if $R{e}_X$ converges to 0, the variational tightness of $b\left(z\right|y)$ to the optimal $p\left(z\right|y)$ depends only on the tightness of $c\left(y\right|z)$ to the optimal $p\left(y\right|z)$.

2.6. Model Variants

We introduce some variants on the basic variational CEB classification model that we will use in Section 3.1.6.

2.6.1. Bidirectional CEB

We can learn a shared representation Z that can be used to predict both X and Y with the following bidirectional CEB model: $Z_X\leftarrow X\leftrightarrow Y\to Z_Y$. This corresponds to the following joint: $p(x,y,z_X,z_Y)\equiv p(x,y)e\left(z_X\right|x)b\left(z_Y\right|y)$. The main CEB objective can then be applied in both directions:

$$\begin{array}{cc}\hfill {\mathrm{CEB}}_{\mathrm{bidir}}\equiv min& -H\left(Z_X\right|X)+H\left(Z_X\right|Y)+{\gamma }_{X}H\left(Y\right|Z_X)\hfill \\ \hfill & -H\left(Z_Y\right|Y)+H\left(Z_Y\right|X)+{\gamma }_{Y}H\left(X\right|Z_Y)\hfill \end{array}$$

(25)

For the two latent representations to be useful, we want them to be consistent with each other (minimally, they must have the same parametric form). Fortunately, that consistency is trivial to encourage by making the natural variational substitutions: $p\left(z_Y\right|x)\to e\left(z_Y\right|x)$ and $p\left(z_X\right|y)\to b\left(z_X\right|y)$. This gives variational CEB${}_{\mathrm{bidir}}$:

$$\begin{array}{cc}\hfill min& 〈loge\left(z_X\right|x)〉-〈logb\left(z_X\right|y)〉-{\gamma }_X〈logc\left(y\right|z_X)〉\hfill \\ \hfill +& 〈logb\left(z_Y\right|y)〉-〈loge\left(z_Y\right|x)〉-{\gamma }_Y〈logd\left(x\right|z_Y)〉\hfill \end{array}$$

(26)

where $d\left(x\right|z)$ is a decoder distribution. At convergence, we learn a unified Z that is consistent with both $Z_X$ and $Z_Y$, permitting generation of either output given either input in the trained model, in the same spirit as Vedantam et al. [33], but without needing to train a joint encoder $q\left(z\right|x,y)$.

2.6.2. Consistent Classifier

We can reuse the backwards encoder as a classifier: $c\left(y\right|z)\propto b(z\left|y\right)p\left(y\right)$. We refer to this as the Consistent Classifier: $c\left(y\right|z)\equiv softmaxb(z\left|y\right)p\left(y\right)$. If the labels are uniformly distributed, the $p\left(y\right)$ factor can be dropped; otherwise, it suffices to use the empirical $p\left(y\right)$. Using the consistent classifier for classification problems results in a model that only needs parameters for the two encoders, $e\left(z\right|x)$ and $b\left(z\right|y)$. This classifier differs from the more common maximum a posteriori (MAP) classifier because $b\left(z\right|y)$ is not the sampling distribution of either Z or Y.

2.6.3. CatGen Decoder

We can further generalize the idea of the consistent classifier to arbitrary prediction tasks by relaxing the requirement that we perfectly marginalize Y in the softmax. Instead, we can marginalize Y over any minibatch of size K we see at training time, under an assumption of a uniform distribution over the training examples we sampled:

$$\begin{array}{cc}\hfill p\left(y\right|z)=& \frac{p\left(z\right|y\left)p\right(y)}{\int d{y}^{\prime }p\left(z\right|y^{\prime })p\left(y^{\prime }\right)}\hfill \end{array}$$

(27)

$$\begin{array}{cc}\hfill \approx & \frac{p\left(z\right|y)\frac{1}{K}}{{\sum }_{k=1}^{K}p\left(z\right|y_k)\frac{1}{K}}=\frac{p\left(z\right|y)}{{\sum }_{k=1}^{K}p\left(z\right|y_k)}\hfill \end{array}$$

(28)

$$\begin{array}{cc}\hfill \approx & \frac{b\left(z\right|y)}{{\sum }_{k=1}^{K}b\left(z\right|y_k)}\equiv c\left(y\right|z)\hfill \end{array}$$

(29)

We can immediately see that this definition of $c\left(y\right|z)$ gives a valid distribution, as it is just a softmax over the minibatch. That means it can be directly used in the original objective without violating the variational bound. We call this decoder CatGen, for Categorical Generative Model because it can trivially “generate” Y: the softmax defines a categorical distribution over the batch; sampling from it gives indices of $Y=y_j$ that most closely correspond to $Z=z_i$.

Maximizing $I(Y;Z)$ in this manner is a universal task, in that it can be applied to any paired data $X,Y$. This includes images and labels – the CatGen model may be used in place of both $c\left(y\right|z_X)$ and $d\left(x\right|z_Y)$ in the CEB${}_{\mathrm{bidir}}$ model (using $e\left(z\right|x)$ for $d\left(x\right|z_Y)$). This avoids a common concern when dealing with multivariate predictions: if predicting X is disproportionately harder than predicting Y, it can be difficult to balance the model [33,34]. For CatGen models, predicting X is never any harder than predicting Y, since in both cases we are just trying to choose the correct example out of K possibilities.

It turns out that CatGen is mathematically equivalent to Contrastive Predictive Coding (CPC) [35] after an offset of $logK$. We can see this using the proof from Poole et al. [36], and substituting $logb\left(z\right|y)$ for $f(y,z)$:

$$\begin{array}{cc}\hfill I(Y;Z)& \ge \frac{1}{K}\sum _{k=1}^K{\mathbb{E}}_{{\prod }_j{y}_k,z\sim p\left(y_j\right)p\left(x_k\right|y_k)e\left(z\right|x_k)}\left[log\frac{e^{f(y_k,z)}}{\frac{1}{K}{\sum }_{i=1}^K{e}^{f(y_i,z)}}\right]\hfill \end{array}$$

(30)

$$\begin{array}{cc}\hfill & =\frac{1}{K}\sum _{k=1}^K{\mathbb{E}}_{{\prod }_j{y}_k,z\sim p\left(y_j\right)p\left(x_k\right|y_k)e\left(z\right|x_k)}\left[log\frac{b\left(z\right|y_k)}{\frac{1}{K}{\sum }_{i=1}^{K}b\left(z\right|y_i)}\right]\hfill \end{array}$$

(31)

The advantage of the CatGen approach over CPC in the CEB setting is that we already have parameterized the forward and backward encoders to compute $I(X;Z|Y)$, so we don’t need to introduce any new parameters when using CatGen to maximize the $I(Y;Z)$ term.

As with CPC, the CatGen bound is constrained by $logK$, but when targeting the MNI, it is more likely that we can train with $logK\ge I(X;Y)$. This is trivially the case for the datasets we explore here, where $I(X;Y)\le log10$. It is also practical for larger datasets like ImageNet, where models are routinely trained with batch sizes in the thousands (e.g., Goyal et al. [37]), and $I(X;Y)\le log1000$.

3. Results

We evaluate deterministic, VIB, and CEB models on Fashion MNIST [38] and CIFAR10 [39]. Our experiments focus on comparing the performance of otherwise identical models when we change only the objective function and vary $\rho$. Thus, we are interested in relative differences in performance that can be directly attributed to the difference in objective and $\rho$. These experiments cover the three aspects of Robust Generalization (Section 2.1): (RG1) (classical generalization) in Section 3.1 and Section 3.1.6; (RG2) (adversarial robustness) in Section 3.1 and Section 3.1.6; and (RG3) (detection) in Section 3.1.

3.1. (RG1), (RG2), and (RG3): Fashion MNIST

Fashion MNIST [38] is an interesting dataset in that it is visually complex and challenging, but small enough to train in a reasonable amount of time. We trained 60 different models on Fashion MNIST, four each for the following 15 types: a deterministic model (Determ); seven VIB models (VIB₋₁, ..., VIB₅); and seven CEB models (CEB₋₁, ..., CEB₅). Subscripts indicate $\rho$. All 60 models share the same inference architecture and are trained with otherwise identical hyperparameters. See Appendix A for details.

3.1.1. (RG1): Accuracy and Compression

In Figure 4 we see that both VIB and CEB have improved accuracy over the deterministic baseline, consistent with compressed representations generalizing better. Also, CEB outperforms VIB at every $\rho$, which we can attribute to the tighter variational bound given by minimizing $R{e}_X$ rather than R. In the case of a simple classification problem with a uniform distribution over classes in the training set (like Fashion MNIST), we can directly compute $I(X;Y)=logC$, where C is the number of classes. In order to compare the relative complexity of the learned representations for the VIB and CEB models, in the second panel of Figure 4 we show the maximum rate lower bound seen during training: $R_X\equiv 〈log\frac{e\left(z\right|x)}{\frac{1}{K}{\sum }_k^{K}e\left(z\right|x_k)}〉\le I(X;Z)$ using the encoder’s minibatch marginal for both VIB and CEB. This lower bound on $I(X;Z)$ is the “InfoNCE with a tractable encoder” bound from Poole et al. [36]. The two sets of models show nearly the same $R_X$ at each value of $\rho$. Both models converge to exactly $I(X;Y)=log10\approx 2.3$ nats at $\rho =0$, as predicted by the derivation of CEB.

3.1.2. (RG2): Adversarial Robustness

The bottom two panels of Figure 4 show robustness to targeted Projected Gradient Descent (PGD) L₂ and L_∞ attacks [14]. All of the attacks are targeting the trouser class of Fashion MNIST, as that is the most distinctive class. Targeting a less distinctive class, such as one of the shirt classes, would confuse the difficulty of classifying the different shirts and the robustness of the model to adversaries. To measure robustness to the targeted attacks, we count the number of predictions that changed from a correct prediction on the clean image to an incorrect prediction of the target class on the adversarial image, and divide by the original number of correct predictions. Consistent with testing (RG2), these adversaries are completely unknown to the models at training time – none of these models see any adversarial examples during training. CEB again outperforms VIB at every $\rho$, and the deterministic baseline at all but the least-compressed model ($\rho =5$). We also see for both models that as $\rho$ decreases, the robustness to both attacks increases, indicating that more compressed models are more robust.

Consistent with the MNI hypothesis, at $\rho =0$ we end up with CEB models that have hit exactly 2.3 nats for the rate lower bound, have maintained high accuracy, and have strong robustness to both attacks. Moving to $\rho =-1$ gives only a small improvement to robustness, at the cost of a large decrease in accuracy.

3.1.3. (RG3): Out-of-Distribution Detection

We compare the ability of Determ, CEB₀, VIB₀, and VIB₄ to detect four different out-of-distribution (OoD) detection datasets. $U(0,1)$ is uniform noise in the image domain. MNIST uses the MNIST test set. Vertical Flip is the most challenging, using vertically flipped Fashion MNIST test images, as originally proposed in Alemi et al. [18]. CW is the Carlini-Wagner L₂ attack [40] at the default settings found in Papernot et al. [41], and additionally includes the adversarial attack success rate against each model.

We use two different metrics for thresholding, proposed in Alemi et al. [18]. H is the classifier entropy. R is the rate, defined in Section 2.5. These two threshold scores are used with the standard suite of proper scoring rules [42]: False Positive Rate at 95% True Positive Rate (FPR 95% TPR), Area Under the ROC Curve (AUROC), and Area Under the Precision-Recall Curve (AUPR).

Table 1. Results for out-of-distribution detection (OoD). Thrsh. is the threshold score used: H is the entropy of the classifier; R is the rate. Determ cannot compute R, so only H is shown. For VIB and CEB models, H is always inferior to R, similar to findings in Alemi et al. [18], so we omit it. Adv. Success is attack success of the CW adversary (bottom four rows). Arrows denote whether higher or lower scores are better. Bold indicates the best score in that column for that OoD dataset.

OoD	Model	Thrsh.	FPR @ 95% TPR ↓	AUROC ↑	AUPR In ↑	Adv. Success ↓
U(0,1)	Determ	H	35.8	93.5	97.1	N/A
	VIB${}_4$	R	0.0	100.0	100.0	N/A
	VIB${}_0$	R	80.6	57.1	51.4	N/A
	CEB${}_0$	R	0.0	100.0	100.0	N/A
MNIST	Determ	H	59.0	88.4	90.0	N/A
	VIB${}_4$	R	0.0	100.0	100.0	N/A
	VIB${}_0$	R	12.3	66.7	91.1	N/A
	CEB${}_0$	R	0.1	94.4	99.9	N/A
Vertical Flip	Determ	H	66.8	88.6	90.2	N/A
	VIB${}_4$	R	0.0	100.0	100.0	N/A
	VIB${}_0$	R	17.3	52.7	91.3	N/A
	CEB${}_0$	R	0.0	90.7	100.0	N/A
CW	Determ	H	15.4	90.7	86.0	100.0%
	VIB${}_4$	R	0.0	100.0	100.0	55.2%
	VIB${}_0$	R	0.0	98.7	100.0	35.8%
	CEB${}_0$	R	0.0	99.7	100.0	35.8%

shows that using R to detect OoD examples can be much more effective than using classifier-based approaches. The deterministic baseline model is far weaker at detection using H than either of the high-performing stochastic models (CEB₀ and VIB₄). Those models both saturate detection performance, providing reliable signals for all four OoD datasets. However, as VIB₀ demonstrates, simply having R available as a signal does not guarantee good detection. As we saw above, the VIB₀ models had noticeably worse classification performance, indicating that they had not achieved the MNI point: $I(Y;Z)<I(X;Z)$ for those models. These results indicate that for detection, violating the MNI criterion by having $I(X;Z)>I(X;Y)$ may not be harmful, but violating the criterion in the opposite direction is harmful.

3.1.4. (RG3): Calibration

A well-calibrated model is correct half of the time it gives a confidence of 50% for its prediction. In Figure 5, we show calibration plots at various points during training for four models. Calibration curves help analyze whether models are underconfident or overconfident. Each point in the plots corresponds to a 5% confidence bin. Accuracy is averaged for each bin. All four networks move from under- to overconfidence during training. However, CEB₀ and VIB₀ end up only slightly overconfident, while $\rho =2$ is already sufficient to make VIB and CEB (not shown) nearly as overconfident as the deterministic model.

3.1.5. (RG1): Overfitting Experiments

We replicate the basic experiment from Zhang et al. [10] by using the images from Fashion MNIST, but replacing the training labels with fixed random labels. This dataset is information-free because $I(X;Y)=0$. We use that dataset to train multiple deterministic models, as well as CEB and VIB models at $\rho$ from 0 through 7. We find that the CEB and VIB models with $\rho <6$ never learn, even after 100 epochs of training, but the deterministic models always learn. After about 40 epochs of training they begin to memorize the random labels, indicating severe overfitting and a perfect failure to generalize.

3.1.6. (RG1) and (RG2): CIFAR10 Experiments

For CIFAR10 [39] we trained the largest Wide ResNet [43] we could fit on a single GPU with a batch size of 250. This was a 62 × 7 model trained using AutoAugment [44]. We trained 3 CatGen CEB${}_{\mathrm{bidir}}$ models each of CEB₀ and CEB₅ and then selected the two models with the highest test accuracy for the adversarial robustness experiments. We evaluated the CatGen models using the consistent classifier, since CatGen models only train $e\left(z\right|x)$ and $b\left(z\right|y)$. CEB₀ reached 97.51% accuracy. This result is better than the 28 × 10 Wide ResNet from AutoAugment by 0.19 percentage points, although it is still worse than the Shake-Drop model from that paper. We additionally tested the model on the CIFAR-10.1 test set [45], getting accuracy of 93.6%. This is a gap of only 3.9 percentage points, which is better than all of the results reported in that paper, and substantially better than the Wide ResNet results (but still inferior to the Shake-Drop AutoAugment results). The CEB₅ model reached 97.06% accuracy on the normal test set and 91.9% on the CIFAR-10.1 test set, showing that increased $\rho$ gave substantially worse generalization.

To test robustness of these models, we swept $ϵ$ for both PGD attacks (Figure 6). The CEB₀ model not only has substantially higher accuracy than the adversarially-trained Wide ResNet from Madry et al. [14] (Madry), it also beats the Madry model on both the L₂ and the L_∞ attacks at almost all values of $ϵ$. We also show that this model is even more robust to two transfer attacks, where we used the CEB₅ model and the Madry model to generate PGD attacks, and then test them on the CEB₀ model. This result indicates that these models are not doing “gradient masking”, a failure mode of some attempts at adversarial defense [2], since these are black-box attacks that do not rely on taking gradients through the target model.

4. Conclusions

We have presented the Conditional Entropy Bottleneck (CEB), motivated by the Minimum Necessary Information (MNI) criterion and the hypothesis that failures of robust generalization are due in part to learning models that retain too much information about the training data. We have shown empirically that simply by switching to CEB, models may substantially improve their robust generalization, including (RG1) higher accuracy, (RG2) better adversarial robustness, and (RG3) stronger OoD detection. We believe that the MNI criterion and CEB offer a promising path forward for many tasks in machine learning by permitting fast amortized inference in an easy-to-implement framework that improves robust generalization.