Feature Protection For Out-of-distribution Generalization

TL;DR

Fine-tuning large pre-trained models on small downstream datasets often improves ID accuracy but harms OOD generalization. The authors diagnose this by showing that fine-tuning distorts semantic representations in $f=[\boldsymbol{\Phi}, \boldsymbol{v}]$, and they compare several feature-preserving strategies, including $L1$/$L2$ regularization, LoRA, knowledge distillation, and WiSE-FT (model averaging). The key finding is that regularization and KD can mitigate forgetting, while WiSE-FT provides the strongest OOD gains, sometimes outperforming the original pre-trained model and full fine-tuning on both ImageNet- and DomainNet-based tasks. This work offers practical guidance for robust deployment of fine-tuned models under distribution shifts by preserving pre-trained features during adaptation.

Abstract

With the availability of large pre-trained models, a modern workflow for building real-world machine learning solutions is to fine-tune such models on a downstream task with a relatively small domain-specific dataset. In such applications, one major challenge is that the small fine-tuning dataset does not have sufficient coverage of the distribution encountered when the model is deployed. It is thus important to design fine-tuning methods that are robust to out-of-distribution (OOD) data that are under-represented by the training data. This paper compares common fine-tuning methods to investigate their OOD performance and demonstrates that standard methods will result in a significant change to the pre-trained model so that the fine-tuned features overfit the fine-tuning dataset. However, this causes deteriorated OOD performance. To overcome this issue, we show that protecting pre-trained features leads to a fine-tuned model more robust to OOD generalization. We validate the feature protection methods with extensive experiments of fine-tuning CLIP on ImageNet and DomainNet.

Figures (5)

• Figure 1: Illustration of the workflow: Pre-train, fine-tune, and evaluation. Nowadays, large neural networks are often pre-trained on large-scale datasets (e.g., LAION). Then, the pre-trained model is fine-tuned on a small dataset from a fine-tuning (source) domain(e.g., ImageNet). Finally, we evaluate the model on target domains. For example, we fine-tune the pre-trained CLIP model on the ImageNet dataset with photo style, but the testing samples are from other styles such as sketch, art, cartoon, etc.
• Figure 2: Accuracy and linear probing accuracy of sketch domain in DomainNet.
• Figure 3: ID and OOD trade-off. (Left): fine-tune on ImageNet and evaluate the OOD performance by the average performance of ImageNet variants, i.e., ImageNet-V2, ImageNet-R, ImageNet-Sketch, ObjectNet; (Right) fine-tune on the 'real' domain of DomainNet and evaluate the generality by the average performance on the Clipart”, “Infograph”, “Painting”, “Quickdraw”, and “Sketch” domains. Different points with the same color mean different parameters.
• Figure 4: Accuracy and linear probing accuracy of 5 OOD domains in DomainNet.
• Figure 5: (Left) Illustration samples of the class lemon from ImageNet and 5 variants; (Right) Illustration samples of the class apple from DomainNet.