DAFT: Distilling Adversarially Fine-tuned Teachers for OOD Generalization
Abstract
We consider the problem of OOD generalization, where the goal is to train a model that performs well on test distributions that are different from the training distribution.
Deep learning models are known to be fragile to such shifts and can suffer large accuracy drops even for slightly different test distributions (Hendrycks & Dietterich, 2019).
We propose a new method – DAFT – based on the intuition that adversarially robust combination of a large number of rich features should provide
OOD robustness.
Our method carefully distills the model from a powerful teacher that learns several discriminative features using standard training while combining them using adversarial training.
The standard adversarial training procedure is modiﬁed to produce teachers which can guide the student better.
We evaluate DAFT on standard benchmarks in the DomainBed framework (Gulrajani & Lopez-Paz, 2020), and ﬁnd that DAFT consistently out-performs well-tuned ERM and distillation baselines by up to 6%, with more pronounced gains for smaller networks.
1. Introduction
Several recent works have shown that standard deep learning models trained with stochastic gradient descent (SGD) style methods can be fragile and might suffer a large drop in accuracy if the test data distribution (also known as target domain) is even slightly different compared to the training data distribution (also known as source domain)
(Hendrycks & Dietterich, 2019; Gulrajani & Lopez-Paz, 2020).
However, in practice, it is quite challenging to obtain training data that exactly matches the test distribution.
For example, due to privacy restrictions we may not be able to access the data of the actual customers of a web-application.
Instead, training data is generated using crowd workers or by seeking volunteers who are willing to donate their data for training.
This clearly leads to distribution shift between the training and test data.
Consequently, it is crucial to design models and training mechanisms that are robust to distribution shifts and can perform well on out-of-distribution (OOD) data.
Even in settings where some amount of data can be collected from the ﬁnal deployment setting, it should be relatively easier to adapt models with good OOD generalization .
OOD problems have been studied in a variety of settings where different amounts of source/ target information might be available.
We consider the OOD generalization setting which is one of the weakest settings, and requires only a labeled training dataset, without any information about the target dataset or even about the sources present in the training dataset.
Note that this setting is slightly different and more challenging than the popular domain generalization setting which requires the identity of source domain for each training point.
Even, in the domain generalization setting, several recent works (Gulrajani & Lopez-Paz, 2020; Vedantam et al., 2021) show that a well-tuned ERM model is still competitive with
SOTA methods.
We propose DAFT to mitigate the problem of OOD generalization ; DAFT is motivated by three key observations as discussed in Sec 3.
Our contributions : In summary, we introduce a novel method – Distillation with Adversarially Finetuned Teacher (DAFT) – that uses the above algorithmic insights to design a robust technique for OOD generalization.
DAFT trains a student model by distilling with a standard trained as well as an adversarially ﬁnetuned teacher model.
Finetuning uses a smooth KL divergence based loss for all logits/prediction values.
See Figure 1 for an overview of DAFT.
We conduct extensive experiments in the standard DomainBed framework (Gulrajani & Lopez-Paz, 2020) – using the prescribed methodology for evaluation, hyperparameter tuning – and compare DAFT against various baselines on the ﬁve OOD datasets in the testbed.
Recall that even for stronger Domain Generalization setting, (Gulrajani &
Lopez-Paz, 2020) showed that well-trained ERM is nearly
SOTA.
In contrast, we demonstrate that DAFT trained in the weaker OOD generalization setting, still consistently out- [Figure] performs ERM (trained according to DomainBed approach) and other baselines.
DAFT is particularly effective for smaller network architectures.
For example, on DomainBed datasets, DAFT trained ResNet-50 models are on an average 4% more accurate than ERM as well as other baselines.
In fact, DAFT+ResNet-50 models are more accurate than an
ERM trained ResNet-152 as well; see Table 1.
Limitations :
In general, different application areas/problems might lead to different forms of distribution shift between training and test distributions.
Consequently, there may not be a single algorithm that is optimal under all distribution shifts.
To address this, we perform experiments on multiple OOD datasets with different characteristics such as satellite images with temporal/geographical shift in TerraIncognita, and object images with varying sources in OfﬁceHome, VLCS, PACS and DomainNet datasets.
In all of these settings, DAFT gives signiﬁcant improvements over
ERM.
However, there may certainly be other forms of distribution shift, beyond those captured in the above datasets, where new ideas might be required to obtain superior performance.
2. Problem Deﬁnition
Problem setting : We are given a training dataset of labeled examples D S
= { ( x i , y i ) : i ∈ [ n ] } , each example drawn independently from a source distribution S .
Given a model
M , we also assume that we can evaluate its accuracy on validation distribution V , which may be different from S .
Note that we do not explicitly access examples from V , but rather use it only for evaluating the accuracy of any given model.
Our goal is to train a model M ∗ that performs well on data drawn from a target distribution T .
Primarily we study the setting of T (cid:54) =
V — especially for the experiments reported in the main paper — but for generality, we report results for T
= V also in the appendix.
We would like to stress that we do not have access to T during model training (and also model/hyper-parameter selection if T (cid:54) = V ).
Motivation : In several settings, there are privacy and proprietary reasons for not having access to or not using validation or target data (i.e., V and T respectively) during model training.
A common strategy to deal with this is to generate training data through other ways such as using crowd workers, requesting some users to volunteer their data (in the case of user data), using simulation models etc.
However, once we train a model, we often have the ability to deploy/evaluate the model for its accuracy on the validation domain V , and in some cases, even on the target domain T .
3. Methodology
Motivation and high level description of algorithm :
We begin by describing the key insights that led to our algorithm, and provide a high-level description.
Recent works have demonstrated that adversarial training can learn more robust features compared to standard training (Ilyas et al., 2019;
Yi et al., 2021).
However, we ﬁnd that vanilla adversarial training does not provide substantial improvements on larger datasets.
Hence, we hypothesize that standard ERM trained models might already learn features which are good for domain generalization (Kirichenko et al., 2022; Kumar et al., 2022; Rosenfeld et al., 2022), but the ﬁnal layer is not able to combine these features in a manner robust to domain shifts.
To further illustrate this point, we perform an experiment using a modiﬁcation of a Fashion-MNIST subset with only two-class: shoe and top .
We superimpose images onto coloured backgrounds, where the colour varies linearly between red (255 , 0 , 0) and green (0 , 255 , 0) .
Training images have a strong correlation between the background colour and label, i.e. color of tops images range from (255 , 0 , 0) to (123 , 132 , 0) , while that of shoes images range between (132 , 123 , 0) and (0 , 255 , 0) .
In general, color can easily distinguish between the classes, but there is a small region between (123 , 132 , 0) and (132 , 123 , 0) , where color cannot distinguish between the classes.
During test time, there is no such correlation with colour i.e. the data is OOD w.r.t.
the train data.
For models trained on this data, we compute the correlation of each neuron at the output of the feature extractor with the shape and colour of the images.
Note that color is a non-robust spurious feature while shape is a robust feature that is strongly correlated with the labels and is useful for prediction despite OOD shifts.
Now, an ERM trained model has an in-domain (ID) test accuracy of 99.9%, and OOD test accuracy of 60%.
Furthermore, only 2 of the 32 features are highly correlated with the robust shape feature, while the rest are correlated with color.
The ﬁnal class output is dominated by the color features.
In contrast, an adversarially trained model has a higher number of shape features (8 out of 32).
But the in-domain accuracy is only 98.3%, and the OOD accuracy is 58%, lower than
ERM.
A probable reason is that even though adversarial training learns more shape-correlated robust features, but the average correlation with shape is much smaller (around 0 .
than the similar correlation of features from standard trained model (around 0 . 8 ).
This is possibly because the shape features learned by adversarially trained models are more suited to the goal of adversarial robustness, while the features learned by standard ERM models are better correlated with the standard classiﬁcation task.
Hence, we introduce the method of adversarial ﬁne-tuning of the last layer after standard ERM based pre-training.
This encourages the model to give a lower prediction weight to color features, and higher weight to the robust shape features learned by the ERM model.
With adversarial ﬁne-tuning, a small ID accuracy drop 99.5% occurs, but the OOD accuracy jumps to 64%.
Next, motivated by the observation that distillation from larger models often helps in-domain performance, we trained our ﬁnal model through distillation of a larger model which was itself trained using adversarial ﬁne-tuning.
Surprisingly, vanilla distillation failed to transfer the superior performance of adversarially teacher model to the student model.
We identify the main reason behind the failure of distillation in transfering the superior OOD performance of teacher to the student to be the following: while the Cross-Entropy loss on adversarial samples ensures that the logit corresponding to the correct class is “robust”, it does not put any constraints on the remaining logits.
Consequently, the remaining logits do not provide useful information for distillation.
In order to tackle this, we add an additional KL divergence term to ensure that all the logits of the teacher model are smooth in the neighborhood of the given input, and are aligned to the logits of the clean image.
The loss to be minimized the teacher is hence
L ( W, ( x, y ))
=
max ˆ x ∈ B (cid:15) ( x ) KL ( z ||
ˆ z )
+ CE
(ˆ z, y ) , where z and ˆ z are the logits of the teacher for x and ˆ x , and CE is the cross-entropy loss.
We use l 2 constrained perturbations with (cid:15) being a hyper-parameter.
The ﬁnal ingredient of our approach is to use two teacher models: the adversarially ﬁne-tuned teacher on inputs where it predicts correctly and the standard trained model on the remaining inputs.
We call the resulting algorithm Distillation of Adversarially Fine-Tuned teacher (DAFT).
4. Experimental Results
In this section, we detail our experimental setup, datasets, baselines and results.
4.1. Experimental Setup
Our experimental setup follows the approach and recommendations of (Gulrajani & Lopez-Paz, 2020).
For all our experiments, we train models of different sizes from the
ResNet (He et al., 2015) family.
Hyper-parameters for all the methods are tuned using the leave-one-domain-out approach descibed in (Gulrajani & Lopez-Paz, 2020).
We use ImageNet pretrained models for comparisons on DomainBed.
We report the mean and std deviation of the metrics across ﬁve random restarts.
Datasets We report OOD accuracy results on 5 different datasets, and the average accuracy across them.
We use all the datasets from the DomainBed (Gulrajani & Lopez-Paz, 2020) benchmark.
Baselines We compare DAFT against the standard ERM method trained on training data D S .
We also compare against AT which trains ( not ﬁnetunes ) models using PGD for l 2 -norm constrained input adversarial perturbations (Madry et al., 2019), as well as TRADES (Zhang et al., 2019) which is a variant.
Since DAFT uses larger models to train a smaller student network, we also compare it against the performance of distilling directly from an ERM trained teacher model.
The teacher in all cases is ResNet-152.
4.2. Results
We compare the OOD accuracy of DAFT against baselines in Table 1.
We notice that our method provides signiﬁcant improvements over standard ERM across all datasets and model sizes.
For example, on OfﬁceHome, we show gains of almost 5% over ERM on all model sizes.
We also note that smaller models trained with DAFT outperform ERM trained larger models; ResNet-34 trained with DAFT is close in performance to ResNet-152 on an average, while
ResNet-50 can beat it.
We also notice that our method outperforms standard logit distillation on all benchmark datasets.
This demonstrates that our method leverages the information provided by larger models in a more efﬁcient manner.
Furthermore, the KLregularization of teachers in DAFT helps improve the transfer, as we demonstrate in the ablation experiments (sec 4.3).
[Table]
ES N
[Table] [Table]
performs better than ERM, while AT is often worse than or similar to ERM due to poor in-domain accuracy of AT. We also notice that DAFT is able to outperform standard baselines by larger margins for datasets like TerraIncognita where the domains are signiﬁcantly different from ImageNet.
This means that features learnt from ImageNet pretraining would be less useful in this scenario.
The better performance of DAFT on this dataset implies that it is able to transfer generalizable features better.
We also verify the gains of
DAFT without using ImageNet pretrained models, but do not report them due to lack of space.
We also performed several ablations of our method to verify the effectiveness of each component.
4.3. Ablations
Does adversarial ﬁnetuning work?
To study the effect of our teacher training paradigms, we compare performance of various teacher models on the PACS, VLCS and OfﬁceHome dataset in Table 2.
We show that it is much better to pre-train a model and ﬁnetune the ﬁnal layer adversarially (AF), rather than training the full model adversarially (AT).
Note that AT is competitive to ERM only on the OfﬁceHome dataset, since there is a high inter-domain similarity in three of the four domains of this dataset, and the images are also similar to ImageNet, on which the models were originally pretrained.
This is consistent with the ﬁndings of (Yi et al., 2021).
Effect of L smooth on distillation : To verify the effect of using teachers trained with L smooth , we present results on three datasets in Fig 2.
For each split, we compute the average gain in OOD accuracy for the teacher (which is a
ResNet-152) when trained with adversarial ﬁnetuning with ( ∆ Teach L smooth ) and without ( ∆ TeachAF ) L smooth .
We then compare the gains over standard distillation observed in students distilled from these teachers ( ∆ Dist L smooth and ∆ DistAF respectively).
Note that gain here refers to the difference in the OOD accuracy of the modiﬁed teacher (resp teacher) model over a standard ERM (resp.
distilled student of a standard teacher) model.
We observe that students distilled from a teacher trained with L smooth obtain similar or even better accuracy gains compared to those achieved by the teacher.
In contrast, students of teachers trained without this term do not even consistently achieve similar accuracy gains as their teachers.
5. Conclusion
Summary :
In this paper, we considered the problem of out of distribution (OOD) generalization, where we are given training examples from a source distribution and are required to output a model which will be evaluated on test examples sampled from a different target distribution.
We ﬁrst observed that the non-robustness of standard trained models on OOD data is primarily due to a non-robust combination of learned features in the ﬁnal linear layer and that the features themselves are capable of obtaining high OOD accuracy.
Inspired by this observation, we designed adversarial ﬁnetuning (AF) which ﬁrst trains the model using standard training and then ﬁnetunes the ﬁnal linear layer using adversarial training.
Motivated by the in-domain accuracy improvements obtained by distillation in prior works, we attempted to train a student model by distilling a teacher model that is trained by AF.
However, we observed that standard distillation does not yield large improvements for AF trained teachers.
We identiﬁed the reason for this to be the instability of logit values around the input and to address this, we incorporated an additional loss term in AF to encourage the logit values of teacher to be smooth.
Finally, to tackle the suboptimal in-domain accuracy of AF trained teacher, we distilled both standard trained and AF trained teachers into the student giving our ﬁnal algorithm DAFT.
On ﬁve benchmark datasets, with diverse kinds of distribution shifts, we showed that DAFT provides signiﬁcantly higher OOD accuracy when compared to ERM as well as standard baselines like adversarial training.
We also presented ablation studies showing the importance of various components of DAFT.
Limitations & Future work : An avenue for future work is devising the optimal way of performing AF – in this work, we only consider ﬁnetuning the ﬁnal linear layer but have not explored if this can be further improved by ﬁnetuning last few layers or other subsets of parameters.
Finally, models that are fragile to OOD shifts and depend on spurious correlation can signiﬁcantly amplify biases in data.
So, further investigation of DAFT for mitigating biases in data is highly interesting.
A. Experimental Setup A.1. Hyperparameters
We use the Adam optimizer for all our experiments, with a batch size of 64.
We run ERM, adversarial training, and distillation for 10000 steps each, while ﬁne-tuning is run for 5000 steps.
We tune the following hyper-parameters for our methods and baselines -
• Learning Rate - Selected from the range [10 − 6 , 10 − 3 ] • Norm of adversarial perturbation (cid:15) -
Selected from the range [0 . 05 , 0 .
5] • Number of adversarial perturbation steps k - An integer selected from the range [3 , 7] • LR for PGD - Selected from the range [10 − 3 ,
10 − 1 ] • Distillation temperature τ - An integer selected from the range [2 , 8] • Weight α for L smooth - Selected from the range [10 − 6 , 10 − 2 ]
The hyperparameters were tuned using random search over the intervals, with 32 conﬁgurations being considered for each algorithm.
A.2. Datasets
The data can be downloaded using the DomainBed repo
A.3. Hardware Setup
We conducted all experiments on a single A100 GPU.
The experimental code was using the DomainBed framework, in
PyTorch.
B. Additional results B.1. Results on WILDS benchmark
We perform experiments on the WILDS benchmark (Koh et al., 2021) with non-ImageNet pre-trained models to verify the efﬁcacy of our approach in settings where pre-training on large datasets is not possible.
We present results on iWildCams and FMoW-WILDS datasets in tables ?? .
We notice that DAFT consistently gives gains over most of the baselines
[Table] [Table]
and
[Table]
methods.
[Table]
We also experiment with a variant of adversarial ﬁne-tuning where we perform perturbations in the feature space of the model rather than the input space.
The comparison with adversarial ﬁnetuning on four datasets is reported in table 5.
We notice that the difference in the performance obtained is not consistent across datasets or sizes.
While the performance of ﬁnetuning in feature space is slightly more for ImageNet pre-trained models, we notice that the range over which (cid:15) needs to be ﬁne-tuned is larger for this variant, and the obtained best (cid:15) differs quite a bit between different models.
On the contrary, for input space perturbations, using the same (cid:15) across models does not degrade performance noticeably.
Note that the last column contains results using non-ImageNet pretrained models, where perturbing in the input space seems to have an edge.
[Table]
C.2. Fine-tuning multiple layers
In table 6, we ﬁne-tune the last three layers instead of just the ﬁnal layer.
We ﬁnd that this leads to slightly improved performance on two datasets, while slightly degrading performance on one.
The optimal parameters to ﬁne-tune remains an open question.
[Table]
C.3. Do adversarial perturbations lead to unstable logits?
In order to verify the effect of adversarial ﬁnetuning with and without the L smooth loss, we compute the logits of an ERM trained model, an adversarially ﬁnetuned model and a KL-regularized ﬁnetuned model on the ”Clipart” split of the OfﬁceHome dataset.
The training data were the ”Real”, ”Product” and ”Painting” splits of the dataset.
We ﬁnd that the mean rank correlation of the KLregularized model with the ERM trained model is higher (0.52 v/s 0.49).
In table 7, we list the average precision@k for k between 1-5 of the logits with the ERM model.
Here prec@k is deﬁned as | topk predictions ( model ) ∩ topk predictions ( ERMModel ) | /k .
As we can see,
L smooth encourages the order of logits to be maintained, while having a similar target accuracy as the AF model.
[Table]
C.4. Additional results on Colored-FashionMNIST
In ﬁg 3, we show examples of images from the FashionMNIST dataset, as well as the l 2 -norm constrained adversarial perturbations.
We ﬁnd that the perturbations mainly change the colour of the images.
For each feature, we compute the relative average perturbation (RAP, i.e. E [ | f i ( x + δ ) − f i ( x ) | | f i ( x ) | ] , where δ is the adversarial perturbation, and f i denotes the i th feature) when the input is perturbed adversarially.
We call this RAP-input.
We also compute the relative average perturbation when the adversarial perturbations are in the feature space, denoted as RAP-feature.
We notice that the maximum perturbation for colour features is much more than that of shape features (11x v/s 0.4x).
This is expected since the adversarially perturbations only change the colour of the image.
Further, we also notice that there is a high correlation between RAP-input and RAP-feature.
In fact,
RAP-feature also follows a similar trend, with the maximum
RAP-feature for shape features being 0.3x, while the maximum RAP-feature for colour features is 13x.
This means that ﬁne-tuning the last layer with either feature perturbations or input perturbations would lead the model to similar classiﬁcation weights.