AutoFT: Learning an Objective for Robust Fine-Tuning

TL;DR

The proposed AutoFT uses bi-level optimization to search for an objective function and hyperparameters that maximize post-adaptation performance on a small OOD validation set, and achieves a new state-of-the-art on the WILDS iWildCam and FMoW benchmarks.

Abstract

Foundation models encode rich representations that can be adapted to downstream tasks by fine-tuning. However, fine-tuning a model on one data distribution often degrades performance under distribution shifts. Current approaches to robust fine-tuning use hand-crafted regularization techniques to constrain the fine-tuning process towards the pretrained model. Yet, it is hard to specify how to adapt relevant characteristics of the foundation model during fine-tuning, as this depends on how the pre-training, fine-tuning, and test data distributions relate to each other. We propose AutoFT, a data-driven approach for robust fine-tuning. Given a task, AutoFT searches for a fine-tuning procedure that enhances out-of-distribution (OOD) generalization. Specifically, AutoFT uses bi-level optimization to search for an objective function and hyperparameters that maximize post-adaptation performance on a small OOD validation set. We evaluate AutoFT on nine natural distribution shifts. Our experiments show that AutoFT significantly improves generalization to OOD inputs, outperforming existing robust fine-tuning methods. Notably, AutoFT achieves a new state-of-the-art on the WILDS iWildCam and FMoW benchmarks, outperforming the previous best methods by $6.0\%$ and $1.5\%$, respectively.

Figures (7)

• Figure 1: Overview of AutoFT: a method for robustly fine-tuning foundation models. While fine-tuning with in-distribution (ID) data (blue), AutoFT searches for a fine-tuning objective that maximizes performance on a small out-of-distribution validation set (red). This validation set serves as a proxy for performance on different distributions (green and purple), allowing AutoFT to learn a robust fine-tuning procedure.
• Figure 2: A summary of our data assumptions and evaluation protocol. The standard approach is to optimize hyperparameters on a validation dataset drawn from the same distribution as the training data. In contrast, AutoFT employs a small out-of-distribution (OOD) validation set for hyperparameter optimization, enhancing the generalizability of the final model. We evaluate all fine-tuned models on data from unseen distribution shifts (green and purple).
• Figure 3: AutoFT outperforms existing methods, both with and without weight ensembling wortsman2022robust. Here, we show the ID-OOD performance curves obtained by linearly interpolating the fine-tuned model weights with the zero-shot weights.
• Figure 4: In binary few-shot classification, AutoFT outperforms existing robust fine-tuning methods. AutoFT outperforms FLYP by $3.1\%$ and full fine-tuning by $4.3\%$ in 32-shot classification on PatchCamelyon.
• Figure 5: Effect of val set size. Performance of AutoFT on iWildCam with varying validation set sizes. Increasing size beyond $1000$ shows minimal additional benefits.