A Unified Contrastive Energy-based Model for Understanding the Generative Ability of Adversarial Training

1 INTRODUCTION . Adversarial Training ( AT ) is one of the most effective approaches developed so far to improve the robustness of deep neural networks ( DNNs ) ( Madry et al. , 2018 ) . AT solves a minimax optimization problem , with the inner maximization generating adversarial examples by maximizing the classification loss , and the outer minimization finding model parameters by minimizing the loss on adversarial examples generated from the inner maximization . Recently , researchers have noticed that such robust classifiers obtained by AT are able to extract features that are perceptually aligned with humans ( Engstrom et al. , 2019 ) . Furthermore , they are able to synthesize realistic images on par with stateof-the-art generative models ( Santurkar et al. , 2019 ) . Nevertheless , it is still a mystery why AT is able to learn more semantically meaningful features and turn classifiers into generators . Besides , AT needs the labeled data { ( xi , yi ) } for training while canonical deep generative models do not , e.g. , VAE ( Kingma & Welling , 2014 ) and GAN ( Goodfellow et al. , 2015 ) only require { xi } . Thus , it is worth exploring if it is possible to train a robust model without labeled data . Several recent works ( Jiang et al. , 2020 ; Kim et al. , 2020 ; Ho & Vasconcelos , 2020 ) have proposed unsupervised AT by adversarially attacking the InfoNCE loss ( Oord et al. , 2018 ) ( a widely used objective in unsupervised contrastive learning ) , which indeed improves the robustness of contrastive encoders . However , a depth investigation and understanding for unsupervised AT is still missing . To address the above issues , in this work , we propose a unified probabilistic framework , Contrastive Energy-based Models ( CEM ) , that provides a principled understanding on the robustness and the generative ability of different training paradigms . Specifically , we make the following contributions : • Demystifying adversarial training and sampling . We firstly propose a probabilistic interpretation for AT , that is , it is inherently a ( biased ) maximum likelihood training of the corresponding energy-based model , which explains the generative ability of robust models learned by AT . Inspired by this , we propose some novel sampling algorithms with better sample quality than previous methods . • A unified probabilistic framework . Based on the understanding above , we propose Contrastive Energy-based Model ( CEM ) that incorporates both supervised and unsupervised learning paradigms . Our CEM provides a unified probabilistic understanding of previous standard and adversarial training methods in both supervised and unsupervised learning . • Principled unsupervised adversarial training and sampling . Specifically , under our proposed CEM framework , we establish the equivalence between the importance sampling of CEM and the InfoNCE loss of contrastive learning , which enables us to design principled adversarial sampling for unsupervised learning . Notably , we show that the sampling methods derived from our framework achieve state-of-the-art sample quality ( 9.61 Inception score ) with unsupervised robust models , which is comparable to both the supervised counterparts and other state-of-the-art generative models . 2 RELATED WORK . Robust generative models . Researchers recently notice that features extracted by robust classifiers are perceptually aligned with humans , while standard classifiers are not ( Engstrom et al. , 2019 ; Kaur et al. , 2019 ) . Furthermore , Santurkar et al . ( 2019 ) show that we can also generate images of high quality with robust classifiers by iteratively updating from a randomly sampled noise . They show the sample quality of robust classifiers is comparable to state-of-the-art generative models like BigGAN ( Brock et al. , 2018 ) . Contrastive learning . Oord et al . ( 2018 ) firstly propose unsupervised contrastive learning by maximizing a tractable lower bound on mutual information ( MI ) , i.e. , the negative InfoNCE loss . However , later works find that the lower bounds degrade a lot with a large MI , and the success of these methods can not be attributed to the properties of MI alone ( Poole et al. , 2019 ; Tschannen et al. , 2020 ) . Our work provides an alternative understanding of unsupervised contrastive learning as importance sampling of an energy-based model , which also enables us to characterize the limitations of existing methods from a new perspective . In fact , contrastive learning can also be seen as a general learning framework beyond the unsupervised scenarios . For example , SupContrast ( Khosla et al. , 2020 ) extends contrastive learning to supervised scenarios . Our work further bridges supervised , unsupervised and adversarial contrastive learning with a unified probabilistic framework . 3 CEM : A UNIFIED PROBABILISTIC FRAMEWORK . Inspired by previous work that bridges discriminative models with energy-based models ( Grathwohl et al. , 2019 ) , in this work , we propose a unified framework , called Contrastive Energy-based Model ( CEM ) , that incorporates both supervised and unsupervised scenarios . Our proposed CEM is a special kind of Energy-based Models ( EBMs ) that models the joint distribution pθ ( u , v ) over two variables ( u , v ) with a similarity function fθ ( u , v ) defined in a contrastive form , pθ ( u , v ) = exp ( fθ ( u , v ) ) Z ( θ ) , ( 1 ) where Z ( θ ) = ∫ exp ( fθ ( u , v ) ) dudv is the corresponding partition function . In other words , in CEM , a pair of samples ( u , v ) has higher probability if they are more alike . In particular , it can be instantiated into the two following variants under different learning scenarios . Parametric CEM . In the supervised scenario , we specify the Parametric CEM ( P-CEM ) that models the joint distribution pθ ( x , y ) of data x and label y in the following form , pθ ( x , y ) = exp ( fθ ( x , y ) ) Z ( θ ) = exp ( gθ ( x ) > wy ) Z ( θ ) , ( 2 ) where gθ : Rn → Rm denotes the encoder , g ( x ) ∈ Rm is the representation of x , and wk ∈ Rm refers to the parametric cluster center of the k-th class . Denote the linear classification weight as W = [ w1 , · · · , wK ] and the logit vector as h ( x ) = g ( x ) > W , we can see the equivalence between P-CEM and JEM ( Grathwohl et al. , 2019 ) as fθ ( x , y ) = gθ ( x ) > wy = hθ ( x ) [ y ] . ( 3 ) Non-Parametric CEM . In the unsupervised scenario , we do not have access to labels , thus we instead model the joint distribution between data x and representation z as pθ ( x , z ) = exp ( fθ ( x , z ) ) Z ( θ ) = exp ( gθ ( x ) > z ) Z ( θ ) , ( 4 ) and the corresponding likelihood gradient is ∇θEpd ( x , z ) log pθ ( x , z ) = Epd ( x , z ) ∇θfθ ( x , z ) −Epθ ( x̂ , ẑ ) ∇θfθ ( x̂ , ẑ ) . ( 5 ) In contrastive to P-CEM that incorporates parametric cluster centers , the joint distribution of NPCEM ( Non-Parametric CEM ) is directly defined based on the similarity between instances . We define the joint data distribution pd ( x , z ) = pd ( x ) pd ( z|x ) through re-parameterization , z = fθ ( t ( x ) ) , t u.a.r.∼ T , x ∼ pd ( x ) , ( 6 ) where u.a.r . denotes sampling uniformly at random and T refers to the user-define pretext , e.g. , a set of data augmentation operators T = { t : Rn → Rn } 1 . 4 SUPERVISED SCENARIO : REDISCOVERING ADVERSARIAL TRAINING AS MAXIMUM LIKELIHOOD TRAINING . In this section , we investigate why robust models have a good generative ability . The objective of AT is to solve the following minimax optimization problem : min θ Epd ( x , y ) [ max ‖x̂−x‖p≤ε ` CE ( x̂ , y ; θ ) ] , where ` CE ( x̂ , y ; θ ) = − log pθ ( y|x̂ ) . ( 7 ) The inner maximization problem is to find an adversarial example x̂ within the ` p-norm ε-ball around the natural example x that maximizes the CE loss . While the outer minimization problem is to find model parameters that minimize the loss on the adversarial examples x̂ . 4.1 MAXIMIZATION PROCESS . For the inner maximization problem , Projected Gradient Descent ( PGD ) ( Madry et al. , 2018 ) is the commonly used method , which generates the adversarial example x̂ by maximizing the CE loss2 ( i.e. , minimizing the log conditional probability ) starting from x̂0 = x : x̂n+1 = x̂n + α∇x̂n ` ( x̂n , y ; θ ) = x̂n − α∇x̂n log pθ ( y|x̂n ) = x̂n + α∇x̂n [ log K∑ k=1 exp ( fθ ( x̂n , k ) ) ] − α∇x̂nfθ ( x̂n , y ) , ( 8 ) while the Langevin dynamics for sampling P-CEM starts from random noise x̂0 = δ and updates with x̂n+1 = x̂n + α∇x̂ log pθ ( x̂n ) + √ 2α · ε ( 9 ) = x̂n + α∇x̂n [ log K∑ k=1 exp ( fθ ( x̂n , k ) ) ] + √ 2α · ε. Eqns . 8 and 9 both have a positive logsumexp gradient ( the second term ) to push up the marginal probability pθ ( x̂ ) . As for the third term , PGD starts from a data point ( x , y ) such that it requires the repulsive gradient to be away from the original data point and do the exploration in a local region . Langevin dynamics instead starts from a random noise and an additive noise ε is injected for exploration . Comparing PGD and Langevin . Following the above analysis , the maximization process in AT can be seen as a ( biased ) sampling method that draws samples from the corresponding probabilistic 1For the data pair , we detach the gradient of z w.r.t . θ and assume that x and t ( x ) have the same marginal distribution pd ( x ) . 2Note that we omit the projection operation and the gradient re-normalization steps . model pθ ( x̂ ) . Compared to Langevin dynamics , PGD imposes specific inductive bias for sampling . With the additional repulsive gradient and ε-ball constraint , it explicitly encourages the samples to be misclassified around the original data points . In practice , adversarial training with such adversarial examples is generally more stable than training JEM with Langevin samples , which indicates that PGD attack is a competitive alternative for the negative sampling method for JEM training . 4.2 MINIMIZATION PROCESS . To begin with , the gradient of the joint log likelihood for P-CEM can be written as follows : ∇θEpd ( x , y ) log pθ ( x , y ) =Epd ( x , y ) ∇θfθ ( x , y ) −Epθ ( x̂ , ŷ ) ∇θfθ ( x̂ , ŷ ) =Epd ( x , y ) ∇θfθ ( x , y ) −Epθ ( x̂ ) pθ ( ŷ|x̂ ) ∇θfθ ( x̂ , ŷ ) , ( 10 ) where ( x , y ) ∼ pd ( x , y ) denotes the positive data pair , and ( x̂ , ŷ ) ∼ pθ ( x̂ , ŷ ) denotes the negative sample pair . As discussed above , the adversarial examples x̂ generated by the maximization process can be regarded as negative samples , and ŷ ∼ pθ ( ŷ|x̂ ) denotes the predicted label of x̂ . To see how the maximum likelihood training of P-CEM is related to the minimization process of AT , we add an interpolated adversarial pair ( x̂ , y ) into Eq . 10 and decompose it as the consistency gradient and the contrastive gradient : ∇θEpd ( x , y ) log pθ ( x , y ) = Epd ( x , y ) ⊗ pθ ( x̂ , ŷ ) [ ∇θfθ ( x , y ) −∇θfθ ( x̂ , ŷ ) ] =Epd ( x , y ) ⊗ pθ ( x̂ , ŷ ) [ ∇θfθ ( x , y ) −∇θfθ ( x̂ , y ) ︸ ︷︷ ︸ consistency gradient +∇θfθ ( x̂ , y ) −∇θfθ ( x̂ , ŷ ) ︸ ︷︷ ︸ contrastive gradient ] . ( 11 ) Next , we show that the two parts correspond to two effective mechanisms developed in the adversarial training literature . AT loss . As the two sample pairs in the contrastive gradient share the same input x̂ , we can see that the contrastive gradient can be written equivalently as Epd ( x , y ) ⊗ pθ ( x̂ , ŷ ) [ ∇θfθ ( x̂ , y ) −∇θfθ ( x̂ , ŷ ) ] =Epd ( x , y ) ⊗ pθ ( x̂ ) [ ∇θfθ ( x̂ , y ) − Epθ ( ŷ|x̂ ) ∇θfθ ( x̂ , ŷ ) ] =Epd ( x , y ) ⊗ pθ ( x̂ ) ∇θ log pθ ( y|x̂ ) , ( 12 ) which is exactly the negative gradient of the robust CE loss ( AT loss ) in Eq . 7 , in other words , gradient ascent with the contrastive gradient is equivalent to gradient descent w.r.t . the AT loss . Regularization . As for the consistency gradient , original AT ( Madry et al. , 2018 ) simply ignores it . Its variant TRADES ( Zhang et al. , 2019 ) instead proposes the KL regularization KL ( p ( ·|x̂ ) ‖p ( ·|x ) ) that regularizes the consistency of the predicted probabilities on all classes , whose optimum implies that p ( ·|x̂ ) = p ( ·|x ) → fθ ( x , y ) = fθ ( x̂ , y ) . Comparing AT and JEM training paradigms . The above analysis indicates that the minimization objective of AT is closely related to the maximum likelihood training of JEM ( Grathwohl et al. , 2019 ) . Compared to JEM that decomposes the joint likelihood into an unconditional model pθ ( x ) and a discriminative model pθ ( y|x ) , the decomposition of AT in Eq . 10 instead stabilizes training by introducing an intermediate adversarial pair ( x̂ , y ) that bridges the positive pair ( x , y ) and the negative pair ( x̂ , ŷ ) . Besides , it can inject the adversarial robustness bias by regularizing the consistency gradient . Together with our analysis on the maximization process , we show that AT is a competitive alternative for training JEM ( a generative model ) with more stable training behaviors . That explains why robust models with AT are also generative .

Summary: The submission proposes a unified probabilistic framework to illustrate the generative capability of adversarial training. It also offers a unified perspective of adversarial training and sampling from both a supervised learning setting and an unsupervised learning setting. The proposed adversarial sampling strategy from the method is extensively demonstrated on different benchmarks, showing better quality compared with existing related works.