SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Andy Zhou $^{1,3}$ , Jindong Wang $^{2}$ , Yu-Xiong Wang $^{1}$ , and Haohan Wang $^{1}$
$^{3}$ AI@UIUC
{andyz3, yxw, haohanw}@illinois.edu, jindong.wang@microsoft.com
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.
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 large-scale 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  $\sim 13.7\mathrm{x}$ .
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
2. Providing a theoretical framework in KD for the choice of a teacher based on the diversity of the data augmentation
3. Proposing a novel data augmentation DAD that outperforms both adversarial training and distillation techniques on natural distribution shifts
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.
$^{1}$ code at https://github.com/andyz245/DiscreteAdversarialDistillation
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 adversarily 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.
We consider a standard dataset  $\{(x,y)\}_{n=1}^{N} \sim 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  $\theta$ . We also consider a discretizer,  $Q$ , and a teacher model  $\phi$  and we use  $\phi(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 \in A$ , the class of all such functions.
Invariance is a desirable property where the model will have the same representation of an input after a transformation is applied. A model,  $\theta$  is said to be invariant if  $\theta(x) = \theta(x')$  for all  $x \in U$ , where  $U$  is the set of all images that can be obtained by a transformation of  $x$  by  $a \in 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  $\hat{a}$  and  $\hat{A}$ . We can consider  $\hat{a}$  to represent an individual distribution  $P$  and  $\hat{A}$  to be drawn from the distribution of distributions  $\hat{P}$ . Our ultimate goal is to train  $\theta$  to be invariant to transformations in  $\hat{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
$$
r _ {\text {w o r s t}} (P, \epsilon) = \max  _ {P ^ {\prime}: w \left(P ^ {\prime}, P\right) \leq \epsilon} \mathbb {E} _ {(x, y) \sim P ^ {\prime}} l (\theta (x), y), \tag {1}
$$
where  $l$  is the loss function and  $r$  depends on an anchor distribution  $P$ , and  $\epsilon$ , 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.
$$
r (P, \epsilon) = \mathbb {E} _ {P ^ {\prime} \sim \hat {P}: w \left(P ^ {\prime}, P\right) \leq \epsilon} \mathbb {E} _ {(x, y) \sim P ^ {\prime}} l (\theta (x), y), \tag {2}
$$
$$
r (P) = \mathbb {E} _ {(x, y) \sim P} l (\theta (x), y). \tag {3}
$$
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
$$
\min  \mathbb {E} _ {(x, y) \sim P} [ l (\theta (x), y) + \max  l (\theta \left(x ^ {\prime}\right), y) ], \tag {4}
$$
where  $x' = x + \epsilon$ ,  $\epsilon$  is the perturbation, and  $l$  is the cross-entropy loss. This achieves adversarial robustness, but cannot generalize well to real-world domain shifts in  $\hat{A}$ . To address this, we consider a generative model,  $Q$ , trained on  $P$  and can model  $\hat{A}$ . Passing an input through  $Q$  in the maximization objective applies a worst-case transformation from  $\hat{A}$ . This modifies (4) to minimize the empirical semantic adversarial risk,
$$
\min  \mathbb {E} _ {(x, y) \sim P} [ l (\theta (x), y) + \max  l (\theta (Q (x ^ {\prime})), y) ]. \tag {5}
$$
3.2 Distillation from a robust teacher
Next, we consider a setting where we also have access to a pretrained  $\phi$  invariant to distribution shifts in  $\hat{A}$ . This enables us to leverage knowledge distillation (KD). This modifies (3) in the following
$$
\min  \mathbb {E} _ {(x, y) \sim P} [ l _ {1} (\theta (x), y) + l _ {2} (\theta (x), \phi (x)) ], \tag {6}
$$
Figure 1: The overall framework of discrete adversarial distillation (DAD). We leverage a foundation model to generate and distill adversarial examples after discretization by a VQGAN.
where  $l_{1}$  is the classification loss, the cross-entropy loss, and  $l_{2}$  is a distance loss between  $\theta (x)$  and  $\phi (x)$ , the KL divergence. (6) can be approximated by the empirical risk where
$$
\min  l _ {1} (\theta (x), y) + l _ {2} (\theta (x), \phi (x)). \tag {7}
$$
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  $\phi$  on augmented samples,
$$
\min  l _ {1} (\theta (x), y) + l _ {2} (\theta (x), \phi (x)) + l _ {2} \left(\theta \left(Q \left(x ^ {\prime}\right)\right), \phi \left(Q \left(x ^ {\prime}\right)\right)\right). \tag {8}
$$
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 cross-entropy 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 vision-language model using only ImageNet [13] data. In the practical setting, we use approximations of an ideal discretizer and robust teacher. For  $\phi$  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  $\epsilon$  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,
$$
\min  l _ {1} (\theta (x), y) + l _ {2} (\theta (x), \phi (x)) + l _ {2} (\theta (Q \left(x ^ {\prime}\right)), \phi (Q \left(x ^ {\prime}\right))). \tag {9}
$$
$$
\text{where} x^{\prime} = \max_{||x^{\prime} - x||_{p}\leq \epsilon}l_{1}(\phi (x),y),  \phi (Q(x^{\prime})) = y.
$$
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  $\hat{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  $\alpha$  depending on the smoothness of  $P$  and  $P'$  where  $cw(P, P')^\alpha$  and  $c > 0$  and is only dependent on  $\alpha$ .
Assumption 3. For function  $\gamma (|\Theta |,n,\beta)$  parameterized by hypothesis space  $|\Theta |$  , number of samples  $n$  , and the probability when the bound holds  $\beta$  , if the samples are i.i.d,  $\gamma (|\Theta |,n,\beta) = 2\mathcal{R}(\mathcal{L})+$ $\sqrt{(\log 1 / \beta) / 2n}$  , where  $\mathcal{R}(\mathcal{L})$  stands for Rademacher complexity and  $\mathcal{L} = \{l_{\theta}\mid \theta \in \Theta \}$  , where  $l_{\theta}$  is the loss function corresponding to  $\theta$  . Additionally, if  $\Theta$  is finite,  $l(\cdot ,\cdot)$  is a zero-one loss, and samples are i.i.d, then  $\gamma (|\Theta |,n,\beta) = \sqrt{(\log(|\Theta|) + \log(1 / \beta)) / 2n}$
Lemma 3.1. Given Assumptions 1 and 2 and variational divergence  $tv$ , for two arbitrary distributions  $P$  and  $P'$  with corresponding density functions  $\delta$  and  $\delta'$ ,  $r(P') \leq r(P) + w(P', P)$ .
Lemma 3.2. Given Assumption 3, Lemma 3.1, and probability at least  $1 - \Gamma$
$$
r \left(P ^ {\prime}\right) \leq r \left(\left(X, Y\right) _ {P}\right) + w \left(P ^ {\prime}, P\right) + \xi \left(n _ {(X, Y) _ {P}}, \Theta , \Gamma\right)
$$
where  $n_{(X,Y)_P}$  denotes the number of sample sizes in the finite dataset  $(X,Y)_P$ , and  $\xi$  is a vanilla term that connects  $n_{(X,Y)_P}$  and  $\Theta$  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.  $\square$
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  $\xi$  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 \sim \hat{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  $\hat{A}$ .
Therefore, we use  $\mathbf{P} \in \hat{A}$  to denote a set of  $m$  distributions, i.e.,  $\mathbf{P} = \{P_1, P_2, \dots, 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', \mathbf{P})$  as the average divergence between distributions within the set. Thus,  $w(P', \mathbf{P}) := \sum_{i}^{m} w(P', P_i) / m$ ,  $\forall P_i \in \mathbf{P}$ .
Lemma 3.3. Given a distribution  $P$ , we generate a new distribution  $P^{*} \in \hat{A}$  using the discretized worst-case adversarial samples of a model  $\theta$ . Training  $\theta$  with adversarial training is equivalent to training  $\theta$  with empirical risk minimization on  $P^{*}$  where  $w(P, P^{*}) \leq \epsilon$ .
We leave the full proof of Lemma 3.3 in Sec. E of the Appendix. Finally, let's denote the model  $\phi$  trained over distribution  $P$  as  $\phi(P)$  and the adversarial data augmentation process that results in a new semantic data distribution as  $D$ . We aim to compare  $w(D(\phi(\mathbf{P})), P')$  and  $w(D(\phi(P)), P')$ . In this context,  $P'$  is any arbitrary testing distribution,  $P$  is a specific training distribution, and  $\mathbf{P}$  represents the set of distributions used for training foundation models.
Lemma 3.4. Assuming  $\hat{P}$  is continuous and has a finite expected value, for two sets of distributions  $\mathbf{P}_1$  and  $\mathbf{P}_2$ , assuming there is at least one distribution in the intersection of  $\mathbf{P}_1$  and  $\mathbf{P}_2$ , for a fixed testing distribution  $P'$ , we have
$$
\mathbb {E} _ {\hat {P}} \Big [ \big | w (D (\phi (\mathbf {P} _ {1})), P ^ {\prime}) - w (D (\phi (\mathbf {P} _ {2})), P ^ {\prime}) \big | \Big ] \leq 2 \mathbb {E} _ {\hat {P}} \big [ \sup _ {P \in \mathbf {P} _ {1}} w (P, P ^ {\prime}) + \sup _ {P \in \mathbf {P} _ {2}} w (P, P ^ {\prime}) \big ]
$$
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  $\phi$ , each trained with distinct data sets  $(\mathbf{P}_1$  and  $\mathbf{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  $\mathbf{P}_1$ , with DAT [41] and the traditional approach of adversarial training based on the student model, referred to as  $\mathbf{P}_2$ . Our results suggest two key implications: