AutoFT: Learning an Objective for Robust Fine-Tuning

TL;DR

AutoFT is a data-driven method that automatically searches for optimal fine-tuning objectives and hyperparameters to improve out-of-distribution generalization of foundation models, outperforming existing methods on multiple benchmarks.

Contribution

AutoFT introduces a bi-level optimization framework to automatically discover fine-tuning objectives tailored for robustness against distribution shifts.

Findings

AutoFT significantly improves OOD generalization across nine benchmarks.

Achieves state-of-the-art results on WILDS iWildCam and FMoW datasets.

Outperforms existing robust fine-tuning techniques by notable margins.

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.