Feature Protection For Out-of-distribution Generalization

TL;DR

This paper investigates how protecting pre-trained features during fine-tuning improves out-of-distribution generalization, demonstrating that feature protection leads to more robust models on unseen data.

Contribution

It introduces feature protection techniques for fine-tuning pre-trained models, enhancing their robustness to out-of-distribution data compared to standard methods.

Findings

Standard fine-tuning overfits features, harming OOD performance.

Feature protection maintains pre-trained features, improving OOD robustness.

Experimental validation on CLIP shows significant OOD performance gains.

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.