CEB Improves Model Robustness

1 INTRODUCTION . 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 . 1.1 CONTRIBUTIONS . 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 . 2 BACKGROUND . 2.1 INFORMATION BOTTLENECKS . 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 . 2.2 IMPLEMENTING A CEB MODEL . 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 . 2.3 ADVERSARIAL ATTACKS AND DEFENSES . 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 .

Summary: 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.