Protect the weak: Class focused online learning for adversarial training

1 INTRODUCTION . The susceptibility of neural networks to adversarial noise ( Goodfellow et al. , 2014 ; Szegedy et al. , 2013 ) has been a grave concern over the launch of such systems in real-world applications . Several techniques defending such attacks that optimize the average performance have been proposed ( Papernot et al. , 2016 ; Raghunathan et al. , 2018 ; Guo et al. , 2017 ; Madry et al. , 2017 ; Zhang et al. , 2019 ) . In response , even stronger attacks have been proposed ( Carlini & Wagner , 2016 ; Engstrom et al. , 2018 ; Carlini , 2019 ) . Indeed , recent studies demonstrate that regardless of the defense there exists an attack that can lower the average performance of the system ( Shafahi et al. , 2018 ) . In this work , we argue that the average performance is not the only criterion that is of interest for real-world applications . For classification , in particular , optimizing the average performance provides no guarantees for the “ weakest ” class . This is critical in scenarios where an attacker can pick the class adversarially in addition to the adversarial perturbation . It turns out that the worst performing class can indeed be much worse than the average in adversarial training . This difference is already present in clean training but we critically observe , that the gap between the average and the worst , is greatly exacerbated in adversarial training . This gap can already be observed on CIFAR10 where the accuracy across classes is far from uniform with 47 % average robust accuracy while the worst class is 14 % ( see Figure 1 ) . The effect is even more prevalent when more classes are present as in CIFAR100 where we observe that the worst class has zero accuracy while the best has 70 % ( see Appendix C where we include multiple other datasets ) . Despite the focus on adverarial training , we note that the same effect can be observed for robust evaluation after clean training ( see Figure 4 ) . This dramatic drop in accuracy for the weakest classes begs for different approaches than the classical empirical risk minimization ( ERM ) , which focuses squarely on the average loss . We suggest a simple alternative , which we dub class focused online learning ( CFOL ) , that can be plugged into existing adversarial training procedures . Instead of minimizing the average performance over the dataset we sample from an adversarial distribution over classes that is learned jointly with the model parameters . In this way we aim at becoming robust to an attacker that can adversarially choose the class . The focus of this paper is thus on the robust accuracy of the weakest classes instead of the average robust accuracy . Concretely , we make the following contributions : • We propose a simple solution which relies on the classical bandit algorithm from the online learning literature , namely the Exponential-weight algorithm for Exploration and Exploitation ( Exp3 ) ( Auer et al. , 2002 ) . The method is directly compatible with standard adversarial training procedures ( Madry et al. , 2017 ) , by replacing the empirical distribution with an adaptively learned adversarial distribution over classes . • We carry out extensive experiments comparing CFOL against three strong baselines across three datasets , where we consistently observe that CFOL improves the weakest classes . • We support the empirical results with high probability convergence guarantees for the worst class accuracy and establish direct connection with the conditional value at risk ( CVaR ) ( Rockafellar et al. , 2000 ) uncertainty set from the distributional robust optimization . Overview We define our new objective in Section 2 , followed by the description of the proposed method in Section 3 . In Section 3.1 we prove convergence guarantees . We then move to the empirical work in Section 4 , where we observe and resolve the issue of robust overfitting in the context of the worst class . Our main empirical findings are covered in Section 5 and we conclude in Section 6 . 1.1 RELATED WORK . Adversarial examples Goodfellow et al . ( 2014 ) ; Szegedy et al . ( 2013 ) are the first to make the important observation that deep neural networks are vulnerable to small adversarially perturbation of the input . Since then , there has been a growing body of literature addressing this safety critical issue , spanning from certified robust model ( Raghunathan et al. , 2018 ) , distillation ( Papernot et al. , 2016 ) , input augmentation ( Guo et al. , 2017 ) , to adversarial training ( Madry et al. , 2017 ; Zhang et al. , 2019 ) . We focus on adversarial training in this paper . While certified robustness is desirable , adversarial training remains one of the most successful defenses in practice . In a concurrent work Tian et al . ( 2021 ) independently observe the non-uniform accuracy over classes in adversarial training , further strengthening the case that lack of class-wise robustness is indeed an issue . However , they mainly focus on constructing an attack that can enlarge this disparity . Minimizing the maximum Connected to our method is that of focused online learning ( FOL ) ( Shalev-Shwartz & Wexler , 2016 ) , which similarly takes a bandit approach , but instead re-weights the distribution over the N training examples , independent of the class label . This naturally leads to a convergence rate in terms of the number of examples N instead of the number of classes k for which usually k N . We compare in more detail theoretically and empirically in Section 3.1 and Section 5 respectively . Interpolations between average and maximum loss have been considered in various other settings : for class imbalanced datasets ( Lin et al. , 2017 ) , in federated learning ( Li et al. , 2019 ) , and more generally the tilted empirical risk minimization ( Li et al. , 2020 ; Lee et al. , 2020 ) . Distributional robust optimization The accuracy over the worst class can be seen as a particular re-weighing the data distribution which adversarially assigns all weight to a single class . Worst case perturbation of the data distribution have more generally been studied under the framework of distributional robust stochastic optimization ( DRO ) ( Ben-Tal et al. , 2013 ; Shapiro , 2017 ) . Instead of attempting to minimizing the empirical risk on a training distribution P0 , this framework considers some uncertainty set around the training distribution U ( P0 ) and seeks to minimize the worst case risk within this set , supQ∈U ( P0 ) Ex∼Q [ ` ( x ) ] . A choice of uncertainty set , which has been given significant attention in the community , is conditional value at risk ( CVaR ) , which aims at minimizing the weighted average of the tail risk ( Rockafellar et al. , 2000 ; Levy et al. , 2020 ; Kawaguchi & Lu , 2020 ; Fan et al. , 2017 ; Curi et al. , 2019 ) . CVaR has been specialized to a re-weighting over class labels , namely labeled conditional value at risk ( LCVaR ) ( Xu et al. , 2020 ) . This was originally derived in the context of imbalanced dataset to re-balance the classes . It is still applicable in our setting and we thus provide a comparison . The original empirical work of ( Xu et al. , 2020 ) only considers the full-batch setting . We complement this by demonstrating LCVaR in a stochastic setting . In Duchi et al . ( 2019 ) ; Duchi & Namkoong ( 2018 ) they are , similarly to our setting , interested in uniform performance over various groups . However , these groups are assumed to be latent subpopulations , which introduces significant complications . It is thus concerned with a different setting , an example being training on a dataset implicitly consisting of multiple text corpora . CFOL can also be formulated in the framework of DRO by choosing an uncertainty set that can re-weight the k class-conditional risks . The precise definition is given in Appendix B.1 . We further establish a direct connection between the uncertainty sets of CFOL and CVaR that we make precise in Appendix B.1 , which also contains a summary of the most relevant related methods in Table 4 . 2 PROBLEM FORMULATION AND PRELIMINARIES . Notation The underlying data distribution is denoted by D with examples x ∈ Rd and classes y ∈ [ k ] . A given iteration is characterized by t ∈ [ T ] , while pyt indicates the yth index of the tth iterate . The indicator function is denoted with 1 { boolean } and unif ( n ) indicates the uniform distribution over n elements . An overview of the notation is provided in Appendix A . In classification , we are normally interested in minimizing the population risk E ( x , y ) ∼D [ ` ( θ , x , y ) ] over our model parameters θ ∈ Rp , where ` is some loss function of θ and example x ∈ Rd with an associated class y ∈ [ k ] . Madry et al . ( 2017 ) formalized adversarially training by modifying this objective with an adversarial perturbation to each example . That is , we instead want to find a parameterization θ of our predictive model which solves the following optimization problem : min θ L ( θ ) : = E ( x , y ) ∼D [ max δ∈S ` ( θ , x+ δ , y ) ] , ( 1 ) where each x is now perturbed by adversarial noise δ ∈ S ⊆ Rd . Common choices of S include norm-ball constraints as in Madry et al . ( 2017 ) or bounding some notion of perceptual distance ( Laidlaw et al. , 2020 ) . When the distribution over classes is uniform this is implicitly minimizing the average loss over all class . This does not guarantee high accuracy for the worst class as illustrated in Figure 1 , since we only know with certainty that max ≥ avg . Instead , we will focus on a different objective , namely minimizing the worst class-conditioned risk : min θ max y∈ [ k ] { Ly ( θ ) : = Ex∼pD ( ·|y ) [ max δ∈S ` ( θ , x+ δ , y ) ] } . ( 2 ) This follows the philosophy that “ a chain is only as strong as its weakest link ” . In our particular setting it models a scenario where the attacker can adversarially choose what class the model is evaluated on . In safety critical application , such as autonomous driving , modeling even just a single class wrong can still have catastrophic consequences . As the maximum in Equation 2 is a discrete maximization problem its treatment requires more care . We will take a common approach and construct a convex relaxation to this problem in Section 3 . 3 METHOD . Since we do not have access to the true distribution D , we will instead minimize over the provided empirical distribution . Let Ny be the set of data point indices for class y such that the total number of examples is N = ∑k y=1 |Ny| . Then , we are interested in minimizing the maximum empirical class-conditioned risk , max y∈ [ k ] L̂y ( θ ) : = 1 |Ny| ∑ i∈Ny max δ∈S ` ( θ , xi + δ , y ) . ( 3 ) We relax this discrete problem to a continuous problem over the simplex ∆k , max y∈ [ k ] L̂y ( θ ) ≤ max p∈∆k k∑ y=1 pyL̂y ( θ ) . ( 4 ) Note that equality is attained when p is a dirac on the argmax over classes . Equation 4 leaves us with a two-player zero-sum game between the model parameters θ and the class distribution p. From the perspective of p it is just a linear objective , albeit adversarially picked by the model . This immediately let us use a no-regret algorithm under simplex constraint , namely Hedge ( Freund & Schapire , 1997 ) : wty = w t−1 y − ηL̂y ( θt ) , pty = exp ( wty ) / k∑ y=1 exp ( wty ) . ( Hedge ) To show convergence for ( Hedge ) the loss needs to satisfy certain assumptions . In our case of classification , the loss is the zero-one loss ` ( θ , x , y ) = 1 [ hθ ( x ) 6= y ] , where hθ ( · ) is the predictive model . Hence , the loss is bounded , which is a sufficient requirement . Note that ( Hedge ) relies on zero-order information of the loss , which we indeed have available . However , in the current form , ( Hedge ) requires so called full information over the k dimensional loss vector . In other words , we need to compute L̂y ( θ ) for all y ∈ [ k ] at every iteration . When accessing the data through mini-batches , it is strictly not possible to compute even an estimate for each y if the number of classes is larger than the batch size . Following the seminal work of Auer et al . ( 2002 ) , we instead construct an unbiased estimator of the k dimensional loss vector L̂ ( θt ) : = ( L̂1 ( θt ) , ... , L̂k ( θt ) ) > based on a sampled class yt from some distribution yt ∼ pt . This stochastic formulation further lets us estimate the class conditioned risk L̂yt ( θ t ) with an unbiased sample i ∼ unif ( |Nyt | ) . This leaves us with the following estimator , L̃ty = { Ly , i ( θ t ) /pty y = y t 0 otherwise ( 5 ) Algorithm 1 : Class focused online learning ( CFOL ) Algorithm parameters : a step rule ModelUpdate for the model satisfying Assumption 1 , adversarial step-size η > 0 , uniform mixing parameter γ = 1/2 , and the loss Ly , i ( θ ) . Initialization : Set w0 = 0 such that q0 and p0 are uniform . foreach t in 0 .. T do yt ∼ pt ; // sample class it ∼ unif ( |Nyt | ) ; // sample uniformly from class θt+1 = ModelUpdate ( θt , Lyt , it ( θ t ) ) ; // update model parameters L̃ty = 1 { y=yt } Ly , it ( θ t ) /pty ∀y ; // construct estimator wt+1y = w t y − ηL̃ty ∀y ; // update the adv . class distribution qt+1y = exp ( wt+1y ) / ∑k y=1 exp ( wt+1y ) ∀y ; pt+1y = γ 1 k + ( 1− γ ) q t+1 y ∀y ; end foreach where Ly , i ( θ ) : = maxδ∈S ` ( θ , xi + δ , y ) . It is easy to verify that this estimator is unbiased , Ey∼p [ L̃ty ] = ( 1− pty ) · 0 + pty · L̂y ( θ t ) pty = L̂y ( θ t ) . ( 6 ) For ease of presentation the estimator only uses a single sample but this can trivially be extended to a mini-batch where classes are drawn i.i.d . from pt . We could pick pt to be the learned adversarial distribution , but it is well known that this can lead to unbounded regret if some pty is small ( see Appendix A for definition of regret ) . Auer et al . ( 2002 ) resolves this problem with Exp3 , which instead learns a distribution qt , then mixes it with a uniform distribution , to eventually sample from pty = γ 1 k + ( 1 − γ ) q t y where γ ∈ ( 0 , 1 ) . Intuitively , this enforces exploration . In the general case this γ needs to be picked carefully and small enough , but similarly to Shalev-Shwartz & Wexler ( 2016 ) , we show that γ = 1/2 suffice in our setting . Our algorithm thus updates qt+1 with ( Hedge ) using the estimator L̃t with sample drawn from pt and subsequently plays pt+1 . CFOL in conjunction with the simultaneous update of the minimization player can be found in algorithm 1 with an example of a learned distribution pt in Figure 2 . Practically , the scheme bears negligible computational overhead over ERM since the softmax required to sample is of the same dimensionality as the softmax used in the forward pass through the model . This computation is negligible in comparison with backpropagating through the entire model . For further details on the implementation we refer to Appendix C.4

Summary: This paper points out that there could be applications where the worst class accuracy could be critical -- for example, the class with the worst class accuracy can be the main target for the attacks and that the difference between the worst robust test accuracy of a class and the average robust test accuracy can be large (also pointed out independently in Tian et al. (2021)).   Motivated by this, the paper proposes a solution to use the Exponential-weight algorithm (Auer et al., 2002), where, as the empirical distribution is taken the adversarial distribution over classes, resulting in a method called Class Focused Online Learning (CFOL). The proposed CFOL method is based on the Focused Online Learning (FOL) (Shalev-Shwartz & Wexler, 2016) method, which does not take the class label into account.  The authors show empirically that CFOL consistently improves the worst class accuracy on CIFAR-10, CIFAR-100, and STL10. Moreover, the authors provide high probability convergence guarantees for the worst class accuracy.