Fast Trainable Projection for Robust Fine-Tuning

TL;DR

This paper introduces Fast Trainable Projection (FTP), a scalable and efficient projection-based fine-tuning method that enhances robustness and speed in transfer learning tasks across various vision benchmarks.

Contribution

FTP provides a novel, computationally efficient projection-based fine-tuning algorithm that improves robustness and speed, and can be integrated with existing optimizers in a plug-and-play manner.

Findings

FTP achieves 35% speedup over prior methods.

FTP improves robustness on OOD datasets across multiple vision tasks.

FTP is adaptable to low-label and continual learning scenarios.

Abstract

Robust fine-tuning aims to achieve competitive in-distribution (ID) performance while maintaining the out-of-distribution (OOD) robustness of a pre-trained model when transferring it to a downstream task. Recently, projected gradient descent has been successfully used in robust fine-tuning by constraining the deviation from the initialization of the fine-tuned model explicitly through projection. However, algorithmically, two limitations prevent this method from being adopted more widely, scalability and efficiency. In this paper, we propose a new projection-based fine-tuning algorithm, Fast Trainable Projection (FTP) for computationally efficient learning of per-layer projection constraints, resulting in an average $35\%$ speedup on our benchmarks compared to prior works. FTP can be combined with existing optimizers such as AdamW, and be used in a plug-and-play fashion. Finally, we show that FTP is a special instance of hyper-optimizers that tune the hyper-parameters of optimizers in a learnable manner through nested differentiation. Empirically, we show superior robustness on OOD datasets, including domain shifts and natural corruptions, across four different vision tasks with five different pre-trained models. Additionally, we demonstrate that FTP is broadly applicable and beneficial to other learning scenarios such as low-label and continual learning settings thanks to its easy adaptability. The code will be available at https://github.com/GT-RIPL/FTP.git.