A C LOSER L OOK AT M ODEL A DAPTATION USING F EATURE D ISTORTION AND S IMPLICITY B IAS
Anonymous authors
Paper under double-blind review
A BSTRACT
Advances in the expressivity of pretrained models have increased interest in the design of adaptation protocols which enable safe and effective transfer learning.
Going beyond conventional linear probing (LP) and fine tuning (FT) strategies, protocols that can effectively control feature distortion, i e., the failure to update features orthogonal to the in-distribution, during FT have been found to achieve improved out-of-distribution generalization.
A popular example is the recent LP+FT protocol which first learns a linear probe and then uses that initialization during FT.
However, in this paper, we find that when adaptation protocols are also evaluated on a variety of safety objectives (e.g., calibration, robustness etc.), that a complementary perspective to feature distortion is required explain protocol behavior.
To this end, we study the susceptibility of protocols to simplicity bias (SB), i e. the well-known propensity of neural networks to rely upon simple features, as SB has recently been shown to underlie several problems in robust generalization.
Using a synthetic dataset, we demonstrate the susceptibility of existing protocols to SB.
Given the strong effectiveness of LP+FT, we propose incorporating hardness-promoting perturbations during LP to obtain initializations for FT that further decrease SB.
We verify the effectiveness of these modified LP+FT protocols by decreasing SB in a controlled setting, and jointly improving OOD generalization and safety on standard adaptation benchmarks
I NTRODUCTION
Through the use of larger datasets (Yalniz et al., 2019), better architectures (Zhai et al., 2022; Chen et al., 2022; Steiner et al., 2022; Tolstikhin et al., 2021), and different self-supervised learning (SSL) approaches (He et al., 2020; Chen et al., 2020; Grill et al., 2020; Caron et al., 2020), the quality of pretrained representations available for transfer learning tasks has dramatically improved.
Indeed, representations from such high-quality
SSL models have been found to be more robust (Hendrycks et al., 2019; Liu et al., 2021), transferable (Ericsson et al., 2021) and semantically consistent (Caron et al., 2021) than their supervised counterparts.
In this regard, there is growing need for adaptation protocols that explicitly capitalize on these improved pretrained features to induce similar beneficial properties, e.g., inducing more than just high accuracy on the target task, after models have been trained on the downstream task. However, standard adaptation protocols that rely upon finetuning ( FT ) all model parameters or training only a linear probe ( LP ) while freezing the network parameters do not maximize the potential of high-quality representations.
For example, while high quality, pre-trained models have sufficiently expressive features to perform well on both in-distribution (ID) and out-of-distribution (OOD) data, LP and FT are not able effectively induce this property in adapted models.
Recently, however, Kumar et al.
(2022) proved that by modifying features only in the ID representation subspace, FT can lead to higher OOD error as it distorts directions outside the ID subspace that are needed OOD generalization.
As both ID and OOD subspaces are represented by the pretrained model, Kumar et al. demonstrate that limiting feature distortion , or controlling updates towards the ID subspace, can lead to improved
ID and OOD performance. To this end, they propose a new protocol which performs LP prior to FT (abbrev. LP + FT ).
By first performing LP , this two-step process ensures that subsequent FT will remain in the vicinity of the original LP solution. This reduce the overall distortion towards the ID distribution subspace and improves performance.
While strong ID and OOD generalization on the target task is indeed an important aspect of transfer learning, practical, high-risk applications, such as healthcare, require that models are also safe (Hendrycks et al., 2021). For example, adapted models should also be well-calibrated, robust to corruptions or adversaries and able to reliably detect anomalous samples (see Figure 1). Given that existing adaptation protocols are primarily focused on improving generalization, it is unclear how existing protocols utilize high-quality pretrained features to promote safe adaptation, and if current protocol design perspectives, such as mitigating feature distortion, will also enable safe generalization.
Our work fills in this gap by evaluating the efficacy of popular adaptation protocols ( LP , FT , LP+ FT ) for both safety and generalization, as well as studying how to improve protocols to induce both these objectives. In particular, we first demonstrate the short-comings of designing protocols to solely limit feature distortion when attempting to induce safe adaptation. We then argue that a complementary perspective, namely avoiding simplicity bias, is needed to ensure both safe and effective adaptation. We focus on simplicity bias, i.e. , the problematic tendency of deep neural networks (DNNs) to prefer simple, potentially brittle features over complex features (Soudry et al., 2018; Gunasekar et al., 2018;
Geirhos et al., 2019; Hermann et al., 2020), as it has been recently shown to underlie many problems in machine learning safety (Shah et al., 2020) but has not been consider the context of adaptation. Using these dual perspectives, we then propose three LP+ FT variants that are able to mitigate both feature distortion and simplicity bias, leading to improved safety and generalization.
Proposed Work.
In this work, we seek to understand the factors relevant to the design of adaption protocols that promote effective and safe generalization.
We take the first step towards this aim by demonstrating limitations in existing LP , FT ,and LP+ FT protocols through an extensive, joint evaluation, and studying adaptation protocols through the complementary lens of avoiding simplicity bias.
Our contributions can be summarized as follows:
Sec. 3: Joint Analysis of Adaptation Protocol Safety and Generalization. We show that when adaptation protocols are evaluated with respect to both ID/OOD generalization and safety, LP+ FT trails behind LP or FT on several safety metrics.
This demonstrates that solely mitigating feature distortion is not sufficient for safe generalization.
We also observe that keeping subsequent FT close to
LP solution is crucial for the improved OOD generalization of LP+ FT . This motivates us to focus on improving the LP initialization as a mechanism for improving both safety and OOD performance.
Sec. 4: Role of Simplicity Bias in (Unsafe) Adaptation. To understand how protocols may induce safe adaptation, we study how different protocols avoid simplicity bias. While simplicity bias (Shah et al., 2020) has been shown to underlie several problems in machine learning safety, to the best of our knowledge, we are the first to consider its role in adaptation settings. We demonstrate that protocols must not only reduce distortion, but that avoiding simplicity bias is also necessary for effective adaptation.
Sec. 5: Improved Protocols for Mitigating Simplicity Bias and Distortion. We propose three, simple modified LP+ FT protocols that help mitigate both simplicity bias and distortion (Sec. 4.1).
In particular, we consider modifying the LP step with uncertainty-driven perturbations (Pagliardini et al., 2022), virtual adversarial training (Miyato et al., 2017) and model-soups (Wortsman et al., 2022), as they are simple and effective strategies. Across synthetic and real datasets, the modified protocols help improve safety and generalization (Sec. 5).  R ELATED W ORK AND B ACKGROUND
Adaptation Protocols.
For a comprehensive overview of transfer learning, please see the surveys of Zhuang et al.
and Pan & Yang.
Here, we discuss the works that are most relevant to our own.
Kirichenko et al.
recently demonstrated that models are able to learn both core features and spurious features. However, classifiers can rely upon spurious features, harming performance on minority groups. To reduce the reliance on spurious features, they propose to retrain the classifier on a small amount of “re-weighting” data, which allows the model to leverage the core features instead of the spurious features.
Other modifications and heuristics have also been proposed to improve FT ’s performance, including side-tuning (Zhang et al., 2019), which tunes a small secondary network that is then combined with the original model, using larger/smaller learning rates for the classifier, as well as regularization-based methods (Jiang et al., 2020). In this work, we focus on two popular and effective protocols, LP and FT .
We additionally study the LP+ FT protocol as it is theoretically-grounded, does not require re-weighting data, designed to exploit high-quality pre-trained representations and achieves SOTA OOD performance during adaptation.
Simplicity Bias.
It is well-known that DNNs demonstrate a bias toward simple, potentially less expressive features (Brutzkus et al., 2017; Soudry et al., 2018; Gunasekar et al., 2018; Geirhos et al., 2019; Hermann et al., 2020), such as textures and backgrounds, and that this bias can lead to shortcuts that limit the generalization of DNNs.
Indeed, recently Shah et al.
(2020) formalized this intuition by more precisely defining simplicity bias, based on the number of linear components to define a decision boundary, and showed that SB leads to non-robust decision boundaries that effects a model’s sensitivity to distribution shifts and adversarial perturbations.
In brief, by learning simple features first, models become invariant to complex features, potentially leading to narrow decision boundaries which can fail to generalize under data shifts.
Notably, DNNs exhibit this bias even when complex features are more expressive and necessary for fitting the distribution.
While various techniques have recently been proposed to mitigate simplicity bias when training from scratch or in the context of pretraining (Teney et al., 2021), we are the first to rigorously study the role of simplicity in the context of model adaptation.
J OINT A NALYSIS OF A DAPTATION P ROTOCOL S AFETY AND G ENERALIZATION
In this section, we evaluate the performance of adaptation protocols across several additional safety objectives (Hendrycks et al., 2021), as practical transfer learning applications require both strong and safe generalization. Through this expanded evaluation, we find that no single protocol is optimal across all safety objectives.
Indeed, the inability of LP+ FT to induce safe adaptation indicates that a complementary perspective to feature distortion, namely simplicity bias, is necessary when designing generalizable and safe protocols (see Sec. 4). We further argue that by constraining models around the LP initialization during FT , LP+ FT may inadvertently harm safety performance by hampering models’ abilities to learn complex, task-specific features needed for robust generalization. While we expand upon the role of LP initialization in Secs.
4 and 5, we begin, here, by introducing the expanded evaluation and experimental set-up.
Experimental Setup.
Three downstream adaptation tasks (and their respective OOD distributions) are considered: CIFAR-10 (ID) → { STL10, CIFAR10.1 } (OOD), Domainnet-Sketch → { DomainnetClipArt, Domainnet-Painting, Domainnet-Real } and Living17 (Source) → Living17 (Target).
These datasets are selected as they correspond to two different types of distribution shifts (standard domain adaptation and subpopulation) and three levels of distortion (low, medium, high).
A MoCo-V ResNet-50 (He et al., 2020) pretrained on ImageNet-1K is used as the base-feature extractor for
CIFAR10 and Living17 experiments, and the CLIP ResNet-50 image encoder pretrained onmillion (image,text) pairs is used for Domainnet-Sketch. These models are selected as they provide sufficiently high quality representations capable of generalizing to both ID and OOD downstream data (Kumar et al., 2022). We perform grid-search to find the best hyper-parameters, and average over 3 seeds.
See supplementary for additional details. Expanded Evaluation. In addition to OOD accuracy on the aforementioned distribution shifts, we report performance on the following metrics in order to evaluate adapted models on key problems in machine learning safety (Hendrycks et al., 2022; 2021):
• Mean Corruption Accuracy (mCA/m
¯ C A): We consider two sets of corruptions: the 15 naturalistic corruptions ( C
• Calibration Error (RMSE) :
It is important that models are well-calibrated so that practitioners may trust the provided predictions in high-risk applications.
We measure the root mean square error of calibration as follows: (cid:114) E C (cid:104) ( P ( Y =
ˆ Y | C = c) − c) 2 (cid:105) , where C indicates the confidence scores, while ˆ Y and Y denote the model’s predictions and ground-truth labels respectively.
• Anomaly Detection Performance (AUROC):
Recognizing when samples are anomalous allows models to abstain from making uninformed and inapplicable predictions.
We consider samples from Blobs, Gaussian, LSUN, Places69, Rademacher, Textures, and SVHN datasets as anomalies and report the AUROC (area under the ROC curve) of the binary classification problem of detecting such samples as anomalies.
• Adversarial Accuracy: DNNs are well-known to be fooled by imperceptible distortions (Ilyas et al., 2019).
We use a 2/225, 10-step PGD attack to measure the robustness of models to such perturbations.
We make the following observations regarding the behavior of different protocols using this expanded evaluation.
In brief, we find that no single protocol is effective across all datasets in jointly obtaining strong and safe adaptation, and that, on low distortion adaptation tasks, the quality of the
LP initialization is critical as pre-trained feature extractor is not substantially updated during LP+ FT .
O BS 1: M ITIGATING F EATURE D ISTORTION DOES NOT INDUCE SAFE ADAPTATION .
We protocol safety generalization met We see that no single protocols achieves top rank across all metrics.
Here, we ask how protocols perform when we consider both safety and generalization objectives to better understand the feature distortion perspective.
In particular, if LP+ FT is able to outperform LP and FT in this expanded evaluation, then it suggests that solely mitigating feature distortion during FT may be sufficient to induce robust adaptation. To test this claim, we rank protocol performance for each safety metric, where ranks are first computed for each dataset separately, and then averaged. Results are shown in Fig.
Smaller ranks correspond to better performance. Results.
LP+ FT obtains the best rank for ID and OOD accuracy as expected, as well as C and C accuracy. However, we also see that FT is better ranked for Adversarial Accuracy and OOD calibration, while LP is better ranked for ID calibration and C-BAR calibration. By design, LP+ FT limits feature distortion.
However, given that LP+ FT trails behind protocols that are not explicitly designed to limit distortion on some safety metrics, it is clear that a complementary perspective is needed to better understand protocol behavior. Indeed, LP+ FT has the best average rank, indicating that it is a good starting point to improve upon. The above results are aggregated across different types of distribution shifts; we extend this analysis next by considering the interplay between individual datasets and protocol performance.
These detailed results are presented in Table 1. O BS 2: L INEAR P ROBING S OLUTIONS M ATTER . Naturally, the amount of distortion required to effectively adapt a pretrained model to a downstream task will vary in accordance to the similarity of the downstream and pretraining data. Here, we seek to understand how protocols behave under different levels of distortion.
In particular, we hypothesize that the LP initialization becomes more influential for LP+ FT in low distortion settings, as subsequent FT remains in the vicinity of initialization. To this end, we compute the batched centered kernel alignment (CKA) score (Nguyen et al., 2021) with respect to the adapted and pretrained models, and take a closer look at performance across metrics. We note that while CKA is better suited for measuring distortion than the L2 norm as used by Kumar et al.
(2022), other neural representation metrics can also be used (Ding et al., 2021; Anonymous, 2023). Results. As shown in Fig.
3, we see that minimal distortion (CKA ≥ 0.9) is required to obtain competitive LP+ FT performance on DomainNet and Living17. However, on CIFAR10, which requires the most distortion as evidenced by lower CKA scores, FT is the most effective protocol for safety measures and is very comparable on generalization performance (see Table 1).
We plot the
CKA similarity between adapted and pretrained models.
DomainNet and Living17 require low distortion, as seen by performance of LP+ FT across metrics, while achieving high CKA ( > 0 . 9 ).
The effectiveness of LP and LP+ FT on Living17 in improving OOD generalization over FT is hardly surprising, as Living17 is a subset of ImageNet, on which the base feature-encoder was already trained.
In contrast, on DomainNet, the difficulty of FT in matching the ID test task performance, despite achieving high training accuracy, suggests FT may learn a solution that over-fits to the training data or relies upon shortcuts (or simple features) that do not generalize to even the test set or OOD data, despite using an expressive CLIP pretrained model.
We emphasize that LP+ FT greatly benefits from strong LP initializations on these low-distortion datasets as corresponding CKA scores show that very limited updates are made during FT . While LP+ FT does induce some meaningful improvements over LP on Living17 and performs comparably to LP on DomainNet, we stress that these benefits arise from keeping the model close to the LP initialization during FT .
Indeed, to obtain acceptable LP+ FT performance, we found that the failure to (i) use small learning rates (3e-7,1e-5) and (ii) freeze batch-norm parameters during FT significantly deteriorate performances. Summary.
Taken jointly, these results suggest that while solely mitigating feature distortion is not sufficient to ensure that adapted models perform well on safety metrics across different levels of shift, improving the LP initialization may be a viable solution to obtaining strong and safe generalization. Indeed, the effectiveness of LP+ FT on low distortion datasets and its high average ranking indicates that it is a promising protocol to build upon.
To understand how to build better protocols, we next introduce simplicity bias as a complementary perspective to feature distortion [Table]
M ITIGATING S IMPLICITY B IAS & F EATURE D ISTORTION FOR S AFE A DAPTATION
Given that simplicity bias underlies various safety issues in machine learning Shah et al. we argue that mitigating feature distortion in a way that minimizes this bias can be a valid mechanism to improve both safety and generalization performance.
Therefore, in this section we first measure the propensity of different protocols to simplicity bias in a controlled setting. Moreover, given our previous observation that LP+ FT models remain in close vicinity of the LP solution after FT we focus on improving the performance of this initial LP initialization so that we may capitalize upon
LP+ FT strong OOD performance, while simultaneously improving safety.
In particular, we propose three light-weight LP+ FT variants that are able to both reduce distortion and SB.
We begin by introducing our synthetic dataset and experimental set-up.
Dataset. As shown in Fig.
4, we create “dominoes” of complex and simple features by pairing each class from CIFAR10 (complex) with the corresponding “digit” class in MNIST (simple), e.g., “bird” samples are paired with digit “2” samples, where the label for each domino is determined by the complex, CIFAR10 sample.
Datasets with three levels of correlation (95%, 99%, 100%) between the simple and complex features are constructed for training.
While 100% correlation allows models to only learn the simple feature for perfect generalization, the more realistic lower correlation settings require models learn at least some aspect of the complex features.
dominoes that enables us to study the effect of simplicity bias on safety and OOD generalization.
Experimental Setup. For evaluation, we also construct a randomized (10% correlation) variant, where simple features are randomly paired with complex features.
(See Fig. 4, panels 3/4).
To assess OOD generalization, we create a variant where complex features are sampled from STL10, instead of CIFAR10.
(See 4, panels 1/2.)
Metrics. We assess the reliance simple features using the following metrics: Randomized Accuracy , e.g., the accuracy on the variant where samples contain random pairings between simple and complex features, and Correlated Accuracy , where pairings between simple and complex features remain correlated.
Models that (are)/(are not) susceptible to simplicity bias will have high Correlated
Accuracy and (low/high) Randomized Accuracy .
Training Details. We use a MoCo-V2 ResNet-50 (He et al., 2020) pretrained on ImageNet-1K as the base-feature extractor.
See supplementary for additional details.
We performed grid-search to find the best parameters, andd report all results over 3 seeds and 3 correlation strengths.
[Table]
Given the above experimental setup, we report the performance of different adaptation protocols in
Table.
Across all correlation strengths FT has the lowest Rand.
OOD Acc.
and high Corr.
OOD Acc.
This clearly indicates that FT has learned to rely upon simple features, effectively disregarding the expressive features of the pretrained model, and is easily susceptible to simplicity bias.
In contrast, by preserving the expressive features of the underlying feature encoder, LP best mitigates simplicity bias in high correlation (0.99,1.0) settings as evidenced by the highest Rand OOD Acc (though Corr.
ID/OOD Acc does slightly suffer).
However, in moderate correlation (0.95), LP+ FT improves upon LP to achieve better Corr.
OOD Acc than LP and the best Rand.
OOD Acc across protocols.
This suggests that when the correlation is not extreme, that moderate distortion, given a suitable LP , is in fact beneficial in mitigating simplicity bias.
At higher correlation strengths (0.99,1.0), however, LP+ FT has lower Rand OOD Acc, while improving the Corr.
OOD Acc relative to LP , indicating in such extreme settings the distortion incurred by subsequent FT is not beneficial and has increased reliance upon simple features.
4. I MPROVED L INEAR P ROBES FOR M ITIGATING S IMPLICITY B IAS
We introduce a variant of the dominoes dataset that enables us to study the effect of simplicity bias on both safety and OOD generalization.
As discussed earlier, adaptation protocols have varying susceptibility to simplicity bias, and mitigating this susceptibility can help improve generalization and safety.
In particular, we observe that LP+ FT and LP are effective protocols for reducing reliance upon simple features on both synthetic datasets, and on low-distortion real datasets (Sec 3).
However, as some level of distortion is typically required when adapting to downstream tasks to obtain sufficient ID task performance, we propose new variants of the LP+ FT protocol that attempt to enable the subsequent FT step to distort features without compromising generalization or increasing simplicity bias.
We note that, while it is certainly possible to modify the FT step as well, modifications to LP are inexpensive as the feature-encoder is not updated.
Moreover, as discussed in Sec. 3, the FT step of LP+ FT often relies upon small learning-rates and frozen batch-norm parameters.
This implies that the fine-tuned solution remains in close vicinity of the initial LP initialization, further motivating a strong starting solution.
To this end, we define the following common notation and introduce modifications to LP step of LP+ FT below.
Let h be the hidden representations, θ model parameters, y labels, ˆ y predicted classes, C the number of classes, g the classifier, and δ the perturbation.
See Supp.
C for additional discussion on the choice of these mitigation strategies.
• LP (VAT) : Virtual Adversarial Training (VAT) (Miyato et al., 2017) enforces local distribution smoothness by minimizing the KL-divergence between the predictions of perturbed pairs of examples.
Since we are using expressive pretrained models, such perturbations maybe meaningful in the inverted latent space as well, and resulting classifiers become robust in some ϵ -neighborhood around each latent-space input.
Formally, let ϵ be some perturbation budget, and α a hyperparameter weighting distributional label smoothness, we minimize the following loss:
[Equation]
where δ := arg max ∥ δ ∥ 2 ≤ ϵ KL [ p ( y | g θ ( h ))
, p ( y | g θ ( h + δ ))] . [Equation] • LP (UDP) : Instead of maximizing the loss to when training with adversarial perturbations, uncertainty-driven perturbations (UDP) (Pagliardini et al., 2022) adversarially maximize a model’s estimated uncertainty and have been shown to be effective in decreasing simplicity bias and improving generalization in non-adaptation settings.
Formally, UDPs are defined as: δ u = arg max ∥ δ ∥ 2 ≤ ϵ H
( g θ ( h ) + δ ) , where H ( g θ ( h ))
= − (cid:80)
c ∈ C
ˆ y c log ˆ y c , (e.g., entropy of predictions).
• LP (Soup) :
Inspired by Wortsman et al.
, we train multiple, sparse, linear probes jointly and then take the average of their weights (aka soup) as the learned LP for subsequent FT .
While soups of large models improve generalization by combining models from the same low-error basin, we consider sparse classifiers soups as an alternative strategy which seeks to average of diverse decision rules, to avoid relying upon a single set of simple features.
Given k classifiers, we minimize min θ 1 ...k 1 k (cid:80) ki
=1 L CE ( g θ i ( h ) , y ) and let g ¯ θ = 1 k (cid:80)
ki
θ i for the FT step.
Empirical Evaluation of Hardness Promoting Augmentations.
We evaluate the effectiveness of the above LP variants, which we collectively refer to as “hardness promoting”, in mitigating the simplicity of bias of LP+ FT .
We make the following observations (see Fig 5).
Across all correlation strengths, we find that using the modified hardness promoting LP s during
LP+ FT (aka hp-
LP+ FT ) improves the Rand.
OOD Accuracy over vanilla LP+ FT ( ≥ 2%) and
FT ( > 20% ).
This clearly indicates that hp-
LP+ FT is indeed effective in decreasing reliance on [Table] [Table] simple features, potentially also leading to improved safety.
Furthermore, with the exception of
LP (Soup) + FT , that hp-
LP+ FT also performs better than vanilla LP+ FT on Corr.
ODD accuracy.
Vanilla FT does out-perform all LP+ FT protocols in this setting, but this is due to reliance upon simple features.
Lastly, we observe that with respect to Corr.
ID Accuracy that hp-
LP+ FT improves performance at low correlation strength, but slightly lose performance at higher correlation strengths.
This is not entirely unexpected as FT ’s reliance upon simple will be useful in the correlated setting.
Given that hp-
LP+ FT is able to reduce reliance upon simple features in this controlled setting, we next evaluate whether these modified protocols are beneficial in improving the performance of LP+ FT on real datasets.
E VALUATING G ENERALIZATION AND S AFETY OF LP+ FT F AMILY
Given the effectiveness of incorporating hardness promoting (hp) augmentations with the family of
LP+ FT protocols (hp- LP+ FT ) in avoiding shortcuts, we also evaluate the modified protocols on the three real-world datasets (Living17, DomainNet, and CIFAR10) with respect to the generalization and safety metrics introduced in Sec. 3.
We summarize our results in Tables 3,4,5 and our observations below.
We tune any method-specific hyperparameters (e.g., epsilon) using ID validation data and report all results over three seeds.
(See supplementary for additional results.)
As discussed in Sec. 3, these three datasets represent scenarios where different levels of distortion are necessary when adapting the pretrained model.
On Living17, a setting which requires minimal distortion during adaptation, we see that vanilla LP+ FT is quite effective with respect to both generalization and safety metrics and is a difficult baseline to surpass.
Indeed, while hp- LP+ FT variants do not lead to significant benefits, they generally perform comparably to vanilla LP+ FT .
On
DomainNet, a setting where fairly low distortion is required for LP+ FT but FT struggles to find [Table] a good solution, hp-
LP+ FT induces some slight benefits with respect to ID/OOD generalization and robustness, though vanilla LP and hp- LP have better calibration performance.
In contrast on CIFAR10, which requires more distortion to obtain an acceptable solution, we see that hp-
LP+ FT leads to improved generalization and a noticeable boost in corruption robustness.
LP (VAT) + FT is particularly effective in this regard.
Lastly, across all datasets, we observe that hp-
LP+ FT protocols lead to similar distortion to vanilla LP+ FT , which suggests that any additional benefits of hp- LP+ FT should not be attributed to only reducing feature distortion.
In summary, we find that while vanilla LP+ FT is already an effective protocol, especially in settings where low distortion is required,
hp-
LP+ FT can provide some benefits and performs competitively.
We suspect that the performance of these modified protocols can further be improved if more sophisticated simplicity bias mitigation strategies are used. Indeed, our central claim, that adaptation protocols should mitigate feature distortion and simplicity, is not dependent on a specific strategy. We further note that while such mitigation strategies may optionally also be used during FT , they cannot solely be used in FT . Indeed, in the case of extreme simplicity, if the LP classifier relies upon simple features to find a low-loss solution, during the subsequent FT step, gradients may not be back propagated in directions that contain complex features.
This entails that the decision boundary continues to rely upon simple features and is at risk of reduced safety performance.
To this end, we recommend incorporating hardness promoting augmentations during LP as a potential safe-guard to simplicity bias.
C ONCLUSIONS
In this work, we took a closer look at the behavior of protocols designed for adapting large-scale pretrained models to downstream datasets.
While it is argued that adaptation protocols should be designed to mitigate feature distortion (e.g., LP+ FT ) in order to improve ID and OOD generalization, we found that when additional aspects of safe generalization are evaluated (e.g., prediction calibration error, adversarial robustness etc.)
mitigating feature distortion alone is not sufficient.
To this end, we considered the complementary perspective, that adaptation protocols should also mitigate simplicity bias.
Using a synthetic dominoes dataset that allows for control over the correlation between simple and complex features, we found that protocols have varying levels of effectiveness in reducing reliance upon simple features.
While, as expected, FT , is most susceptible to simplicity bias, we see that LP+ FT is able to balance both distortion and simplicity bias in settings where the correlation between simple and complex features is not too extreme.
Motivated by the benefits of LP+ FT and given the known relationship between simplicity bias and sub-optimal generalization, we used “hardness promoting” LP initializations to further improve LP+ FT ’s performance.
We createed such initializations by using virtual adversarial or uncertainty driven perturbations during LP , or creating a soup of sparse linear probes.
These modifications further reduced LP+ FT ’s reliance upon simple features on the synthetic dataset.
On three real-world datasets, these improved protocols led to better safety and generalization performance, further validating the need to consider both distortion and simplicity bias when design adaptation protocols.
A A PPENDIX B A DDITIONAL R ELATED W ORK
For a comprehensive overview of transfer learning, please see the surveys of Zhuang et al.
and Pan &
Yang.
Here, we discuss a few directly works directly relevant to our own.
Recently, Kumar et al. demonstrated that learning probing prior to fine-tuning (e.g., LP+ FT ) can improve both in-distribution and out-of-distribution performance when transferring to a downstream task given a highly expressive, pretrained model.
They demonstrated that FT only modifies features in the ID representation subspace and not in other directions, which can lead higher OOD error as direction outside the ID subspace are necessary for OOD generalization.
However, by initializing
FT with a trained linear probe, feature distortion can be decreased since this initialization is closer to optimal model, and thus requires less distortion in ID subspace, preserving the expressiveness of the original model.
Concurrently, Kirichenko et al. demonstrated that models are able to learn both core features and spurious features.
However, classifiers can rely upon spurious features, harming performance on minority groups.
To reduce the reliance on spurious features, they propose to retrain the classifier on a small amount of “re-weighting” data, that allows the model to leverage the core features instead of the spurious features.
Other modifications and heuristics have also been proposed to improve fine-tuning, including sidetuning Zhang et al.
(2019), which tunes a small secondary network that is then combined with the original model, using larger/smaller learning rates for the classifier, as well as regularization-based methods Jiang et al.
We focus on the LP+ FT protocol, as it is principled and achieves strong
OOD performance.
Additionally, several works have studied properties of the model that influence the effectiveness of transfer learning Azizpour et al.
; Huh et al.
; Kornblith et al.
(2019), including the robustness of pretrained features Salman et al.
(2020); Utrera et al.
While the connection between adversarial training and improved feature representations Allen-Zhu & Li (2021); Kaur et al.
(2019) has been studied, we use virtual adversarial training during LP to learn a better classifier that is less reliant upon simple features, and we do not use an adversarially trained feature extractor.
Finally, we note that while we are, to the best of our knowledge, the first to consider this holistic evaluation of safety and generalization in the context of transfer learning with highly expressive pretrained models,
Hendrycks et al.
have considered the trade-offs induced by different data augmentation strategies Yun et al.
; Devries & Taylor (2017); Hendrycks et al.
; Cubuk et al.
(2019; 2020) on safety metrics in supervised learning.
We emphasize that while our evaluation is similar, that our work focuses on a different context and contains an additional layer of complexity as we consider the interaction between adaptation protocols, generalization behavior and safety performance.
C E XPERIMENTAL D ETAILS
Please see the Anonymous Git repository for training details.
In brief, we performed grid-search to find the best parameters, which are as follows.
For CIFAR-10 and CIFAR-100, we train only the classifier for 200 epochs with LR=30 during LP .
For FT , the entire model is trained for 20 epochs with LR=1e-5.
For LP+ FT , the model’s classifier is initialized with the solution found by LP , and then it is fine-tuned for 20 epochs.
A grid-search was conducted to determine the LR for LP and
FT .
For Domain-Net Experiments, we use 200 epochs with LR=30 during LP .
For FT , the entire model is trained for 20 epochs with LR=3e-4.
For LP+ FT , the model’s classifier is initialized with the solution found by LP , and then it is fine-tuned for 20 epochs, using LR=3e-7.
Furthermore, following Kumar et al.
, we freeze the batchnorm layers during LP+ FT .
A CLIP Radford et al.
(2021) pretrained ResNet-50 is used for the DomainNet experiments, while a MoCoV2He et al.
(2020) is used for all CIFAR experiments.
We use augmentation functions from timmWightman (2019) and compute CKA scores using the packaged provided by torch-cka.
When using augmented protocols, the same LRs are used.
Note, all results were obtained by averaging over 3 seeds.
We consider model soups of sizes 5,10,20, tune ϵ in 0.005, 0.01, 0.02 and 0.1 for UDP, and α in 0.001, 0.01, 0.1 for VAT.
For CIFAR-MNIST results, LP is done for 100 epochs, and FT is done for 20 epochs.
Motivation for Hardness Promoting Variants.
We selected UDP (Pagliardini et al., 2022), VAT (Miyato et al., 2017), and model-soups (Wortsman et al., 2022) as simplicity bias mitigation strategies due to their effectiveness and ease of use. We emphasize, however, that our findings are not specific to the choice of a given mitigation strategy and we expect that advancements in such strategies will further improve the effectiveness of our proposed LP+ FT variants. At present, the selected strategies are strong, representative mitigations that we have confirmed are effective at mitigating simplicity bias in the adaptation context using the synthetic dominoes dataset in Sec. 4.
We conceptually justify each strategy here: UDP is designed to help mitigate simplicity bias by learning by a large margin classifier, opposed to a narrow margin classifier that relies upon simple features.
As noted by Shah et al.
(2020), such narrow margin classifiers are sensitive to small perturbations and the simple features supporting the decision boundary may not be discriminative under distribution shifts.
By maximizing uncertainty (instead of loss) to create adversarial perturbations, UDP is able to learn a maximum-margin classifier that is better able to handle such shifts.
Notably, to create such a maximum-margin classifier, the model will necessarily learn more complex features.
We use virtual adversarial training to help avoid reliance upon simple features, as VAT enforces distribution smoothness so that classifiers become robust in some epsilon neighborhood around the input. We note that we are performing this training in the hidden representation space, so perturbations correspond may be altering high-level semantics. To maintain strong performance under such high-level perturbations, the model should learn to rely upon more complex features, and learn a better margin classifier; and lastly, we use model-soups so that we may learn a set of classifiers that rely upon disjoint sets of features. By learning a set of diverse classifiers, we are able to average classifiers that have learned to rely upon different features, instead of becoming overly reliant upon a single simple feature.
In future work, we intend to build a theoretical framework that helps us better justify these interventions and create new ones
[Table]
D A DDITIONAL R ESULTS
[Table] [Table] [Table] [Table]