THAT: Two Head Adversarial Training for Improving Robustness at Scale

Convolutional neural networks have demonstrated remarkable performance in a multitude of computer vision tasks like image classification [16, 48], object detection [34, 3, 24], etc. Despite the power of these state-of-the-art models, they are found to be extremely unstable to input perturbations [17, 22, 12]. This fragility can be exploited by crafting adversarial examples, which are optimized to manipulate networks while appearing innocuous to humans.

One popular method to mitigate the brittleness of neural networks is adversarial training [25, 14, 20], in which network parameters are updated using adversarially perturbed images. This produces a hardened network that is robust to adversarial perturbations in the pixel space. While adversarial training is able to increase the robustness of classifiers, it often reduces accuracy on clean images at test time [51, 40, 2]. This accuracy loss is believed to be in part because of fundamental tradeoffs between accuracy and robustness [35, 40, 2], and in part because of domain shift between clean and adversarial image distributions [45, 27, 10].

In light of this, we propose to align features of clean images with their adversarially perturbed counterparts for improved clean accuracy and robustness. Our approach is inspired by recent self-supervised learning frameworks [5, 15], in which two random crops of the same image are considered as a positive pair and their distance in feature space is minimized to learn generic feature representations. Here, a clean image and its adversarial copy form a positive pair, and by contrasting it with other pairs, we modulate standard adversarial training to prevent the adversarial feature distribution from drifting away from the natural distribution.

With this in mind, we introduce Two Head Adversarial Training (THAT), an adversarial training framework that uses multiple training objectives to boost robustness. THAT adversarially trains a robust encoder using two parallel heads, a feature head and a classification head, both sitting on a shared ResNet backbone. For each training step, clean images are passed through a naturally trained clean network to produce features. Then, adversarial examples are made by attacking the classification head of the robust encoder. A loss function is then computed that contains two terms: a contrastive loss that promotes similarity between the natural features from the clean encoder and the adversarial features from the robust encoder, and a classification loss from the classification head. The former term enforces that feature representations of adversarial images are aligned with natural feature representations (i.e., domain shift is minimized), while the second term ensures that the feature representations contain the information needed for accurate classification. The robust encoder is then updated to minimize the combined loss.

At test time, we consider two different classification modes for defense. In addition to using a standard classification head, we also consider a nearest-neighbor-based classification procedure that relies on the feature extraction head.

We conduct extensive experiments on ImageNet [8] using standard adversarial training updates, and also using accelerated (a.k.a, “free” [36]) adversarial updates. In both settings, we demonstrate that THAT outperforms other adversarial training frameworks in terms of both clean accuracy and robustness with different backbones.

Early adversarial training methods crafted adversarial examples with the Fast Gradient Sign method (FGSM) [14], but this approach results in networks that are easily broken by multi-step attacks. Madry et al. use project gradient descent (PGD) to generate adversarial examples to improve robustness for multi-step attacks [25]. PGD performs gradient ascent on input images in the signed gradient direction with respect to the classification loss, and then clips the perturbation to enforce an $\ell_{\infty}$-norm constraint. Kannan et al. further regularize adversarial training by penalizing the difference between logits from clean images and their adversarial variants [20]. Xie et al. introduce denoising blocks that use self-attention in feature maps to improve adversarial training [46]. Balaji et al. improve adversarial training for ImageNet by imposing different robustness criteria for each training sample [2]. The TRADES method achieves high levels of robustness by training with a loss that promotes similarity between the predicted label scores for natural and adversarial images [51]. Due to the huge computational overhead of adversarial training methods, Shafahi et al. introduce a “free” adversarial training strategy that updates model weights and generates adversarial examples by replaying mini-batches [36].

In our work, we build upon standard and free adversarial training settings and demonstrate that explicitly contrasting features improves both clean accuracy and robustness. Dubey [10] et al. show that nearest neighbor search can be used for deflecting strong attacks, but this requires an extremely large gallery of reference images ($\sim 1$ billion) to achieve good performance. Instead, we show that our network is able achieve strong defense via nearest neighbors using the much smaller ImageNet dataset. In addition, Mao et al. use a triplet loss for better adversarial training [27], which requires manual hard negative mining. Our contrastive loss automatically performs hard negative mining with a number of negative samples.

TRADES is based on the decomposition $\mathcal{R}_{rob}=\mathcal{R}_{nat}+\mathcal{R}_{bnd},$ which represents the robust accuracy of a model as the sum of the natural risk (i.e., the 0-1 loss), and the “boundary” risk. The boundary risk is the probability that, for an input $X$, there exists an input point $X^{\prime}$ within a ball around $X$ such that $X$ and $X^{\prime}$ are assigned different labels. These terms are then upper bounded by a smooth loss and minimized. The multi-class TRADES objective is:

where $\mathcal{L}$ is a loss function that measures the discrepancy between class label distributions predicted for $X$ and its adversarial example $X+\mathbf{\delta}$, and $\lambda$ is a scalar parameter.

For 2-class problems, the loss (2) is a tight convex relaxation of the non-convex 0-1 loss whenever $\mathcal{L}$ is a calibrated loss function that satisfies several weak assumptions [51]. However, the tightness result for TRADES is specific to binary classification; the relaxation is no longer tight for multi-class problems, and becomes extremely pathological for many-class problems like ImageNet.

To understand why this is, observe that TRADES uses the KL divergence for $\mathcal{L}(\cdot,\cdot)$ when solving multi-class problems. If $p$ and $q$ are the class distributions predicted by the network for $X$ and its adversarial example, respectively, then the TRADES loss contains the KL-divergence term:

where the sum in Eqn. (4) is over class indices. For typical ImageNet images, the distributions $p$ and $q$ contain just a few relevant/large class probabilities, and then a “long tail” with nearly 1000 small probabilities that are generally discarded when making classification decisions. The KL loss used by TRADES sums over all of these class labels, and its value becomes highly dominated by the many contributions from classes with small probabilities. Furthermore, the smallest values of $q_{i},$ which appear in the denominator of Eqn. (4), make the loss most sensitive to unlikely classes. In practice, we often ignore such class labels with near-zero probability, and yet they account for a dominant share of the TRADES objective.

Our proposed method is similar in spirit to TRADES —it enforces similarity between network outputs for natural and adversarial images. But unlike TRADES, our proposed method uses a contrastive loss that is computed using the inner product between feature vectors, and is dominated by large entries in feature representations. This enables THAT to do contrastive learning while avoiding the pathological behaviors that TRADES suffers in the many-class regime.

In this section, we introduce THAT, a contrastive framework with a naturally trained clean encoder and an adversarially trained robust encoder. The method learns feature representations that are robust to attacks while simultaneously achieving high classification accuracy. Each of these objectives is achieved with a loss function on a separate head of the network. We describe each component of our framework here.

where ${\bm{\hat{z}}}$ is the normalized logits based on ${\bm{z}}$, ${\bm{\hat{w}}}_{i}$ is the $i$ the column in $\bm{W}_{ac}$ representing the normalized weights for the $i$-th class, $y_{i}\in\mathbb{R}^{\{0,1\}}$ is the label for the $i$-th class, and $C$ is the total number of classes in the dataset. $\eta$ is a learnable parameter to control the sharpness of the distribution. Here, ${\bm{\hat{w}}}_{i}$ is considered as the class prototype, and we are forcing the logit vectors to lie close to it. This is similar in spirit to Eqn (5), in which features from an adversarial example are mapped to be close to features of its clean twin.

Finally, the combined training objective of THAT can be written as:

The nearest-neighbor classifier computes feature representations for all training samples with the clean encoder and stores them in a memory bank $\bm{M}_{\text{train}}$. Given a test image $\hat{X}$, we first compute its feature ${\bm{u}}_{t}$ through the robust feature head. Then, the embedding is compared with those of all training samples to retrieve the top-$k$ nearest neighbors. $N_{k}$ similarity scores are computed based on dot products, and a weighted average of neighbor labels is computed. More precisely, the confidence of $\hat{X}$ belonging to the $c$-th class is defined as:

where ${\bm{u}}^{i}$ is the feature for $i$-th neighbor sample, and ${\bm{y}}^{i}$ is the corresponding one-hot label.

We presented THAT, a two-stream contrastive learning framework for improved robustness and clean accuracy. THAT is trained using two loss functions; one to align the feature distributions between natural and adversarial images, and one to promote good classification accuracy. The resulting model is able to defend against strong adversarial attacks at test time not only using the hardened classifier but also using a KNN search. Through extensive experiments, we demonstrate THAT achieves better results than alternative methods for ImageNet under a wide range of settings.