Paper: How to Fine-Tune Vision Models with SGD (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
> 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
1. We present results for **CIFAR-10** in Table 6. 
2. We do not test on ImageNet because some of the models are pretrained using supervised labels on ImageNet.
3. We picked distribution shift datasets because we wanted to examine the accuracy **both in-distribution and out-of-distribution** (a critical problem in modern deep learning). We note that the datasets we use are **popular evaluation benchmarks**: as one metric, DomainNet has 1300 cites, WILDS has 900 cites, Waterbirds has 1100 cites, which is comparable to iNaturalist and VTAB. For some recent examples of transfer learning papers that focus on these datasets, see Kumar et al. 2022, Goyal et al., 2022, Ghosal et al., 2022, Wortsman et al. 2022. We also examine real world datasets in satellite data (FMoW) and tumor detection (Camelyon). 

That said, we agree that testing on more datasets will solidify our claims and will add this to the future work section.

Reviewer comment (Experiments):
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.

Author response:
> 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

Just to clarify—our claim in Section 4 is not based on CLIP. We examined all the models pretrained on ImageNet (multiple models), and compared their performance across multiple datasets. We noticed that SGD performs well on the dataset that is most similar to ImageNet (Living-17, which is a subset of ImageNet). We agree that more evidence would be good—we have softened the phrasing to "This suggests that SGD **may** work better when fine-tuning and pretraining data are similar—we leave a more thorough analysis to future work". Happy to make more changes—how does that sound?

Reviewer comment (Evaluation):
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%.

Author response:
> 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%

Just to clarify, for a ViT-L, SGD got **99.0%** accuracy, and AdamW and SGD (freeze-embed) got **99.3%**. We calculated **error reduction** as (99.3% - 99.0%) / (100% - 99.0%) = **30%**. The reason we examine the error reduction is that **at such high accuracy levels**, it is much more **difficult to improve the accuracy**. For example, it is often not too hard to improve the accuracy from 80.0% to 82.0% (a 2% improvement), but improving the accuracy from 99.0% to 101% is impossible. In other words the 0.3% accuracy improvement should be contextualized with the maximum possible accuracy improvement of 1.0%. However, we acknowledge your point and have **edited to note the exact accuracy** in the text. Happy to make additional changes as you like.

Reviewer comment (Writing):
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).

Author response:
> "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)."

Thank you for pointing this out. We agree and apologize for the mistake. To clarify:
1. SGD (no momentum) stores 2 floats per parameter
2. SGD (with momentum) stores 3 floats per parameter
3. AdamW stores 4 floats per parameter

We have fixed the text to say that AdamW maintains 1.33$\times$ to 2$\times$ more states as SGD in the introduction.\*

\* Calculation: 4 / 3 = 1.33 and 4 / 2 = 2

Reviewer comment (Evaluation):
Although the experimental results achieve SOTA performance and practical memory savings, the improvement of the proposed approach, compared to the vanilla approach is marginal.

Author response:
We note that the improvements of SGD (freeze embed) are quite significant:
- Comparison to SGD: SGD (freeze-embed) achieves **5% higher accuracy OOD**, averaged across all datasets and models, and also performs better on the vast majority of datasets.
- Comparison to AdamW: SGD (freeze-embed) uses substantially less memory. For example, **AdamW consumes 50% more memory** than SGD (freeze-embed, no momentum) on a ViT-L---on top of the significant memory savings, SGD (freeze-embed, no momentum) still achieves higher OOD accuracy.

Reviewer comment (Novelty):
Reprogramming [1, 2] and Visual Prompting [3, 4, 5] is also a parameter-efficient finetuning approach, please discuss these methods in the related work.

Author response:
Thank you for the useful pointers—we have **added these to the related works**.

Reviewer comment (Evaluation):
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?

Author response:
1. **SGD-freeze embed works consistently well**, so we recommend using it as a general guideline. SGD freeze-embed does better than SGD in 80% of the 70 settings (7 models $\times$ 5 datasets $\times$ {ID, OOD}).  In table 1 (out-of-distribution), we show 35 model-dataset pairs (7 models, and 5 datasets). SGD freeze-embed does better than SGD in 27/35 of the comparisons. In table 2 (in-distribution), SGD freeze-embed does better than SGD in 29/35 comparisons. 
2. In modern deep learning it's rarely the case that a single optimizer dominates in every case—our paper gives a detailed analysis as to when and why each of the optimizers can do better. To get further improvements, practitioners can look at the **embedding layer gradients** at the start of training (which is very cheap---a single backward pass). If these are similar to (or smaller than) the gradients of other layers, then we expect SGD (even without freeze embed) to work fairly well.
3. In addition, our experiments and analysis in Appendix E show that SGD works better when the model was pretrained with SGD, or when the fine-tuning data is similar to pretraining (e.g., Living-17, which is a subset of ImageNet). So researchers and practitioners can also use their understanding of the pretraining optimization process, and the relationship between the pretraining and fine-tuning data, to get further improvements.

Reviewer comment (Evaluation):
Please report the mean and standard deviation of the results across multiple rounds (different seeds).

Author response:
We **reported mean and std-dev for CLIP** (the most advanced of the models we consider for distribution shift) across multiple seeds in Table 5. 
Running multiple seeds for all of table 2 was too expensive—we already did a sweep over 4 optimizers x 5 datasets x 7 models x 6 learning rates (and also other optimizers such as LARS and LAMB, ablations on weight decay, etc)—multiplying all results by multiple seeds was computationally infeasible. However, we hope that the multiple seeds for CLIP, and breadth of results across datasets and models, is helpful. If there are any specific results you are particularly interested in, we can run additional seeds for them!

Reviewer comment (Experiments):
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.

Author response:
This is a good idea! We wanted to be thorough and prioritize covering different types of datasets, distributions shifts, and models for this work, but examining segmentation and detection is a good idea for future work. We note that follow-up work has shown that SGD (freeze-embed) works well for NLP models such as RoBERTa—so it is more general than image classification.

Reviewer comment (Reproducibility):
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.

Author response:
We apologize for the confusion. Gradual unfreezing is another ablation we tried—it is not a main contender in the paper, but we give results in Appendix B. The idea is to unfreeze layers from top to bottom over the course of training. We have added a description to Appendix B.

Reviewer comment (Experiments):
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.

Author response:
This is an interesting idea! We note that learning rate warmup is generally helpful for AdamW (because it leads to better estimates for the moments), and generally has a smaller benefit for SGD. Also, we did sweep over many learning rates including some very small learning rates and selected the best, which hopefully partially accounts for this. In this work we focused our additional experiments on different optimizers (e.g., AdamW, LARS, LAMB), lower learning rates on the embedding layer, linear probing, SGD with weight decay, etc.

We agree that the learning rate schedule is good to explore in future work, and have mentioned this in the revised submission.

Reviewer comment (Experiments):
Will such modified SGD still achieve performance gain on NLP model such as Llama2?

Author response:
We note that follow-up work has shown that SGD (freeze-embed) **works well for NLP models** such as RoBERTa. We think examining whether it helps Llama2 is a great avenue for future work—based on our results in Appendix E we believe this is potentially promising because Llama2 was pretrained using AdamW, and that is when our method tends to outperform SGD.

Reviewer comment (Experiments):
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?

Author response:
In Appendix E, we run controlled experiments where we pretrain a ResNet-50 model with (a) SGD, or (b) AdamW. We then fine-tune each of these models with (1) SGD, (2) AdamW, and (3) SGD Freeze-Embed. We find that **when the model is pretrained with SGD, then it is better to fine-tune with SGD than with AdamW.**

Reviewer comment (Novelty):
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.

Author response:
> hinges on models which have an embedding layer  
> recommendation models… embedding layer contains most of the model parameters  

Some recent follow-ups to our paper have found that freezing the embedding layer can lead to improvements in NLP models such as RoBERTa. 

Conceptually, we can think of the embedding layer as an abstract interface between the raw input and the later networks in the model, in which case our kind of analysis may be useful for a more general class of models.

Recommender systems is another interesting domain—we have noted this and the cite in the future work section.

Reviewer comment (Experiments):
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.

Author response:
> hinges on models which have an embedding layer  
> recommendation models… embedding layer contains most of the model parameters  

Some recent follow-ups to our paper have found that freezing the embedding layer can lead to improvements in NLP models such as RoBERTa. 

Conceptually, we can think of the embedding layer as an abstract interface between the raw input and the later networks in the model, in which case our kind of analysis may be useful for a more general class of models.

Recommender systems is another interesting domain—we have noted this and the cite in the future work section.

Reviewer comment (Evaluation):
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?

Author response:
> From the practitioners point of view, how should one choose which method to use?

1. **SGD-freeze embed works consistently well**, so we recommend using it as a general guideline. SGD freeze-embed does better than SGD in 80% of the 70 settings (7 models $\times$ 5 datasets $\times$ {ID, OOD}).
2. In modern deep learning it's rarely the case that a single optimizer dominates in every case—our paper gives a detailed analysis as to when and why each of the optimizers can do better. To get further improvements, practitioners can look at the **embedding layer gradients** at the start of training (which is computationally very cheap---a single backward pass). If these are similar to (or smaller than) the gradients of other layers, then we expect SGD (even without freeze embed) to work fairly well.
3. In addition, our experiments and analysis in Appendix E show that SGD works better when the model was pretrained with SGD, or when the fine-tuning data is similar to pretraining (e.g., Living-17, which is a subset of ImageNet). So researchers and practitioners can also use their understanding of the pretraining optimization process, and the relationship between the pretraining and fine-tuning data, to get further improvements.

Reviewer comment (Reproducibility):
For Fig. 2, are the gradient norms normalized by number of elements?

Author response:
> For Fig. 2, are the gradient norms normalized by number of elements?

Figure 2 does not normalize by number of elements, however in Appendix F we show figures where they are normalized by the parameter norm, and the number of parameters.

Reviewer comment (Evaluation):
Why do you think SGD outperforms the other alternatives across so many experiments?

Author response:
> Why do you think SGD outperforms the other alternatives across so many experiments?

Could we please clarify which alternative you are referring to here? SGD (freeze-embed) and AdamW generally perform better than SGD in the majority of experiments