CEB IMPROVES MODEL ROBUSTNESS
Year: 2019

Abstract:
We demonstrate that the Conditional Entropy Bottleneck (CEB) can improve model robustness. CEB is an easy strategy to implement and works in tandem with data augmentation procedures. We report results of a large scale adversarial robustness study on CIFAR-10, as well as the IMAGENET-C Common Corruptions Benchmark, IMAGENET-A, and PGD attacks.

heading: 1 INTRODUCTION
text: We aim to make models that make meaningful predictions beyond the data they were trained on. Generally we want our models to be robust. Broadly, robustness is the ability of a model to continue making valid predictions as the distribution the model is tested on moves away from the empirical training set distribution. The most commonly reported robustness metric is simply test set performance, where we verify that our model continues to make valid predictions on what we hope represents valid draws from the exact same data generating procedure.
Adversarial robustness tests robustness in a worst case setting, where an attacker (Szegedy et al., 2013) makes limited targeted modifications to the input that are as fooling as possible. Many adversarial attacks have been proposed and studied (Szegedy et al., 2013; Carlini & Wagner, 2017b;a; Kurakin et al., 2016a; Madry et al., 2017). Most machine-learned systems are currently believed to be vulnerable to adversarial examples. Many defenses have been proposed, but very few have demonstrated robustness against a powerful, general-purpose adversary (Carlini & Wagner, 2017a; Athalye et al., 2018). While robustness to adversarial attacks continues to attract interest, recent discussions have emphasized the need to consider other forms of robustness as well (Engstrom et al., 2019). The Common Corruptions Benchmark (Hendrycks & Dietterich, 2019) measures image models robustness to more mild but real world sorts of perturbations. Even these modest perturbations can be very fooling for traditional architectures.
One of the few general purpose strategies that demonstrably improves model robustness is Data Augmentation (Cubuk et al., 2018; Lopes et al., 2019; Yin et al., 2019). However, it would be nice to identify loss-based solutions that can work in tandem with the data augmentation approaches. Intuitively, by performing modifications of the inputs at training time, the model is prevented from being too sensitive to particular features of the inputs that don’t survive the augmentation procedure.
Alternatively, we can try to make our models more robust by making them less sensitive to the inputs in the first place. The goal of this work is to experimentally investigate whether, by systematically limiting the complexity of the extracted representation using the Conditional Entropy Bottleneck (CEB), we can make our models more robust in all three of these senses: test set generalization (e.g., classification accuracy on “clean” test inputs), worst-case robustness, and typical-case robustness.

heading: 1.1 CONTRIBUTIONS
text: This paper is primarily empirical. We demonstrate:
• CEB models are easy to implement and train. • CEB models demonstrate improved generalization performance over deterministic base-
lines on CIFAR-10 and ImageNet. • CEB models show improved robustness to adversarial attacks on CIFAR-10. • CEB models show improved robustness on the IMAGENET-C Common Corruptions
Benchmark, the IMAGENET-A Benchmark, and targeted PGD attacks.
Additionally, we show that adversarially-trained models fail to generalize to attacks they weren’t trained on, by comparing the results on L2 PGD attacks from Madry et al. (2017) to our results on the same baseline architecture. This result underscores the importance of finding ways to make models robust that do not rely on knowing the form of the attack ahead of time.

heading: 2 BACKGROUND

heading: 2.1 INFORMATION BOTTLENECKS
text: The Information Bottleneck (IB) objective (Tishby et al., 2000) aims to learn a stochastic representation Z ∼ p(z|x) that retains as much information about a target variable Y while being as compressed as possible. The objective:1
max I(Z;Y )− σ(−ρ)I(Z;X), (1)
uses a Lagrange multiplier σ(−ρ) to trade off between the relevant information (I(Z;Y )) and complexity of the representation (I(Z;X)). Because Z depends only on X (Z ← X ↔ Y ): Z and Y are conditionally independent given Z:
I(Z;X,Y ) = I(Z;X) + I(Z;Y |X) = I(Z;Y ) + I(Z;X|Y ). (2)
This allows us to write the information bottleneck of Equation (1) in an equivalent form:
max I(Z;Y )− e−ρI(Z;X|Y ). (3)
Just as the original Information Bottleneck objective (Equation (1)) admits a natural variational lower bound (Alemi et al., 2017), so does this form. We can variationally lower bound the mutual information between our representation and the targets with a variational decoder q(y|z):
I(Z;Y ) = Ep(x,y)p(z|x) [ log
p(y|z) p(y)
] ≥ H(Y ) + Ep(x,y)p(z|x) [log q(y|z)] . (4)
While we may not know H(Y ) exactly for real world datasets, in the information bottleneck formulation it is a constant outside of our control and so can be dropped in our objective. We can variationally upper bound our residual information:
I(Z;X|Y ) = Ep(x,y)p(z|x) [ log
p(z|x, y) p(z|y)
] ≤ Ep(x,y)p(z|x) [ log
p(z|x) q(z|y)
] , (5)
with a variational class conditional marginal q(z|y) that approximates ∫ dx p(z|x)p(x|y). Putting both bounds together gives us the Conditional Entropy Bottleneck objective (Fischer, 2018):
min p(z|x)
Ep(x,y)p(z|x) [ log q(y|z)− e−ρ log p(z|x)
q(z|y)
] (6)
Compare this with the Variational Information Bottleneck (VIB) objective (Alemi et al., 2017):
min p(z|x)
Ep(x,y)p(z|x) [ log q(y|z)− σ(−ρ) log p(z|x)
q(z)
] . (7)
The difference between CEB and VIB is the presence of a class conditional versus unconditional variational marginal. As can be seen in Equation (5): using an unconditional marginal provides a looser variational upper bound on I(Z;X|Y ). CEB (Equation (6)) can be thought of as a tighter variational approximation than VIB (Equation (7)) to Equation (3). Since Equation (3) is equivalent to the IB objective (Equation (1)), CEB can be thought of as a tighter variational approximation to the IB objective than VIB.
1 The IB objective is ordinarily written with a Lagrange multiplier β ≡ σ(−ρ) with a natural range from 0 to 1. Here we use the sigmoid function: σ(−ρ) ≡ 1
1+eρ to reparameterize in terms of a control parameter ρ on
the whole real line. As ρ→∞ the bottleneck turns off.

heading: 2.2 IMPLEMENTING A CEB MODEL
text: In practice, turning an existing classifier architecture into a CEB model is very simple. For the stochastic representation p(z|x) we simply use the original architecture, replacing the final softmax layer with a dense layer with d outputs. These outputs are then used to specify the means of a d-dimensional Gaussian distribution with unit diagonal covariance. That is, to form the stochastic representation, independent standard normal noise is simply added to the output of the network (z = x + ). For every input, this stochastic encoder will generate a random d-dimensional output vector. For the variational classifier q(y|z) any classifier network can be used, including just a linear softmax classifier as done in these experiments. For the variational conditional marginal q(z|y) it helps to use the same distribution as output by the classifier. For the simple unit variance Gaussian encoding we used in these experiments, this requires learning just d parameters per class. For ease of implementation, this can be represented as single dense linear layer mapping from a one-hot representation of the labels to the d-dimensional output, interpreted as the mean of the corresponding class marginal.
In this setup the CEB loss takes a particularly simple form:
E wy · (f(x) + )− log∑ y′ ewy′ ·(f(x)+ ) − e −ρ 2 (f(x)− µy) (f(x)− µy + 2 )  . (8) Here the first term is the usual softmax classifier loss, but acting on our stochastic representation z = f(x) + , which is simply the output of our encoder network f(x) with additive Gaussian noise. The wy is the yth row of weights in the final linear layer outputing the logits. µy are the learned class conditional means for our marginal. are standard normal draws from an isotropic unit variance Gaussian with the same dimension as our encoding f(x). The second term in the loss is a stochastic sampling of the KL divergence between our encoder likelihood and the class conditional marginal likelihood. ρ controls the strength of the bottleneck and can vary on the whole real line. As ρ → ∞ the bottleneck is turned off. In practice we find that ρ values near but above 0 tend to work best for modest size models, with the tendency for the best ρ to approach 0 as the model capacity increases. Notice that in expectation the second term in the loss is (f(x) − µy)2, which encourages the learned means µy to converge to the average of the representations of each element in the class. During testing we use the mean encodings and remove the stochasticity.
In its simplest form, CEB training a classifier amounts to injecting Gaussian random noise in the penultimate layer and learning estimates of the class averaged output of that layer with the stochastic regularization shown. In Appendix B we show simple modifications to the TPU-compatible ResNet implementation available on GitHub from the Google TensorFlow Team that produce the same core ResNet-50 models we use for our ImageNet experiments.

heading: 2.3 ADVERSARIAL ATTACKS AND DEFENSES
text: Attacks. The first adversarial attacks were proposed in Szegedy et al. (2013); Goodfellow et al. (2015). Since those seminal works, an enormous variety of attacks has been proposed (Kurakin et al. (2016a;b); Moosavi-Dezfooli et al. (2016); Carlini & Wagner (2017b); Madry et al. (2017); Eykholt et al. (2017); Baluja & Fischer (2017), etc.). In this work, we will primarily consider the Projected Gradient Descent (PGD) attack (Madry et al., 2017), which is a multi-step variant of the early Fast Gradient Method (Goodfellow et al., 2015). The attack can be viewed as having four parameters: p, the norm of the attack (typically 2 or∞), , the radius the the p-norm ball within which the attack is permitted to make changes to an input, n, the number of gradient steps the adversary is permitted to take, and i, the per-step limit to modifications of the current input. In this work, we consider L2 and L∞ attacks of varying and n, and with i = 43 n .
Defenses. A common defense for adversarial examples is adversarial training. Adversarial training was originally proposed in Szegedy et al. (2013), but was not practical until the Fast Gradient Method was introduced. It has been studied in detail, with varied techniques (Kurakin et al., 2016b; Madry et al., 2017; Ilyas et al., 2019; Xie et al., 2019). Adversarial training can clearly be viewed as a form of data augmentation (Tsipras et al., 2018), where instead of using some fixed set of functions to modify the training examples, we use the model itself in combination with one or more
adversarial attacks to modify the training examples. As the model changes, the distribution of modifications changes as well. However, unlike with non-adversarial data augmentation techniques, such as AUTOAUG, adversarial training techniques considered in the literature so far cause substantial reductions in accuracy on clean test sets. For example, the CIFAR-10 model described in Madry et al. (2017) gets 95.5% accuracy when trained normally, but only 87.3% when trained on L∞ adversarial examples. More recently, Xie et al. (2019) adversarially trains ImageNet models with impressive robustness to targeted PGD L∞ attacks, but at only 62.32% accuracy on the non-adversarial test set, compared to 78.81% accuracy for the same model trained only on clean images.

heading: 2.4 COMMON CORRUPTIONS
text: The Common Corruptions Benchmark (Hendrycks & Dietterich, 2019) offers a real world test of model robustness in light of common image processing pipeline corruptions. Figure 4 shows examples of the 15 corruptions present in the benchmark. IMAGENET-C is a modified test set of Imagenet images with the 15 corruptions applied at five different strengths. Within each corruption type we evaluated the average error at each of the five levels (Ec = 15 ∑5 s=1Ecs). To summarize the performance across all corruptions we report not only the average corruption error across all 15 tasks (avg = 115 ∑ cEc), but also the commonly reported Mean Corruption Error (mCE), which reweights the errors on each task according to the performance of a baseline ALEXNET model (Hendrycks & Dietterich, 2019):
mCE = 1
15 ∑ c ∑5 s=1Ecs∑5 s=1E ALEXNET cs . (9)
There are slightly different pipelines that have been used in the literature for the IMAGENET-C task (Lopes et al., 2019). In this work we used the ALEXNET normalization numbers and data formulation as in Yin et al. (2019).

heading: 2.5 NATURAL ADVERSARIAL EXAMPLES
text: The IMAGENET-A Benchmark (Hendrycks et al., 2019) is a dataset of 7,500 naturally-occurring “adversarial” examples across 200 ImageNet classes. The images exploit commonly-occurring weaknesses in ImageNet models, such as relying on textures often seen with certain class labels.

heading: 3 EXPERIMENTS

heading: 3.1 FASHION-MNIST EXPERIMENTS
text: As a warm up, we consider Wide ResNet (Zagoruyko & Komodakis, 2016) models trained on Fashion MNIST (Xiao et al., 2017), and evaluated on targeted PGD L2 and L∞ attacks. 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. Results are shown in Figure 1.
In this experiment we wanted to compare the performance of VIB, CEB, and a deterministic baseline. In Figure 1 (left) we see that both VIB and CEB have improved accuracy over the deterministic baseline. In order to compare the relative complexity of the learned representations for the two models, in the second panel we show the maximum lower bound seen during training on the rate: E [ log p(z|x)1
K ∑K k p(z|xk)
] ≤ I(Z;X) using the encoder’s minibatch marginal for both VIB and CEB.2
The two sets of models show nearly the same rate lower bound at each value of ρ.
The right two panels of Figure 1 show robustness to the targeted PGD L2 and L∞ attacks. Here CEB outperforms VIB. We also see for both models that as ρ decreases, the robustness to both attacks increases. In line with the proposed Minimum Necessary Information criterion from Fischer (2018), at ρ = 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 ρ = −1 gives only a small improvement to robustness, at the cost of a large decrease in accuracy.

heading: 3.2 CIFAR-10 EXPERIMENTS
text: 28×10 Wide ResNet Experiments We trained a set of 25 28×10 Wide ResNet (WRN) CEB models on CIFAR-10 at ρ ∈ [−1,−0.75, ..., 5], as well as a deterministic baseline. They trained for 1500 epochs, lowering the learning rate by a factor of 0.3 after 500, 1000, and 1250 epochs. This long training regime was due to our use of the original AUTOAUG policies, which requires longer training. The only additional modification we made to the basic 28×10 WRN architecture was the removal of all Batch Normalization (Ioffe & Szegedy, 2015) layers. Every small CIFAR-10 model we have trained with Batch Normalization enabled has had substantially worse robustness to L∞ PGD adversaries, even though typically the accuracy is much higher. For example, 28×10 WRN CEB models rarely exceeded more than 10% adversarial accuracy. However, it was always still the case that lower values of ρ gave higher robustness. As a baseline comparison, a deterministic 28×10 WRN with BatchNorm, trained with AUTOAUG reaches 97.3% accuracy on clean images, but 0% accuracy on L∞ PGD attacks at = 8 and n = 20. Interestingly, that model was noticeably more robust to L2 PGD attacks than the deterministic baseline without BatchNorm, getting 73% accuracy compared to 66%. However, it was still much weaker than the CEB models, which get over 80% accuracy on the same attack (Figure 2). Additional training details are in Appendix A.1.
Figure 2 demonstrates the adversarial robustness of CEB models to both targeted L2 and L∞ attacks. The CEB models show a marked improvement in robustness to L2 attacks compared to an adversarially-trained baseline from Madry et al. (2017) (denoted Madry). Figure 3 shows the robustness of five of those models to PGD attacks as is varied. We selected the four CEB models to represent the most robust models across most of the range of ρ we trained. Note that of the 25 CEB models we trained, only the models with ρ ≥ 1 succesfully trained. The remainder collapsed to chance performance. This is something we observe on all datasets when training models that are too low capacity. Only by increasing model capacity does it become possible to train at low ρ. Note that this result is predicted by the theory of the onset of learning in the Information Bottleneck and its relationship to model capacity from Wu et al. (2019).
62×7 Wide ResNet Experiments. In order to explore the effect of model size on training, and to train at lower ρ, we trained the largest Wide ResNet we could fit on a single GPU with a batch size of 250. This was a 62×7 model similar to the ones above, including the use of AUTOAUG, but we additionally enabled BatchNorm. We were able to train at ρ = 0 with this larger model, which reached 97.51% accuracy. This result is better than the 28×10 Wide ResNet from AUTOAUG 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 new CIFAR-10.1 test set (Recht et al., 2018), getting accuracy of 93.6%. This is a gap of 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
2This lower bound on I(X;Z) is the “InfoNCE with a tractable encoder” bound from Poole et al. (2019).
AUTOAUG results). The same model at ρ = 5 reached 97.05% accuracy on the normal test set and 91.9% on the CIFAR-10.1 test set, showing that increased ρ gave substantially worse generalization.
To test robustness of these models, we swept for both PGD attacks, which we show in Figure 3. The main result is that the 62× 7 CEB0 model not only has substantially higher accuracy than baseline Wide ResNets trained with AUTOAUG, it also beats the adversarially-trained model on both the L2 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 62×7 CEB5 model and the adversariallytrained model to generate PGD attacks, and then test them on the CEB0 model. This result helps to counter possible claims that these models are doing “gradient masking” (the more compelling evidence against gradient masking is that the robustness of the model is strongly correlated with the hyperparameter ρ, whose only effect is to constrain the amount of information the model captures).
We additionally tested both models on the CIFAR-10 Common Corruptions test sets. At the time of training, we were unaware that AUTOAUG’s default policies for CIFAR-10 contain brightness and contrast augmentations that amount to training on those two corruptions from Common Corruptions (as mentioned in Yin et al. (2019)), so our results are not appropriate for direct comparison with other results in the literature. However, they still allow us to compare the effect of bottlenecking the information between these two large models. The ρ = 5 model reached an mCE3 of 61.2. The ρ = 0 model reached an mCE of 52.0, which is a dramatic relative improvement.

heading: 3.3 IMAGENET EXPERIMENTS
text: To demonstrate CEB’s ability to improve robustness to real world data shifts, we trained six different types of networks on ImageNet at 224×224 resolution and two different sizes of RESNET, RESNET-50 and RESNET-152, and then tested them on IMAGENET-C, IMAGENET-A, and targeted PGD attacks. As a simple baseline we trained RESNET-50 networks with no data augmentation. We then trained the same networks but as CEB networks at ten different values of ρ = (1, 2, . . . , 10). AUTOAUG (Cubuk et al., 2018) has previously been demonstrated to improve robustness markedly on IMAGENET-C so next we trained our baseline RESNET-50 model with AUTOAUG. We similarly trained these AUTOAUG models as CEB models with ten different values of ρ. IMAGENET-C numbers are also sensitive to the model capacity. To assess whether CEB can benefit larger models we repeated the experiments with a modified RESNET-50 network where every layer was made twice as wide. Finally, we repeated the above six model types with RESNET-152 baselines and CEB models without AUTOAUG, with AUTOAUG, and with AUTOAUG and twice as wide. All other hyperparameters (learning rate schedule, L2 weight decay scale, etc.) remained the same across all models. In total we trained 66 ImageNet models – 6 deterministic baselines varying augmentation, width, and depth, and 60 CEB models additionally varying ρ. The results for the RESNET-50 models are summarized in Figure 4 and Table 1. For RESNET-152, see Figure 5 and Table 2.
The CEB models highlighted in Figures 4 and 5 and Tables 1 and 2 were selected by cross validation. These were values of ρ that gave the best clean test set accuracy. Despite being selected for classical generalization, these models also demonstrate a high degree of robustness on both average- and worst-case perturbations. In the case that more than one model gets the same test set accuracy, we choose the model with the lower ρ, since we know that lower ρ correlates with higher robustness. The only model where we had to make this decision was for RESNET-152 with AUTOAUG, where five models all were within 0.1% of each other, so we chose the ρ = 3 model, rather than ρ ∈ {5...8}.
IMAGENET-C and IMAGENET-A. Both data augmentation and increasing model capacity have positive effects on robustness to both IMAGENET-C and IMAGENET-A, but for all three classes of models CEB gives substantial additional improvements.
Targeted PGD Attacks. We tested on the random-target version of the PGD L2 and L∞ attacks (Kurakin et al., 2016a), both at = 16, n = 20, and i = 2, which is considered to be a strong attack still (Xie et al., 2019). Similar to our results on CIFAR-10, model capacity makes a substantial difference to whitebox adversarial attacks. In particular, none of the ResNet-50 models perform well, getting less than 1% top-1 accuracy. However, the Resnet-152 CEB models show a dramatic improvement over the deterministic baseline models, with top-1 accuracy increasing from
3The mCE is computed relative to a baseline model. We use the baseline model from Yin et al. (2019).
0.09% to 17.09% between the deterministic baseline and the ρ = 1 models without AUTOAUG, a relative increase of 187 times, and increases nearly as large for the AUTOAUG and wide AUTOAUG models. Interestingly, for the PGD attacks, AUTOAUG was detrimental – the RESNET-152 models without AUTOAUG were more robust than those with AUTOAUG. Only the wide RESNET-152 models with AUTOAUG exceeded the robustness of the narrow RESNET-152 without AUTOAUG.

heading: 4 CONCLUSION
text: The Conditional Entropy Bottleneck (CEB) provides a simple mechanism to improve robustness of image classifiers. We have shown that CEB gives a tighter variational bound on the IB objective than the closely-related VIB, while also having consistently better test accuracy and robustness. We have shown a strong trend toward increased robustness as ρ decreases in the standard 28×10 Wide RESNET model on CIFAR-10, and that this increased robustness does not come at the expense of accuracy relative to the deterministic baseline. We have shown that CEB models at a range of ρ essentially dominate an adversarially-trained baseline model, even on the attack the adversarial model was trained on, and have incidentally shown that the adversarially-trained model generalizes to at least one other attack less well than a deterministic baseline. Finally, we have shown that on ImageNet, CEB provides substantial gains over deterministic baselines in both validation accuracy and robustness to Common Corruptions, Natural Adversarial Examples, and targeted Projected Gradient Descent attacks. We hope these empirical demonstrations inspire further theoretical and practical study of the use of bottlenecking techniques to encourage improvements to both classical generalization and robustness.

heading: A EXPERIMENT DETAILS
text: Here we present additional technical details for the CIFAR-10 and ImageNet experiments.
A.1 CIFAR-10 EXPERIMENT DETAILS
We trained all of the models using Adam (Kingma & Ba, 2015) at a base learning rate of 10−3. We lowered the learning rate three times by a factor of 0.3 each time. The only additional trick to train the CIFAR-10 models was to start with ρ = 100, anneal down to ρ = 10 over 2 epochs, and then anneal to the target ρ over one epoch once training exceeded a threshold of 20%. This jump-start method is inspired by experiments on VIB in Wu et al. (2019). It makes it much easier to train models at low ρ, and appears to not negatively impact final performance.
For the 62×7 models, we used the data augmentation policies for CIFAR-10 found by AUTOAUG and trained the models for 800 epochs, lowering the learning rate by a factor of 10 at 400 and 600 epochs.
A.2 IMAGENET EXPERIMENT DETAILS
We follow the learning rate schedule for the ResNet 50 from Cubuk et al. (2018), which has a top learning rate of 1.6, trains for 270 epochs, and drops the learning rate by a factor of 10 at 90, 180, and 240 epochs. The only difference for all of our models is that we train at a batch size of 8192 rather than 4096. Similar to the CIFAR-10 models, in order to ensure that the ImageNet models train at low ρ, we employ a simple jump-start. We start at ρ = 100 and anneal down to the target ρ over 12,000 steps. The first learning rate drop occurs a bit after 14,000 steps. Also similar to the CIFAR 28×10 WRN experiments, none of the models we trained at ρ = 0 succeeded, indicating that RESNET-50 and WRN 50×2 both have insufficient capacity to fully learn ImageNet. We were able to train RESNET-152 at ρ = 0, but only by disabling L2 weight decay and using a slightly lower learning rate. Since that involved additional hyperparameter tuning, we don’t report those results here, beyond noting that it is possible, and that those models reached top-1 accuracy around 72%.

heading: B CEB EXAMPLE CODE
text: Here we give the core changes needed to make ResNet CEB models, based on the TPU-compatible ResNet implementation from the Google TensorFlow Team.
# In model.py: def resnet_v1_generator(block_fn, layers, num_classes, ...):
def model(inputs, is_training): # Build the ResNet model as normal up to the following lines: inputs = tf.reshape(
inputs, [-1, 2048 if block_fn is bottleneck_block else 512]) # Now, instead of the final dense layer, just return inputs, # which for ResNet50 models is a [batch_size, 2048] tensor. return inputs
Listing 1: Modifications to the model.py function.
# In resnet_main.py add the following imports and functions: import tensorflow_probability as tfp tfd = tfp.distributions
def ezx_dist(x): """Builds the encoder distribution, e(z|x).""" dist = tfd.MultivariateNormalDiag(loc=x) return dist
def bzy_dist(y, num_classes=1000, z_dims=2048): """Builds the backwards distribution, b(z|y).""" y_onehot = tf.one_hot(y, num_classes) mus = tf.layers.dense(y_onehot, z_dims, activation=None) dist = tfd.MultivariateNormalDiag(loc=mus) return dist
def cyz_dist(z, num_classes=1000): """Builds the classifier distribution, c(y|z).""" # For the classifier, we are using exactly the same dense layer # initialization as was used for the final layer that we removed # from model.py. logits = tf.layers.dense(
z, num_classes, activation=None, kernel_initializer=tf.random_normal_initializer(stddev=.01))
return tfd.Categorical(logits=logits)
def lerp(global_step, start_step, end_step, start_val, end_val): """Utility function to linearly interpolate two values.""" interp = (tf.cast(global_step - start_step, tf.float32)
/ tf.cast(end_step - start_step, tf.float32)) interp = tf.maximum(0.0, tf.minimum(1.0, interp)) return start_val * (1.0 - interp) + end_val * interp
Listing 2: Modification to the head of resnet main.py.
# Still in resnet_main.py, modify resnet_model_fn as follows: def resnet_model_fn(features, labels, mode, params): # Nothing changes until after the definition of build_network: def build_network(): # Elided, unchanged implementation of build_network.
if params['precision'] == 'bfloat16': # build_network now returns the pre-logits, so we'll change # the variable name from logits to net. with tf.contrib.tpu.bfloat16_scope():
net = build_network() net = tf.cast(net, tf.float32)
elif params['precision'] == 'float32': net = build_network()
# Get the encoder, e(z|x): with tf.variable_scope('ezx', reuse=tf.AUTO_REUSE): ezx = ezx_dist(net) # Get the backwards encoder, b(z|y): with tf.variable_scope('bzy', reuse=tf.AUTO_REUSE): bzy = bzy_dist(labels)
# Only sample z during training. Otherwise, just pass through # the mean value of the encoder. if mode == tf.estimator.ModeKeys.TRAIN: z = ezx.sample() else: z = ezx.mean()
# Get the classifier, c(y|z): with tf.variable_scope('cyz', reuse=tf.AUTO_REUSE): cyz = cyz_dist(z, params)
# cyz.logits is the same as what the unmodified ResNet model # would return. logits = cyz.logits
# Compute the individual conditional entropies: hzx = -ezx.log_prob(z) # H(Z|X) hzy = -bzy.log_prob(z) # H(Z|Y) (upper bound) hyz = -cyz.log_prob(labels) # H(Y|Z) (upper bound)
# I(X;Z|Y) = -H(Z|X) + H(Z|Y) # >= -hzx + hzy =: Rex, the residual information. rex = -hzx + hzy
rho = 3.0 # You should make this a hyperparameter. rho_to_gamma = lambda rho: 1.0 / np.exp(rho) gamma = tf.cast(rho_to_gamma(rho), tf.float32)
# Get the global step now, so that we can adjust rho dynamically. global_step = tf.train.get_global_step()
anneal_rho = 12000 # You should make this a hyperparameter. if anneal_rho > 0: # Anneal rho from 100 down to the target rho # over the first anneal_rho steps. gamma = lerp(global_step, 0, aneal_rho,
rho_to_gamma(100.0), gamma)
# Replace all the softmax cross-entropy loss computation with # the following line: loss = tf.reduce_mean(gamma * rex + hyz) # The rest of resnet_model_fn can remain unchanged.
Listing 3: Modifications to resnet model fn in resnet main.py.

Review summaries:
The paper modifies existing classifier architectures and training objective, in order to minimize "conditional entropy bottleneck" (CEB) objective, in attempts to force the representation to maximize the information bottleneck objective. Consequently, the paper claims that this CEB model improves general test accuracy and robustness against adversarial attacks and common corruptions, compared to the softmax + cross entropy counterpart. This claim is supported by experimental results on CIFAR-10 and ImageNet-C datasets.

This paper studied the effectiveness of Conditional Entropy Bottleneck (CEB) on improving model robustness. Three tasks are considered to demonstrate its effectiveness; generalization performance over clean test images, adversarially perturbed images, and images corrupted by various synthetic noises. The experiment results demonstrated that CEB improves the model robustness on all considered tasks over the deterministic baseline and adversarially-trained classifiers.