How to Fine-Tune Vision Models with SGD
Authors: Ananya Kumar
Ruoqi Shen
Sebastien Bubeck
Suriya Gunasekar
Venue: ICLR 2024 poster (2024)
Keywords: fine-tuning
SGD
freezing layers
distribution shift
Primary area: transfer learning, meta learning, and lifelong learning
TL;DR: SGD can do worse than AdamW under distribution shifts, but simple changes make SGD competitive (and more memory efficient)

Abstract: SGD and AdamW are the two most used optimizers for fine-tuning large neural networks in computer vision. When the two methods perform the same, SGD is preferable because it uses less memory (12 bytes/parameter with momentum and 8 bytes/parameter without) than AdamW (16 bytes/parameter). However, on a suite of downstream tasks, especially those with distribution shifts, we find that fine-tuning with AdamW performs substantially better than SGD on modern Vision Transformer and ConvNeXt models. We find that large gaps in performance between SGD and AdamW occur when the fine-tuning gradients in the first "embedding" layer are much larger than in the rest of the model. Our analysis suggests an easy fix that works consistently across datasets and models: freezing the embedding layer (less than 1% of the parameters) leads to SGD with or without momentum performing slightly better than AdamW while using less memory (e.g., on ViT-L, SGD uses 33% less GPU memory). Our insights result in state-of-the-art accuracies on five popular distribution shift benchmarks: WILDS-FMoW, WILDS-Camelyon, BREEDS-Living-17, Waterbirds, and DomainNet.

Published as a conference paper at ICLR 2024
How TO FINE-TUNE VISION MODELS WITH SGD
Ananya Kumar
Ruoqi Shen,
ananya@cs. stanford.edu
shenr3ecs wash wessiinge ington
Sébastien Bubeck
Suriya Gunasekar suriyaggmmicrosoft.ccm
www.bbeccciccoooof...oooo com
ABSTRACT
SGD and AdamW are the two most used optimizers for fine-tuning large neural networks in computer vision. When the two methods perform the same, SGD is preferable because it uses less memory (12 bytes/parameter with momentum and 8 bytes/parameter without) than AdamW (16 bytes/parameter), However, on a suite of downstream tasks, especially those with distribution shifts, we find that fine-tuning with AdamW performs substantially better than SGD modern Vision Transformer and ConvNeXt models. We find that large gaps in perfor- mance between SGD and AdamW occur when the fine-tuning gradients the first "embedding" layer are much larger than in the rest of the model. Our analysis suggests an easy fix that works consistently across datasets and models: freezing the embedding layer (less than 1 % of the parameters) leads to SGD with or without momentum performing slightly better than AdamW while using less memory (e.g., VIT-L, SGD uses 33% less GPU memory). Our insights result in state-of- the-art accuracies on five popular distribution shift benchmarks: WILDS-Camelye Camelyon, BREEDS-Living 17, Waterbirds, and DomainNet.
INTRODUCTION
Fine-tuning large pretrained models on downstream tasks has become dominant approach in deep learning (Korbllth et al., 2019; Chen et al. 2020 Zhai et al. 2020). The two most commonly used optimizers in current practice
SGD and AdamW (Kingma & Ba 2015 Loshchilov & Hutter 2019 [ While most modern vision architectures (ViTs, ConvNeXts, and variants) increasingly use A for pretraining, it is still common
SGD for fine-tuning. Part of the appeal is that SGD is more memory and compute efficient: AdamW uses 1.33x and 2 as much memory per parameter as SGD with and without momentum, respectively (Ginsburg al. 2019 Dettmers et al. 2022).At the same time, in terms of fine-tuning accuracies, prior work (Dosovitskiy et al. I Steiner et al., 2021 Kumar et al. 2022) report similar performance between AdamW W and SGD on ImageNet like domains that are closer to pretraining data. In contrast, we reach different conclusions when fine-tuning on datasets that are far from pretraining data or have substantial distribution shifts.
Caron et al. 2021 Radford et al.) 202T), ConvNeXts (Liu et al. 2022), and ResNets (Kolesnikov et al.] 2020 He al. 2016), of different sizes and pretraining modalities. When pretrained on a large corpus and then fine-tuned, these models achieve near state-of-the-ar performance on downstream benchmarks. In addition to good transfer learning, we also want fine-tuned models to handle practical distribution shifts gracefully. So we focus on distribution shift datasets that have both in-distribution (ID) and dftto--disttibution (OOD) evaluations: WILDS-FMoXS WILDS- WILDS-Camelyee Camelyon, Waterbirds, BREEDS- Living-17, DemainNet. These were selected to capture different types of data shifts (subpopulation shifts, spurious correlations, style shifts), including real world shifts in medical imaging and satellite remote sensing from the WILDS benchmark (Koh et al. 2021).
We examine 7 popular models, including vision transformers (Dosovitskiy et al. 2021
We find that on newer models like ViTs and ConvNeXt, Adam W can significantly outperform SGD, especially OOD. Averaged across the datasets, fine-tuning a CLIP ViT-B/16 model with AdamW W gets 2.1 higher accuracy ID and 8.1 1 higher accuracy OOD compared to SGD (Figure Tbp. These gains
By default, the deep learning usage of SGD minihatch stochastic gradient descent with momentum.

---

Published as a conference paper at ICLR 2024
Hear
ID accuracy
OOD accuracy
pectrained
SGD
90.0%
67.9%
Irnn, 504++11
ITC
AdamW
(+2.1%)
(+8.1%)
SGD (frezze-embed) SGD (freeze-embed no mom.
(+2.1%)
(+8.0%)
freeze
Embbediing
(+2.2%
(+9.0%)
(a) Simplified schematic of ViT illus- trating how we do freeze-embed.
(b) Performance of different fine-tuning methods on CLIP ViT-B/16 averaged over distribution shift datasets.
Figure 1: We fine-tune models including ViTs, DINO, CLIP, ConvNeXt, ResNet, on 5 distribution shift datasets (Living-17, Waterbirds, DomainNet, WILDS-Camelyee WILDS- Fine-tuning with SGD gets lower accuracies than AdamW on modern pretrained models (Vision Transformers and ConyNeXt), especially OOD. Interestingly, minor tweak to SGD where we freeze the first "embedding" layer (< 1% of parameters see Figure Tay is competitive with A while using lower GPU memory. Further, dropping the momentum state in SGD gives additional gains in accuracy at even lower memory cost.
consistent across models too averaged across all models and datasets, AdamW gets 1.2% higher accuracy ID and 4.0% higher accuracy OOD (Tables 12
A key difference between AdamW and SGD, is that AdamW normalizes the gradient update of each parameter using an estimate of their second moments. Thus, parameters with consistently high gradients will change less when using Adam W than with SGD. Towards understanding these dynamics are better, we examine the gradients at each layer of our pretrained models. We find that for the models where AdamW significantly outperforms SGD, the gradients at pretrained initialization of the first "embedding" layer are much larger than the gradients of the other layers.
To test if over-training of the embedding layer in fact why SGD performs worse than Adam W we consider a minor modification where we freeze the embedding layer and tune the rest of the model with SGD (Figure Taa- -e call this SGD (freeze-embed), embedding layers are only small fraction (around 0.7% for ViT-B/16) of the total parameters of the model, so a priori we might not expect a substantial difference in accuracies. However, surprisingly this simple freezing of the embedding layer consistently improves SGD performance across most models and datasets and achieved ID and OOD accuracies that are competitive with or better than AdamW (Figure [b). Averaged across all datasets and models, SGD ffreeze-embed gets 76.7% accuracy OOD (vs. 72.0% for SGD and 76.0% for Adam The analogous Adam (freeze-embed) gets 76.5%, which does not improve over SGD (freeze-embed), supporting that freeze-embed may be the reason AdamW outperforms SGD (it is not an independent axis of improvement). We also tried a more memory efficient variation. SGD (freeze-embed, no momentum), which drops the momentum state in "GD-interestinggl. this gets even slightly better OOD accuracy than the other methods (76.9%) despite using even less memory.
vision transformer models, the
In terms of memory usage, our profiling (Table 4 shows that on a ViT-B/16, AdamW uses 16% and 36% more memory than SGD (freeze-embed) and SGD (freeze-embed, no momentum), respectively. The memory overhead of AdamW increases with the model size. On ViT-L/14 the overheads of A are 18%, and 49%, respectively.
These methods and insights, while simple, lead to state-of-the-art accuracies on all five datasets: WILDS Cameel,, WILDS-FMoW, DomainNet, Waterbirds, and BREEDS Living-17, while being more memory efficient than AdamW.
2 SCOPE AND SETUP
We use the following notation: a network map fo : X' + y is represented composition of layers (head) r(L)
as fo 1 f .... o f())) /(embed)) where (fhead; oz,,.., 01, fmmbed) denote all the parameters of the model. We use /(embed) and /(head) to denote blocks that can conceptually be considered the "embedding" layer and the "head", respectively.
Fine-tuning. Consider networks that have been pretrained to get an initialization Opretrain We focus on fine-tuning on labeled dataset D Prinetune from a new task. Concretely, given loss
2

---

Published as a conference paper at ICLR 2024
function y x Rzo. we minimize the training loss L(0) Periil Lizg)) ((fo(z).))
using iterative optimization algorithms starting from the initialization Opretrain
We evaluate the accuracy of fine-tuned models on held-out in-distribution (ID) and out-of-distribution (OOD) test datasets. For ID evaluation, we use samples Diest Pfinetune from the same distribution as the fine-tuning training dataset. To examine whether we have learned robust model, we consider benchmarks that also provide OOD test examples Deed Poods which differs from the fine-tuning distribution Prinetune in practically meaningful ways.
2.1 OPTIMIZERS: SGD AND ADAMW
The two most common optimizers for minimizing the fine-tuning loss L(0) from pretrained initial- ization are SGD (with/no momentum) and AdamW. We will introduce other optimizers as needed. Compared to vanilla SGD (no momentum), which only stores the parameters and gradients as opti- mizer states, SGD (with momentum) stores extra state per parameter (to track the first moment), while Adam W stores 2 extra states per parameter (to track the first and second moments) see Appendix
A for exact updates. This corresponds to difference between AdamW and SGD and SGD (no momentum) of 4GB and 8GB GPU memory per IB parameter model during training With the current scale of the models 100s of billions parameters and increasing, such memory overheads are very costly (Dettmers et al. 2022). Thus, understanding when and how compared to AdamW can significantly improve training of large scale models.
use the cheaper SGD
2.2 DATASETS AND MODEL ARCHITECTURES
Datasets. We choose five fine-tuning benchmarks that capture different types of data shifts (sub- population shifts, spurious correlations, style shifts), including two real world shifts.
1. Living-17 (Santurkar et al. 2020) is sub-population shift dataset from the BREEDS benchmark. The goal is to classify an image as one of 17 animal categories with ID and OOD data from different sub-categories. For example, in the "bear" category, the ID dataset contains black bears and sloth bears and the OOD dataset has brown bears and polar bears.
2. Waterbirds (Sagawa al., 2020) is spurious correlation dataset where the goal is to classify an image as a "waterbird" or "landbird". In the ID dataset, "water" backgrounds are typically correlated with "waterbird" labels, but are uncorrelated in the OOD dataset.
3. DomainNet (Peng et al.) 2019) is a domain adaptation benchmark. ID contains sketch images, and the OOD contains real images of the same categories. We use the version the dataset from Tan et al. (2020).
4. WILDS-FMoW (Christic et al., 2018 Koh et al. 2021) consists of remote sensing satellite images. The goal is to classify satellite image into one of 62 geographical categories. The ID dataset contains satellite images from across the world between 2002 and 2012, and the OOD dataset contains images from Africa in 2017.
5. WILDS DaLDD--Cameyon (Bandi et al. 2018 Koh et al. 2021) is a medical images dataset for detecting tumors in tissue slides. The II and OOD datasets contain slides from different hospitals.
Model Architectures We consider seven popular pretrained models that span different architectures (vision transformers and convolutional networks), sizes, and pretraining objectives (multi modal, supervised, self-supervised).
(1-2) CLIP ViT-B/16 and CLIP ViT-L/14 (Radford et al.) 202T): CLIP vision transformers of two
sizes pretrained on mulli-modee WebImageTTexx dataset.
(3) ViT-B/16 (Dosovitskiy et al., 2021): : vision transformer pretrained on Imagenet-2 -21..
DINO ViT-B/16 (Caron al. (202 self-supervised ViT pretrained on ImageNet-1
ConvNeXt-B (Liu et al.) 2022): modernized convnet pretrained on ImageNet-21k using
advanced data augmentations and Mix as in Touvron et al. (202I).
"The bottleneck for memory in older ResNe(X)ts is typically the number of activations. However, in modem large transformer training, memory requirements are of the same scale as the number of parameters. Further, techniques such as gradient accumulation and gradient checkpointing can make the activation memory small leaving the optimizer states as the main bottleneck.
3

---

Published as a conference paper at ICLR 2024
(6-7) BiT ResNet-50 and BiT ResNet-101 101 (Kolesnikov et al. 2020): ResNetV2 models of two sizes pretrained on ImageNet-21k.
((embed) and { For vision transformer models, we consider the patch-to-toker embedding fontes
Jorai
layer along with its layer norm if present as the embedding layer. For convolutional networks, the embedding layer refers to the "stem" block along with the first stage: the "stem" in ResNerV2 is a 7 x convolution with stride 2 followed by a 2 x 2 MaxPool; while in ConvNeXt it is a non- overlapping 4 x 4 convolution with stride 4. For each downstream task, we replace the final layer of all the pretrained models with randomly initialized classifier head g(head)
3 SGD, ADAMW, AND LAYER GRADIENTS
Modern deep learning models increasingly use AdamW for pretraining, where it has been repeatedly shown to produce better features for downstream tasks than SGD (Dosovitskiy et al. 202T Lu
et al.
2021 2022). For fine-tuning, on the other hand, there are no systematic studies or a definitive
answer as to whether AdamW SGD should be used (Dosovitskiy et al.) 2021 Touvron et al. 2021).
The ablation study in TTuvron et al. (2021) even found that there is no difference in performance between AdamW and SGD when fine-tuning on ImageNet-IK. ConvNext (Liu et al. 2022) and Swin transformers (Liu et al. a comparison with SGD. In this work we focus on better understanding the dynamics of Adam and SGD during the fine-tuning phase. Detailed results are discussed in Section 4 We first highlight some initial observations below.
(02T) papers report using AdamW for fine-tuning, but they do not mention
3.1 ADAMW VS SGD
Adam W outperforms SGD. We find that, generally speaking, fine-tuning with AdamW produces better results than with SGD, especially for more recent models like ViT variants and ConvNeXt. The gaps are more dramatic on out-of-distribution (OOD) test accuracies compared to in-distribution (ID). See Table and Table 2 for the full OOD and ID results, respectively. Averaged across the 7 models and 5 datasets, Adam W gets an OOD accuracy of 76.0% (vs. 72.0% for SGD), and an ID accuracy of 91.5% (vs. 90.3% for SGD). We emphasize that this happens even though we sweep over 6 learning rates and early stop.
AdamW ~ SGD on BiT ResNets. Breaking down the results by model type, we find that Adam W and SGD perform comparably for older convolutional networks, namely BiT ResNet-50 and BiT ResNet- 101. For example, on a BiT ResNet-101, AdamW gets an average OOD accuracy of 74.5% (vs. 74.6% for SGD). However, for newer models including vision transformers and the modernized convolutional network (Con,NeXt), AdamW gets much higher accuracies.
3.2
EXAMINING LAYER GRADIENTS
A key operational difference between AdamW and SGD is that AdamW divides the gradient update for each parameter by a weighted running average of its second moments. So parameters with consistently high gradients will change less when using AdamW than with SGD. This suggests examining the gradients across different components of the neural network (at pretrained initialization) they vary lot, then AdamW and SGD will behave very differently.
We measure the average gradient norm at each layer, across minibatches of the training set. More formally, recall our notation for network layers as fo = / minibatch B- {x11, vI), ...,, (Io, yb)} of random samples from the training data, the stochastic gra- dient of layer E (embed, 1.2..... . L. head) is given by: (B) [ >!' Vf, (x:;) y:). The norm of the stochastic gradient, (Be(B))l2, roughly measures how much layer will change after one step of SGD from pretrained initialization. We use the average gradient norm across all minibatches BI.... Bm in the training set Dirain as measure of movement in the first SGD step.
p(embed) ((head) p(L) o o... I S'gentes Given
1 Σ Ig (B)ll, where E E {embed, 1 - L, head).
Ginit
(3.1)
4

---

Published as a conference paper at ICLR 2024
tid Bed wed laver
DERINRED
Figure 2: We visualize the layer-wise gradient norms our models on DomainNet at the pretrained initialization. For each layer, we plot the average minibatch gradient norm G7 computed in equation 3.1 We highlight special layers: gradient norms of the "embedding" layer parameters G shown red-squares, and those the classifier "head"'s G are show as green-triangle. The middle layer gradient norms are shown as black-cireles. For transformer models, we have separate (black) points for MLP and the attention layers.
Note that we measure all the gradients at the pretrained initialization geretrain before performing any gradient updates. In Figure 2 we plot the mayer-wise gradient norms Gipit on DomainNet. Similar plots for other datasets and with alternative normalization are provided in Appendix
First-layer gradient is an outlier. Apriori, we expect the gradients of the classifier "head" (green- triangles) to be large as they are randomly initialized while the rest of the model is pretrained (Kumar) al., 2022). Interestingly, we see Figure 2 (see also Appendix F) that for the recent models (ViTs stand out as outliers with
and ConyNeXt), gradient norms of the "embedding" layers much larger gradients than the other layers. These are also the models where see big gaps between AdamW W and SGD performance in Tables 1
This suggests that the embedding layer plays a distinctive role newer models. Since SGD uses the same learning rate for all layers, we will end up make substantially larger updates to the embedding layer compared to other layers-leading to either over-tuning the embedding layer or under-tuning the remaining layers. Over-training of the embedding layer is undesirable given the common wisdom that lower layers ought to be tuned less as they learn more transferable features (Kumar et al., (022). AdamWo the other hand adaptively normalizes the movement of each parameter. On the other hand, computationally. SGD is preferable over Adam due to lower memory footprint.
How avoid the above issues with SGD? One possibility is to different learning rates for different layers, but this might end up requiring extensive hyperparameter tuning. Another option is use other low-memory footprint optimizers with layerwise normalization techniques like LARS (You et al. 2017a) and LAMB (You et al. 2020). In our initial experiments (see Table while these methods improved over SGD, they did not close the gap with AdamW. Instead, we found much simpler modification to SGD that consistently leads to accuracies competitive with or better than AdamW W while using lower memory (Section
3.3 WHY FIRST-LAYER GRADIENTS ARE HIGHER?
Why are the first-layer gradients higher for modern vision models such Vision Transformers and ConvNeXt? There are two plausible hypotheses: (a) architectural differences from ResNet models, or (b) optimization differences between pretraining and fine-tuning- -iTT ResNet models were pretrained with SGD while the Vision Transformer and ConvNeXt models need to be pretrained with AdamW (Steiner et al. 2021). In Appendix [E we run controlled experiments and find that the optimizer mismatch appears to be the key factor. To control for architecture, we use the same ResNet architecture and pretrain using either (1) SGD or (2) A. For the SGD pretrained ResNet the first layer has smaller gradients than the other layers, while for the AdamW pretrained ResNet the first layer has higher gradients than the other layers (Figure 33. In line with our previous results, fine-tuning an AdamW W pretrained model with SGD leads to worse OOD accuracy (Table 13), but fine-tuning an SGD pretrained model with SGD performs well. On the other hand, architectural changes lead to smaller changes.
3.4 FREEZE-EEMEEDING
While the presence of large embedding layer gradients in Figure correlates with the observed performance gaps between AdamW and SGD, it is not clear that this observation its potential
5

---

Published as a conference paper at ICLR 2024
issues discussed above are definitive causes for SGD performing worse than AdamW. To further test the intuition, in the next section, we consider a "freeze- treeze--embdd variation of fine-tuning, where we simply freeze the embedding layer to its pretrained initialization, and then fine-tune the rest of the model as usual. The embedding layer only consists of a small fraction of the model parameters (e.g.. 0.7% in base vision transformer), so apriori freezing the embedding layer is very tiny tweak-ane and we would not expect large change in accuracy. However, if the hypothesis above holds merit, then we would expect the modification to aid SGD but not AdamW.
DETAILED EXPERIMENTS ON FREEZE-EMBEDDIN
We consider a simple variation of SGD fine-tuning where we freeze the embedding layer and only perform SGD updates on the rest of the network we call this method "SGD (freeze-embed)" For further memory gains, we also consider SGD without momentum "SGD (freeze-embed, no momentum)" this section we discuss the results of fine-tuning our models on distribution shift benchmarks mentioned in Section 2 We use the implementations of SGD and Adam W in Py Torch. For each method, we train for the same number of epochs using a cosine learning rate schedule, and sweep over 6 starting learning rates (ensuring that the optimal learning rate is in the middle of the sweep). For all datasets we follow prior work (Kumar et al., 2022) and pick the best learning rate and early stop based the ID validation accuracy. See Appendix A for additional details.
Table and Thhl the OOD and ID accuracies, respectively, of the four methods across models and datasets. For each model and dataset, we also highlight the relative gain/loss of AdamW and SGD (freeze-embed) from regular SGD in green/red. We discuss the main observations below.
1. AdamW outperforms SGD. We see that AdamW largely outperforms SGD, often by substantial margins on ViT and ConvNeXt models. Particularly in OOD evaluation, the gaps are remarkable. In the few instances where SGD performs better the gaps much smaller. The differences between SGD and Adam are generally more modest for older ResNet models.
SGD (freeze-embed) as well as SGD (freeze-embed, no momentum) are competitive with or better than Adam W. For each individual model, averaged across the datasets SGD (freeze-embed) variants are consistently the best or tied-best method on OOD accuracy, and only minimally worse than AdamW on ID) accuracy (right columns of Table 12. Averaged across all the models and datasets, SGD (free7e-embed) performs the best OOD, getting an average OOD accuracy of 76.7% (vs. 71.9% for SGD and 76.0% for AdamW). On ID data, SGD (Freeze-embed) closes 85% of the gap between SGD and AdamW, getting an average accuracy of 91.3% (vs. 90.2% for SGD and 91.5% for AdamW), SGD (freeze-embed, no momentum) gets the highest average OOD accuracy of 76.9% (vs. 71.9% for SGD and 76.0% for Adam W) and competitive ID accuracy of 91.2% (vs. 90.2% for SGD and 91.5% for AdamW), while saving additional memory.
3. Larger CLIP model gains more from SGD (freeze-embed) and Adam On a CLIP-ViT-B/16 we see that SGD
(eeez) and Adam get about 8% higher OOD accuracy than SGD. Upon scaling to a larger CLIP-ViT-L/14 which is also our best model for OOD performance, we see even higher gains. On accuracy than SGD, and a 0.7% higher OOD accuracy than A This suggests that our findings might be even more relevant for larger models.
SGD (rrezze-embed) gets a 14.3% higher OOD
4. AdamW is red when SGD (freez)-embed) is red. Across all our models and datasets, see that in instances where SGD (freeze-embed) is worse than SGD (i.e.. highlighted as red), Adam is also worse than SGD. This is not always the case the other way around. For example, on ViT B-/16 fine-tuned on NomainNet and Camelyon, OOD performance of AdamW is worse than SGD, but SGD (frecze-embed) is competitive the best of the two. At the same time, on Living- 17 OOD evaluation, multiple models have both A and SGD (freeze-embed) perform significantly worse than SGD.
5. SGD performs well when fine-tuning data is closer to pretraining. Breaking down the results by datasets, the main trend we see is that with models that were pretrained on ImageNet-21 l (all models here except CLIP) and fine-tuned on Living-17. 17, the performance of SGD is typically higher than Adam W and SGD (freeze-embed)) Images in Living-17 are derived from ImageNet- 1K and hence arguably, in this case the fine-tuning distribution is closest to pretraining. This suggests that SGD may work better when tine-tuning and pretraining data are similar- we leave a more thorough analysis to future work.

---

Published as a conference paper at ICLR 2024
laving-17
Waterbirds
DomainNet
FMoW
Camelyon
Avg.
CLIP 8/16
SGD Adam
80.0
62.5
72.8
37.3
96.8
67.9
CLIP VIT-8/16
8288 (-2.)
71.9 (49.4)
89,2 (+16.4)
40.7 (-3.4)
95.7 (48.9)
76.0 (+8.1)
CLIP V/T-8/16
SGD (freeze embed)
83,2 (+3.2)
73.7 (x11.2)
88.2 (+15.4)
40.2 (+2.)
94.3 (47.5)
75.9 (+8.0)
CLIP
SGD (frezee embed, momentuum)
83.1 (+3.3)
90.4 (+17.9)
(+16.2)
38.8 (-1.5
93.3 (+6.5)
76.9 (+9.0)
CLI 21T-1/14
SGD
84.2
65.0
60.8
41.0
83.2
66,8
CLIP MTT-//14
Adam
58.0 (-3.8)
85.3 (+20.2)
93.8 (+33.0) 0)
18.3 (-7.3)
95.9 (+12.7)
82.2 (+15.0)
CLIP VIT-1/14
SGD (frezze-embel)
90,5 (-6.3
84.7 (+19.7)
93.1 (+32.33
49,9 (-8 9)
96.5 (+13.3)
83.0 (+16
CLIP VIT-1/144
SGID (freeze--cnin ceooeedemddd..om.... momentum)
89.2 (-5.0)
(421.8) 8)
46.8
(413.5)
(+15.8)
Sup VIT-R/16
SGD
89.5
77:4
N6.3
33.5
92.6
75.8
Sup VIT-R/16
Adam
88.3 (-1.2)
81 6. (14.2)
84.4 (-1.9)
35.9 (-2.4)
87.9 (-4.7)
75.6 (-0.2)
Sup VIT-B/16
SGD ceezz)-
88.0 (-1.5)
82.4 (45.0)
86.3 (10.0)
34.4 (+0.9)
93.7 (+1.1)
77.0 (+:1.2
Sup VIT B/16
SGD drreze embed. momentumm
$8.4 (-1.3)
76.6 (-0.8)
$6.1 (.0.2)
34.6 (-1.3)
(:5.1)
74.6 (-1.2)
DINO
33.6
68.2
DINO VIT-1116
Adam W
$7.4 (-0.8)
61.2 (+5.1)
77.4 (+1.4)
35,8 (-2.2)
91.9 (+5.0)
70.7 (+2.5)
DINO VIT-B/16
SGD (freese-emine embee))
86.7 (-1.5)
67.9 (+11.8)
78A (+2.4)
35,9 (-2.3)
90.6 (+3.7)
71.9 (+3.7)
DINO VIT-R/16
SGD (freeze embod, momentum)
88.2 (-0.0)
68.5 (412.4)
79.9 (43.9)
33.4 (-0.2)
93,9 (47.0)
72,8 (+4.6)
CoevNext-Burt Base
39.7
CrevNeet- Rose
Adam
90.3 (-3.7)
89.8 (49.6)
89.5 (-0:3)
38.4 (-1:3)
89.5 (46.5)
79.5 (+222)
Conv/Next-Bos
SGD (freeze-emhad)
92.6 (-1.4)
86.9 (46.7)
91.2 (+1.4)
38.2 (-1.5)
S3.1 (45.1)
79.4 (+2.1)
ConvNest- Base
SGD direeze embed, no moceentum)
93.3 (-0.7)
83.6 (43.4)
91.5 (+1.7)
38.1 (-1.6)
96.0 (+13.0)
80.5 (+3.2)
ResNes-50
SGD
84.3
76.5
80.0
34,1
90.4
73.1
ResNee-50
Adam'
83.1 (-1.2)
74.8 (-1.7)
(+4.0)
33.8 (-0.3)
(+2.0)
73.6 (+0.5)
B:T ResNee-50
SGD (freeze embrd)
84.1 (-0.2)
75.5
(-1.0)
82.3 (+2.3)
35,0 (-0.9)
95.2 (+1.8)
74.4 (+1.3)
BIT ResNee-50
SGD (frex ot embed. no momentum)
83.8 (-0.5)
76.0
(40.5)
83.8 (+3.8)
33.6 (-0.5)
95.6 (45.2)
74.5 (+1.4)
B:T ResNex-101 101
SGD
82,8
76.9
86.2
36,0
89.3
74.6
B/T ResNes-101
Adam
82,9 (-0-3)
79.4 (42.5)
37.0 0g
(40.4)
(.0.1)
B/T ResNes-101
SGD (frezze-cmbad)
83.1 (-0.3)
77.3
(40.4)
(-0:2)
36.0 (:2.0)
95.5
(46.2)
75.6 (+1.0)
B:T ResNet-101 101
SGD (frucae-embod no momentum)
82,9 (-0.3)
79.1
(42.2)
86.0
(-0.2)
36.4 (-1.6) (.66
95,9 (+5.6)
76.1 (+1.5)
Table 1: Out-of-distribution (OOD) accuracies including SGD embbed) without momentum.
Living117
Waterbinds
DomminNet
FMoW
Camelyon
Avg.
CLIP
SGD
97.8
88.8
67.0
99.4
90.0
CLIP VIT-B/15
Adam
98.1 (+0.3)
97.7 (-0.5)
95.0 (+6.2)
70.1 (+3.1)
99.5 (+0.1)
92.1 (+2.1)
CLIP VIT-R/16
SGD (freene- cahhed)
98.2 (40.4)
97.8 (10.6)
94.9 (16.1)
70.0 (13.0)
99.5 (10.1)
92.1 (42.1)
CLIP VIT-8/16
SGD (fresze- embed, momenturs)
98.2 (+0.4)
97.9 (+0.7)
95.2 (16.4)
70.1 (x3.1)
99.5 (s0.1)
92.2 (+2.2)
CLIP VIT-L/14
SGD
98.4
97.3
84.3
69.0
99.4
89.7
CLIP VIT-I/14
Adam
98.9 (40.5)
9M.8 (-1.5)
96.9 (+12.6)
74.5 (15-5)
99.6 (40.2)
93.7 (44.0)
CLIP VIT-L/M4
SGD (ireca)-mmbed)
98.7 (+0.3)
98.9 (-1.6)
97.1 (+12.8)
74.5 (+5.5)
99.6 (40.2)
93.7 (+1.0)
CLIP VIT-L/14
(fiezzeeebee, momenturn)
(40.4)
(-1.4)
17.3 (+13.0)
74.3 (+5:3)
99.5 (40.1)
(+4.0)
Sup WTTBB116
SGD
98.6
99.1
91.7
64.1
99.4
90.6
Sup VTT-B/16
A
98.7 (+0.1)
99.0 (:0.1)
91.7 (-0.0)
66.4 (+2.3)
99.5 (+0.1)
91.1 (+0.5)
Sup WTT-/116
SGD (freeze-ene -erreeeeeeedd
98.7 (10.1)
99.2 (10.1)
91.5 (-0.2)
65.0 (10.9)
99.6 (10.2)
90.8 (10.2)
Sup WTT-B/16
(freeze embed, momentura)
(-0.1)
98.9 (40.2)
90.6 (-1.1)
65.7 (+1.6)
99.5 (+0.1)
90.6 (+0.0)
DINO VIT-8/16
SGD
98.4
97.0
88.2
62.4
4
89.1
DINO VIT-R/16
Adam W
98.5 (40.1)
97.9 (-0.9)
89.4 (41.2)
66.0 (43.6)
99.6 (40.2)
90.3 (+1.2)
DINO V/T-B/16
SGD (reeca-embbll)
98.4
97.5 (-0.5)
89.0 (40.8)
63.5 (+1.1)
99.5 (+0.1)
(+0.5)
DINO VIT-B/16
SGD (frezz-embel, no momenturn)
VM.5 (40:1)
97.5 (-0.5)
89.2 (+1.0)
63.8 (4)
99.5 (40.1)
89.7 (40.6)
Com-Nex- Base
SGD
98.7
99.0
94.8
66.3
99.4
91.6
Com-Next- Base
Adum
98.6 (-0.1)
99.5 (+0.5)
91.5 (-0.3)
68,8 (+2.5)
99.7 (+0.3)
92.2 (+0.6)
Com Base
SGD (freeze-cme ceneeee
(+0.4)
95.1 (+0.3)
67.4 (+1..)
99.5 (10.8)
92.0 (10.4)
Com-NextBBas Base
SGD (frezze embed, no momentur))
98.8 (+0.1)
99.2 (+0.2)
95.0 (+0.2)
66.0 (40.3)
99.4 (:0.0)
91.7 (+0.1)
BIT ResNet-50
SGD
97.4
98.4
89.3
64.6
99.5
89.8
BIT ResNes-50
Adam W
97.2 (-0.2)
98.5 (40.1)
89.2
65.1 (40.5)
99.5 (10.0)
89.9 (10.1)
BiT ResNct-0
SGD (frocac-emibd))
97.6 (+0.2)
98.5 (+0.1)
89.2
(-0.1)
64.8 (40.2)
99.5 (40.0)
89,9 (40.1)
BIT ResNet-50
SGD (frezze-emeed, no momenturn)
97.5 (40.1)
9X.4 (-0.0)
89.2 (40.1)
63.6 (-1.0)
99.4 (-0.1)
89.6 (-0.2)
BiT ResNet-100 101
SGD
98.3
98.9
92.0
66.0
99.4
90.9
BiT ResNet-
Adam
98.4 (+0.1)
(-0.3)
91.1 (-0.9)
67.0 (+1.0)
99.6 (+0.2)
90,9 (-0.0)
BiT ResNes-101
SGD (freeze embed)
98.4 (+0.1)
98.8
(.0.1)
91.5
(-0.5)
65.9 (60.1)
99.5 (.0.1)
90.8 (-0.1)
BiT ResNee- 101
SGD (freeze embed, no momentura)
98.1
(-0.2)
98.9 (-0.0)
91.5
(-0.5)
66.3 (+0.3)
99.3 (-0.1)
90.8 (-0.1)
Table 2: In-distribution (ID) accuracies including SGD (freeze-embed) without momentum.
4.1 SOTA EXPERIMENTS AND RESULTS
Our experiments get new state-of- the art results for OOD accuracy on all 5 datasets. On 3/5 datasets (Living- 17, WILDS-FMOW and WILDS-Camelyee best, while in other 2, AdamW has small edge. Here, state-of-the-arr means that the numbers we get are better than, to our knowledge, any reported number and all numbers and are better than standard full fine-tuning with SGD at the time of the preprint of the paper. We show the best results from our paper in Table with comparison to the previous wate-of-the-art
our proposed SGD (freeze-embed) does the
the official leaderboard,
As a final point, we mention that if our hypothesis that AdamW would inherently avoid over-tuning of the embedding layer were true, unlike for SGD, freezing the embedding for AdamW would not be beneficial. In the Appendix, we expand the Tables [ to include Adam V (freeze-embed) and indeed see that freeze-embed does not provide complementary gains on top of AdamW.
7

---

Published as a conference paper at ICLR 2024
Living-17
Waterbirds
DomainNet
FMoW
Camelyon
Best prior result
87.6
89.3
87.2
47.6
93.3
Best result from our paper
90.5
89.8
93.8
49,9
96.5
SGD
SGD
SGD
Optimizer for best result (freeze-embed) embed))
Adam
AdamW
ftrezze-emeedd
Model for best result
CLIPVIT-L/14
ConvNeXI-B B
CLIP VIT-/114
CLIP VIT-1/14
CLJP VIT-1/11
SGD fireze-embedd result
90.5
86.9
93.1
49,9
96.5
Table 3: Our OOD accuracy results compared with the best reported numbers in prior work on these datasets.
We restrict to methods that do not use OOD data for hyperparameter selection early stopping. To our knowledge, the previous ofttc-o--hhe-at results are from Wortsman et al. (2022) for FMoW, RRoee al. (202 for Camelyon, Kumar al. (2022) for Living-17 17 and the version DomainNet introduced by Tan et al. (2020), and Ghosal et al. (2022) for Waterbirds. The WILDS numbers and references are taken from the official WILDS
leaderboard (as of 28 Sep 2022), and for Waterbirds we consider all methods that do not use group labels. For Living-17, we omit models pretrained with ImageNet- IK as Living-17 17 is subset of ImageNet-IK.
DETAILED ANALYSIS OF CLIP.
CLIP models have strong transfer learning performance and robustness among our 7 models, CLIP models had the best ID and OOD accuracies averaged across datasets. So we did a more detailed analysis of the CLIP ViT-B/16 with other optimizer which are summarized in Table 5
CLIP ViT-B/16
CLIP VYT-L/14
CLIP ViT-H/14
AdamW SGD
2.7 GB
7.1 GB
Out-of-
2.3 GB
6.1 GB
10.1 GB
SGD (freeze-embed)
2.3 GB
6.1 GB
10.1 GB
SGD (freeze-embed, no momentum)
2.0 GB
4.8 GB
7.6 GB
Table 4: We profile the GPU memory consumption on three CLIP models of varying sizes, on a Titan-X GPU, SGD (freeze-embed) gives practical gains over AdamW, especially if we drop momentum. The largest CLIP VIT-H/14 does not fit GPU memory when using AdamW, but fits memory with other optimizers. Note that the treeze-embed methods perform competitively or better than Adam on accuracy, as shown Table
GPU memory profiling. We profiled the GPU memory consumption of our 4 fine-tuning methods, on 3 models, namely CLIP-ViT ViT B/16, CLIP ViT-L/14, and OpenCLIP ViT-H/14 the original CLIP model only scales up to ViT-L/14, so we used the OpenCLIP (Ilbarco al.. 2021) for ViT-H/14. The profiling was done using Weights and Biases on a Titan-X GPU with micro-batch size of 1.
In Table 4 we ViT-B/16, AdamW uses 16% and 36% more memory than SGD (freeze-embed) and SGD (freeze-embed, no momentum), respectively. The gains are better for larger models: on VIT-L/14, AdamW uses 18% and 48% more memory respectively. On ViT-H/I4, AdamW runs out of memory, while SGD (frezze-embed) and SGD (freeze-embed, no momentum) are able to run, showing that the gains are at least 20% and 60% respectively.
For large models, it is common to use additional tricks like gradient checkpointing to reduce the activation memory. That is, for speed penalty (at most 2x), we only need to store activations in vL of T. layers when doing backpropagation. Gradient checkpointing would further increase our gains over AdamW W since they do not change the memory consumed by the weights but can substantially decrease the memory consumed by the model's activations.
SGD (freeze-embed) and AdamW outperform LAMB and LARS. We also ran two alternative adaptive gradient methods, LARS YYou et al.) 201 7017b) and LAMB (You et al. (200)) also sweeping over 6 learning rates and early stopping. These are alternate methods with layerwise normalization that can avoid over-training of large gradient layers. Moreover, like SGD and SGD (freeze-embed),
LARS also has a lower memory footprint than Adam W. In Table 5 we see that while LARS and LAMB get higher accuracies than SGD, they do worse than SGD (freeze-embed) and Adam W both ID and OOD. In this case, our modification with freeze-embed appears to be more effective.
Generating Table involved over 600 fine-tuning runs, so we weren't 't able repeat this for every model.
8

---

Published as a conference paper at ICLR 2024
Laving-17
Woterhirss
DemoinnNe
FMol
Camelyon
Avg:
00D
00D
(0.2)
80.0
(0.2(0.1)
62.3 (5.0)
(8.1 (2.8
(0.0
27.3(1.3)
99.4 (0.0)
(1.I)
90.0
Adam'W
98.1 (0.1)
(1.2)
(0.0)
71.9 (2:4)
(5,0
(1.1))
(0.2)
(0.7
(0.0)
(freeze embed)
98.2 (0.3)
83.2 (0.8)
(0.1)
73.7 (1.1)
(0.3)
85.2 (0.7)
(0.0
(0.7
99.5 (0.0)
94.3 (0.3)
92.1
75.9
SGD
(froorescm, no monn
98.2
83.1
80.4
95.2
38.8
99.5
93.3
92.2
76.9
LAMB
79.5
64.0
99.5
91.7
73.2
LARS
97.7
83.9
97.1
48.6
93.3
67.0
38.6
99.3
93.3
90.9
69.6
Table 5: CLIP VIT-B/16 performance with new optimizers and confidence intervals. For addition to SGD, Adam W, and SGD (freeze-embed), provide 90% confidence intervals based on 3 runs of each hyperparameter configuration. In addition, we show accuracies for SGD (freeze-embed, momentum); J; well comparison two other optimizers, LARS (You et al. 2017b) and LAMB (You et al. 2020), which use layer-wise normalization in their updates.
CIFAR-10 results. As proof of concept for standard (non-OOD) transfer learning, we fine-tune CLIP ViT models on CIFAR- 10. Even at high accuracies, AdamW and SGD (freezo-embed) improve per- formance. SGD (freeze-embed) gets 20% and 30% lower error than SGD CLIP ViT-B/16 and ViT-L/14, respectively. For example, on a CLIP ViT-L/14, SGD gets 99.0% accuracy but SGD (freeze-embed) gets 99.3% accuracy.
CLIP VIE-Bis
CLIP WE--14
Scil (hiesse cnibed)
98.4
99.3
Table 6: CIFAR-10 10 accuracy
Composes with other fine-tuning methods. In Appendix B we show that SGD (freeze- can compose with other fine-tuning improvements such as LP-FT (Kumar et al. 2022 for further gains.
6 ADDITIONAL RELATED WORKS
Many works in transfer learning propose freezing parameters while fine-tuning to preserve pretrained information. For example, linear probing, which freezes the entire model except the head (Kumar et al. 2022 Wortsman et al.) 2021) ) and zero-shot models (Radford et al.) 2021) have been shown to
improve OOD performance over standard fine-tuning. In NLP, methods such prefix-tuning (Li & Liang] 2021) and prompt-tuning (Lester et al. 2021 > have been shown to improve OOD accuracy,
There are many other parameter ellicient
methods well (Elsayed et al. 2019 Chen 2023 Bahng et al. 2022 Tsao et al. 2023 Jia et al., (022). Other works propose regularizing
parameters towards initial initialization, freezing the first several layers, using different learning rates for
different layers, or tuning different layers for different examples (Long et al.) 2013 Ge & Yu 2017
Howard & Ruder
2018 Guo et al., 2019 Zhang et al., 2020; /hhu et al. 2020 Jiang et al. 2021
Aghajanyan et al. 2021). Typically, large fraction of the model is frozen to preserve the pretrained
information a key difference in this work is that we find that freezing a very small fraction of the
model (<1% of the parameters) can lead to substantial and consistent improvements in accuracy.
Other optimizers have been proposed to reduce Adam W's memory footprint, including LARS YYou et al. (017b) and AdaFactor (Shazeer & Stern 2018). Our method is simpler and achieves better accuracies with same or better memory gains. A complementary line of work Dettmers et al.) (2022) study quantization mechanisms for optimizer states. These tools, although developed for AdamW. can also be used with SGD (freeze-embed) to get additional gains in memory.
7 DISCUSSION.
We note that the methods we consider are not complex. We showed that minor tweak of freezing the embedding layer overwhelmingly improves the performance across the board when fine-tuning with SGD. We clarify that we do not claim that SGD (freeze-embed) is a substantially better method than AdamW in terms of accuracy. Rather, it is remarkable that with its simplicity, we can already achieve comparable or even slightly better accuracy than AdamW across a wide range of models and benchmarks, while using much less memory. We hope future work can focus on extending these methods to other applications such as NLP, image segmentation, and recommender systems (Naumov) et al. 2019), models without embedding layers, and more sophisticated learning rate schedules.

---

Published as a conference paper at ICLR 2024
8 ACKNOWLEDGEMENTS
We thank Percy Liang, Tengyu Ma, Yuanzhi Li, and Zhiyuan Li for helpful comments.
REFERENCES
Armen Aghajanyan, Akshat Shrivastava, Anchit Gupta, Naman Goyal, Luke Zeuteeoyer, and Sonal Gupta. Better fine-tuning by reducing representational collapse. In International Conference on Learning Representations (ICLR), 2021.
Hyojin Bahng, Ali Jahanian, Swami Sankaranarayanan, and Phillip Isola. Exploring visual prompts for adapting large-scale models, 2022.
Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, al. From detection of individual metastases to classification of lymph node status the patient level: the CAMELYON17 challenge. IEEE Transactions on Medical Imaging, 38(2)5550-560, 2018.
Mathilde Caron, Hugo Touvron, Ishan Misra, Herve Jegou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties self-supervised vision transformers. In International Conference Computer Vision (ICCV), 2021. 2021.
Pin-Yu Chen. Model reprogramming: Resource-efficien cross-domain machine learning, 2023.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International Conference on Machine Learning (ICML), pp. 1597-1607, 2020.
Gordon Christie, Neil Fendley, James Wilson, and Ryan Mukherjee. Functional map of the world. In Computer Vision and Pattern Recognition (CVPR), 2018
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via biock-wise quantization. 9th International Conference on Learning Representations, ICLR, 2022.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby, An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021.
Gamaleldin F. Elsayed, Ian Goodfellow, and Jascha Sohl-Dickstein. Adversarial reprogramming of neural networks. International Conference on Learning Representations (ICLR), 2019.
Weifeng Ge and Yizhou Yu. Borrowing treasures from the wealthy: Deep transfer learning through selective joint fine-tuning. In Computer Vision and Pattern Recognition (CVPR), 2017.
Soumya Suvra Ghosal, Yifei Ming, and Yixuan Li. Are vision transformers robust to spurious correlations? arXiv preprint 1oXXv:2008.027900 2022.
Boris Ginsburg. Patrice Castonguay, Oleksii Hrinchuk, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Huyen Nguyen, Yang Zhang, and Jonathan M. Cohen. Training deep networks with stochastic gradient normalized by layerwise adaptive second moments. 2019.
Yunhui Guo, Honghui Shi, Abhishek Kumar, Kristen Grauman, Tajana Rosing, and Rogerio Feris. Spottune: Transfer learning through adaptive fine-tuning. Computer Vision and Pallern Recog- nition (CVPR), 2019.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Computer Vision and Pattern Recognition (CVPR), 2016.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. In Association for Computational Linguistics (ACL.), 2018.
10

---

Published as a conference paper at ICLR 2024
Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt. Openclip, July 2021. URL httph:ht//////... /oi. http:///oo.rrg/....333 If you use this software, please cite it as below.
zenodo. 5143773
Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie. Serge Belongie, Bharath Hariharan, and Ser-Nam Lim. Visual prompt tuning. In European Conference Computer Vision (ECCV), 2022.
Haoming Jiang. Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Tuo Zhao. Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regu- larized optimization. In International Conference on Learning Representations (ICLR), 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsub- ramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton A. Earnshaw, Imran Haque, Sara Beery, Jure Leskovee, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. WILDS: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning (ICML), 2021.
Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. In ECCV, 2020.
Simon Kornblith, Jonathon Shlens, and Quoc V. Le. Do better imagenet models transfer better? In Computer Vision and Pattern Recognition (CVPR), 2019.
Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma. and Percy Liang, Fine-tuning can distort pretrained features and underperform out-of-distriuttion. In International Conference on Learning Representations (ICLR), 2022.
Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning, arXiv preprint vrXiv:2104.086%0 2021.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. In Association for Computational Linguistics (ACL), 2021.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021.
Zhuang Liu, Hanzi Mao, Chaozheng Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Computer Vision and Puttern Recognition (CVPR), 2022.
Mingsheng Long, Jianmin Wang, Guiguang Ding, Jiaguang Sun, and Philip Yu. Transfer feature learning with joint distribution adaptation. In Proceedings of the IEEE international conference on computer vision, pp. 2200-2207, 2013.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR, 2019.
Subhabrata Mukherjee and Ahmed Hassan Awadallah. Distilling transformers into simple neural networks with unlabeled transfer data. arXiv preprint arXiv:1910.01769, 2019.
Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sun- daraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G. Azzolini, Dmytro Dzhulgakov, Andrey Mallevich, Ila Cherniavskii, Yinghai Lu, Raghuraman Krish- namoorthi, Ansha Yu, Volodymyr Kondratenko, Stephanie Pereira, Xianjie Chen, Wenlin Chen, Vijay Rao. Bill Jia. Liang Xiong, and Misha Smelyanskiy. Deep learning recommendation model for personalization and recommendation systems. CoRR, abss/990..eaa 2019. URL /ttp://axiiiii wr.. http:///awvvv..or/a/////////

---

Published as a conference paper at ICLR 2024
Xingchao Peng, Qinxun Bai, Xide Xia, Zijun Huang, Kate Saenko, and Bo Wang. Moment matching for multi-source domain adaptation. In International Conference on Computer Vision (ICCV), 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning (ICML), volume 139, pp. 8748-8763, 2021.
Alexander Robey, George J. Pappas, and Hamed Hassani. Model-based domain generalization. In NeurlPS, 2021.
Miguel Romero, Yannet Interian, Timothy Solberg, and Gilmer Valdes. Targeted transfer learning to improve performance in small medical physics datasets. Medical Physics, 47(12):6246-66 6256,
2020.
Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance regularization for worst-case generalization. In International Conference on Learning Representations (ICLR), 2020.
Shibani Santurkar, Dimitris Tsipras, and Aleksander Madry. Breeds: Benchmarks for subpopulation shift. arXiv, 2020.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning (ICML), 2018.
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. ArXiv, abs/2106. 10270, 2021.
Shuhan Tan, Xingchao Peng, and Kate Saenko. Class-imbalanced domain adaptation: An empirical odyssey. arXiv preprint arXiv:1910. 10320, 2020.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herv'e J' J'egou. Training data-efficient image transformers & distillation through attention. In /CML, 2021.
Hsi-Ai Tsao, Lei Hsiung, Pin- Chen, Sijia Liu, and Tsung-Yi Yi Ho. Autovp: An automated visual prompting framework and benchmark, 2023.
Mitchell Wortsman, Gabriel Ilharco, Mike Li, Jong Wook Kim, Hannanch Hajishirzi, Ali Farhadi, Hongseok Namkoong, and Ludwig Schmidt. Robust fine-tuning of zero-shot models. arXiv preprint arXiv:2109.0/900 2021.
Mitchell Wortsman, Gabriel Ilhareo, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, and Ludwig Schmidt. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time. In International Conference on Machine Learning (ICML), 2022.
Jiaxuan You, Xiaocheng Li, Melvin Low, David Lobell, and Stefano Ermon. Deep gaussian process for crop yield prediction based on remote sensing data. In Association for the Advancement of Artificial Intelligence (AAAI), 2017a.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv: Computer Vision and Pattern Recognition, 2017b.
Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song. James Demmel, Kurt Keutzer, and Cho-Jui Hsich. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations (ICLR), 2020.
12

---

Published as a conference paper at ICLR 2024
Xiaohua Zhai, Joan Puigcerver, Alexander Kolesnikov, Pierre Ruyssen, Carlos Riquelme, Mario Lucic, Josip Djolonga, Andre Susano Pinto, Maxim Neumann, Alexey Dosovitskiy, Lucas Beyer, Olivier Bachem, Michael Tschannen, Marcin Michalski, Olivier Bousquet, Sylvain Gelly, and Neil Houlsby. A large-scale study of representation learning with the visual task adaptation benchmark. arXiv, 2020.
Jeffrey O Zhang, Alexander Sax, Amir Zamir, Leonidas Guibas, and Jitendra Malik. Side-tuning: baseline for network adaptation via additive side networks, In European Conference on Computer Vision (ECCV), 2020.
Chen Zhu, Cheng. Zhe Gan, Siqi Sun, Tom Goldstein, and Jingjing Liu. FreeLB: Enhanced adversarial training for natural language understanding. In International Conference on Learning Representations (ICLR), 2020.
13

---

Published as a conference paper at ICLR 2024
A ADDITIONAL TRAINING DETAILS
SGD and Adam W updates. We start with initialization 8(0: @pretrain Given hyperparam- eters, minibatch size batch_size and number of epochs num_epochss our algorithms run for T num epocks. Doain steps. At steps & = o.,,...,7? we select randomly shuffled minibatch batch size
B, from Derain (reshuffled at end of each epoch) and compute the minibatch stochastic gradient [B,] ((.y)BB Vol(fo(z), The parameters g(+) are updated as follows.
For SGD, in addition to gradients dt and weights ((0), we maintain first order momentum estimate my optimizer state, initialized m-1 0. The SGD(1), (1,;)) update with ne learning rate, / momentum, and À weight decay, is given by
9t 9t + AB(1)
my pmm+11 St (first moment)
((+1) g( m
For AdamW, we maintain two additional optimizer states: the first moment estimate m, and second moment estimate vg initialized as m- = v-1 = 0. The AdamW(ne,.5 31, 32, X) update with % learning rate, (31, .2)) betas, and À weight decay is given by
m = fmm-1 1 + - -3119
(first moment)
U/ = 82 + - B2)gi
(second moment)
m
Ut
m; (1 - 891
(1
g((+1) = (1 - T/e ()(())
m+
17t Vo +e
Gradual unfreezing. The success of freeze-embed for SGD is inline with our hypothesis that it beneficial to change the lower layers less during fine-tuning. We could further refine this idea- instead of just freezing the first "embedding" layer, could start by freezing the entire pretrained model (except for the head) and gradually unfreeze the rest of the model over training. Strategies like this have been explored by ULMFIT (Howard & Ruder 2018) and others (Mukherjee & Awadallah 2019 Romero et al., 2020).
found that naively using single learning rate throughout this process or even a cosine learning rate schedule is not sufficient to achieve good performance. Instead, fit an exponential learning rate decay law for our experiments (see details on the learning rate schedule below). The performance of -unfuullunfreezing on CLIP ViT-B/16 is shown in Table 7 performed particularly well. For example, got an OOD accuracy of 93.2% on DomainNet and 96.5% on Camelyon- wwich are both competitive with SGD (freeze-embed) and A on much larger CLIP ViT-L/14 model. On average, performed slightly worse than SGD (freeze-embed, no momentum). Overall, we find that the most consistent improvements come from simply freezing the embedding layer, which suggests that the embedding layer plays a particularly prominent role in modern vision architectures.
For gradual-unfreezing
8 On some datasets, cradual-unfreezznn
Hyperparameter details
Learning rate. We sweep over 6 learning rates for SGD ([3e-5, le-4, 3e-4, 1e-3, 3e-3, 1e-2]) and for AdamW ([3-77, le-6, 3e-6, le-5, 3e-5, le-4)). As standard, we use smaller learning rates for AdamW because they work better. The best learning rate is chosen based on ID validation accuracy. For the LP-FT experiments which we run in Section B we divide all the learning rates by 10, since have trained the linear probe, so the learning rate required fine-tune the entire model is smaller.
Other optimizer parameters. We use all the default options for SGD and Adam in pytorch except set the SGD momentum parameter to 0.9. The other default options are: for SGD, the weight decay is 0 and we use momentum without dampening, and for AdamW, the weight decay is 0.01, betas are (0.9,0.999) and epsilon is le-08.
14

---

Published as a conference paper at ICLR 2024
Number of epochs. We train for 20 epochs Living- 17, 20 epochs Waterbirds, 50 epochs DomainNet, 5 epochs on HILDS-FMoW, and 3 epochs on Camelyon. We use the same number of training epochs for all the optimizer methods we report results
Learning rate schedule. With the exception Gradual-urge sunfuelluffreezing Table and we use a cosine learning schedule and decay the starting learning rate over the course of training steps. Learning rates schedule for Gradual Gradual-unfreezing unfreezing is described below.
Learning rate schedule for gradual unfreezing. For gradual unfreezing, do not use cosine learning rate scheduler. Instead, at epoch I. (0-indexed) out of T, we multiply the base learning rate by exp(3.73 3373+ (1.0 1/(T- 1))). This means we multiply the learning rate by exp(3.73) 41.7 in epoch 0, and by 1 in the last epoch. The intuition is that when we are tuning smaller number of layers we need higher learning rate than for full fine-tuning- example, the optimal learning rates for head tuning is higher than the optimal learning rate for full fine-tuning. The exact constant (3.73) was selected by comparing the optimal learning rate for head tuning with full finettuning on Waterbirds (the optimal learning rate for head-tuning was approximately exp(3.73) 73) larger for head-tuning). Without this decay schedule, for example using a vanilla cosine learning rate scheduler, gradual unfreezing worked similarly worse than vanilla full fine-tuning.
Stopping criteria. The results presented in the paper are from models early stopped based on ID validation accuracy. We sanity checked that the conclusions are similar if use the last checkpoint.
Data augmentation and input preprocessing. Additionally, we the following preprocessing and augmentations on our input images. We use very basic augmentations (e.g.. only horizontal flips for WILDS, and standard augmentations from past work), including for state-of-the-ar results:
1. For WILDS-FMoW the images are 224 x 224, random horizontal flip augmentation. We do not perform any augmentations at test time.
do not resize, and only perform
2. For WILDS -DDD--Cyll we resize the images to 224 x 224 with a bilinear interpolation (standard in pyorch), and only perform a random horizontal flip augmentation. We do not perform any augmentations at test time, just the resize.
3. For Living-17 17, follow (Kumar et al.,
2022) and perform RandomResizedCropp to 224 x 224 sized images (using the default options in pyorch), and then a random horizontal flip augmentation while training. At test-time resize the image 256 256 and then take a centercrop of size 224 x 224.
4. For DomainNet, we follow (Kumar et al. 2022) and first resize the image to 256 x 256 with bicubic interpolation, then take a RandomCrop of size 224 224 (using the default options in pyorch). and then random horizontal flip augmentation while training. At test-time we simply resize the image to 224 224 with bicubic interpolation.
5. For Waterbirds, we resize the image to 224 x 224 with bicubic interpolation and then take a centercrop of size 224 224. We apply the same transformation at time.
Embedding layer. For SGD (freeze-embed), the exact layers we freeze are follows:
1. CLIP ViTs: We freeze the patch-to-token embedding layer and layernorm.
2. Supervised and DINO VITs:: We freeze the patch-to-token embedding layer (there is no layernorm after this)
3. BiT-ResNets: We freeze the 'stem' and the first convolution block of the model. We tried freezing less and more of the model in our initial experiments, but it did not seem to help.
4. ConvNeXt-B: We freeze the 'stem' and the first stage of the model.
B ABLATIONS ON CLIP
In Tables we show additional ablations for the CLIP ViT-B/16 model on all datasets We tried:
Running each ablation for all models all datascts too computationally expensive.
15

---

Published as a conference paper at ICLR 2024
1. SGD (freeze-embed, not layer-norm): For the CLIP model our freeze-embed variation freezes the bottom embedding layer along with the layernorm right after that. We ran an ablation where we only freeze the bottom linear embedding layer, but not the layer norm. This performs comparably with SGID (freeze-embed), which suggests that freezing the input layer is what's important, and the layer norm does not matter much.
2. SGD (5x lower LR on embed layer): Another idea from analysis-where found that the embedding layer seems to be why AdamW does better than SGD-is to use a smaller learning rate on embedding layer. Typically, this would involve additional hyperparameter tuning hence undesirable. However, as a heuristic, we ran SGD with 5x smaller learning rate for the embedding layer (since the gradients in the embedding layer are about 5x larger) compared other layers. As expected, this improves over SGD, but does not do as well as SGD (freeze-embed).
SGD (no momentum): Since SGD (freeze-ember no momentum) performed very well in our experiments, also tried fine-tuning with full SGD (no freezing), but without momentum, We found that SGD (no momentum) and SGD perform comparable.
SGD (weight decay): Vanilla SGD is done without weight decay, but AdamW incorporates weight decay. We ran this ablation to confirm that the gains of AdamW are not because of weight decay. We used the torch SGD optimizer, and set the weight decay argument to 0.01. Indeed, we found that SGD and SGD (weight decay) perform comparably, which suggests that weight decay is not the reason for the improved performance of AdamW.
Linear probing: We freeze the pretrained model, and only train a linear probe on the features of the CLIP ViT-B/16 model. We train a logistic regression classifier using the sklearn library, sweeping over 50 regularization values in np.logspacc(-7,22,
LP-FT: Kumar et al. 2022) show that first linear probing, and then full fine-tuning the entire model often works better, especially nut-of-distribution We run LP-FT well, and for the full fine-tuning step we use SGD, AdamW, or SGD (freeze-embed), to test if our conclusions still hold with better fine-tuning method. Note that the test accuracy Waterbirds was bit unstable early on so for this dataset we use the last epoch instead of early stopping on ID validation accuracy. Indeed, even with LP-FT, find that Adam W slightly outperforms SGD out-of- offdistribution. and SGD treeze-embedl outperforms both methods with an average accuracy of 77.5%. The accuracies with LP-FT higher than regular fine-tuning, in line
Kumar al. (0022).
Algorithms
Living-17
Waterbirds
DomainNet
FMoW
Camelyon
Avg.
Baselines
SGD
80.0
62,5
72.8
37.3
86.8
67.9
AdamW
SGD (freeze-embed)
83,2
73.7
75.9
Our methods
SGD (freeze-embed, no momentum)
83.1
80.4
89.0
38.8
93.3
76.9
Gradual-unfrezzin
81,9
69.1
93.2
40.5
96.5
76.2
Variations of SGD (frecze-emeed)
SGD (5x lower embed layer)
SGD (freeze-ebbed, not layer-norun)
83.6
83.3
74,3
89.1
39.3
95.7
92.9
75.0
75.9
Other freezing: Linear prohing
86,2
60.4
89.1
29.0
92.6
71.5
Other layerwise normal-
LAMB LARS
79.5
64.0
90.4
38.8
93.4
73.2
methods
83,9
48.6
83.8
38.6
93.3
69.6
Variations of SGD withous any freezing
SGD (no momentum)
$1.4
59.2
76.7
37.9
84.3
67.9
SGD (weight decay)
83.9
SGD
86.7
67.3
89.2
37.9
94.1
75.0
Variations of LP-PT
Adama
84.5
68.2
90.1
39.7
95.8
75.7
SGD (frecze-embed)
83.1
75.9
90.8
41.8
96.0
77.5
Table -uitof-distributton (OOD) accuracy of more optimizers on CLIP VIT-B/16. We weight decay, momentum, and unfreezing the layer norm at the bottom of the model do not make much a difference.
16

---

Published as a conference paper at ICLR 2024
Algorithms
Living-17
Waterbinds
DomainNet
Camelyon
Avg.
SGD
97.8
67.0
90.0
AdamW
98.1
97.7
95.0
70.1
99.5
92.1
SGD (fresze-embed)
Our methods
SGD (freeze-embed, momentum)
98.2
97.9
95.2
70.1
99.5
92.2
Gradual Gradual-unfreezing
Voriations of SGD (freaz-embed)
SGD (freeze embed, not layer nome-n
98.0
98.0
95.4
70.2
99.5
92.2
SGD (5x lower LR on embed layer)
98.0
97.5
94,8
68.7
99.5
91.7
Other freezing: Linear prohing
97.8
96.6
94.5
47.2
96.1
86.4
Other fayerwise normal. Ecation methods
LAMB
98.2
97.8
95.1
67.9
99.5
91.7
Voriations of SGD without any freezing
SGD momentum)
98.0
97.1
89.5
66.4
99.3
90.1
SGD (weight decay)
97.6
97.2
87.9
66.4
99.3
89.7
Variations of LP-PT
AdamW (frezz-embeed)
98.2
98.2
95.7
69.2
99.5
92.2
Table 8: In-distribution (ID) accuracy of more optimizers on CLIP ViT-B/16. We find that weight decay, momentum, and unfreezing the layer norm at the bottom of the model do not make much of a difference.
17

---

Published as a conference paper at ICLR 2024
C ABLATIONS ON FREEZING CONVNEXT
The convNeXt- Base model consists of a stem (6.3k parameters) then 4 stages (415 thousand, 1.7 million. 58 million, 27 million parameters respectively), followed by layernorm (2k parameters). In the main paper, we freeze the stem and the first stage which consists of 0.5% of the parameters of the entire model. This is similar the fraction of parameters the patch embedding layer of the vision transformers, which we freeze.
An alternative choice is to freeze only the stem of the ConvNeXt layer, which is 0.007% of the parameters. We run an ablation where we try this choice of freezing, which we call SGD (freeze-stem). For our original approach of freezing the stem and the first block, we call it SGD (freeze-stem- blookk-l)... Results for OOD are in Table 9 and for ID are in Table I0
Living-17 17
Waterbirds
DomainNet
FMoW
Camelyon
Avg.
SGD AdamW
94.0
80.2
89.8
39.7
83.0
77.3
90.3 (-3.7)
89.8 (+9.6)
89.5 (-0.3)
38.4 (-1.3)
89,5 (+6.5)
79.5 (+2.2)
SGD (freere-stem--locke
92.6 (-1.4)
86.9 (+6.7)
91.2 (+1.4)
38.2 (-1.5)
88.1 (45.1)
79.4 (+2.1)
SGD (frecze-stem)
91.4 (-2.6)
84.7 (+4.5)
91.0 (+1.2)
40.3 (+0.6)
86.2 (43.2)
78.7 (+1.4)
Table 9: Out-of-distribution (OOD) accuracies for ConvNeXt where either freezing just the stem layer (0.5% of model parameters), the stem and the first block (0.007% of model parameters).
Living-17
Waterbirds
DomainNet
FMoW
Camelyon
Avg.
SGD
98.7
99.0
94.8
66.3
99.4
91.6
AdamW
98.6 (+0.1)
99.5 (+0.5)
94.5 (-0.3)
68.8 (+2.5)
99.7 (+0.3)
92.2 (+0.6)
SGD firezze-stem-blockk]
98.6
(-0.1)
99.4 (40.4)
95.1 (40.3)
67.4 (+1.1)
99.5 (+0.1)
92.0 (+0.4)
SGD (frecze-stem)
98.8 (+0.1)
99.3 (40.3)
95.1 (+0.3)
67.2 (+0.9)
99.5 (+0.1)
92.0 (40.4)
Table 10: In-distribution (ID) accuracies for ConvNeXt where we try either freezing just stem layer (0.5%
of model parameters), or the stem and the first block (0.007% of model parameters).
D RESULTS FOR ADAMW (FREEZE-EMBED)
In our paper, we hypothesized that AdamW and SGD (freeze-embed) improve on SGD for the same reason-the they change the embedding layer less. Based on this hypothesis, we would expect that the gains of AdamW and freeze-embed are not complementary. Indeed, we find that the AdamW (freeze-embed) variation performs similarly to Adam W and SGD (freeze-embed).
Tables 11 & 12 expand on the OOD and ID accuracy results in Tables IT & 12 respectively, to include Adam W (freeze-embed), Overall, Adam (freeze-embed) and SGD (freeze-embed) per- form comparably for all the recent vision models, both fairly close to Adam AdamW-although there differences in some models and datasets. Averaged across all the datasets and models, AdamW, AdamW (freeze-embed), and SGD (freeze-embed) get 76%, 76.5%. and 76.7% accuracy, respec- tively, compared to 72.0% for SGD. Averaged across all the in-distribution datasets and models, A, AdamW (freeze-embed), and SGD (frecze-embed) get 91.5%. 91.5%, and 91.3% accuracy, respectively, compared to 90.3% for SGD. Overall the performances of the three models are similar enough to suggest that these methods work well for the same reason that they tune the embedding layers less, and that freeze-embed is not an independent axis of improvement.
18

---

Published as a conference paper at ICLR 2024
Living-17
Waterbirds
DomainNet
FMoW
Camelyon
Avg:
CLIP VIT-8/16
SGD
80.0
62.5
72.8
37.3
86.8
67.9
CLP VIT-R/16
AdamW
8288 (42.8)
71.9 (49.4)
89.2 (416.4)
40.7 (43.4)
95.7 (48.9)
76.0 (48.1)
CLIP VIT-8/16
SGD (frezz-eemedd)
83.2 (13.2)
73.7 (:11.2)
88.2 (:154)
40.2
94.3 (47.5)
75,9 (+8.0)
CLIP V/T-8/16
A (feeze-emmee)
82.4 (+2.4)
69.5 (+7.0)
88.2 (+15.4)
40.7 (+3.4)
96.7 (+9.9)
75.5 (+7.6)
CLIP VIT-L/14
SGD
84.2
65.0
60.8
41.0
83.2
66.8
CLIP VIT-L/14
AdamW
68.0 (+3.8)
85.2 (-20.2)
93.8 (+33.0)
48.3 (+7.3)
95.9 (+12.7)
82.2 (+15,4)
VIT-L/14
SGD (mreze-embee))
90.5 (+6.3)
81.7 (-19.7)
(+32.3)
49.9 (+8.9)
(413,3)
83.0 (+16.2)
CLIP VIT-/11
Adam fieeze-eebook
88.0 (+3.8)
84.5 (+ (+19.6)
93.8 (+33.0)
44.6 (+3.6)
96.4 (+13.2)
81.5 (+14.7)
Sup VIT-B/16
SGD
89.5
77.4
86.3
33.5
92.6
75.8
Sup VIT-R/16
Adam
88.3 (-1.2)
81.6 (44.2)
84.4 (-1.9)
38.9 (42.4)
87.9 (-4.7)
75.6 (-0.2)
Sup VIT-B/16
SGD (rreze-embel))
88.0 (-1.5)
82,4 (+5.0)
86.3 (40.0)
34.4 (+0.9)
93.7 (+1.1)
77.0 (+1.2)
Sup VIT-R/16
AdamW (freee-cmhol)
88.1 (-1.4)
82,4 (+5.0)
82.3 (-400
35.7 (42.2)
93.0 (40.44
76.3
(40.5)
DINO ViT B/16
SGD
88.2
56.1
76.0
33.6
86.9
68.2
DINO VIT-B/16
AdumW
87.4 (-8.8)
61.2 (65.1)
77.4 (+114)
35.8 (12.2)
91.9 (45.0)
70:7 (42.5)
DINO V/T-
SGD (freeze- embed))
86.7 (-1.5)
67.9 (=11.8)
78.4 (+2.4)
35.9 (x2.3)
90.6 (+3.7)
71.9 (+3.7)
DINO B/16
AdamW (freese embed)
86.8 (-1.4)
64.5 (+8.4)
76.6 (+0.6)
35.4 (+1.8)
93.5 (+6.6)
71.4
(+3.2)
Com-Next-BBB Base
SGD
94,0
80.2
89.8
39.7
83.0
17.3
Next-Baae
Adam
90.3 (-3.7)
89.8 (+9.6)
89.5 (-0.3)
38.1 (-1.3)
89.5 (+6.5)
79.5 (+2.2)
ConYNext Base
SGD (freeze embed)
92.6 (-1.4)
86.9 (+6.7)
91.2 (+1.4)
38.2 (-1.5)
88.1 (+5.1)
79.4 (+2.1)
ConvNext-Base
Adam (frecze-embed
88.1 (-5.9)
91.6 (+11,4)
89.8 (+0.0)
41.6
87.7 (+4.7)
79.8
(+2.5)
ResNes-50
SGD
84,3
76.5
34.1
90.4
73.1
BiT ResNet-50
Adami
83.1 (-1.2)
74.8 (-1.7)
84.0 (+1.0)
33.8 (-0.3)
92.4 (+2.0)
73.6 (+0.5)
BiT ResNet-50
SGD (freaze-embed)
84,1 (-0.2)
75,5 (-1,0)
82.3 (42.3)
35.0 (40.9)
95.2 (++.8)
74.4 (+1.3)
ResNes-50
Adam W (freeeeemmed)
82.9 (-1.4)
77.3 (40.8)
83.3 (43.3)
36.3
(43.3)
74.7 (41.6)
BIT 101
SGD
82.8
76.9
86.2
38.0
89.3
74.6
BIT ResNer- 101
AdamWW
82.9 (10.1)
79.4 (+2.5)
83.5 (-2.7)
37.0 (-1.0)
89.7 (10.4)
74.5 (-0.0)
BTT ResNeet-10 101
SGD (freeze embed)
83.1 (+0.3)
77.3 (40.4)
86.0 (-0.2)
36.0 (-2.0)
95.5 (+6.2)
75.6 (+1.))
ResNes-101
Adam'W (freze-emhod)
84.8 (12.0)
78.3 (41.4)
84.3 (-1.9)
38.2 (10.2)
95.0 (45.7)
76.1 (41.5)
Table 11: Out-of-distribuutio (OOD) accuracies with Adam (freeze-emhed). This is an expansion of Table | to include Adam W (freeze-emhed) OOD results for fine-tuning 7 popular models across 5 benchmark datasets. On OOD performance averaged across all models and datasets, AdamW (freeze-embed) gets slightly hetter accuracy than Adam but slightly worse than SGD (frezze-embed).
Living-17
Waterhinds
DomainNet
FMoW
Camelyon
Avg.
VIT./116
SGD
97.8
97.2
67.0
99.4
CLIP WIT-/166
Adamm
98.1 (40.3)
97.7 (40.5)
95.0 [46.2)
70.1 (43.1)
99.5 (40 1)
92.1 (42.1)
CLIP VIT-B/16
SGD (freze-cmbbd)
98.2 (+0.4)
97.8 (+0.6)
91.9
(+6.1)
70,0 (+3.0)
99.5 (40.1)
92.1 (+2.1)
VIT-8/16
Adam (frozo-embxd))
98.3 (+0.5)
(40.6)
(46.5)
70.5 (43.5)
99.6 (40.2)
(42.3)
CLIP VIT-L/14
SGD
98.4
97.3
84.3
69.0
99.4
89.7
CLIP VIT-L/14
AdamW
95.9 (+0.5)
98.8 (+1.5)
96.9 (+12.6)
74,5 (+5.5)
99.6 (+0.2)
93.7 (+1,0)
VIT.L/14
SGD (freeze-embed)
98.7 (+0.3)
(+1.6)
(+12.8)
74.5 (+5.5)
99.6 (10.2)
(+4.0)
CLIP WIT-1/14
AdamW (frezzoembed))
98.7 (40.3)
98.88 (+1.5)
96.7 (412.4)
75.1 (46.1)
99.6 (40.2)
93.8 (44.1)
Sup VIT-R116
SGD
98.6
99.1
64.1
99.4
Sup VIT-B/16
Adam'
98.7 (+0.1)
99.0 (-0.1)
91.7
(-0.0)
66.4 (+2.3)
99.5 (40.1)
91.1 (+0.5)
Sup VIT-B/16
SGD (freeze embed)
98.7 (+0.1)
99.2 (+0.1)
91.5
(.0.2)
65.0 (+0.9)
99.6 (+0.2)
90.8 (+0.2)
VIT-R/16
Adam (frezzeeemed))
98.6 (10.0)
99.0
901
(-0.)
66.88 (v2.7)
99.5 (10.1)
91.0
DINO VIT-8/16
SGD
98.4
97.0
88.2
62.4
99.4
89.1
DINO VIT-B/16
Adami
98,5 (+0.1)
97.9 (40.9)
89.4 (+1.2)
66.00 (43.6)
99.6 (+0.2)
90.3 (+1
VIT-B/16
SGD (frezze-cmbed)
98.4 (+0.0)
(+0.5)
(+0.8)
63.5 (+1.1)
99.5 (+0.1)
(+0.5)
DINO VTT-B/16
AdamW embed)
98.4 (10.0)
97.8 (10.8)
89.7 (+1.5)
65.7 (13.3)
99.6 (10.2)
90.3 (+1.2)
ConvNext-Be Base
SGD
98.7
99.0
94.8
66.3
99.4
91.6
ConsNext-Base
Adam'
98.6 (:0:1)
99.5 (40.5)
94.5 (403)
688 (42.5)
99.7 (40.3)
92.2 (40.6)
ConyNext-Bae Basc
SGD (frecze-cmibd)
98.6 (-0.1)
99.4 (+0.4)
95.1 (+0.3)
67.1 (+1.1)
99.5 (+0. 1)
92.0 (+0-4)
ConyNext-Base
Adamm (frocce-mibbc))
98.7 (+0.0)
99.5 (40.5)
94.7
(-0.1)
69.2 (+2.9)
99.6 (+0.2)
92.4 (+0.8)
BiT ResNet-50
SGD
97.4
98.4
89.3
64.6
99.5
89.8
B/T ResNet-50
Adamm
97.2 (.0.2)
98.5 (10.1)
89.2 (.0.1)
65.1 (10.5)
99.5 (10.0)
89.9 (10.1)
BIT ResNet.50
SGD (freeze- emhed)
97.6 (+0.2)
98.5 (10.1)
89.2
(.0.1)
64.8 (10.2)
99,5 (10.0)
89.9 (+0.1)
BIT ResNet-50
Adamm (frezzecmbbed)
97.4 (40.0)
98,4 (-0.0)
(-0.2)
64.9 (40.3)
99.6 (40.1)
89.9 (40.1)
BIT ResNet-101
SGD
98_3
98.9
92.01
66.0
99.4
90.9
ResNet-101 101
Adam'
98.4 (+0.1)
(-0.3)
(-0.9)
67.0 (+1.0)
99.6 (+0.2)
BIT ResNet-101
SGD (frezze-cmbed))
98.4 (+0.1)
98.8
(.0.1)
91.5
(.0.5)
65.9 (:0.1)
99.5 (+0.1)
90.8
(40.1)
BIT ResNet-101
AdamW (freeze-emeed)
98.2
(-0.1)
98.7
(-0:2)
91.1
(-0.9)
66.6 (10.6)
99.6 (10.2)
90.9 (40.0)
Table 12: In-distribution (ID) accuracies with AdamW (freeze-embed). This is expansion of Table 2 to include AdamW (frezz--emeed) ID results for our 7 models and datasets. AdamW, AdamW (frezze-embed), and SGD (freeze-embed) all perform comparably on ID accuracies.
19

---

Published as a conference paper at ICLR 2024
E
WHAT COULD CAUSE LARGE "EMBEDDING" LAYER GRADIENTS?
Generally speaking, we find that Adam W fine-tuning leads to better performing and more robust models than SGD fine-tuning, specially in modemn pretrained models. Our algorithms to close this gap were inspired by the observation that on modern vision pretrained models, the embedding layers have substantially larger gradients at pretrained initialization compared other layers. This could lead to over-training of the embedding layer when using SGD, which we hypothesized would be bad for robust fine-tumie tunettnnig..-- The success of SGD ffreeze-embede in our experiments adds further evidence to our hypothesis. But, why are the gradients of embedding layers at pretrained initialization high in the first place? More generally, why does Adam W do better than SGD during fine-tuning? We discuss two plausible hypotheses and then test these out in a controlled experiment.
Algorithmic aspects: pretraining algorithm? Among our 7 models, the more recent models like vision transformers and ConyNeXt, which were pretrained with AdamW are also the ones with largest gaps in performance between Adam and SGD fine-tuning. BiT ResNets that were pretrained with SGD had much smaller differences between AdamW and SGD. This strongly suggests that the discrepancy between the algorithms in pretraining and fine-tuning might be cause for the performance gaps. For example, it is possible that pretraining with AdamW implicitly biases towards contigurations that most benefit from Adam W updates. This would also explain why other adaptive algorithms like LARS and LAMB are not competitive even though they perform some form of layer-we layer-wise normalization. On the other hand it does not explain why such effects would distinctively impact the "embedding" layer and not the other layers. In a related manner, is also unclear why SGD (freeze-embed) would be able to overcome such implicit biases from Adam W pretraining.
Architectural aspects: "patchifying" first layer? There are substantial architectural differences between the newer transformer and ConvNeXt models and the older ResNets which could also contribute to the newer models working better with AdamW. Most notably, vision transformer is fundamentally different architecture from convolutional networks. At the same time, the biggest differences between the architectures- like self-attention, softmax lomeliieartt, or fully connected layers-ane are not the likely contributors to the differences between Adam W and SGD in fine-tuning. This is because, we also see gaps between these methods with the ConvNeXt models which lack the major transformer components. Rather, is the possible that some of the designs that were adapted from transformers to ConvNeXt contributes to the differences between Adam and SGD fine-tuning. Among these, many primarily affect the higher layers such as heavy third stage, depthwise convolution, inverted bottleneck, and larger kernels, and are unlikely to cause the lower layer gradients to be high. The key design change we believe is important here the use of "patchify stem" that could cause distinctive changes in the gradients for lower blocks.
The "stem" layer primarily controls how the input is processed for the rest of the network. VITs and ConvNext down-samples the input images from non-overlapping patch tokens, compared to ResNets that use denser-overlapped convolution followed by max-pool, The coarser non-overlap might lead the embedding layer to attune more closely to the patch distribution in pretraining. This might not be an issue by itself as pixel/patch level information is noisy anyways and the higher layers can extract more robust features from across the patches. A possible dynamics in fine-tuning is as follows: On new datasets where the patch distributions are very different from pretraining, the embedding layers might have large gradients. In standard SGD, this might cause the embedding layer moving too quickly to fit the new patches before the higher layers can adapt to the changes. Instead, freezing the embedding layer would pass along the raw patches with minimal processing and lets the model adapt to the new distribution based on higher level features, which are likely more robust.
Summary of hypotheses. In conclusion, we hypothesize that the differences between AdamW and SGD fine-tuning arise from a combination of the above ceasons-pretraining with AdamW and large "patchify" embedding layer. Additionally, the use of GeLU activation and layer normalization might also change the optimizer dynamics although these minor changes from ReLU activation and group normalization used in BiT ResNets. It is of interest to systematically explore these reasons further in future work.
Controlled experiments to test hypotheses. We test these hypotheses out in controlled experi- ments and find that the mismatch between the pretraining and fine-tuning optimization method seems
20

---

Published as a conference paper at ICLR 2024
to be the key factor (for both why SGD underperform during fine-tuning, and why the first-layer gradients of modern pretrained models is higher).
To test this out, we pretrained a BiT-ResNet-50 model using either (a) SGD or (b) Adam For
each of these pretraining optimizer choices, we also tried using either (i) the original BiT-ResNet-50 architecture or (ii) modifying it by "patchifying" the first layer. This gives us 4 pretrained models. For each of these pretrained models, we tried fine-tuning it with (1) SGD, (2) AdamW, and (3) OiD-Freezz--Emee.NNote that for ResNet models we consider the stem and the first stage as part of the 'embedding' layer.
For the AdamW pretrained model, fine-tuning with SGD does worse than AdamW on the OOD test sets (with or without the patchify architecture change), but for the SGD pretrained models fine-tuning with SGD does slightly better than AdamW. For example, for the standard BiT ResNet-50 architecture pretrained with A. fine-tuning with Adam W gets an average 1.3% higher accuracy OOD. For the BiT ResNet-50 (patchify) architecture pretrained with AdamW, fine-tuning with AdamW gets an average 0.9% higher accuracy OOD, On the other hand, for both ResNet models pretrained with SGD, fine-tuning with AdamW does slightly worse than SGD. This suggests that the optimizer mismatch (between pretraining and fine-tuning) is a key reason for why fine-tuning with SGD can do worse than fine-tuning with Adam W. Nonetheless, fine-tuning with SGD ffreeze-embed consistently performs well, getting the highest average OOD accuracy in 3/4 cases (its OOD accuracy is in between SGD and Adam W on a ResNet patchify model pretrained with SGD).
The full results for OOD are in Table 13 and for ID are in Table 14 As in the main paper, the ID) accuracies of all methods are fairly close.
Pretraining
Fine-tuning
Living-17
Waterbirds
DomainNet
FMoW
Camelyon
Avg.
SGD
SGD
83,9
63.4
77.9
34.2
88.3
69.6
SGD
Adam
77.5 (-6.4)
66.4 (43.0)
75.5 (-2.4)
33.4 (-0.8)
94.6 (+6.3)
69.5 (-0.1)
SGD
SGD (freeze-embed)
84,2 (+0,3)
65.3 (+1.9)
79.3 (+1.4)
32.2 (-2.0)
95.6 (47.3)
71.3 (+1.7)
AdamW
SGD
81.4
56,9
74.4
30.7
86.6
66.0
AdamW
AdamW
80.1 (-1.3)
65.6 (48.7)
73.6 (-0.8)
30,6 (-0.1)
86.6 (+0.0)
67.3 (+1.3)
AdamW
SGD (freeze-emhed)
81.1 (-0.3)
60,9 (+4.0)
77.5 (+3.1)
30,3 (40.4)
89.6 (+3.0)
67.9 (+1.9)
SGD (Patchify)
SGD
82.5
67.0
78.6
35.2
91.0
70.8
SGD (Patchiffy)
Adam
79.8 (-2.7)
64.3 (-2.7)
75.2 (-3.4)
34.6 (40.6)
90.9 (:0.1)
69.0 (-1.8)
SGD (Patchifyy)
SGD (freeze-embed)
82.5 (40.0)
62,9 (-4.1)
78.5 (-0.1)
34.9 (-0.3)
91.2 (+0.2)
70.0 (-0.8)
AdamW (Patchify)
SGD
81.3
61.7
74.5
31.7
93.3
68.5
AdanW (Patchify)
Adam
79.3 (-2.0)
69.5 (+7.8)
73.1 (-1.4)
33.2 (+1.5)
91.9 (-1.)
69.4 (+0.9)
AdamW (Patchify)
SGD (freeze-emhed)
82.0 (+0.7)
64.0 (42.3)
77.0 (+2.5)
31.1 (40.6
95.0 (+1.7)
69.8 (+1.3)
Table 13: Out-of--istribution (OOD) accuracies for the different BiT-ResNet pretrained models.
Pretruining
Fine-tuning
Living-17
Waterbirds
DomainNet
FMoW
Camelyon
Avg.
SGD
SGD
97.6
97.6
88.3
63.1
99.5
89,2
SGD
AdamW
97.8 (+0.2)
97.6 (-0.0)
87.2 (-1.1)
62.4 (-0.7)
99.5 (+0.0)
88.9 (-0.3)
SGD
SGD (freeze-embed)
97.6 (-0.0)
97.6 (-0.0)
87.3 (-1.0)
62.7 (-0.4)
99,5 (40.0)
89.0 (-0.2)
AdamW
SGD
97.2
97.2
87.5
58,9
99.3
88.0
AdamW
AdamW
97.4 (-0.2)
97.3 (40.1)
86.5 (-1.0)
58.1 (-0.8)
99.5 (+0.2)
87.8 (-0.2)
AdamW
SGD (frezze-embed)
97.4 (+0.2)
97.1 (40.1)
86.9 (-0.6)
58.2 (-0.7)
99.3 (+0.0)
87.8 (-0.2)
SGD (Patchify)
SGD
97.9
97.5
88.3
62.8
99.4
89.2
SGD (Patchify)
AdumW
97.5 (-0.4)
97.3 (40.2)
87.0 (-1.3)
62.2 (-0.6)
99.5 (+0.1)
88.7 (-0.5)
SGD (Patchify)
SGD (freeze embed)
97.6 (:0.3)
97.5 (.0.0)
87.6 (:0.7)
63.1 (+0.3)
99.4 (+0.0)
89.0 (:0.2)
AdanW (Patchify)
SGD
97.5
97,3
87.7
59.6
99.3
88.3
AdamW (Patchify)
AdamW
(+0.1)
97.2 (00.)
86.2 (-1.5)
58.4 (-1.2)
99.5 (+0.2)
87.8 (.0.5)
AdamW (Putchify)
SGD (frecze-embed)
97.6 (+0.1)
97.2 (40.1)
87.0 (-0.7)
59.4 (-0.2)
99.2 (40.1)
88.1 (-0.2)
Table 14: In-distribution (ID) accuracies for the different BiT-ResNet pretrained models.
As in Section 3.2 we also plot the average gradient norm at each layer, across minibatches of the dataset, as described in Equation 3.I] We show the plot for DomainNet, Waterbirds, and Living- 17 in Figure 3 We also show plots for alternative ways of normalizing the gradients in Figure 4 and Figure 5] For SGD pretrained ResNet, the embedding layer has lower gradient than the other layers. For AdamW W pretrained ResNet, the embedding layer has higher gradient than the other layers. In
21

---

Published as a conference paper at ICLR 2024
other words, pretraining with SGD versus with AdamW leads to substantial difference in the embedding layer. Using a 'patchify' stem does not substantially change these results.
These results provide further evidence that (1) SGD does worse than AdamW when fine-tuning Vision Transformer and ConvNeXt models because these models are pretrained using AdamW, (2) the embedding layer plays key role here, since pretraining with different optimizers leads to very different behavior at the embedding layer, (3) SGD (frezze-embed) can potentially improve fine- tuning, by freezing the embedding layer. Without this, SGD either over-optimizers the embedding layer (if the learning rate is large), or under-optimizers the other layers (if the learning rate is small).
F ADDITIONAL PLOTS ON GRADIENT NORMS AT THE PRETRAINED INITIALIZATION
In Figure 6 we show plots for the gradients norms different layers for Living-17 and Waterbirds. These are analogous to Figure 2 for DomainNet in the main paper. We again see that, among the pretrained layers. the embedding layer has much higher gradients the modern architectures.
We also consider two ways of normalizing the gradient magnitude. In the first method, we divide the norm of the gradient by the norm of the parameters in order to capture the relative "movement" in the first SGD step as opposed to the absolute "movement" which is captured by gradient norm itself. In the second method, we divide the norm of the gradient by the square root of the number of parameters. This is to check that layer does not simply have a larger gradient because it has more parameters. The reason we use square root as follows: suppose each parameter has gradient c, then the layerwise gradient norm scales with the square root of the number of parameters. Also, the first step of Adam update essentially a signed gradient descent step. wherein we ignore weight decay, the per-layer "movement" is the square root of the number of parameters. So this normalization can be thought of as relative size SGD update compared to Adam in each layer at initialization. For visualization purposes, we exclude the 'head' layer gradient in these plots they often much larger than the others so the plots become hard to see we include the 'head' layer. Note that we expect the head layer to have higher gradients because is randomly initialized (Kumar et 2022). For ViT models, we omit gradients of the cls token, position embedding, and layer norm after the embedding layer.
22

---

Published as a conference paper at ICLR 2024
middle layers
embedding layer
3-
2
1
SGD
AdamW
SGD (Pstchify)
AdamiW (Patchify)
(a) Layer-wis wise gradient norms on Living- 17 at pretrained initialization
middle layers
4-
embedding layer
3-
SGD
AdamW
SGD (Patchify)
Adamiw (Patchify)
(b) Layer-wise gradient norms on Living-17 pretrained initialization
middle layers
embedding layer
3.0
2.5
15
10
0.5
SGD
AdamW
SGD (Patchify)
AdamW (Patchify)
(c) Layer-wise wise gradient norms on Waterbirds at pretrained initialization
Figure 3: We visualize the layer-wise gradient norms of the four Bit-ResNet models (a) DomainNet, (b) Living-17 17 and (e) Waterhirds, the pretrained initialization. For better visualization, we omit the head from the plot, which predictably has much larger gradients than others (since it randomly initialized). the same as Figure 2 gradient norms of "embedding" and "middle" layers are shown as red-squares and black-circles. respectively. We see that the "embedding" layer has higher gradient (than the other layers) for models pretrained with AdamW, but lower gradient (than the other layers) for models pretrained with SGD, which supports the hypotheses that the "embedding' layer plays key role, and that pretraining with AdamW vs. SGD leads to very different models and is responsible for this bchavior.
format
23

---

middle for Published as a conference paper at ICLR 2024
layers embedding layer
0.04
0.03
0.02
0.01
0.00
AdamW
SGD (Patchity)
AdamW (Patchify)
(a) Layer-wise gradient norms divided by parameter norm, on DomainNet at pretrained initialization
middle layers
embedding layer
0.030
0.025
0.020
0.015
0,010
0.000
SGD
AdamW
SGD (Patchify)
Adamw (Patchify)
(b) Layer-wise gradient norms divided by parameter norm, on Living-17 17 at pretrained initialization
middle layers
embedding layer
0.025
0.020
0.015
0.010
0.005
0.000
5GD
AdamW
SGD (Patchify)
AdamiW (Patchify)
(c) Layer-wise gradient norms divided by parameter norm, on Waterbirds at pretrained initialization
Figure 4: We visualize the layer-wise gradient norm, divided by the
of the parameters on DomainNet. pretrained initialization. For better visualization, we omit the head
(b) Living-17, 17, and (c) Waterbirds, from the plot, which predictably has much larger gradients than the others (since it is randomly initialized). The format is the same as Figure gradient norms of "embedding" and "middle" layers are shown as red-squares and black-circles, respectively. Under this normalization scheme, the "embedding" layer has much higher gradients when pretrained with AdamW, but "embedding" layer gradient somewhere models pretrained with SGD.
24

---

Published as a conference paper at ICLR 2024
0.0030
middle layers
embeddiing layer
0.0025
0.0015
0.0005
0.0000
AdomW
SGD (Patchify)
Adami (Patchify)
(a) Layer-wise gradient norms divided by #paaameterss DomainNet at pretrained initialization
middle layers
embediiing layer
0,0015
0,0010
0.0005
0.0000
SGD
AdamW
SGD (Patchify)
AdamW (Potchify)
(b) Layer-wise gradient norms divided by v Appramettrrs, on Living- 17 at pretrained initialization
0.00175
middle layers
embedding layer
0.00150
0,00125
0,00100
0.00075
0.00050
0.0C025
SGD
Adomw
SGD |Patchify)
AdamW (Patchify)
(c) I Layer-wise gradient norms divided by /#parameters, on Waterbirds at pretrained initialization
Figure 5: We visualize the layer-wise gradient norm, divided by the square root of the number of parameters (a) DomainNet, (b) Living- 17, and (c) Waterbirds. at the pretrained initialization. For better visualization, omit the head from the plot which The format the same as Figure 2 gradient norms of "embedding" and "middle" layers are shown as red- squares and black-circles, respectively. Under this normalization scheme, the "embedding" layer has higher gradients in all cases. However, the embedding layer gradient about 2-3 times larger (than other layers) for models pretrained with SGD, but over 10 times larger (than other layers) for models pretrained with AdamW.
predictably much larger than the others (since is randomly initialized).
25

---

Published as a conference paper at ICLR 2024
middle layers
head layer
25
embedding layer
15
clip_b
clip_1
dino_b
convnext_p
resnet50
resnet101
(a) Layer-wise gradient norms Living-17 17 pretrained initialization
60
middle layers
head layer
embedding layer
40
10
clip_b
clip_
vit_b
dino_b
convnext_
resne:50
resnet101
(b) Layer-wise gradient norms on Waterbirds at pretrained initialization
Figure 6: We visualize the layer-wise gradient norms our models on (a) Living-17 17 and (b) Waterbirds, at the pretrained initialization. The format is the same Figure 2 gradient norms of "embedding", "head", and "middle" layers are shown red-squares, green-triangles and black-circles, respectively.
26

---

Published as a conference paper at ICLR 2024
middle layers
2.0
embedding layer
1.5
1.0
0.0
clip
clip_l
vit
dino_h
convnext
resiee50p
resnet101
(a) Layer-wise gradient norms divided by parameter norm, on DomainNet at pretrained initialization
middle layers
embedding layer
2.0
1.0
0.5
0.0
clip
clip
vit.
dino__
convnext
resine500
resnet101
(b) Layer-wise gradient norms divided by parameter norm, on Living-17 17 at pretrained initialization
middle layers embedding layer
2.0
1.5
1.0
0.5
0.0
clip_b
clip_]
vit
dino_b
convnext_b
resnet50
resnet101
(c) Layer-wise gradient norms divided by parameter norm, on Waterbirds at pretrained initialization
Figure 7: We visualize the layer-wise gradient norm, divided by the norm of the parameters on (a) DomainNet. (b) Living-17, and (c) Waterbirds, the pretrained initialization. For better visualization, omit the head from plot, which predictably has much larger gradients than others (since it randomly initialized). The format is the same as Figure gradient norms of "embedding" and "middle" layers are shown as red-squares and black-circles. respectively. Under this normalization scheme, the embedding layer has higher gradients than the other layers in all models. However, the gradient is only slightly larger for ResNet models, and substantially also provides support for why freezing the embedding layer in
larger for the Vision Transformer Vision Transformers might make a larger difference.
27

---

Published as a conference paper at ICLR 2024
middle layers
0.030
embedding layer
0.025
0.020
0 0.015
0.000
clip_b
clip_l
vit_b
dinoo_t
convnext_
resnet50
resnet101
(a) Layer-wise gradient norms divided by Aparameters, DomainNet at pretrained initialization
0.035
middle layers
embedding layer
0.030
0.025
0.020
0.015
0.005
0.000
clip_1 b
clip_l
vit
no
convnext
resnet50
resnet101
(b) Layer-wise gradient norms divided by ##parameters, on Living-17 17 at pretrained initialization
middle layers
embedding layer
0.04
0.03
0.02
0.01
dip_b
vit
dino_b
convnext_e
b
resnet50
(c) Layer-wise gradient norms divided by ##parameters, on Waterbirds pretrained initialization
Figure 8: We visualize the layer-wise gradient norm, divided by the square root of the number of parameters on (a) DomainNet, (b) Living-17, 17, and (c) Waterbirds, at the pretrained initialization. For better visualization, omit the head from the plot which has predictably much larger than the others (since it is randomly initialized). format the same as Figure 2 gradient norms of "embedding" and "middle" layers are shown red- squares and black-cireles, respectively. Under this normalization, we see that the gradients of the embedding layer are much larger than the other layers in all models, including ResNets.
28


Reviews:
[{"review_id": "VB1TPgOXfh", "reviewer": "Reviewer_S1Jb", "summary": "This paper analyzes why AdamW often outperforms SGD when finetuning large pretrained vision models. It turns out that the embedding layer gradients are quite large and AdamW is able to suppress the huge swings in embedding layer values caused by these gradients. A simple method is proposed whereby the embedding layer is frozen during finetuning. This method is shown to perform on par or better than AdamW in both in-distribution and out of distribution finetuning. Moreover, SGD with frozen embedding layer uses less memory than AdamW (roughly 1.3x less).", "questions": "1) Why does the memory overhead of adamw relative to sgd increase with model size? It seems like the relative difference in memory overhead should be constant\n2) For Fig. 2, are the gradient norms normalized by number of elements?\n3) Why do you think SGD outperforms the other alternatives across so many experiments?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "I found the paper to be extremely well written. All claims were backed up by experimental results. The method is simple, yet effective. Very good paper.", "weaknesses": "1) It seems that the entire paper hinges on models which have an embedding layer. While these models are popular now, they may not be popular forever, which limits the long term impact of this method.\n2) In models beyond vision, such as recommendation models (i.e. DLRM https://arxiv.org/abs/1906.00091), the embedding layer contains most of the model parameters. In such cases, it is also unclear whether this kind of method could work. Of course, the authors are explicit that this paper is about vision models in particular, so I am not penalizing them for this (possible) shortcoming\n3) In some cases, SGD performed as well or better than the other variants. From the practitioners point of view, how should one choose which method to use?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "74D5Wz1kZ1", "reviewer": "Reviewer_RNj8", "summary": "This paper found that Adam achieves better OOD and ID accuracy than SGD, mainly because the embedding layer of some currently popular models yields larger gradients. Such larger gradients can be well treated by Adam, but may cause over-training of the embedding layer by SGD. Therefore, this work proposes two variants of SGD (freezing the embedding layer; freezing the embedding layer and no momentum), which are found to boost the accuracy significantly.", "questions": "1. Will such modified SGD still achieve performance gain on NLP model such as Llama2?\n2. As you mentioned, the reason why the early embedding layer has larger gradients is that such models are typically pre-trained by Adam. What will the result be if they are models are pre-trained by SGD? Will it be worse than with Adam?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Very extensive experiments are conducted and thorough analysis is provided. \n2. The presentation is good and motivation is clear and strong.\n3. The minor yet effective modification on SGD shows good performance improvement, with less memory requirement than Adam.", "weaknesses": "Since I am not researching around the optimization domain, I can not clearly point out what is the weakness of this paper. See my questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g2nC21J2kM", "reviewer": "Reviewer_6FtM", "summary": "The paper proposes a method to fine-tune pre-trained models using SGD instead of AdamW to reduce the amount of memory the optimizer states take. The main observation of the paper is that when fine-tuning on out of distribution data the gradient magnitudes on the first embedding layer for architectures pre-trained with AdamW and fine-tuned with SGD are large. The authors hypothesize that the large gradients result in overfitting on OOD data. To address the paper freezes the embedding layer when fine-tuning with SGD and shows that the models can be fine-tuned as well as or better than using AdamW with lower memory footprint. The experimental evaluation and the results validate the hypothesis and the effectiveness of the proposed freeze SGD fine-tuning method.", "questions": "* Given an optimizer mismatch can the overfitting of early layers be avoided by low learning rates in the early part of fine-tuning? Overall the learning rate schedule seems to have an important role to play when there is an optimizer mismatch. Have the authors tried using different learning rate schedules and are the large gradient problems persistent throughout the optimization or confined to initial steps when fine-tuning.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* As the models are become larger reducing the memory footprint to train/fine-tune models becomes increasingly important. One way to reduce the memory footprint is to use optimizers which do not use additional much state. The paper shows that by freezing a the initial embedding layer SGD with momentum or just plain SGD can match performance of fine-tuning the full model with AdamW. \n* The ablation studies show interesting observations on the role of the optimizer difference when pre-training and fine-tuning. Some of these observations can help practitioners in choosing the optimizer and hyper parameters when fine-tuning.", "weaknesses": "* All the down stream fine-tuning experiments are focused on classification. It would be good to see if the observations hold for other tasks like detection, segmentation etc. \n* Gradual unfreezing is mentioned in Table 5 but the method is not clearly described in the paper. There is some mention of the learning rate schedule with gradual unfreezing in the appendix but does not fully explain what was done.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mK5knklJcB", "reviewer": "Reviewer_zcmx", "summary": "This paper unveils a captivating strategy for enhancing the fine-tuning of large vision models (e.g., Vision Transformer, ConvNeXt models, etc.). While SGD and AdamW are both widely used, SGD is favored for its memory efficiency when they perform equally well. However, the authors reveal that especially in scenarios with *distribution shifts*, AdamW excels in fine-tuning modern Vision Transformer and ConvNeXt models.\n\nThe key observation is that this performance gap between SGD and AdamW is most pronounced when the gradients in the first \"embedding\" layer are significantly larger than those in the rest of the model. To address this issue, the authors propose a simple solution: freezing the embedding layer, which accounts for less than 1% of the parameters. This approach not only matches AdamW's performance but also conserves memory, with up to a 33% reduction in GPU usage on ViT-L.\n\nThe proposed approach leads to SOTA results on five popular distribution shift benchmarks, such as WILDS-FMoW, WILDS-Camelyon, BREEDS-Living-17, Waterbirds, and DomainNet. In essence, this paper offers an enchanting recipe to narrow the performance gap between SGD with momentum and AdamW. Their findings shed light on optimizing large vision models.", "questions": "Please refer to the weakness. In addition, please address the following questions:\n\n1. It is clear that the proposed approach is not always the best choice, how does the developer or engineer decide which approach to use when it comes to different datasets?\n\n2.  Please report the mean and standard deviation of the results across multiple rounds (different seeds).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is well-written, the motivation is sound, and the problem has been addressed well.\n2. The paper addresses the crucial issue of improving out-of-distribution (OOD) accuracy when fine-tuning vision models.\n3. The paper achieves state-of-the-art results on multiple benchmarks, demonstrating its efficacy in narrowing the performance gap between SGD and AdamW.", "weaknesses": "1. Although the experimental results achieve SOTA performance and practical memory savings, the improvement of the proposed approach, compared to the vanilla approach is marginal.\n\n2. The experimental results show that using SGD (freeze-embed) may also result in poor performance in some cases, how does the author explain this phenomenon?\n\n3. Reprogramming [1, 2] and Visual Prompting [3, 4, 5] is also a parameter-efficient finetuning approach, please discuss these methods in the related work.\n\n> [1] Elsayed et al. Adversarial Reprogramming of Neural Networks (ICLR 2019)\n> \n> [2] Chen. Resource-Efficient Cross-Domain Machine Learning \n> \n> [3] Bahng et al. Exploring Visual Prompts for Adapting Large-Scale Models\n> \n> [4] Tsao et al. AutoVP: An Automated Visual Prompting Framework and Benchmark\n> \n> [5] Jia et al. Visual Prompt Tuning (ECCV 2022)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xXksg7J42P", "reviewer": "Reviewer_qVNV", "summary": "- The paper investigates the differences in utilizing AdamW vs. SGD for finetuning (not pretraining) modern ViT and Conv nets\n- While usually AdamW outperforms SGD for finetuning, it finds that freezing the embedding layer and then applying SGD (with or without momentum) results in performance which is at par or better than AdamW on distribution shift benchmarks (WILDS-FMoW, WILDS-Camelyon, BREEDS-Living-17, Waterbirds, and DomainNet)\n- The paper points out that models pretrained with AdamW do not finetune well with SGD, and that in such situations the embedding layer has higher grad norms than the rest of the parameters\n- The paper produces SOTA results for the datasets when finetuning CLIP and ConvNeXt models with either SGD + freeze-embed or AdamW", "questions": "It is not clear (1) when researchers should use SGD (freeze-embed) as compared to SGD and AdamW and (2) in which situations the high grad norms for the embedding layer happen. Is this paper only relevant for tasks with a domain shift from pretraining, or for OOD evaluations? How does it work for more popular and / or in-domain transfer learning tasks, such as ImageNet or iNaturalist?\n\nThe paper does cite other works which can get similar performance with AdamW and SGD on ImageNet or iNaturalist, but it would still be worthwhile to show results with the hyperparameters in the paper and with SGD (freeze-embed). \n\nA better understanding outside of the distribution shift datasets will help me bump up my rating.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper looks into the well known phenomenon of SGD finetuning not matching up to AdamW for certain situations in vision models and finds some interesting observations\n  - SGD generally only outperforms on the considered evaluations when pretraining is performed with AdamW, when SGD pretraining is used then there isn't much of a gap\n  - The embedding layer has a very high grad norm when pretraining with AdamW which potentially results in SDG updating its weights too quickly, which is why the freezing operation improves SGD's results significantly\n- The paper is well written and easy to read, and the experiments are convincing. The additional experiments in the appendix are also very useful to confirm the hypotheses presented in the paper.", "weaknesses": "- While the paper's focus is distribution shift evaluations, it doesn't touch upon conventional downstream evaluation tasks such as ImageNet-1k, iNaturalist, datasets in VTAB (https://github.com/google-research/task_adaptation), etc. Even if some of the conclusions hold only for these distribution shift datasets, it is still important to contrast the observations around grad norms and SGD vs. AdamW vs. SGD (freeze-embed) on other common visual transfer tasks. \n- In Section 4 it is mentioned \"SGD performs well when fine-tuning data is closer to pretraining.\" and the conclusion is drawn from just a single data point, CLIP, which is trained on very different data and very different losses. There isn't enough data to back this claim. \n  - Minor: It also says that all models except CLIP were trained on ImageNet-21K, whereas DINO is trained on ImageNet-1k.\n- The CIFAR-10 results on CLIP are presented in a misleading manner \"SGD (freeze-embed) gets 20% and 30% lower error than SGD on CLIP ViT-B/16 and ViT-L/14, respectively.\" where the difference across SGD, AdamW and SGD (freeze-embed) is only <=0.4%\n- The paper does a grid search on a single hyperparameter, learning rate, and sweeps over 6 values, but holds the weight decay constant at 0 for SGD and 0.01 for AdamW. Models pretrained differently (SSL, CLIP, supervised) usually require very different finetuning settings, so this search is a bit limited, but this isn't a major weakness since the primary point is showing that SGD (freeze-embed) catches up to AdamW without recipe changes to SGD.\n- Minor: The paper claims contradictory numbers for SGD and AdamW memory consumption (12 bytes vs 16 bytes per parameter in the Abstract, and AdamW maintaining 2x to 3x more states as SGD in the Introduction).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)