Demystifying Adversarial Training via A Uniﬁed Probabilistic Framework
Abstract
Adversarial Training (AT) is known as an effective approach to enhance the robustness of deep neural networks.
Recently researchers notice that robust models with AT have good generative ability and can synthesize realistic images, while the reason behind it is yet under-explored.
In this paper, we demystify this phenomenon by developing a uniﬁed probabilistic framework, called
Contrastive Energy-based Models (CEM).
On the one hand, we provide the ﬁrst probabilistic characterization of AT through a uniﬁed understanding of robustness and generative ability.
On the other hand, our CEM can also naturally generalize AT to the unsupervised scenario and develop principled unsupervised AT methods.
Based on these, we propose principled adversarial sampling algorithms in both supervised and unsupervised scenarios.
Experiments show our sampling algorithms signiﬁcantly improve the sampling quality and achieves an Inception score of 9.61 on
CIFAR-10, which is superior to previous energybased models and comparable to state-of-the-art generative models.
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 classiﬁcation loss and the outer minimization ﬁnding model parameters by minimizing the loss on adversarial examples generated from the inner maximization.
Recently, researchers have noticed that such robust classiﬁers 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 state-of-the-art Accepted by the ICML 2021 workshop on A Blessing in Disguise: The Prospects and Perils of Adversarial Machine Learning. generative models (Santurkar et al., 2019).
Nevertheless, it is still a mystery why AT is able to learn more semantically meaningful features and turn classiﬁers into generators.
Besides, as generative methods, AT requires labeled data { ( x i , y i ) } while canonical generative models only require unlabeled data, e.g. , VAE (Kingma & Welling, 2014) and
GAN (Goodfellow et al., 2015).
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 & Nvasconcelos, 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.
In this paper, we address the above issues by proposing a uniﬁed framework with the following contributions:
• A uniﬁed probabilistic framework.
We propose Contrastive Energy-based Model (CEM) that incorporates both supervised and unsupervised learning paradigms.
Our CEM provides a uniﬁed probabilistic understanding of previous standard and adversarial training methods in both supervised and unsupervised learning.
• Demystifying AT.
Based on CEM, we propose the ﬁrst 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.
• Principled unsupervised AT.
Speciﬁcally, 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 training and sampling for unsupervised learning.
• State-of-the-art adversarial sampling algorithms.
Inspired by this, we propose some novel sampling algorithms with better sample quality than previous methods.
Notably, we show that our sampling methods 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. A Uniﬁed Probabilistic Framework
Our proposed CEM is a special kind of Enery-Based Models (EBMs) (LeCun et al., 2006) that models the joint distribution p θ ( u , v ) with a contrastive similarity function f θ ( u , v ) ,
[Equation]
where u , v are two random variables, θ are model parameters and Z ( θ ) =
(cid:82) exp ( f θ ( u , v ))
d u d v is the corresponding partition function.
In CEM, the log probability log p θ ( u , v ) is proportional to the similarity between u and v , as measured by the similarity function f θ ( u , v ) .
Parametric CEM.
In the supervised scenario, we specify the Parametric CEM (P-CEM) that models the joint distribution p θ ( x , y ) between data x ∈ R n and label y ∈
[ K ] , p θ
[Equation]
x [Equation] [Equation] where g θ : R n → R m denotes the encoder, g ( x ) ∈ R m is the representation of x , and w k ∈ R m refers to the parametric cluster center of the k -th class.
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
[Equation]
and the corresponding likelihood gradient is
[Equation]
In contrastive to P-CEM that incorporates parametric cluster centers, the joint distribution of NP-CEM (Non-Parametric
CEM) is directly deﬁned based on the similarity between instances (Oord et al., 2018).
We deﬁne the joint data distribution p d ( x , z )
=
p d ( x ) p d ( z | x ) through re-parameterization, z = f θ ( t ( x )) , where t ∼T and x ∼ p
d ( x ) , (5) where T refers to the user-deﬁne pretext, e.g., a set of data augmentation operators T = { t : R n → R n } 1 .
3. Rediscovering Adversarial Training
In this section, we investigate why robust models have a good generative ability from the perspective of P-CEM.
In general, AT solves the following minimax problem:
[Equation]
The inner maximization problem is to ﬁnd an adversarial example ˆ x within the (cid:96) p -norm ε -ball around the natural example x that maximizes the CE loss.
While the outer minimization problem is to ﬁnd model parameters that minimize the loss on the adversarial examples ˆ x .
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 loss 2 starting from ˆ [Equation] x , that is,
[Equation]
while the Langevin dynamics for sampling from P-CEM starts from random noise ˆ x 0 = δ and updates with
[Equation]
The two update rules above 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.
Thus, the maximization process in AT can be seen as a (biased) sampling method that draws samples from the corresponding probabilistic model p θ (ˆ x ) .
Compared to
Langevin dynamics, PGD imposes speciﬁc inductive bias for sampling.
With the additional repulsive gradient and ε -ball constraint, it explicitly encourages the samples to be misclassiﬁed around the original data.
In practice, adversarial training with such adversarial examples is generally more stable than training P-CEM with Langevin samples, which indicates that PGD attack is a competitive alternative for the negative sampling method for P-CEM training.
Minimization Process.
To begin with, the gradient of the joint log likelihood for P-CEM can be written as:
[Equation]
where ( x , y ) ∼ p
d ( 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 Eqn.
(3) and decompose it as the consistency gradient and the contrastive gradient:
[Equation]
Next, we show that the two parts correspond to two effective mechanisms developed in the adversarial training literature
[Equation]
AT loss.
As the two sample pairs in the contrastive gradient share the same input ˆ x , we can see that the contrastive gradient (second term) can be written equivalently as
[Equation]
ˆ x )] , which is exactly the negative gradient of the AT loss (cid:96)
CE
(ˆ x , y ; θ ) in Eqn.
(6) (see details in Appendix).
Regularization.
As for the consistency gradient (ﬁrst term), original AT (Madry et al., 2018) simply ignores it.
Its variant TRADES (Zhang et al., 2019) instead proposes the KL regularization KL( p ( ·| ˆ x ) (cid:107) p ( ·| x )) that regularizes the consistency of the predicted probabilities on all classes, whose optimum implies that the consistency gradient vanishes, i.e., p ( ·| ˆ x )
= p ( ·| x ) →
f θ ( x , y )
=
f θ (ˆ x , y ) .
The above analysis indicates that the minimization objective of AT is closely related to the maximum likelihood training of P-CEM, and TRADES further injects adversarial robustness prior by regularizing the consistency gradient.
Together with the analysis on the maximization process, we show that
AT is a competitive alternative for training P-CEM (a generative model) with more stable training behaviors.
That explains why robust models with AT are also generative.
3.1. Discussion on Standard Training
In the above discussion, we have explained why adversarial training is generative from the perspective of P-CEM.
In fact, it can also help characterize why classiﬁers with Standard
Training (ST) are not generative (i.e., poor sample quality).
A key insight is that if we replace the model distribution p θ (ˆ x ) with the data distribution p d ( x ) in Eqn.
(3), we have
[Equation]
∇ θ f θ ( x , ˆ y )
= ∇ θ E p d ( x ,y ) log p θ ( y | x ) , which is the negative gradient of the CE loss.
Thus, ST is equivalent to training P-CEM by simply replacing modelbased negative samples ˆ x ∼ p θ ( x ) with data samples x ∼ p
d ( x ) .
This approximation makes ST computationally efﬁcient with good accuracy on natural data, but signiﬁcantly limits its robustness on adversarial examples (as model-based negative samples).
Similarly, because ST ignores exploring negative samples while training, standard classiﬁers also fail to generate realistic samples.
3.2. Extension to Unsupervised AT
In the above discussion, we establish the connection between AT and CEM in the supervised scenarios.
Moreover, this perspective also enables us to generalize adversarial training to unsupervised scenarios via our uniﬁed framework, CEM.
Speciﬁcally, we can extend the discussion of
P-CEM (supervised) to NP-CEM (unsupervised) and derive the corresponding inner and outer optimization objectives in a principled way.
Due to the limit of space, more details are provided in the Appendix.
4. Principled Adversarial Sampling
The interpretation of AT through our uniﬁed framework also inspires us to design principled sampling algorithms from robust models.
Below, we present our adversarial sampling algorithms for both supervised and unsupervised scenarios.
4.1. Supervised Scenario
Targeted Attack (TA).
Previously, to draw samples from a robust classiﬁer, Santurkar et al.
(2019) utilize a targeted attack that optimizes the input initialized by a random noise ˆ x 0 towards a speciﬁed class
ˆ y ∼ p
d (ˆ y ) :
[Equation]
exp( f θ (ˆ x n
[Equation]
k ))
(cid:35) .
Compared to PGD attack in Eqn.
( ?? ), while pushing ˆ x towards ˆ y , TA has a negative logsumexp gradient that decreases the marginal probability p θ (ˆ x ) .
This could explain why TA is less powerful for adversarial attack and is rarely used for adversarial training.
Conditional Sampling (CS).
To overcome the drawback of targeted attack, a natural idea would be dropping the negative logsumexp gradient.
In fact, we can show that this is equivalent to sampling from the conditional distribution:
[Equation]
and its Langevin dynamics takes the form:
[Equation]
[Table]
Unsupervised
[Table]
Samples drawn this way essentially follow an approximated model distribution, p θ (ˆ x , ˆ y )
≈ p d (ˆ y )
p
θ (ˆ x | ˆ y ) .
Thus, CS can be seen as a debiased targeted attack algorithm
[Table]
4.2. Unsupervised Scenario
Without labels, we can also draw samples from p θ ( x ) from our unsupervised robust models, i.e., NP-CEM, by performing Langevin dynamics with K negative samples { z − k } : {+ˆ x n +
 = ˆ x n + α ∇ ˆ x n log p θ (ˆ x n ) + √+}
[Equation]
α · ε ≈ ˆ x n + α ∇ ˆ x n (cid:34) log
K K (cid:88) [Equation] In comparison, as we do not have access to real data points as our anchor points for adversarial attack, PGD attack may not be effective at sampling for the unsupervised scenario.
5. Experiments
In this section, we validate our proposed method by evaluating the adversarial sampling algorithms derived from our framework.
Speciﬁcally, we evaluate the sample equality with both supervised and unsupervised robust models on
CIFAR-10.
For supervised robust models, we adopt the same pretrained ResNet50 checkpoint on CIFAR-10 as San-
[Table]
[Table] [Table]
turkar et al.
As for the unsupervised case, we are the ﬁrst to consider sampling from unsupervised robust models.
We train ResNet18 and ResNet50 (He et al., 2016) encoders following the setup of an existing unsupervised adversarial training method ACL (Jiang et al., 2020).
The training attack is kept the same as that of the supervised case for a fair comparison.
As for the sampling algorithms, we tune the step size α , noise scale ε , and number of steps K for better sample quality.
More details are provided in Appendix.
Comparison with other generative models .
In Table 1, we compare the sample quality of adversarial sampling methods with different kinds of generative models, where our adversarial sampling methods outperform many deep generative models and obtain state-of-the-art Inception scores on par with StyleGAN2 (Karras et al., 2020).
Comparison among adversarial sampling methods.
In
Table 2, we further compare the sample quality of different adversarial sampling methods.
For supervised models, we can see that indeed TA obtains the lowest IS, while CS can signiﬁcantly reﬁne the sample quality.
For unsupervised models, we can see that Langevin dynamics outperforms PGD consistently by a large margin.
In particular, conditional sampling initialized with class-wise noise can improve a little on the sample quality compared to unconditional sampling.
6. Conclusion
In this paper, we proposed a uniﬁed probabilistic framework, named Contrastive Energy-based Model (CEM), which not only explains the generative ability of adversarial training, but also provides a uniﬁed perspective of adversarial training and sampling in both supervised and unsupervised paradigms.
Extensive experiments show that sampling methods derived from our framework indeed demonstrate better robustness and sample quality than state-of-the-art methods. Acknowledgement
Yisen Wang is supported by the National Natural Science
Foundation of China under Grant No. 62006153 and Project 2020BD006 supported by PKU-Baidu Fund.
Zhouchen
Lin is supported by the National Natural Science Foundation of China (Grant No.s 61625301 and 61731018),
Project 2020BD006 supported by PKU-Baidu Fund, Major Scientiﬁc Research Project of Zhejiang Lab (Grant No.s 2019KB0AC01 and 2019KB0AB02), and Beijing Academy of Artiﬁcial Intelligence. Jiansheng Yang is supported by the National Science Foundation of China under Grant No. A. Principled Unsupervised Adversarial Training with CEM
In this section, we show that with our uniﬁed framework, we can naturally extend the interpretation developed for supervised adversarial training to the unsupervised scenario.
A.1. Understanding Unsupervised Standard Training
Recently, the following InfoNCE loss is widely adopted for unsupervised contrastive learning of data representations (Oord et al., 2018;
Chen et al., 2020; He et al., 2020),
[Equation]
where f θ ( x , ˆ x ) =
g θ ( x ) (cid:62) g θ (ˆ x ) calculates the similarity between the representations of the two data samples, x , ¯ x are generated by two random augmentations (drawn from T ) of the same data example, and
ˆ X = { ˆ x j } Kj =1 consists of K independently drawn negative samples.
In practice, one of the K negative samples is chosen to be the positive sample ¯ x .
Therefore, InfoNCE can be seen as an instance-wise
K -class cross entropy loss for non-parametric classiﬁcation.
Perhaps surprisingly, we show that the InfoNCE loss is equivalent to the importance sampling estimate of our NPCEM by approximating the negative samples from p θ ( x ) with data samples from p d ( x ) , as what we have done in standard supervised training (Section C.3):
[Equation]
which is exactly the negative gradient of the InfoNCE loss.
In the above analysis, for an empirical estimate, we draw
N positive pairs ( x i , z i ) ∼ p
d ( x , z ) , and for each anchor x i , we further draw K negative samples { z − ik } from p d (ˆ z ) through reparameterization (Eqn. (5)).
The negative phase of NP-CEM is supposed to sample from p
θ (ˆ z | ˆ x ) , where samples semantically close to the anchor sample ˆ x , a.k.a.
hard negative samples, should have high probabilities.
However, InfoNCE adopts a non-informative uniform proposal p d (ˆ z ) for importance sampling, which is very sample inefﬁcient because most samples are useless (Kalantidis et al., 2020).
This observation motivates us to design more efﬁcient sampling scheme for contrastive learning by mining hard negatives.
For example, Robinson et al.
(2021) directly replace the plain proposal with
[Equation]
while keeping the reweighing term.
From the perspective of NP-CEM, it will introduce bias into the importance sampling, which should be treated carefully.
In all, NP-CEM provides a principled framework to characterize and develop contrastive learning algorithms.
A.2. Proposed Unsupervised Adversarial Training
AT is initially designed for supervised learning, where adversarial examples can be clearly deﬁned by misclassiﬁcation.
However, it remain unclear what is the right way to do Unsupervised Adversarial Training (UAT) without access to any labels.
Previous works (Jiang et al., 2020; Ho & Nvasconcelos, 2020; Kim et al., 2020) have carried out UAT with the adversarial InfoNCE loss, which works well but lacks theoretical justiﬁcation.
Our uniﬁed CEM framework offers a principled way to generalize adversarial training from supervised to unsupervised scenarios.
Maximization Process.
Sampling from p θ ( x ) can be more difﬁcult than that in supervised scenarios because it does not admit a closed form for variable z .
Thus, we perform
Langevin dynamics with K negative samples { z − k } ,
[Equation]
While the PGD attack of the InfoNCE loss (Eqn. 18),
[Equation]
resembles the Langevin dynamics as they both share the positive logsumexp gradient that pushes up p θ (ˆ x ) , and differs by a repulse negative gradient − f θ (ˆ x , z ) away from the anchor z , which is a direct analogy of the PGD attack in supervised learning (Section 3).
Therefore, we believe that the PGD attack of InfoNCE is a proper way to craft adversarial examples by sampling from p θ ( x ) [Equation]
Minimization Process.
Following the same routine in Section C.2, with the adversarial example ˆ x ∼ p θ (ˆ x )
[Equation]
we can insert an interpolated adversarial pair (ˆ x , z ) and decompose the gradient of NP-CEM into the consistency gradient and the contrastive gradient,
[Equation]
Seed
TA
Langevin
CS
PGD
Langevin
PGD
[Figure]
supervised models (ResNet50), unsupervised models (ResNet18), unsupervised models (ResNet50).
Following the analysis for supervised AT, it is also easy to see that the contrastive gradient is equivalent to the gradient of the Adversarial InfoNCE loss utilized in previous work (Jiang et al., 2020; Ho & Nvasconcelos, 2020;
Kim et al., 2020) with adversarial example ˆ x ,
[Equation]
B. More Details on Adversarial Sampling
In this part, we provide more details of our adversarial sampling experiments.
Models.
For supervised robust models, we adopt the same pretrained ResNet50 checkpoint on CIFAR-10 as Santurkar et al.
(2019) 3 for a fair comparison.
The model is adversarially trained with (cid:96) 2 -norm PGD attack with random start, maximal perturbation norm 0 .
5 , step size 0 .
1 andsteps.
As for the unsupervised case, we are the ﬁrst to consider sampling from unsupervised robust models.
We train
ResNet18 and ResNet50 (He et al., 2016) encoders following the setup of an existing unsupervised adversarial training method ACL (Jiang et al., 2020).
The training attack is kept the same as that of the supervised case for a fair comparison.
More details are provided in Appendix.
Sampling protocol.
In practice, our adversarial sampling methods take the following general form as a mixture of the
PGD and Langevin dynamics,
[Equation]
where g k is the update gradient, ε k is the diffusion noise, Π S is the projector operator, and δ is the (conditional) initial seeds drawn from the multivariate normal distribution whose mean and covariance are calculated from the CIFAR-10 test set following Santurkar et al.
(details in Appendix).
Note that there are four hyper-parameters in our sampling protocol: step size α , (cid:96) 2 -ball size β , noise scale η , and iteration steps K , for which we list our choice in Table ??
(the analysis on these parameters can be found in Appendix).
We evaluate the sample quality quantitatively with Inception
Score (IS) (Salimans et al., 2016) and Fr´echet Inception
Distance (FID) (Heusel et al., 2017).
Visualization.
We also demonstrate our randomly drawn (not cherry-picking) samples in Figure 1.
It can be seen that both our supervised and unsupervised methods achieve signiﬁcantly better sample quality than previous baselines.
C. Omitted Technical Details
For a concise presentation, we have omitted several technical details in the main text.
Here we present a complete description of the derivation process.
C.1. Log Likelihood Gradient of EBM
In Section 3, we have introduced Energy-based Models (EBM) and the gradient of their log likelihood.
We now show how it can be derived out.
For a EBM in the following form,
[Equation]
the gradient of the log likelihood can be derived as follows:
[Equation]
C.2. Equivalence between AT Loss and Contrastive Gradient in Supervised Learning
In Section [Equation] we have claimed that the contrastive gradient equals to the negative gradient of the AT loss) following the same deduction in Eqn.
[Equation]
which is exactly the negative gradient of the canonical AT loss (Madry et al., 2018).
C.3. Connection between Standard Training and JEM
In Section 3.
[Equation]
we have claimed that if we replace the model distribution p θ (ˆ x ) with the data distribution p d ( x ) in Eqn.
(3), the log likelihood gradient of JEM is equivalent to the negative gradient of the CE loss.
Here we give a detailed proof as follows:
[Equation]
C.4. Equivalence between InfoNCE Loss and Non-parametric CEM
In Section A.1, we have claimed that the the log likelihood gradient of NP-CEM equals to exactly the negative gradient of the InfoNCE loss when we approximate p θ ( x ) with p d ( x ) .
The derivation is presented as follows:
[Equation]
C.5. Equivalence between Adversarial InfoNCE and NP-CEM
In Section A.
[Equation]
we have developed the unsupervised analogy of AT loss and regularization.
In particular
[Equation]
we have claimed that contrastive gradient is equivalent to the gradient of the
Adversarial InfoNCE loss (i
[Equation]
e., the InfoNCE loss of the adversarial example ˆ x ) utilized in previous work (Jiang et al., 2020; Ho & Nvasconcelos, 2020; Kim et al., 2020).
It can be derived following Eqn.
[Equation]