A Closer Look at Model Adaptation using Feature Distortion and Simplicity Bias
Abstract
In order to achieve strong in-distribution (ID) and out-of-distribution (OOD) generalization during transfer learning, it was recently argued that adaptation protocols should better leverage the expressivity of high-quality, pretrained models by controlling feature distortion (FD), i.e., the failure to update features orthogonal to the
ID.
However, in addition to OOD generalization, practical applications require that adapted models are also safe.
To this end, we study the susceptibility of common adaptation protocols to simplicity bias (SB), i.e., the well-known propensity of neural networks to rely upon simple features, as this phenomenon has recently been shown to underlie several problems in safe generalization.
Using a controllable, synthetic setting, we demonstrate that solely controlling FD is not sufficient to avoid SB, harming safe generalization.
Given the need to control both SB and FD for improved safety and ID/OOD generalization, we propose modifying a recently proposed protocol with goal of reducing SB.
We verify the effectiveness of these modified protocols in decreasing SB on synthetic settings, and in jointly improving OOD generalization and safety on standard adaptation benchmarks.
Introduction
requires that adapted models not only generalize well to in- and out-of distribution data but also that they do so safely.
Due to rapid improvements in the representation quality of largescale, pretrained self-supervised models (LSPM) [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13], there has been growing interest in developing transfer learning or adaptation protocols which are expressly designed to exploit these expressive features.
However, standard protocols, e.g., fine-tuning ( FT ) all layers or only training a linearprobe ( LP ), do not effectively utilize this expressivity and achieve a sub-optimal in-distribution (ID) vs. out-of-distribution (OOD) generalization trade-off [14, 10, 15, 16].
Kumar et al.
argue that feature distortion (FD), i.e., the phenomenon of exclusively updating only a subset of features aligned with the ID data, leads to decreased OOD generalization and propose a new family of LP+FT protocols to improve this trade-off.
By performing LP prior to FT , LP+FT protocols make fewer updates to the LSPM during FT , and thus reduce FD without compromising on task performance.
However, in addition to strong ID/OOD generalization, practical deployment requires that models are also safe [17], e.g., well-calibrated, robust to anomalous/corrupted/adversarial examples and avoid shortcuts (see Figure 1).
Yet, we find that even the well-studied LP , FT and LP+FT protocols achieve varying levels of success when considering both, these additional safety metrics, and datasets with different degrees of distribution shift from the original, pretraining dataset (ImageNet) (see vanilla protocols in Table.
This observation clearly suggests that a complimentary perspective to FD is needed to understand and improve the behavior of adaptation protocols with respect to safety metrics.
Proposed Work.
To this end, we study the susceptibility of protocols to simplicity bias (SB) [18, 19, 20, 21, 22], i.e. , the tendency of deep neural networks (DNNs) to prefer simple features over complex features and learn thin, non-robust decision boundaries [23].
Using a configurable, synthetic dataset, we find that FT is particularly prone to SB and that LP+FT does somewhat help mitigate both SB and distortion.
However, in settings where both OOD generalization and avoiding SB are required, LP+FT can comprise upon performance by exclusively prioritizing FD mitigation.
Using these insights, we aim to systematically mitigate both SB, which is known to influence the robustness of DNNs, and FD, which influences ID/OOD generalization during model adaptation.
To this end, we propose modifications to LP+FT , where we leverage either optimized perturbations (virtual adversarial training, uncertainty-driven perturbations) or construct soups ( i.e. , average multiple models) of probes to find LP initializations that reduce the SB for the subsequent FT step.
On the synthetic dataset, we indeed find that modified protocols decrease SB and, on real datasets, these protocols improve both OOD generalization and safety metrics.
Simplicity Bias, Feature Distortion and Adaptation
We evaluate the susceptibility to simplicity bias of LP , FT , and LP + FT protocols on a custom variant of the dominoes dataset [23].
As shown in Fig.
2, we create “dominoes" of complex and simple features by pairing each class from
CIFAR10 (complex) with the corresponding “digit" class in
MNIST (simple), and consider 3 levels of correlation (95%, 99%, 100%) between the simple and complex features.
To assess OOD generalization, we create a variant where complex features are sampled from STL10, instead of CIFAR10.
(See 2, panels 1/2.
We consider two 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 SB will have high Correlated Accuracy and (low/high)
Randomized Accuracy .
We use a MoCo-V2 ResNet-50 [6] pretrained on ImageNet-1K as the base-feature extractor and results are averaged over 3 seeds.
See supplementary for additional details.
Results.
We make the following observations based on only on the vanilla protocols, shown in black and gray in Fig.
Here, the Rand OOD setting is equivalent to “safety evaluation" as models must avoid shortcuts to perform well under distribution shifts.
Across all correlation strengths, FT has lower Rand.
OOD Acc.
and higher Corr.
OOD Acc.
than LP+FT .
This clearly indicates that FT is highly susceptible to SB.
In contrast, given that LP+FT has higher Rand.
OOD Acc.
and comparable
Corr.
ID Acc., LP+FT more effectively decreases SB in order to do well on the OOD dataset.
In the appendix, we include additional results which demonstrate that under high correlation (0.99,1.0),
LP is more effective at decreasing SB, as any additional distortion is harmful.
However, in moderate correlation (0.95), additional distortion is in fact beneficial to LP+FT .
2.1 Improved Linear Probes for Mitigating Simplicity Bias
While our results on the dominoes dataset indicate that LP+FT and LP are effective protocols for reducing SB, we found that additional distortion can helpful in moderate correlation settings.
This indicates that decreasing FD alone is unlikely to achieve to optimal safety performance.
To that end, we propose new variants to the LP step of the LP+FT protocol which attempt to enable the subsequent
FT step to distort features without compromising generalization or increasing SB.
While its possible to modify the FT step as well, modifications to LP are inexpensive as the feature-encoder is not updated and, given that the fine-tuned solution remains in close vicinity of the initial LP initialization, strong starting solution is well-motivated.
To this end, we introduce the following modifications.
• LP(VAT) :
Virtual Adversarial Training (VAT) [24] 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.
• LP(UDP) : Instead of maximizing the loss when training with adversarial perturbations, uncertaintydriven perturbations (UDP) [25] maximize a model’s estimated uncertainty and have been shown to be effective in decreasing SB and improving generalization in non-adaptation settings.
• LP(Soup) :
Inspired by [26], 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.
Empirical Evaluation of Hardness Promoting Augmentations.
We evaluate the effectiveness of the above LP variants, which we collectively refer to as “hardness promoting”, in reducing SB of LP+FT and summarize the results in Fig.
We make the following observations.
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 simple features, potentially also leading to improved safety
Furthermore, with the exception of LP(Soup) + FT , hp- LP+FT also performs better than vanilla LP+FT on Corr.
ODD accuracy.
Lastly, we observe that with respect to Corr.
ID Accuracy that hp-
LP+FT improves performance at low correlation strength, but slightly loses performance at higher correlation strengths.
This is not entirely unexpected as FT ’s reliance upon simple will be useful in the correlated setting.
Evaluating Generalization and Safety of LP+FT Family
Given the effectiveness of incorporating hardness promoting (hp) augmentations with the family of
LP+FT protocols (hp- LP+FT ) in avoiding shortcuts/SB, we also evaluate on the modified protocols on the three real-world datasets (Living17, DomainNet, and CIFAR10) with respect to ID/OOD generalization and safety metrics.
We summarize our results in Table 1 and our observations below. (See supplementary for additional results and details.)
These three datasets represent scenarios where different levels of distortion, measured using CKA scores [27, 28], 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 a good solution, we see that 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.
To this end [Table] we recommend incorporating hardness promoting augmentations during LP as a potential safe-guard to simplicity bias.
[Table]
Discussion
In this work
[Table]
we considered factors important to the design of adaptation protocols which can induce both strong generalization and strong safety performance when performing transfer learning using high-quality pretrained models.
We find that while the recently proposed LP + FT protocol does achieve impressive OOD accuracy by mitigating feature distortion [Table] simple FT or LP can outperform it with respect safety objectives, such as robustness to corruptions, calibration error [Table]
and anomaly detection performance. We argue that feature distortion alone is not sufficient to understand this generalization vs. safety trade-off and that protocols should also consider susceptibility to simplicity bias.
To this end, we propose using optimized perturbations (virtual adversarial training, uncertainty driven perturbations) or constructing soups to find better LP initialization, which will better enable subsequent FT to decrease SB and improve safety performance.
We verify the benefits of modified protocols empirically on synthetic and real datasets, where, respectively, SB is decreased and safety performance improved. Acknowledgements
This work was performed under the auspices of the U.S. Department of Energy by the Lawrence
Livermore National Laboratory under Contract No. DE-AC52-07NA27344, Lawrence Livermore
National Security, LLC.and was supported by the LLNL-LDRD Program under Project No. 21-ERD012.
PT was an intern at Lawrence Livermore National Labs while working on this project.
A Appendix
In this section, we provide a brief overview of directly relevant work.
adaptation protocols used with pre-trained representations, widely adopted augmentation strategies and popular metrics used for evaluating the safety of deep neural networks.
For a comprehensive overview of transfer learning, please see the surveys of Zhuang et al.
[29] and
Pan et al.
Here, we discuss a few directly works directly relevant to our own.
Adaptation Protocols.
Given a pre-trained model, common practices for adapting it to a target task of interest are to either fine-tune ( FT ) all model parameters or train only a linear probe ( LP ) while freezing the network parameters.
In general, large-scale, pre-trained models have sufficiently expressive features to perform well on both ID and OOD data.
However, in practice, LP can result in lower classification accuracies due to the limited expressivity of the linear probe, while end-to-end
FT can distort pretrained features toward ID data, thus harming OOD performance.
Given these inherent limitations of both LP and FT , there is a need to design adaptation protocols that can lead to models with improved generalization performance and safety characteristics.
In this spirit, Kumar et al.
[16] recently proposed to perform LP prior to FT (abbrev.
LP + FT ) and demonstrated that this protocol improves OOD performance without comprising ID generalization of the target task by limiting feature distortion.
Namely, 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.
Concurrently, Kirichenko et al.
[31] also found that retraining the last-layer of a model on a minority group or a simple data re-weighting can significantly improve robustness to spurious correlations and argue that classifiers poorly utilize the expressive features learnt by the model, instead relying upon spurious (instead of core) features.
Notably, the model is able to learn both spurious and core features, so only simple re-weighting on minority data is needed, if we assume disentangled features.
In this paper, we focus on the generic protocols such as LP , FT , and LP + FT , since they are effective, inexpensive and do not to perform any additional re-weighting data.
Notably, we find that feature distortion can explain the ID vs. OOD generalization behavior of FT and LP , it cannot be straightforwardly extended to understand the trade-off between OOD and safety performance, necessitating a complementary perspective (simplicity bias).
Other modifications and heuristics have also been proposed to improve fine-tuning, including sidetuning [32], 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 [33].
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 [34, 35, 14], including the robustness of pretrained features [36, 37].
While the connection between adversarial training and improved feature representations [38, 39] 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, while we consider a holistic evaluation of safety and generalization in the context of transfer learning with highly expressive pretrained models, Hendrycks et al [40] have considered the trade-offs induced by different data augmentation strategies [41, 42, 43, 44, 45] 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.
Simplicity Bias.
It is well-known that deep neural networks demonstrate a bias toward simple, potentially less expressive features [18, 19, 20, 21, 22], such as textures and backgrounds, and that this bias can lead to shortcuts that limit the generalization of DNNs.
Indeed, recently Shah et al. [23] 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.
While various methods have recently been proposed to mitigate simplicity bias when training from scratch or in the context of pretraining [46], we focus the susceptibility of existing adaptation protocols to simplicity bias as a tool for gaining insights into their safety behavior.
B Experimental Details
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 [47] pretrained ResNet-50 is used for the DomainNet experiments, while a MoCoV2[6] is used for all CIFAR experiments.
We use augmentation functions from timm[48] 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. and 0.1 for UDP, and α in 0.001, 0.01, 0.1 for VAT.
For CIFAR-MNIST results, LP is done forepochs, and FT is done for 20 epochs.
B.1 Safety Evaluation
LP , FT , and LP + FT protocols are evaluated for generalization and safety performance on three downstream adaptation tasks: CIFAR-10, Living17, and Domainnet-Sketch, where we report the following additional metrics for each dataset.
We select these datasets as they correspond to two different types of distribution shifts (standard domain adaptation and subpopulation) and 3 levels of distortion (low, medium, high).
Our safety evaluation protocol is similar to [40].
• OOD Accuracy: Models are expected to generalize well under the following distribution shifts: Living17(Source) → Living17(Target), CIFAR-10 → {STL10, CIFAR10.1} and Domainnet-Sketch → {Domainnet-ClipArt, Domainnet-Painting, Domainnet-Real}.
CIFAR10 and Domainnet-Sketch shifts are popular domain-adaptation datasets, while Living17 is a recently proposed sub-population shift benchmark [49]. • Mean corruption accuracy (mCA/m ¯ C A):
We consider two sets of corruptions that a model should be robust to: the 15 naturalistic corruptions (C) proposed by [50], and the 10 perceptually dissimilar corruptions ( ¯ C ) proposed by [51].
Corruptions are applied to each ID test dataset and the average accuracy over each set is reported. • 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 [52].
We use a 2/225, 10-step PGD attack to measure the robustness of models to such perturbations. • Representational Similarity (CKA) : Kumar et al [16] claim that feature distortion harms OOD performance of adapted models by only updating representations in the span of the training data.
We measure this distortion by computing the batched centered kernel alignment (CKA) score [27] with respect to the representations of the ID test dataset obtained from the pretrained and adapted models.
C Additional Results
[Table]
[Table] [Table]