Title: Distilling Out-of-Distribution Robustness from Vision-Language Foundation Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: robustness

knowledge distillation

adversarial training

data augmentation

generalization
Abstract: We propose a conceptually simple and lightweight framework for improving the robustness of vision models through the combination of knowledge distillation and data augmentation. We address the conjecture that larger models do not make for better teachers by showing strong gains in out-of-distribution robustness when distilling from pretrained foundation models. Following this finding, we propose Discrete Adversarial Distillation (DAD), which leverages a robust teacher to generate adversarial examples and a VQGAN to discretize them, creating more informative samples than standard data augmentation techniques. We provide a theoretical framework for the use of a robust teacher in the knowledge distillation with data augmentation setting and demonstrate strong gains in out-of-distribution robustness and clean accuracy across different student architectures. Notably, our method adds minor computational overhead compared to similar techniques and can be easily combined with other data augmentations for further improvements.

Distilling Out-of-Distribution Robustness from Vision-Language Foundation Models

Abstract

We propose a conceptually simple and lightweight framework for improving the robustness of vision models through the combination of knowledge distillation and data augmentation. We address the conjecture that larger models do not make for better teachers by showing strong gains in out-of-distribution robustness when distilling from pretrained foundation models. Following this finding, we propose Discrete Adversarial Distillation (DAD), which leverages a robust teacher to generate adversarial examples and a VQGAN to discretize them, creating more informative samples than standard data augmentation techniques. We provide a theoretical framework for the use of a robust teacher in the knowledge distillation with data augmentation setting and demonstrate strong gains in out-of-distribution robustness and clean accuracy across different student architectures. Notably, our method adds minor computational overhead compared to similar techniques and can be easily combined with other data augmentations for further improvements.

1 Introduction

One of the goals of machine learning is to develop systems that can generalize effectively across diverse populations and environments, much like human intelligence. Despite the impressive advancements in neural networks that surpass human performance in various tasks in computer vision, their generalization capabilities remain inadequate when faced with out-of-distribution data, such as adversarial perturbations [60], unusual colors and textures [17, 58, 68], or challenging contexts [29].

One major line of research addresses this issue with more sophisticated training strategies [37], including adversarial training [39], data augmentation [18, 28, 79], or other regularizations [35, 46, 61, 73, 69]. In this paper, we focus on adversarial-training-based data augmentation [75], which can enhance the quantity and diversity of training data. In addition, theoretical work suggests achieving high robustness requires significantly more samples than clean accuracy [56]. This has also been shown empirically [3, 24], most recently with transformers [14] which achieve robustness on a variety of computer vision tasks. In addition to weak inductive bias, powerful model capacity, and grounding with language, these models are often trained with large-scale datasets [12, 14, 48, 76], up to billions of images [57] that encompass many real-world distribution shifts. As a result, these foundation models [1] exhibit remarkable zero-shot generalization, especially on natural distribution shifts such as artistic renderings, but require large amounts of compute and heavily parameterized models.

In this paper, we aim to connect these two lines of work. We investigate if it is possible to improve robustness by introducing a foundation model as a teacher to distill robust representations and help generate a diverse data augmentation. We conduct our analysis without requiring the teacher's largescale dataset and focus on out-of-distribution robustness by introducing an image-to-image generative model to discretize optimized perturbations. We aim to leverage in-distribution data to a greater

extent and conduct our investigation with CLIP [48]. This marks a departure from existing work in knowledge distillation (KD), which tends to focus on smaller models and datasets. In fact, prior work [8] has called into question the utility of distilling from stronger teachers over training from scratch altogether. However, we find that although this model capacity gap can impair improvements in clean accuracy, distilling from robust teachers improves out-of-distribution robustness, even when leveraging only in-distribution data. Surprisingly, distilling on clean ImageNet images from CLIP with the original KD objective [32] results in a more robust ResNet34 than training with state-of-the-art regularization methods, despite a parameter difference of ∼ 13.7x.

However, it is currently unclear in what settings a teacher's robustness can reliably transfer to a student and how to best combine distillation with data augmentation. We aim to answer this question both theoretically and empirically. We view adversarial training and data augmentation in the lens of domain generalization and prove that the diversity of the augmented samples leads to improved robustness. Our findings further suggest that foundation models make for strong teachers due to their diverse training distribution.

Building upon these findings, we introduce discrete adversarial distillation (DAD) , a KD framework that further distills the robustness of a teacher model by leveraging the adversarial examples of the teacher discretized by a VQGAN [15] as data augmentation. Notably, these samples are generated in an offline fashion, adding minor computational overhead compared to standard adversarial training. Intuitively, a foundation model will produce more diverse adversarial samples than a teacher trained on the same distribution, and we provide a theoretical framework using Wasserstein distance to formalize this proposition. Empirically, when distilling CLIP to a ViT-B, we achieve robust accuracy of 46.1% on ImageNet-Sketch [68] and 65.1% on ImageNet-Rendition [26], improving on the state of the art by 17.8% and 11.3% respectively. DAD can also be freely combined with existing regularization techniques, resulting in improvements in clean accuracy. In summary, our contributions are 1

Establishing the KD for out-of-distribution robustness setting and a proposing a novel KD objective based on data augmentation

Providing a theoretical framework in KD for the choice of a teacher based on the diversity of the data augmentation

Proposing a novel data augmentation DAD that outperforms both adversarial training and distillation techniques on natural distribution shifts

2 Related Work

We define out-of-distribution accuracy/robustness as a model's performance on non-adversarial distribution shifts, adversarial accuracy/robustness to the case of robustness of adversarial examples, and clean accuracy as evaluation on a dataset drawn from the same distribution.

Data augmentation. Data augmentation is frequently used as regularization [75] by expanding the quantity and diversity of training data. This is often achieved through simple transformations such as rotations or image crops or more advanced techniques such as image mixing [79, 28], reinforcement learning [10, 81] or adversarial training [21, 31, 41] to find the optimal transformation.

Adversarial training (AT) was initially introduced to enhance model robustness by training with adversarial examples [39]. Although effective for defending against adversarial attacks, several works [64, 78, 70] have indicated a trade-off between adversarial clean accuracy in AT, limiting its effectiveness as a general data augmentation. Considerable efforts [47, 49] have been made to minimize this trade-off and directly use adversarial examples as data augmentation [50, 74], but there is still a considerable gap in out-of-distribution performance compared to foundation models like CLIP [48].

One line of work has recently been adapted to this issue. The model-based robustness paradigm [4, 52] leverages the disentangled latent representations of pretrained generative models to improve or validate out-of-distribution robustness [80], and can be used to improve adversarial examples. Most similar to our work is [23, 2, 41], which use a GAN or VAE [15, 34, 67] to discretize or discover adversarial examples. However, we leverage both a pretrained discretizer and foundation model, and adapt the AT objective to a knowledge distillation setting.

Knowledge distillation. Knowledge Distillation (KD) is a technique for training a student model with guidance from a stronger teacher model, widely applied in vision and language tasks [6, 32, 55, 71]. Most works focus on improving the KD objective with different knowledge transfer objectives, such as feature distance [7, 54, 63], attention [77], distribution [45], activation boundary [30], and sample distance [38, 44, 66]. [8] raises the model capacity gap issue, where training suffers when the size of the student and teacher models differ, but we find that there is still benefit to distilling from a robust teacher. Another line of work, defensive distillation, aims to distill adversarial robustness from an adversarially trained teacher [20, 82, 83]. We have a similar goal, but for out-of-distribution robustness and propose a loss objective not previously explored in prior works.

3 Method

We consider a standard dataset { ( x, y ) } N n =1 ∼ P N where instances and their labels ( x n , y n ) are drawn from a distribution P and are used for training the student model θ . We also consider a discretizer, Q , and a teacher model ϕ and we use ϕ ( x ) to denote the output of a model given the sample x . a denotes a function that applies a data augmentation on x , also known as a transformation function. Additionally, a ∈ A , the class of all such functions.

3.1 Setup

Invariance is a desirable property where the model will have the same representation of an input after a transformation is applied. A model, θ is said to be invariant if θ ( x ) = θ ( x ′ ) for all x ∈ U , where U is the set of all images that can be obtained by a transformation of x by a ∈ A , which includes the identity transformation. a ranges from worst-case imperceivable perturbations to real-world distribution shifts like artistic sketches [68]. In this paper, we focus on the latter, denoted as ˆ a and ˆ A . We can consider ˆ a to represent an individual distribution P and ˆ A to be drawn from the distribution of distributions ˆ P . Our ultimate goal is to train θ to be invariant to transformations in ˆ A . A model that maintains the same representation under transformations or distribution shifts of x is said to be robust, which we define as the worst-case expected risk where

where l is the loss function and r depends on an anchor distribution P , and ϵ , the maximum deviation allowed under the Wasserstein's distance metric. Similarly, we can define the expected robustness and expected risk in terms of an arbitrary distribution, including the training distribution.

The robustness of the resulting model is highly dependent on x ′ , P , and the choice of data augmentation. It is also susceptible to adversarial attacks, where x ′ is a worst-case perturbation of x . Adversarial robustness can be improved with adversarial training, which couples the outer minimization objective from (3) with an inner maximization objective in the following

where x ′ = x + ϵ , ϵ is the perturbation, and l is the cross-entropy loss. This achieves adversarial robustness, but cannot generalize well to real-world domain shifts in ˆ A . To address this, we consider a generative model, Q , trained on P and can model ˆ A . Passing an input through Q in the maximization objective applies a worst-case transformation from ˆ A . This modifies (4) to minimize the empirical semantic adversarial risk,

3.2 Distillation from a robust teacher

Next, we consider a setting where we also have access to a pretrained ϕ invariant to distribution shifts in ˆ A . This enables us to leverage knowledge distillation (KD). This modifies (3) in the following

CLIP

Image

CLIP

Text

Text prompts

A photo of

a {object}

VQGAN

L (Q(

𝒙

'),

𝒙

,

𝒚

,

𝜽

, φ)

𝒙

+

δ

),

𝒚

)

𝒙

𝒚

φ

φ

δ

{rabbit}

Q

𝒙

'

𝒙

𝒙

'

𝜽

Distill from

φ

Update

𝜽

parameters

where l 1 is the classification loss, the cross-entropy loss, and l 2 is a distance loss between θ ( x ) and ϕ ( x ) , the KL divergence. (6) can be approximated by the empirical risk where

Following theoretical work [43], distilling from a robust teacher with l 2 improves generalization due to minimizing the population risk, which has lower variance. In this formulation, the output of the teacher acts as a more robust supervisory signal than the label, encompassing the probability distribution of classes. This allows the student to learn the teacher representations on in-distribution data, but our experiments show that this is inadequate for out-of-distribution robustness, even when using a robust teacher. To address this, we combine (5) and (7) to also distill the representations of ϕ on augmented samples,

The teacher is more robust, and is able to "solve" the perturbation for the student through distillation. Like [41] and [72], we train our models with both the original and augmented samples, expanding the size of the dataset and maintaining the original information path for x . We do not use the crossentropy loss or y labels for x ′ , as these labels may be wrong and could limit the expressiveness of the data augmentation. This allows us to use adversarial samples of the teacher in a novel maximization objective and obtain stronger empirical results.

3.3 Discrete Adversarial Distillation

The goal of our method, discrete adversarial distillation (DAD), is to distill from a large-scale visionlanguage model using only ImageNet [13] data. In the practical setting, we use approximations of an ideal discretizer and robust teacher. For ϕ we use CLIP [48], which was trained on a large-scale dataset and achieves impressive zero-shot generalization on a variety of natural distribution shifts.

For Q , we use a pretrained VQGAN [15], following [41], which also finds minimal improvements with a stronger discretizer. The VQGAN consists of an encoder, decoder, and quantization, where the encoder learns a latent representation of the input x , the quantization maps the representation to a visual codebook entry, and the decoder reconstructs x . We denote this process as Q ( x ) . In the adversarial training setting, Q discretizes x , a worst-case perturbation ϵ is added by approximating the maximization objective to obtain x ′ , and the resulting image is then discretized by Q again. To improve the transfer of robustness from the teacher, we hypothesize a better augmentation is more informative and exposes more teacher knowledge. We make this rigorous in the following section.

To generate these examples, we adopt adversarial training and modify the maximization objective of (5) to use the worst-case transformations of the teacher. We hypothesize a teacher trained on more diverse data will have more informative adversarial examples. To ensure the correctness of the perturbation, we only use samples that are still classified correctly by the teacher after generation. We use the teacher as an "oracle", allowing it to distill correct representations of the transformed image. Additionally, we generate these samples in an offline manner asynchronously from the pretrained teacher and add them to the original dataset during the training of the student. These examples

only need to be generated once for each teacher and be reused as additional data. This adds minor additional training costs compared to online adversarial training or DAT [41], which has 11x and 3.5x the cost of standard training, respectively [41]. Our full objective is described as,

3.4 Theoretical Investigation

We aim to investigate how to best distill robustness from a robust teacher trained on a large-scale dataset. We find that robust performance can be connected to the distance between the training and test distributions. A data augmentation can represent a new distribution, and the robustness of a model trained on this distribution can be quantified by its diversity and closeness to the test distributions. Although its representation on in-distribution samples can distill a degree of robustness, we show that due to being closer to the test distribution, it is more effective to leverage the discretized adversarial examples of the teacher than the student as our data augmentation of choice.

We begin with some assumptions.

Assumption 1. For an arbitrary data pair ( x, y ) , transformations in ˆ A do not alter the semantics of the data. We can also say we consider an ideal labeling function where any ( x, y ) pair can be correctly mapped, y = f ( x )

Assumption 2. Any arbitrary distributions P and P ′ we compare possess smooth probability densities controlled by two constants c and α depending on the smoothness of P and P ′ where cw ( P, P ′ ) α and c > 0 and is only dependent on α .

Assumption 3. For function γ ( | Θ | , n, β ) parameterized by hypothesis space | Θ | , number of samples n , and the probability when the bound holds β , if the samples are i.i.d, γ ( | Θ | , n, β ) = 2 R ( L ) + √ (log 1 /β ) / 2 n , where R ( L ) stands for Rademacher complexity and L = { l θ | θ ∈ Θ } , where l θ is the loss function corresponding to θ . Additionally, if Θ is finite, l ( · , · ) is a zero-one loss, and samples are i.i.d, then γ ( | Θ | , n, β ) = √ (log( | Θ | ) + log(1 /β )) / 2 n

Lemma3.1. Given Assumptions 1 and 2 and variational divergence tv , for two arbitrary distributions P and P ′ with corresponding density functions δ and δ ′ , r ( P ′ ) ≤ r ( P ) + w ( P ′ , P ) .

Lemma 3.2. Given Assumption 3, Lemma 3.1, and probability at least 1 -Γ ,

where n ( X,Y ) P denotes the number of sample sizes in the finite dataset ( X,Y ) P , and ξ is a vanilla term that connects n ( X,Y ) P and Θ with the generalization error bound.

Proof. We apply conventional generalization analysis through uniform convergence to Lemma 3.1. We leave the full proof of Lemma 3.1 in Sec. E in the Appendix.

This results in an intuitive conclusion: empirical robustness performances depends on the divergence between the training and testing distributions, as well as two additional elements. The first is the empirical error term on the training distribution, which can be quantified, and the second is a technical term influenced by the sample size and hypothesis space. The exact manner in which ξ depends on these parameters is contingent on the particular convergence analysis employed.

Therefore, the decisive factor for robustness is the degree of deviation between the training and testing distributions. Therefore, using diverse data augmentations close to the testing distribution will lead to the largest gains in robustness. This intuitive understanding suggests that training with distributions generated from the teacher will be more advantageous, as the teacher, having been trained on a large dataset, encapsulates more diverse distributions.

There findings are applicable to any arbitrary distributions P ∼ ˆ P . Nevertheless, this doesn't inherently encapsulate the characteristics of foundation models that trained on data from across the internet, composed of a variety of semantic distributions from ˆ A .

Therefore, we use P ∈ ˆ A to denote a set of m distributions, i.e., P = { P 1 , P 2 , . . . , P m } , and we consider a pretrained foundation model trained with such a set of distributions. To facilitate forthcoming discussions, we extend the notation of w to encompass sets, defining w ( P ′ , P ) as the average divergence between distributions within the set. Thus, w ( P ′ , P ) := ∑ m i w ( P ′ , P i ) /m, ∀ P i ∈ P .

Lemma 3.3. Given a distribution P , we generate a new distribution P ∗ ∈ ˆ A using the discretized worst-case adversarial samples of a model θ . Training θ with adversarial training is equivalent to training θ with empirical risk minimzation on P ∗ where w ( P, P ∗ ) ≤ ϵ .

We leave the full proof of Lemma 3.3 in Sec. E of the Appendix. Finally, let's denote the model ϕ trained over distribution P as ϕ ( P ) and the adversarial data augmentation process that results in a new semantic data distribution as D . We aim to compare w ( D ( ϕ ( P )) , P ′ ) and w ( D ( ϕ ( P )) , P ′ ) . In this context, P ′ is any arbitrary testing distribution, P is a specific training distribution, and P represents the set of distributions used for training foundation models.

Lemma 3.4. Assuming ˆ P is continuous and has a finite expected value, for two sets of distributions P 1 and P 2 , assuming there is at least one distribution in the intersection of P 1 and P 2 , for a fixed testing distribution P ′ , we have

We leave the full proof of Lemma 3.4 in Sec. E of the Appendix. Our findings provide a comparison of the differences in the training mechanisms for various models ϕ , each trained with distinct data sets ( P 1 and P 2 ) and subjected to adversarial training. The methodology can easily be simplified to compare the bounded results between adversarial training based on the teacher model and standalone adversarial training by setting one of the training datasets to consist of a single distribution.

In the scope of our investigation, we compare DAD and discrete adversarial training based on the teacher model, referred to as P 1 , with DAT [41] and the traditional approach of adversarial training based on the student model, referred to as P 2 . Our results suggest two key implications:

Given a fixed P 2 , a more diverse P 1 potentially results in greater variations in performance. We show visualizations that support this in Sec. D in the Appendix. In other words, the use of larger, more diverse pretrained datasets for the teacher model or more diverse data augmentations for the student model increases the likelihood of creating a robust final model. This is also been shown empiricially in prior work investigating the source of robustness in foundation models [16, 62]. However, in practice, the efficacy of distilling from this teacher depends on a variety of factors, including student-teacher architecture, training objective, and student capacity.

For a fixed P 2 , the greater the distance between the testing dataset P ′ and P 2 , the more likely it is that the performance gains will be realized by distilling the teacher model trained on a more extensive set of training data. To put it intuitively, if the testing dataset closely resembles the training dataset (i.e., w ( P , P ′ ) is small), then adversarial training based on the teacher model might not yield significant performance improvements. However, if the testing dataset differs considerably from the training dataset, then adversarial training based on a teacher model that has been trained on a larger dataset is more likely to yield superior performance gains. This observation aligns with our empirical results.

4 Experimental Results

4.1 Experimental Setup

Models. We focus primarily on ResNet50 [25] and ViT-B/16 [14]. We distill from a frozen pretrained CLIP-ViT-L/14 [48], trained on 224 x 224 resolution images with a patch size of 14.

Datasets. We train our models on ImageNet-1K [13]. We use several evaluation datasets. For in-distribution performance, we evaluate on ImageNet and ImageNet-V2 [51], a replication of ImageNet's evaluation set. We focus our study on natural distribution shifts and evaluate on ImageNetA [29], a set of adversarially filtered natural images misclassified by a ResNet50, ImageNet-Sketch [68] which contains artistic sketches of objects, and ImageNet-Rendition [26] which contains abstract

or rendered objects. To observe performance on distributions that are out-of-distribution for the CLIP teacher, we evaluate on synthetic distribution shifts ImageNet-C [27] , which applies corruptions (snow, blur, noise, etc.) to ImageNet, and Stylized-ImageNet [17], which processes ImageNet with style transfer from a source image.

4.2 Baselines

DAD consists of both a data augmentation and knowledge distillation objective. We compare to both types of methods in our experiments.

Common data augmentations. For the simplest baseline, we follow [59] and train with common data augmentations Mixup [79], which combines images and labels, and Randaugment [11], which learns a policy over common transformations such as brightness or shear.

DAT. We also compare against the state-of-the-art data augmentation, DAT [41], which uses a VQGAN [15] to discretize adversarial examples in adversarial training. DAT uses the standard adversarial training objective 5.

Knowledge distillation. We compare against other logit-based knowledge distillation objectives, which only distill the output logits of the teacher. We consider standard knowledge distillation 7 and DIST [33], which aims to address the model capacity gap issue by distilling logit class relationships. Neither method natively supports distillation on augmented samples, so we also compare to defensive distillation objectives ARD [20] and RSLAD [83]. ARD modifies 7 to use the KL divergence between the student logits on the augmented sample with the teacher logits on the normal sample. RSLAD is an extension of ARD that replaces the cross-entropy terms with a KL divergence loss. For a fair comparison, we use DAD as the data augmentation.

4.3 Main Experimental Results on ViT-B/16 and ResNet50

ImageNet-1K. Tab. 1 shows results on ImageNet-1K and its distribution shifts. We compare against ViT-B/16 and ResNet50 models without data augmentation and with the state-of-the-art data augmentation approaches, PyramidAT [31] and DAT [41]. We combine DAD with the data augmentations used in AugReg [59], MixUp [79] and RandAugment [11]. We find that DAD has the best average performance across datasets for both models. For ViT-B we find that DAD has competitive in-distribution performance, but greatly improves performance on natural distributions. Compared to Pyramid AT and DAT, DAD also generalizes well to ResNet50. This suggests that the

DAD data augmentation can be used across student architectures and that due to distillation, DAD is especially effective when training smaller models.

We also demonstrate DAD can be combined with existing approach DAT for stronger in-distribution performance. We add our distillation objective to the DAT objective and train the student on both the teacher's and its own adversarial samples. However, this comes at the cost of lower performance on natural distribution shifts, although we do observe that DAD + DAT still outperforms the prior state-of-the-art on ImageNet-Sketch and ImageNet-Rendition.

Synthetic distribution shifts. We also evaluate our models on synthetic distribution shifts composed of generated transformations in Tab. 2. Since the diverse training distribution of CLIP is mostly composed of natural distribution shifts, it has weaker zero-shot generalization to synthetic distribution shifts, and this performance is inherited in the student model. In fact, zero-shot CLIP is already outperformed by some compared methods on ImageNet-C, and Stylized-ImageNet. However, for ResNet50 DAD also has the best ImageNet-C performance, likely due to compared methods being specialized for certain distribution shifts [36, 74] or architectures [31].

Distillation. In Tab. 3 we compare DAD to knowledge distillation objectives. KD [32] and DIST [33] are vanilla distillation approaches without data augmentation or AT. ARD [20] and RSLAD [83] are defensive distillation objectives that use a adversarially robust teacher to encourage invariance to perturbations. For a fair comparison, we use CLIP-ViT-L/14 as the teacher and discretize the adversarial examples. We find that our distillation objective outperforms vanilla and defensive distillation objectives. We note that all methods can transfer robustness to the student, even methods without data augmentation.

ImageNet-21k. In Tab. 4 we show further gains in robustness from applying DAD to a ViT-B/16 pretrained on ImageNet-21K. We fine-tune this model with our method using only our method. Despite the baseline model performing worse than the variant trained with AugReg, DAD achieves

higher relative and absolute gains in robustness. We hypothesize the larger training distribution equips the student with useful inductive biases that let it more easily learn the more out-of-distribution adversarial examples generated from CLIP. We note that the CLIP training set is still ∼ 28.2x larger so this does not contradict our theory, but it may also be beneficial to train or pretrain the student on a more diverse dataset for a smoother distillation process.

4.4 Ablations

Other student and teacher architectures. Although our method and theory is adapted for foundation models, we investigate its efficacy on other models and teachers in Tab. 5. We consider a different large-scale teacher, CLIP-RN101, as well as teachers trained on ImageNet-1K that achieve out-ofdistribution robustness through methods besides large-scale training, like Discrete ViT [40] or ViT-B [14] trained with DAT [41]. We also consider smaller students like ResNet34 [25] and ViT-S.

We find that distilling robustness in our setting depends on several factors, but most crucially, the robustness of the teacher. Like other distillation techniques, we find that our method can transfer representations between various student/teacher architectures. We find that our method is also susceptible to the model capacity gap, with lower clean accuracy on ResNet34 when distilling from CLIP than training from scratch. However, using CLIP results in the best performance on natural distribution shifts. Despite the more similar architecture, distilling from CLIP-RN101 across students is less effective than distilling from the more robust CLIP-ViT-L. We include similar results with vanilla knowledge distillation in Sec. A of the Appendix.

Pure data augmentation. We study in Tab. 6 the effect of training on the DAD adversarial examples purely as a data augmentation technique, without distillation. Although DAD remains competitive, we find significant drops in performance, suggesting that it is difficult for the student to learn robust representations of these images on its own. However, we continue to observe improvements on natural distribution shifts, suggesting these samples are closer to CLIP's training distribution. However, training with DAD samples is significantly cheaper than DAT and Pyramid AT, making it more efficient in practice.

Computational cost analysis. Since DAD uses adversarial examples generated from a frozen teacher, there is no need to regenerate them during training. This amortizes the otherwise significant cost of adversarial training. We compare the cost of DAD with other adversarial data augmentation approaches in Tab. 7. By avoiding the need to continuously generate new adversarial examples, the only remaining cost for DAD is training on a larger dataset, making it cheaper than similar methods.

Additional ablations on choice of generative model, use of gradients, and transfer to adversarial robustness can be found in Sec. B in the Appendix.

5 Conclusion and limitations

We conduct the first study on distilling out-of-distribution robustness. We develop a framework for the use of foundation models in this setting and empirically and theoretically validate their advantages as a teacher. We propose discrete adversarial distillation (DAD) which uses the discrete adversarial examples of the teacher as a more diverse data augmentation and directly distill its most diverse representations. However, we find that DAD tends to be biased towards the performance of the CLIP teacher, exhibiting improvements mostly on natural distribution shifts. In practice, these shifts tend to be the most useful, and with the small computational cost of using DAD, we encourage practitioners to adopt it when training small models. We hope the development and release of improved foundation models and generative models will further demonstrate the effectiveness of our method.

We encourage further work to understand the limitations of machine vision models in out-ofdistribution settings. More robust models carry the potential risk of automation bias, i.e., an undue trust in vision models. However, even if models are robust against corruptions in finite out-ofdistribution datasets, they might still quickly fail on the massive space of semantic transformations in real-world data. Understanding under what conditions model decisions can be deemed reliable is still an open research question.

References

Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, et al. On the opportunities and risks of foundation models. arXiv:2108.07258 , 2022. 1

Dan Andrei Calian, Florian Stimberg, Olivia Wiles, Sylvestre-Alvise Rebuffi, András György, Timothy A Mann, and Sven Gowal. Defending against image corruptions through adversarial augmentations. In ICLR , 2022. 2

Yair Carmon, Aditi Raghunathan, Ludwig Schmidt, John C Duchi, and Percy S Liang. Unlabeled data improves adversarial robustness. In NeurIPS , 2019. 1

Junbum Cha, Kyungjae Lee, Sungrae Park, and Sanghyuk Chun. Domain generalization by mutualinformation regularization with pre-trained models. In ECCV , 2022. 2

Minwoo Chae and Stephen G. Walker. Wasserstein upper bounds of the total variation for smooth densities. Statistics Probability Letters , 163:108771, 2020. 18

Guobin Chen, Wongun Choi, Xiang Yu, Tony Han, and Manmohan Chandraker. Learning efficient object detection models with knowledge distillation. In NeurIPS , 2017. 3

Liqun Chen, Zhe Gan, Dong Wang, Jingjing Liu, Ricardo Henao, and Lawrence Carin. Wasserstein contrastive representation distillation. CVPR , 2020. 3

Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation. In ICCV , 2019. 2, 3

Francesco Croce and Matthias Hein. Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks. In ICML , 2020. 17

Ekin D. Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V. Le. Autoaugment: Learning augmentation strategies from data. In CVPR , June 2019. 2

Ekin Dogus Cubuk, Barret Zoph, Jon Shlens, and Quoc Le. Randaugment: Practical automated data augmentation with a reduced search space. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, NeurIPS , 2019. 7

Mostafa Dehghani, Josip Djolonga, Basil Mustafa, et al. Scaling vision transformers to 22 billion parameters. In ICML , 2023. 1

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR , 2009. 4, 6

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv:2010.11929 , 2020. 1, 6, 7, 8, 9, 15

Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In CVPR , 2021. 2, 4, 7, 15

Alex Fang, Gabriel Ilharco, Mitchell Wortsman, Yuhao Wan, Vaishaal Shankar, Achal Dave, and Ludwig Schmidt. Data determines distributional robustness in contrastive language image pre-training (clip). In ICML , 2022. 6

Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel. Imagenet-trained cnns are biased towards texture; increasing shape bias improves accuracy and robustness. In ICLR , 2019. 1, 7

Golnaz Ghiasi, Hugo Touvron, Tom DeVries, Tsung-Yi Lin, Yann LeCun, and Roozbeh Mottaghi. Augmax: Adversarial composition of mixtures for robust data augmentation. arXiv:2106.00582 , 2021. 1

Evarist Giné and Richard Nickl. Mathematical foundations of infinite-dimensional statistical models . Cambridge university press, 2021. 18

Micah Goldblum, Liam Fowl, Soheil Feizi, and Tom Goldstein. Adversarially robust distillation. AAAI , 2020. 3, 7, 8

Chengyue Gong, Tongzheng Ren, Mao Ye, and Qiang Liu. Maxup: A simple way to improve generalization of neural network training. In CVPR , 2021. 2

Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv:1412.6572 , 2014. 10, 17

Sven Gowal, Chuan Qin, Po-Sen Huang, Taylan Cemgil, Krishnamurthy Dvijotham, Tim Mann, and Pushmeet Kohli. Achieving robustness in the wild via adversarial mixing with disentangled representations. In CVPR , 2020. 2

Sven Gowal, Sylvestre-Alvise Rebuffi, Olivia Wiles, Florian Stimberg, Dan Andrei Calian, and Timothy A Mann. Improving robustness using generated data. NeurIPS , 2021. 1

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR , 2016. 6, 7, 8, 9

Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. The many faces of robustness: A critical analysis of out-of-distribution generalization. In ICCV , 2021. 2, 6

Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In ICLR , 2019. 7

Dan Hendrycks, Norman Mu, Ekin D. Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshminarayanan. Augmix: A simple data processing method to improve robustness and uncertainty. ICLR , 2019. 1, 2, 9

Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. In CVPR , 2021. 1, 6

Byeongho Heo, Minsik Lee, Sangdoo Yun, and Jin Young Choi. Knowledge transfer via distillation of activation boundaries formed by hidden neurons. In AAAI , 2019. 3

Charles Herrmann, Kyle Sargent, Lu Jiang, Ramin Zabih, Huiwen Chang, Ce Liu, Dilip Krishnan, and Deqing Sun. Pyramid adversarial training improves vit performance. In CVPR , 2022. 2, 7, 8, 10

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. arXiv:1503.02531 , 2015. 2, 3, 8, 15

Tao Huang, Shan You, Fei Wang, Chen Qian, and Chang Xu. Knowledge distillation from a stronger teacher. In NeurIPS , 2022. 7, 8

Tero Karras, Samuli Laine, and Timo Aila. A style-based generator architecture for generative adversarial networks. In CVPR , 2019. 2

Woo Jae Kim, Yoonki Cho, Junsik Jung, and Sung-Eui Yoon. Feature separation and recalibration for adversarial robustness. In CVPR , 2023. 1

Yingwei Li, Qihang Yu, Mingxing Tan, Jieru Mei, Peng Tang, Wei Shen, Alan Yuille, et al. Shape-texture debiased neural network training. In ICLR , 2020. 7, 8

Haoyang Liu, Maheep Chaudhary, and Haohan Wang. Towards trustworthy and aligned machine learning: A data-centric survey with causality perspectives. arXiv:2307.16851 , 2023. 1

Yufan Liu, Jiajiong Cao, Bing Li, Chunfeng Yuan, Weiming Hu, Yangxi Li, and Yunqiang Duan. Knowledge distillation via instance relationship graph. In CVPR , 2019. 3

Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In ICLR , 2018. 1, 2, 17

Chengzhi Mao, Lu Jiang, Mostafa Dehghani, Carl Vondrick, Rahul Sukthankar, and Irfan Essa. Discrete representations strengthen vision transformer robustness. In ICLR , 2022. 9

Xiaofeng Mao, Yuefeng Chen, Ranjie Duan, Yao Zhu, Gege Qi, Shaokai Ye, Xiaodan Li, Rong Zhang, and Hui Xue. Enhance the visual representation via discrete adversarial training. In NeurIPS , 2022. 2, 4, 5, 6, 7, 8, 9, 10, 17

Jieru Mei, Yucheng Han, Yutong Bai, Yixiao Zhang, Yingwei Li, Xianhang Li, Alan Yuille, and Cihang Xie. Fast advprop. In ICLR , 2022. 7, 8, 10

Aditya K Menon, Ankit Singh Rawat, Sashank Reddi, Seungyeon Kim, and Sanjiv Kumar. A statistical perspective on distillation. In ICML , 2021. 4

Wonpyo Park, Dongju Kim, Yan Lu, and Minsu Cho. Relational knowledge distillation. In CVPR , 2019. 3

Nikolaos Passalis and Anastasios Tefas. Learning deep representations with probabilistic knowledge transfer. In ECCV , 2018. 3

C. Qin, J. Martens, S. Gowal, D. Krishnan, K. Dvijotham, A. Fawzi, S. De, R. Stanforth, and P. Kohli. Adversarial robustness through local linearization. In NeurIPS , 2019. 1

Rahul Rade and Seyed-Mohsen Moosavi-Dezfooli. Reducing excessive margin to achieve a better accuracy vs. robustness trade-off. In ICLR , 2021. 2

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICML , 2021. 1, 2, 4, 6, 7, 8, 9, 15

Aditi Raghunathan, Sang Michael Xie, Fanny Yang, John C. Duchi, and Percy Liang. Understanding and mitigating the tradeoff between robustness and accuracy. In ICML , 2020. 2

Sylvestre-Alvise Rebuffi, Francesco Croce, and Sven Gowal. Revisiting adapters with adversarial training. In ICLR , 2023. 2

Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In ICML , 2019. 6

Alexander Robey, George J. Pappas, and Hamed Hassani. Model-based domain generalization. In NeurIPS , 2021. 2

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In CVPR , 2022. 15

Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio. Fitnets: Hints for thin deep nets. In ICLR , 2015. 3

Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv:1910.01108 , 2020. 3

L. Schmidt, S. Santurkar, D. Tsipras, K. Talwar, and A. Madry. Adversarially robust generalization requires more data. In NeurIPS , 2018. 1

Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmarczyk, and Jenia Jitsev. Laion-5b: An open large-scale dataset for training next generation image-text models. In NeurIPS , 2022. 1

Ali Shahin Shamsabadi, Ricardo Sanchez-Matilla, and Andrea Cavallaro. Colorfool: Semantic adversarial colorization. In CVPR , 2020. 1

Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv:2106.10270 , 2022. 7, 8, 10, 15, 16

Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, D. Erhan, Ian J. Goodfellow, and Rob Fergus. Intriguing properties of neural networks. In ICLR , 2014. 1

Jihoon Tack, Sihyun Yu, Jongheon Jeong, Minseon Kim, Sung Ju Hwang, and Jinwoo Shin. Consistency regularization for adversarial robustness. In AAAI , 2022. 1

Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Measuring robustness to natural distribution shifts in image classification. arXiv:2007.00644 , 2020. 6

Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive representation distillation. In ICLR , 2020. 3, 16

Dimitris Tsipras, Shibani Santurkar, Logan Engstrom, Alexander Turner, and Aleksander Madry. Robustness may be at odds with accuracy. In ICML , 2019. 2

Zhuozhuo Tu, Jingwei Zhang, and Dacheng Tao. Theoretical analysis of adversarial learning: A minimax approach. In NeurIPS , 2019. 19

Frederick Tung and Greg Mori. Similarity-preserving knowledge distillation. In CVPR , 2019. 3

Aaron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete representation learning. In NeurIPS , 2017. 2

Haohan Wang, Songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global representations by penalizing local predictive power. In NeurIPS , 2019. 1, 2, 3, 6

Haohan Wang, Zeyi Huang, Xindi Wu, and Eric Xing. Toward learning robust and invariant representations with alignment regularization and data augmentation. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , 2022. 1

Haohan Wang, Xindi Wu, Zeyi Huang, and Eric P Xing. High-frequency component helps explain the generalization of convolutional neural networks. In CVPR , 2020. 2

Huan Wang, Yijun Li, Yuehai Wang, Haoji Hu, and Ming-Hsuan Yang. Collaborative distillation for ultra-resolution universal style transfer. In CVPR , 2020. 3

Huan Wang, Suhas Lohit, Mike Jones, and Yun Fu. What makes a "good" data augmentation in knowledge distillation - a statistical perspective. NeurIPS , 2020. 4

D. Wu, S.-T. Xia, and Y. Wang. Adversarial weight perturbation helps robust generalization. In NeurIPS , 2020. 1

Cihang Xie, Mingxing Tan, Boqing Gong, Jiang Wang, Alan Yuille, and Quoc V. Le. Adversarial examples improve image recognition. In CVPR , 2020. 2, 7, 8, 10

Suorong Yang, Weikang Xiao, Mengcheng Zhang, Suhan Guo, Jian Zhao, and Furao Shen. Image data augmentation for deep learning: A survey. arXiv:2204.08610 , 2022. 1, 2

Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu. Coca: Contrastive captioners are image-text foundation models. arXiv:2205.01917 , 2022. 1

Sergey Zagoruyko and Nikos Komodakis. Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer. In ICLR , 2017. 3

Hongyang Zhang, Yaodong Yu, Jiantao Jiao, Eric Xing, Laurent El Ghaoui, and Michael Jordan. Theoretically principled trade-off between robustness and accuracy. In ICML , 2019. 2

Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. Mixup: Beyond empirical risk minimization. In ICLR , 2018. 1, 2, 7

Peiyan Zhang, Haoyang Liu, Chaozhuo Li, Xing Xie, Sunghun Kim, and Haohan Wang. Foundation model-oriented robustness: Robust image model evaluation with pretrained models. arXiv:2308.10632 , 2023. 2

Xinyu Zhang, Qiang Wang, Jian Zhang, and Zhao Zhong. Adversarial autoaugment. In ICLR , 2019. 2

Shiji Zhao, Jie Yu, Zhenlong Sun, Bo Zhang, and Xingxing Wei. Enhanced accuracy and robustness via multi-teacher adversarial distillation. In ECCV , 2022. 3

Bojia Zi, Shihao Zhao, Xingjun Ma, and Yu-Gang Jiang. Revisiting adversarial robustness distillation: Robust soft labels make student better. ICCV , 2021. 3, 7, 8

6 Appendix

The appendix is organized as follows. First, in Sec. A, we show additional results on using the original knowledge distillation objective. In Sec. B we show the results of additional ablations on the generative model, use of gradients, and transfer to adversarial robustness. In Sec. C we provide additional hyperparameter and implementation details. In Sec. D we show visualizations of DAD and the diversity of its data augmentation to support our theory. In Sec. E we provide full proofs from our theory. Finally, in Sec. F we provide visualizations of generated DAD samples compared to standard and DAT samples.

A Additional results on vanilla knowledge distillation

In Tab. 8 we find that surprisingly, distilling from CLIP on only in-distribution data is able to transfer robust representations, but is generally outperformed by DAD. This works especially well on smaller models, like ResNet34. In fact, it can also improve clean accuracy compared to training from scratch, for all the models we test except ViT-B. In Tab. 9, we find that distilling from CLIP generally results in the highest average robust performance, especially for natural distribution shifts. However, any robust teacher can distill robustness in this setting, including a ResNet50 trained with AugMix as the only robustness intervention. This matches our results for Tab. 5.

B Additional ablations

B.1 Choice of generative model

We use VQGAN [15] as our generative model of choice due to its nature as a image-to-image model, making it suitable as a discretizer. To justify this, we also experiment with Stable Diffusion [53], a

text-to-image generative model. We use the generic prompt "A photo of an object". We observe in Tab. 10 a significant decrease in performance when trained using DAD compared to VQGAN. This suggests the importance of using a discretizer for DAD. Perhaps modifying the text prompt for could boost performance and be an interesting avenue for future work, especially since CLIP also requires a text prompt.

B.2 Use of gradients

DAD is based on adversarial training and uses gradients to find the most diverse and useful data augmentations. To show the importance of using teacher gradients to generate adversarial examples, we implement a sampling-based baseline where we discretize images without the added perturbation. The results in Tab. 11 are significantly worse than DAD, indicating the need to use gradients to discover diverse samples. This is also supported by our theoretical analysis that indicates more diverse adversarial examples are better for robustness. Higher in-distribution performance also suggests the samples are less diverse.

B.3 Transfer to adversarial robustness

Although we center our study on out-of-distribution robustness, it is natural to examine the effect of DAD on adversarial robustness due to the use of adversarial training. In Tab. 12 we attack our trained models with adversarial attacks of various difficulty. We observe a small improvement in adversarial robustness for simpler attacks, but neither DAT or DAD is robust to AutoAttack. This is because the perturbation is discretized and no longer represents the original adversarial example. We observe that DAT is stronger than DAD for ViT-B. Unlike out-of-distribution robustness, since adversarial robustness is based on perturbations generated with gradients from the base model, DAT models are trained on images closer to these perturbations than DAD models (which were trained on perturbations generated with CLIP gradients). However, for ResNet50, DAD is better even for adversarial robustness as distillation is able to help smaller capacity models learn discrete adversarial examples. We also observe higher ResNet50 performance in general in Tab. 1.

C Implementation details

We adopt official hyperparameter settings for a fair comparison for our baselines. For knowledge distillation, we use a temperature of t = 4 for all models and a = 0 . 5 , following [63]. For DAD, we also weight the second KL-divergence term by a . All ViT models are trained with the AugReg [59] hyperparameter and data augmentation configurations.

Following [41], we use the pretrained VQGAN weights from the official GitHub 2 . The VQGAN with f = 8, d = 4 and K = 16384 is used for main experiments.

We use one iteration for the adversarial attack, and an attack learning rate of 0.1.

For the DAT + DAD variant, we add an additional cross entropy loss term with the student adversarial example to the overall training objective, and weight by a .

We conduct all of our experiments on 8 32GB NVIDIA V100 GPUs.

D Wasserstein distance comparisons

To justify our theoretical framework and empirical results we investigate the relationship between Wasserstein distance and model performance. We use a pretrained ResNet50 and calculate Wasserstein distance from batch norm statistics on different distributions using 1000 mini-batches. These distributions are data augmentations generated with the respective methods. In 2a we find that DAD tends to have better performance the larger the distribution shift. This is likely due to the help of distillation letting the model learn robust representations. Additionally, in 2b, we find that relative to the Wasserstein distance between clean ImageNet images and a distribution shift and baseline models, our method has higher performance.

Comparison of Degree of Distribution shift toPerformance

80

ImageNet

DAT

DAD

60

(Ours)

Performance

40

20

0

0.1

0.2

0.3

0.4

0.5

Wassersteindistance

E Proofs

Proof of Lemma 3.1.

Comparison ofDegreeofDistributionShifttoRelative

Performance

0.8

DAT

DAD

(Ours)

0.6

RelativePerformance

0.4

0.2

0

-0.2

0

0.2

0.4

0.6

0.8

Wassersteindistance

Proof. Variational divergence tv measures the divergence between distributions, where B is the set of measurable subsets in P and P ′

Following [5] and from Assumption 2, we can bound total variation with Wasserstein's distance. We take K : R → R , a kernel satisfying a suitable moment condition, so for any coupling of P and P ′ ,

̸

by the Jensen's inequality. Assume that p, p ′ ∈ H α 1 ( R ) . Let { ϕ m } be the orthonormal system in L 2 ([ -1 , 1]) of Legendre polynomials defined by

for x ∈ R and m ∈ N . Define

Then, by Propositions 4.1.5 and 4.1.6 of [19],

Since max x ∈ [ -1 , 1] ( | ϕ m ( x ) | ∨ | ϕ ′ m ( x ) | ) is bounded by a constant depending only on m , where ϕ ′ m is the derivative of ϕ m and a ∨ b is the maximum of a and b ,

Thus,

where ϕ m (0) and α depends only on α . By the triangle inequality, we have

If we take

the proof is complete.

Proof of Lemma 3.3.

Proof. Let ( Z, d Z ) be a metric space where Z = X × Y and d Z is defined as:

where d X and d Y represent the metric in the feature space and label space, respectively. Then we can define the Wasserstein distance between P and P ∗ ,

where Γ( P, P ∗ ) denotes the collection of all measures on Z × Z with marginals P and P ∗ on the first and second factors, respectively.

Following [65], we use the minimax approach, considering the worse-case P ∗ in the Wasserstein ball B w p ε of radius ε centered around P where

Next we define a transport map T : Z → Z to push P to P ∗ as follows:

where x ∗ = arg max x 0 ∈ P ∗ l ( θ ( x 0 ) , y ) . By the definition of d Z , d Z (( x, y ) , ( x ∗ , y )) = d X ( x, x ∗ ) .

Finally, let P ∗ = T θ # P , the pushforward of P by T θ , then we have R ( P, ϵ ) = R ( P ∗ ) . By the definition, we have

Therefore, r ( P, ϵ ) = r ( P ∗ ) . This lets us establish upper bound on the worst-case in the Wasserstein ball and bound the adversarial expected risk. Next we define the radius of the adversary constrained by B as ε B := sup x ∈ B d X ( x, 0) . For any hypothesis h and the corresponding P ∗ = T θ # P , we have w ( P, P ∗ ) ≤ ε B . By the definition of the Wasserstein distance, we have

where the last inequality uses the translation invariant property of d X . Therefore, we have

Proof of Lemma 3.4

Proof. We use P to denote the (at least) one distribution in the intersection of P 1 and P 2 .

where c is a positive constant.

Similarly, we can have

F Visualizations

(a) Clean

(c) DAT

(d) DAD

2IB

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a framework that improves the image classification model's robustness by distilling CLIP models and augmenting adversarial learning with pre-trained generative models. The authors follow classical adversarial learning to generate perturbed examples and then input the examples to VQ-GAN. Lastly, they distill a CLIP model into a smaller network with both normal and augmented adversarial examples. While I have some concerns about the results, extensive experiments verify the effectiveness of the proposed method in many settings.\n\nSTRENGTHS:\n1. The extensive experiments cover various settings, such as different architectures for teacher and student models and variations of ImageNet.\n2. The authors provide some theory explaining that a more diverse distribution could be helpful for distillation.\n\nWEAKNESSES:\n1. The research problem is unclear. In L37, the authors claim to \"train small robust models with a large-scale model as a teacher, but without access to the teacher's training data.\" However, they focus more on out-of-distribution robustness in the later parts and the experiments.\n \n2. The proposed method needs to be better justified. In L118, the authors suddenly introduce generative models to augment adversarial examples and claim that it helps generalizability without any reasoning. If possible, the authors should provide more motivation and insights into it and conduct some experiments.\n\n3. Many notations are not defined. For example, the definition of $\\theta(x, x\\prime)$ is missing. What exactly does $Q(x\\prime)$ work? \n\n4. The novelty is limited. The main contribution seems solely to leverage VQ-GAN to enhance distillation. The reason for choosing VQ-GAN is also missing.\n\n5. The research scope is limited. The authors claim in L45 that distilling large-scale models into smaller models improves adversarial robustness. However, the experiments only consider CLIP models.\n\nQUESTIONS:\n1. The performance could be better, even though it outperforms prior work. In Table 1, neither the baselines nor the proposed method achieves the CLIP baseline performance. Maybe one should look into cross-modality instead of augmenting query datasets for distillation.\n\n2. What does it mean by *DAD (no distill)*? I recommend the authors refer the readers to the corresponding loss terms.\n\n3. The proposed method is similar to leveraging more data to perform distillation (Eq. 8). The author may like to explore using different kinds of data (e.g., ImageNet + CIFAR-10 or ImageNet + ImageNet with augmentation), which could even outline the contribution of the proposed method.\n\nLIMITATIONS:\n1. The authors only consider distilling CLIP models.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper proposes a simple and lightweight framework for improving the robustness of vision models through knowledge distillation. This paper applies pre-trained teacher model to generate adversarial examples and apply VQGAN as a data augmentation method to generate more informative adversarial samples. This paper provides a theoretical framework for applying robust teacher in the knowledge distillation with data augmentation.

STRENGTHS:
+The paper applies a large pre trained model based on knowledge distillation to enhance the model's out-of-distribution robustness with data augmentation, this idea is novel.

+The paper provides a theoretical analysis of how to best distill robustness from a robust teacher trained on a large-scale dataset.

+The experiment of the paper can prove the effectiveness of the method.

WEAKNESSES:
-The motivation for applying VQGAN is not so clear. The article claims that VQGAN can discretize adversarial examples, but can other data augmentation technology such as common data augmentation methods (translation, cropping), other variants of GAN, or Stable Diffusion achieve better performance? This part seems to be a direct application, and I think a discussion of the necessity should be provided.

-Although this paper has a pretty performance of out-of-distribution robustness, I wonder about the performance of adversarial robustness (such as PGD or AutoAttack).

-The frame diagram is somewhat rough and crude (this will not affect the final rating, but I hope the author can make some modifications).

QUESTIONS:
see the weakness

LIMITATIONS:
see the weakness
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a new method to improve the robustness of vision models through knowledge distillation by leveraging a large-scale pre-trained teacher model (CLIP) and a VQ-GAN discretizer. The paper does a great job motivating the method and arguing for knowledge distillation when tackling out of distribution robustness. It also includes extensive theoretical justifications. And finally, the paper introduces and evaluates a novel distillation objective (\"DAD\"). In the proposed method (\"DAD\") the student model is trained to minimize a loss made up of a standard cross-entropy loss, a KL to the teacher's predictions (i.e. standard knowledge distillation) and an additional KL between the teacher and student's predictions on additional (teacher-adversarial) data points (which is novel). The paper also includes a reasonable evaluation section for the method with overall positive results, while highlighting some considerable performance improvements on natural distribution shifts over baselines.\n\nSTRENGTHS:\nThe paper is relatively well written and easy to follow. The paper introduces a relatively novel and interesting knowledge distillation loss which should have increasingly wider applicability and grow in significance as more and more of the ML world is moving towards leveraging pre-trained foundation models.\n\nWEAKNESSES:\n- Missing paragraph detailing DAT: since a large proportion of the strongest results seem to be relying on combining DAD with DAT. The paper should also include a description of DAT's loss and the \"DAT + DAD\" loss written out (in the appendix if no space available in the main body).\n- The paper mentions the method has \"negligible computational overhead\" -- I strongly disagree. Computing adversarial examples for a CLIP model that also passes them through VQGAN is not negligible. One can say that the computation cost is amortized though -- as you only have to do it once and then read it from disk.\n- Missing \"Avg\" column in Table 2 and 3, and best results are not shown in bold in Table 2.\n- Loss is unclear (Eq 9) -- please expand and clarify this to show which variables are being optimized over, at which stage, which are the constraints exactly, and specify the losses (what is l_1 and l_2 precisely - these are only mentioned in passing in line 125-126 being \"typically\"/\"usually\" CE and KL); and also please use or remove \"n\", as currently it is not used to index anything.\n\n- Missing references, e.g., on [1] (on worst case risk over a population of distributions) and [2] (for a related formulation of a similar set of assumptions and similar theoretical conclusion):\n\n[1] Arjovsky et al., Invariant Risk Minimization.\n\n[2] Calian et al., Defending against image corruptions through adversarial augmentations.\n\n- Figure 1 does not seem to be faithful depiction of the method as described in the paper, so it should ideally be improved significantly (the caption could also give more details towards explaining the method).\n\nQUESTIONS:\n- At a high level, looking at equation (9) of the full model loss, I'm trying to understand how each component affects downstream robustness. So, I'd like to able to map the contribution of each of the two \"l_2\" losses and the way in which the adversarial example is computed to specific results. This seems like a pretty straightforward and important ablation. Table 3 and Table 5 present just a slice of this ablation as far as I can tell. Could you please clarify how one can reconstitute this ablation from the results presented in the paper?\n\n- The adversarial examples could also not be computed but rather perturbations could be sampled instead (while still using the VQGAN), or random data augmentation methods (or just an identity function used, i.e. which would correspond to just dropping the last l_2 loss in Eq 9). (From Table 6 it seems that more iterations actually results in worse performance across many metrics.) Have you experimented with this?\n\n- In Table 1, does the row \"DAT + DAD (Ours)\" refer to training with DAT & DAD on top of AugReg-ViT[55]? If yes, please mark it with a leading \"+\" for clarity.\n\n- In line 303, the paper states \"These adversarial examples are generated with DAD\" -- was \"DAT\" intended?\n\nLIMITATIONS:\nThe authors adequately addressed the method's limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel approach called discrete adversarial distillation (DAD) to train small, robust models using large-scale models as teachers. The authors establish the knowledge distillation framework for out-of-distribution robustness and provide a theoretical framework for utilizing large-scale models as teachers. DAD achieves state-of-the-art performance on natural distribution shifts, surpassing adversarial training and traditional distillation techniques.\n\nSTRENGTHS:\nOverall, this paper is also well-written and clearly structured; the authors provide both theoretical analysis and empirical experiments to verify their approach. The reviewer would like to list some strengths of the work.\n\nFirstly, it extends the knowledge distillation framework to the domain of out-of-distribution robustness, opening up new possibilities for training small, robust models that can generalize effectively across diverse populations and environments.\n\nSecondly, this paper highlights the advantages of leveraging large-scale models as teachers, providing a theoretical framework to support this approach. This understanding can inform the development and utilization of robust models in practical applications.\n\nLastly, the experimental results showcase state-of-the-art performance on natural distribution shifts, surpassing existing techniques. This signifies the potential of the proposed DAD method in addressing the challenges of robustness and generalization in ML models.\n\nWEAKNESSES:\nThe reviewer found it hard to understand the main idea at first glance, the paper could improve its clarity in figures, as well as in presenting the assumptions and theoretical analysis, and its connection to the proposed method. Explicitly linking the theoretical framework to the practical implementation of the DAD method would help readers grasp the underlying principles and motivations behind the proposed approach more effectively.\n\nQUESTIONS:\nPlease refer to the weakness. In addition, why the authors did not provide the code for review but answered the *Reproducibility* as \"yes\"? Please provide the code during the rebuttal period, using Anonymous GitHub is a viable option. Also, the authors claim that the proposed only brings negligible computational overhead, but the reviewer could not find the part that addresses computing resources; what exactly are the computational costs for the proposed framework?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper presents a knowledge distillation framework for vision model, by leveraging a teacher model (CLIP). In the setup, a discretizer (VQGAN) model is introduced to the adversarial examples from the teacher model. Adversarial training (AT) objective is adapted in the knowledge distill setting. The proposed discrete adversarial distillation (DAD) demonstrate improved performance in out-of-distribution tasks, comparing to other AT approaches.

STRENGTHS:
This paper present an interesting approach for knowledge distillation for large scale vision model. The experiment is thoroughly conducted with various of SOTA approaches compared, and comprehensive ablation study. The writing is easy to follow and the idea is well established.

WEAKNESSES:
1. One main concern is on comparing the proposed method with DAT. In DAT, the idea of discretizer is introduced and VQGAN was also considered in their framework. In their paper, various tasks were evaluated, including image classification, self-supervised learning, and object detection. To me, DAD seems to be applying the similar methodology, but on a OOD setup. 

2. In the experiment, it seems like DAT has achieved competitive results in ImageNet-1K, and even better results in ImageNet-21K. What is the key advantage of DAD over DAT?

QUESTIONS:
1. How come DAD+DAT performs worse than DAT in some cases?
2. Are there any qualitative results? what are the drawbacks from previous model(s) that are better addressed by the newly proposed DAD?
3. Have you try training the teacher model in parallel with the student model, not just provide adversarial samples from offline manner?

LIMITATIONS:
Please see weakness.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08