Rethinking Weight Decay for Robust Fine-Tuning of Foundation Models

TL;DR

This paper introduces Selective Projection Decay (SPD), a new weight decay method that selectively regularizes model layers to improve robustness and generalization when fine-tuning foundation models.

Contribution

SPD is a novel weight decay technique that selectively applies regularization to certain layers, enhancing robustness and generalization in fine-tuning foundation models.

Findings

SPD improves in-distribution generalization.

SPD enhances out-of-distribution robustness.

Adam with SPD outperforms standard methods on benchmarks.

Abstract

Modern optimizers such as AdamW, equipped with momentum and adaptive learning rate, are designed to escape local minima and explore the vast parameter space. This exploration is beneficial for finding good loss basins when training from scratch. It is not necessarily ideal when resuming from a powerful foundation model because it can lead to large deviations from the pre-trained initialization and, consequently, worse robustness and generalization. At the same time, strong regularization on all parameters can lead to under-fitting. We hypothesize that selectively regularizing the parameter space is the key to fitting and retraining the pre-trained knowledge. This paper proposes a new weight decay technique, Selective Projection Decay (SPD), that selectively imposes a strong penalty on certain layers while allowing others to change freely. Intuitively, SPD expands and contracts the parameter search space for layers with consistent and inconsistent loss reduction, respectively. Experimentally, when equipped with SPD, Adam consistently provides better in-distribution generalization and out-of-distribution robustness performance on multiple popular vision and language benchmarks. Code available at~\url{https://github.com/GT-RIPL/Selective-Projection-Decay.git}