A Distributional Robustness Perspective on Adversarial Training with the $\infty$-Wasserstein Distance

1 INTRODUCTION . We call adversarial example an input which is a human-imperceptible -perturbation1 compared to a real example that results in an incorrect classification from a classifier [ Goodfellow et al. , 2014 ; Sun et al. , 2018 ; Athalye et al. , 2018 ; Santurkar et al. , 2019 ; Nguyen et al. , 2015a ; Kurakin et al. , 2016 ; Moosavi-Dezfooli et al. , 2016 ] . The particularity of these examples , which justifies their study , is the fact that most of the so-called adversarial examples are misclassified by a large majority of state-of-the-art neural networks . Long seen as bugs , Santurkar et al . [ 2019 ] asserts that the existence of adversarial examples is explained by the presence of easy-to-perturb patterns within the data distribution that are not perceptible by humans but useful for the classification task . In response to this discovery , different methods have been developed to train robust classifiers and to craft adversarial attacks by determining optimal perturbations [ Nguyen et al. , 2015b ; Papernot et al. , 2016a ; Goodfellow et al. , 2014 ; Papernot et al. , 2016b ; Tramèr et al. , 2017 ; Madry et al. , 2017 ] . Among these training methods , adversarial training [ Goodfellow et al. , 2014 ; Tramer et al. , 2017 ; Madry et al. , 2018 ] has settled as one of the strongest baselines to train robust classifiers . This method is relatively simple : it consists of training the classifier directly on batches of adversarial examples , leading to the following optimization problem [ Madry et al. , 2018 ] : min θ∈Rd E ( x , y ) ∼pdata [ max‖x̃−x‖∞≤ ` ( fθ ( x̃ , y ) ) ] ( 1 ) where ` is usually the cross-entropy loss and pdata is the dataset distribution . Our goal is to rigorously connect adversarial training and robustness by providing a distributional robustness 1Adversarial example can be defined with respect to any notion of distance that captures the fact that an -perturbation is imperceptible by a human . In this work , for the sake of simplicity , we focus on the ` ∞ norm , which is the most common one . perspective [ Delage and Ye , 2010 ; Ben-Tal and Nemirovski , 1998 ; Sinha et al. , 2018 ] for adversarial training . Distributional robustness is a framework to study predictive models that aim at being robust against distribution shift . Given a set of possible distribution P that a given predictive model fθ could encounter the distributionally robust classification task is : min θ∈Rd max padv∈P Ex̃∼padv [ ` ( fθ ( x̃ , y ) ) ] ( 2 ) Note that while standard adversarial training aims to use the point-wise optimal adversarial x̃j examples ( i.e . for each natural input xi ) , our perspective considers the optimal adversarial distribution . Working at the scale of distributions leads us to draw inspiration from the Kantorovich relaxation in Optimal Transport theory . To do so , we will define the set P using coupling measures between natural and adversarial distributions with the help of∞-∞-Wasserstein distance . Contributions . Our contributions are three-fold : First , we provide a strong connection between adversarial training and adversarial robustness and distributional robustness by using some tools from optimal transport , we call this formulation adversarial transport . Our second contribution is to show that , in this specific setting , there exists , against any given classifier fθ , an optimal distribution of adversarial examples and to provide a closed-form solution for this distribution . Our third contribution is to use these theoretical insights to come up with a practical training method using Langevin Monte-Carlo sampling to jointly find the optimal classifier and the optimal distribution of adversarial examples . By using this new technique for adversarial training we obtain robust classifiers outperforming the standard baseline [ Madry et al. , 2018 ] in terms of robustness and clean accuracy . Moreover , this training technique provides a speed-up of respectively ×200 for MNIST and ×8 for CIFAR in terms of training time . Related Work . Among the various works about defenses and attacks in the context of adversarial examples , Adversarial Training [ Madry et al. , 2018 ] is the most common baseline . However , in this procedure , adversarial examples are generated independently of the others by searching for the optimal perturbation for each one . Thus , AT seeks point-wise optimality and not global optimality . Here , the goal is to make our classifier robust to any unknown adversarial example distribution . It justifies the consideration of the DRO framework . DRO enables a mathematical formulation for dealing with uncertainty in complex systems [ Delage and Ye , 2010 ; Lam , 2018 ; Rahimian and Mehrotra , 2019 ] . In ML , DRO tries to minimize the loss over the worst-case distribution in a neighborhood of the observed training data distribution [ Duchi and Namkoong , 2016 ; 2018 ; Chouzenoux et al. , 2019 ; Duchi et al. , 2016 ] . Thus , DRO is pertinent in the case of adversarial examples and allows to bring a principled distributional perspective on AT . Within the DRO framework , one crucial point is to properly define the uncertainty set which contains the underlying data distribution . The importance of this set lies in the fact that it allows to reduce the computation time and the search area of the underlying distribution [ Delage and Ye , 2010 ; Ben-Tal and Nemirovski , 1998 ; Lam , 2018 ] . Different methods have been proposed to construct uncertainty sets . Among them , one considers the help of OT operators and defines the set with all the probability measures such as the p-Wasserstein distance from the observed distribution is less than a given constant ( p-Wasserstein ball ) [ Blanchet and Murthy , 2019 ] . Some works have started to study p-Wasserstein DRO in the case of adversarial examples . Zheng et al . [ 2019 ] defines Distributionally Adversarial Attacks ( DAA ) which is a generalization of PGD attacks on the space of probability measures with W2-DRO and [ Staib and Jegelka , 2017 ] did the Wp-DRO generalization of AT . E.Wong et al . [ 2019 ] also considers the Wp ball using an explicit formulation of the uncertainty set with coupling matrices . However , while the previous works modified the AT problem—the Wp constraint on the adversarial distribution does not correspond to the distribution of adversarial examples computed in AT—our work aims at overcoming that . The original AT problem stays a the scale of the example and considers OT to deal with the displacement of pixels . This perspective leads to the discovery that the standard AT problem corresponds to an∞-Wasserstein DRO problem with the ` ∞ underlying geometry . In the meantime , as in our paper , entropic regularization is used to find closed forms of the coupling matrix . To extend this understanding of adversarial examples in the case of distributions we generalize this uncertainty set with the ∞-∞-Wasserstein distance and add entropic regularization . 2-∞- Wasserstein DRO has already started to receive consideration in the case of AT [ Gao et al. , 2017 ] . Nevertheless , the use of the ` ∞ within the∞-Wasserstein distance allows both to consider the DRO problem at the dataset scale but also to ensure that the adversarial condition is respected for each element of the adversarial distribution . Together with the use of coupling matrices in our formulation , we found a concave and tractable formulation of our problem which results in a different technique to sample from this distribution ( namely Langevin Monte Carlo ) . Finally , we propose to keep them in memory to reuse the past samples for the update of the robust classifier . Notations . In this work , we will consider a measurable finite dimensional input space X ⊂ Rd . We will note B ( X ) the Borel σ-algebra of X andM ( X ) 1+ the set of all probability measures on X . If not precised , a map T : X → X is considered measurable . We , recall that given a separable metric space X s.t . any probability measure on X is a Radon measure , for T : X → X , the push-forward measure of β = T ] α ∈ M ( X ) 1+ of some α ∈ M ( X ) 1+ satisfies : ∀h ∈ C ( X ) , ∫ X h ( x̃ ) dβ ( x̃ ) =∫ X h ( T ( x ) ) dα ( x ) . 2 ADVERSARIAL TRAINING , TRANSPORT , AND DISTRIBUTIONAL ROBUSTNESS . In this work , we say that x̃ ∈ X is an adversarial example of f at the data point ( x , y ) if ‖x̃−x‖∞ ≤ [ Goodfellow et al. , 2014 ] . Given a loss function ` , an optimal adversarial example x̃ aims at maximizing the following optimization problem : x̃ ∈ arg max ‖x̃−x‖∞≤ ` ( fθ ( x̃ ) , y ) . ( 3 ) In practice , a large value for ` ( fθ ( x̃ ) , y ) will yield a misclassification of the example y . By considering the adversarial training problem ( 1 ) , one can notice that the distribution of input-output pairs seen by the classifier fθ depends on the value of its own parameters . We will call this distribution padv . Such a distribution can be described as the pushforward of a transport map T that respect the adversarial example constraint ‖T ( x ) − x‖∞ ≤ , ∀x ∈ X . In other words we have , padv = T ] pdata : ( x̃ , y ) ∼ padv ⇔ ∃ ( x , y ) ∼ pdata s.t . x̃ = T ( x ) and ‖x̃− x‖∞ ≤ . ( 4 ) This remark leads us to our first proposition that explicitly proposes a distributional robust optimization ( DRO ) formulation of the adversarial training problem . Proposition 2.1 ( DRO formulation of adversarial training ) . The adversarial training problem ( 1 ) can be formulated as a DRO problem ( 2 ) where : P = { padv : ∃T : X → X , padv = T ] pdata and ‖T ( x ) − x‖∞ ≤ , ∀x ∈ X } ( 5 ) Interestingly , one can draw parallels with Monge ’ s formulation of optimal transport ( OT ) , were given an initial measure µ and a target measure ν , the task is to solve the following problem inf T ∫ X c ( x , T ( x ) ) dµ ( x ) subject to ν = T ] µ . ( 6 ) On the one hand , ( 6 ) and the DRO problem defined in Proposition 2.1 are drastically different problems in essence . It can be explained by the difference between the constraints on the transportation map and the fact that while Equation 6 aims at minimizing a function of the transportation map T , Equation 2 aims at maximizing a function of the transported distribution padv . On the other hand , it is insightful to notice how both problems depend on a set of transportation maps to connect the initial and the target distributions . The limitations of Monge ’ s problem are well known in the OT community : due to its non-linear constraint , the problem is often intractable and may not have a solution ( i.e . the infimum may not be achieved ) . That is why Kantorovich [ 1940 ] , proposed to relax ( 6 ) into a problem with linear constraints by considering coupling measures ( an approach that allows a mass splitting from a natural example xi to different locations x̃j ) . Such a relaxation had the advantage to provide a better understanding of the original problem leading to novel theoretical guarantees ( e.g. , the existence of solutions for Kantorovich ’ s problem , duality , and in some situations equality between Monge ’ s problem and Kantorovich ’ s problem ) as well as tractable algorithms [ Sinkhorn , 1964 ; Cuturi , 2013b ] . In a similar vein , we believe our perspective on adversarial examples that combine ideas from distributional robustness and adversarial transport leads to a better understanding of the problem of adversarial training . Thus , we will use that insight to consider coupling measures to relax the adversarial training problem ( 1 ) . It will lead to a new DRO formulation for adversarial training that we call adversarial transport .

Summary: - This paper proposes using $\infty$-Wasserstein distributional robustness for improving model robustness. Specifically, to obtain a tractable solution, it adds an entropic regularization term to the primal form of $\infty$-Wasserstein distributional robustness. The Langevin algorithm is used to sample adversarial examples for adversarial training.  - Although the experimental results show the merit of the proposed method, the experiments are still humble with a lack of experiments on more impactful architectures such as WideResNet and comparison to other SOTA adversarial training methods.