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: This paper justifies why models trained with adversarial training are good generative models by proposing a probabilistic framework. The paper is heavily influenced by the JEM paper and proposes a generalized form of analysis in that paper to include the unsupervised scenario. The generalization defines a joint distribution $p(x, z)$ over data $x$ and representations $z$. For the supervised case, each class is represented by a "center" $w_y$ and for the unsupervised one, $p(x, z) = p(x) p(z|x)$ and $p(z|x)$ is reparametrized with augmentation.  Then the authors use this probabilistic framework to analyze the adversarial training. Specifically, they show PGD is a biased form of the Langevin dynamics in their formulation. Then they show how the maximum likelihood training of their probabilistic framework is similar to adversarial training. As a result, they conclude adversarial training leads to a model with high generative power. Finally, they propose refined sampling strategies from adversarially trained models.  In the last part of the paper, they use their framework to propose an unsupervised adversarial training method followed by sampling algorithms from the produced models.