Title: Fast Trainable Projection for Robust Fine-tuning
Venue: NeurIPS 2023
Accepted: yes
Keywords: fine-tuning

transfer learning

regularization
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.

Fast Trainable Projection for Robust Fine-Tuning

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 .

1 Introduction

With new progress being made in pre-training of foundation models every year, such as selfsupervised [1, 2, 3] or language-supervised training [4], their potential has gone far beyond merely speeding up convergence [5]. They have demonstrated superior transferability to other tasks, reducing the need for data and improving robustness and generalization capabilities [6, 7, 8]. The problem of how to fine-tune (transfer) a foundation model such that we maintain its robustness and generalization capabilities acquired during pre-training on large datasets has therefore become an essential research topic. This problem is hard because the conventional machine learning paradigm of validating on held-out training data does not impose any constraints on robustness and generalization w.r.t. the foundation models. For example, fine-tuning with a slightly large learning rate can easily destroy capabilities that reside in the foundation models [8], while performing well on the target task.

To maintain the robustness and generalization capability of the pre-trained model when fine-tuning, recent projection-based methods explicitly constrain the distance between the fine-tuned and the pre-trained models through projection. For example, MARS-SP [9] specifies a distance constraint shared by all layers in a neural network. However, it is practically intractable to tune a constraint for each layer (poor scalability ). TPGM [10] proposes to automatically learn different constraints

Yen-Cheng Liu

Zsolt Kira

𝛾

!

𝑾

!

𝑾 #

!

1. UGD

3. Projection

2. ProjUpdate

𝑾

"

(a) One Step FTP Diagram

Time (it/s)

Real (ID)

Sketch

Painting

Infograph

1.0

1.5

TPGM

FTP

(b) Image Classification

Time (it/s)

Clean (ID)

Fog

Defocus

Brightness

1.0

1.75

TPGM

FTP

(c) Semantic Segmentation

for each layer, solving the issue of scalability in MARS-SP, however, with increased computational overhead (poor efficiency ). These limitations prevent the method from being adopted more widely.

To achieve scalability and efficiency simultaneously, we propose Fast Trainable Projection (FTP), for learning both the projection constraints and the main model in a single forward pass (Fig. 1a), significantly reducing computation overhead in prior works while achieving competitive performance. Specifically, FTP removes the algorithmic redundancy of extra training procedures required in TPGM [10], which requires sampling a separate batch of data and running a nested training loop. FTP achieves this by 1) utilizing different batches of training data sampled at consecutive steps and 2) re-using gradients calculated for the main model update (Sec. 3.2). This leads to a 35% speedup with comparable performance in fine-tuning (Fig. 1b, 1c). The efficiency improvement and easy adaptability as a drop-in replacement with existing optimizers are essential to making projection-based methods applicable to more fine-tuning problems. For example, we implement SGDP, an SGD variant with built-in FTP. SGDP can be used as a drop-in replacement for SGD (details in Appendix 8.7) as:

To demonstrate this, we test FTP on four different vision tasks, image classification, semantic segmentation, human parts segmentation, and surface normal estimation. FTP shows superior OOD performance under domain shift or natural corruptions on all benchmarks. Moreover, we apply FTP to a continual learning (CL) benchmark and achieve state of the art performance when combined with a simple CL technique.

Finally, we show that FTP is a special instance of hyper-optimizers [11, 12, 13, 14, 15, 16, 17, 18] that aims to reduce the manual tuning of optimization hyper-parameters such as learning rate by learning them automatically through automatic differentiation and nested optimization. Theoretically, to understand why FTP and other projection methods can maintain the robustness of the pre-trained models, we propose to establish a theoretical connection between robustness and projection through the lens of Liptschitz continuity, a widely adopted measure of robustness [19, 20, 21]. In summary, our contributions are:

We present a new fine-tuning algorithm, Fast Trainable Projection, to efficiently learn the projection constraints and fine-tune the model simultaneously, bringing significantly improved computation efficiency w.r.t. prior works [10] in Sec. 3.2.

We show that FTP is a special instance of hyper-optimizers that aims to reduce manual tuning of hyper-parameters through nested optimization in Sec. 3.3.

We discuss a dual perspective of the fine-tuning robustness in the feature space and the weight space of a model to mathematically understand why projection can maintain the robustness of the pre-trained models in Sec. 3.4.

We show superior robustness on OOD datasets on four vision tasks with five pre-trained models and SOTA performance on a continual learning benchmark, all with a 35 % speedup in Sec. 4.

2 Related Works

We summarize related works in (general) robust fine-tuning into three categories: when, where, and how much to fine-tune, depending on their underlying strategy. Moreover, we discuss recent advances in fine-tuning language-image pre-trained models, which have inspired specialized fine-tuning strategies. When to fine-tune: LP-FT [7] discovers that fine-tuning the entire network can distort features in the pre-trained models and proposes to first only fine-tune the last linear layer followed by training the entire network with a small learning rate. We will include LP-FT in our experiments. Where to fine-tune: Instead of fine-tuning the entire network, some methods investigate the choice of weights to fine-tune. SpotTune [22] learns where to fine-tune through an additional policy network. However, SpotTune needs to retain the policy network, the pre-trained model, and the fine-tuned model in memory for inference, adding significant computation at inference time. Recently, SurgicalFT [23] proposes to use the GradientNorm heuristic, the ratio of the gradient norm to the parameter norm, to determine which layer to fine-tune. Parameter-efficient fine-tuning methods are another example of this category. While they aim to minimize the parameters tuned, they have been shown to improve OODgeneralization performance as well in NLP applications [24, 25, 26, 27, 28, 29]. We specifically compare to two recent parameter-efficient methods that only tune the bias terms: Bitfit [30] for Transformers [31] and Partial Fusion [32] for ResNets [33]. How much to fine-tune: Our work belongs to this category where the entire neural network is fine-tuned simultaneously. Specifically, we can split works into two sub-categories: regularization and projections. Regularization: DELTA [34] proposes to regularize the output (feature maps) of a neural network to that of its pre-trained model. This requires two separate passes through the pre-trained model and the fine-tuned model increasing both memory and computation overhead. L2-SP [35] instead regularizes the L2 distance between the fine-tuned model and the pre-trained model, serving as a strong baseline. Projection: Utilizing projection to enforce a close distance to the pre-trained model has been studied in prior works: MARS-SP [9] and TPGM [10]. We dedicate a section to revisit them later in the method section (Sec. 3.1). Language-Image Pre-trained Models. Several recent works have proposed special fine-tuning strategies for language-image pre-trained models with zero-shot capability. WISE-FT [8] achieves SOTA performance by linearly interpolating a fine-tuned model and its initialization at the end of fine-tuning. However, it only applies to a subset of pre-trained models with linear connectivity such as CLIP [4]. FT-Like-Pretrain [36] proposes to use a contrastive fine-tuning strategy, the same strategy used in pre-training for those models, instead of the conventional cross-entropy loss for many vision tasks. The method has demonstrated superior results when combined with WISE-FT, where WISE-FT contributes the most to the improvement. Similarly, we will also combine our method with WISE-FT to show improved OOD performance using CLIP.

3 Method

3.1 Review: Enforcing Projection and Learning Constraints

In this work, we focus on fine-tuning a pre-trained model, where W 0 ∈ R n × m is the weights of a linear layer in the pre-trained model, to a downstream task. We denote W t as the fine-tuned model at training iteration t and w i as the i th row of a matrix W ∈ R n × m . Several prior works [9, 10] have attempted to use projection to improve fine-tuning robustness. The most vanilla formulation in MARS-SP [9] has two steps: unconstrained gradient descent and projection .

Unconstrained Gradient Descent (Abbrev. UGD), the projection-based methods first compute the updated model weights ˜ W t without projection. For example, at iteration t , given a batch of training data D tr t , we first obtain ˜ W t as the following,

where g t is the derivative of the loss function L D tr t ( W t -1 ) calculated on D tr t w.r.t. W t -1 and Opt ( · ) is an existing optimization algorithm, such as SGD, AdamW [37].

Projection. MARS-SP [21] projects the updated model ˜ W t towards its initialization W 0 with a pre-defined projection constraint γ for all layers using the MARS matrix norm (see Appendix . 8.2) as shown below in Eq. 2.

However, MARS-SP has poor scalability because it is practically intractable to hand-tune different constraints for each layer, which results in sub-optimal performance as reported by TPGM [10]. Instead of a pre-defined γ for all layers, TPGM proposes to learn a different constraint γ t for each layer 1 . and updates them iteratively during training. This enables TPGM to customize a different regularization strength for each layer and to have superior performance for both ID and OOD data.

ProjUpdate. Given as input the frozen unconstrained model ˜ W t from UGD (Eq. 1), TPGM adds an intermediate ProjUpdate function before projection , which samples a separate set of data from the validation dataset D val t and uses a standalone training loop to update the projection constraints γ t while keeping the model ˜ W t frozen. Specifically, ProjUpdate creates a temporary projected model W p by projecting ˜ W t towards W 0 based on the previous constraint γ t -1 using Eq. 2, i.e., W p = Π( ˜ W t , W 0 , γ t -1 ) . Therefore, W p ( γ t -1 ) 2 can be viewed as a function of γ t -1 . Then FTP calculates the gradient ∇ γ t by taking a derivative of the loss function L D val t ( W p ( γ t -1 )) w.r.t. γ t -1 :

With the calculated gradient, TPGM uses an existing optimizer Opt ( · ) to update γ t . This procedure, sampling D val t and calculating the derivative ∇L D val t ( W p ( γ t -1 )) , is the key to learning projection constraints because the unconstrained model ˜ W t (highlighted above and calculated in Eq. 1), was updated on the training data D tr t and γ t is now updated on separate data D val t . The discrepancy between D tr t and D val t allows TPGM to find a better projected model W p (projected between ˜ W t and W 0 ) by updating γ t , which balances between fitting the training data D tr t and generalizing to D val . Finally, with an updated γ t , TPGM again projects ˜ W t towards W 0 to obtain the final model W t using Eq. 2, replacing the pre-defined γ with a learned γ t . A flow chart of TPGM is in Fig. 2.

The algorithm demonstrated the capability to automatically learn different constraints for each layer, solving the scalability issue in MARS-SP. However, TPGM introduces extra computation in the additional training loop. In the next section, we propose a scalable and efficient projection algorithm that learns the projection constraints for each layer without separate validation data and loops.

3.2 FTP: Fast Trainable Projection

To inherit the scalability of TPGM while reducing the computational overhead, we propose Fast Trainable Projection (FTP) (Algorithm 1). Similar to TPGM, the algorithm has three components: UGD , ProjUpdate , Projection . The ProjUpdate component is the major contributor to efficient computation. It builds on a key insight: Instead of sampling separate data D val t each time, we use two training data batches sampled independently at consecutive steps, e.g., D tr t -1 and D tr t . Specifically, we use D tr t to update γ t instead of D val t . As a result, the optimization of γ t re-uses most of the computation used for the optimization of the main model.

ProjUpdate. Specifically, instead of taking a derivative of L D val t ( W p ) w.r.t. γ t -1 as in TPGM, FTP calculates the gradient of γ t -1 by the derivative of the loss function on the current training data L D tr t ( W t -1 ) w.r.t. γ t -1 . Note that W t -1 = Π( ˜ W t -1 , W 0 , γ t -1 ) is also a function of the

Algorithm 1 FTP: F ast T rainable P rojection.

Require:

W

0

the pre-trained model

Require:

κ

positive gradient annealing rate

Require:

µ

1

e

-

2

, β

1

, β

2

(0

.

9

,

0

.

999)

fixed parameters for AdamUpdate

for

t

= 1

...T

do

{

g

t

←∇L

D

tr

(

W

t

-

1

)

˜

W

t

Opt

(

W

t

-

1

,

g

t

)

▷

Unconstrained Gradient Descent (Eq. 1)

if

t

= 1

then

γ

t

= 1

e

-

8

▷

Initialize

γ

else









∇

γ

t

∑

i

(

g

i,

⊺

t

( ˜

w

i

t

-

1

-

w

i

0

)

1

∥

˜

w

i

t

-

1

-

w

i

0

∥

1

)

if

∇

γ

t

>

0 :

∇

γ

t

=

κ

∇

γ

t

γ

t

AdamUpdate

(

γ

t

-

1

,

∇

γ

t

, t

)

▷

ProjUpdate (Eq. 4,Eq. 5, Alg. 2)

W

t

= Π( ˜

W

t

,

W

0

, γ

t

)

▷

Projection (Eq. 2)

Unconstrained GD

ProjUpdate

Projection

𝑾

!"#

= Π 𝑾

$

$"#

, 𝑾

%

, 𝛾

!"#

𝛾

!

𝑾 $

$

𝑾

𝒕"𝟏

𝑾 $

$

𝑾

𝒕"𝟏

𝑾

𝒑

= Π 𝑾

$

$

, 𝑾

%

, 𝛾

!"#

𝛾

!

𝑾

$

= Π 𝑾

$

$

, 𝑾

%

, 𝛾

!

TPGM

FTP

𝑾 $

$"#

was updated on

𝒟

!"#

!)

,

𝛾

!

is updated on

𝒟

!

!)

𝑾 $

$

was updated on

𝒟

!

!)

,

𝛾

!

is updated on

𝒟

!

*+,

𝑾

$

= Π 𝑾

$

$

, 𝑾

%

, 𝛾

!

∇ℒ

𝒟

!

"#$

𝑾

/

𝛾

!"#

∇ℒ

𝒟

!

!%

𝑾

!"#

𝛾

!"#

∇ℒ

0

!

!%

𝑾

!"#

∇ℒ

0

!

!%

𝑾

!"#

constraint γ t -1 as a result of projection from the previous step. Hence, by virtue of the chain rule, the gradient of L D tr t ( W t -1 ( γ t -1 )) w.r.t. γ t -1 is,

where the summation loops over each row in the matrix W t -1 because the same constraint γ t -1 is enforced for all rows (see MARS norm in Appendix Eq. 15 and Eq. 2 ) so the final gradient is the summation of all gradients for each row. Similar to TPGM, because the starting point of projection ˜ W t -1 (highlighted above) was updated using the previous training batch D tr t -1 and the gradient ∇ γ t is calculated using the current batch D tr t , the discrepancy between D tr t -1 and D tr t enables FTP to learn meaningful projection constraints. Crucially, we proposed a novel formulation that allows for re-using the gradient g t used for calculating the unconstrained model ˜ W t in the UGD step (Eq. 1).

Gradient Annealing. Prior work [10] noticed that learning projection constraints for each layer can suffer from underfitting because the learned constraints can be too conservative, and used an additional regularization to help reduce this negative effect. For FTP, we introduce a simple technique that uses a single gradient annealing factor for all layers, κ ∈ [0 , 1] to modulate the magnitude of the positive gradient ∇ γ t > 0 , which contributes to the shrinkage of the constraints. When ∇ γ t > 0 ,

For example, when κ = 0 , the projection constraint γ will not receive any positive gradient and is therefore non-decreasing. With the annealed gradients ∇ γ t , we update the constraint using the Adam update rule [38] because Adam is suitable for non-stationary optimization, where the optimal values change over time. Please see Appendix 8.3 for a detailed algorithmic description of AdamUpdate and additional discussion on how FTP saves computation.

Finally, after obtaining the updated γ t from AdamUpdate , FTP applies the learned constraints to project the current unconstrained model ˜ W t towards the pre-trained model W 0 using Eq. 2 with a different constraint for each layer. The complete algorithm is summarized in Alg. 1. For a quick comparison with TPGM, we provide a side-by-side computation flow chart of FTP in Fig. 2.

Implicit Assumption. The algorithmic difference between TPGM and FTP makes an implicit assumption. Specifically, after obtaining the updated constraints γ t after AdamUpdate , if the algorithm were to follow TPGM, the next step would be applying the updated constraints to recalculate the previous model W t -1 since γ t is updated based on W t -1 (Eq. 4). However, instead of rolling back, FTP applies the updated constraints directly to the current unconstrained model ˜ W t to calculate W t . This step assumes smoothness in the update of γ t , i.e., the γ t does not change drastically in consecutive steps. The assumption is valid since γ t is updated by AdampUpdate (Alg. 2 in Appendix 8.3) which uses a moving average update with a momentum of 0 . 9 . So the change of γ t is very smooth because of the high discount factor of 0 . 9 . Importantly, it enables us to re-use the same gradient g t available for computing the current unconstrained model ˜ W t to update γ t . This is the key to saving computation because the separate training loop as a result of 'rolling back' is the main computation bottleneck in TPGM.

3.3 FTP as a Hyper-Optimizer for Fine-Tuning

The FTP algorithm in Alg. 1 bears motivational and algorithmic similarity to a recent resurrection of hyper-optimizers [11, 12, 13, 14, 15, 16, 17, 18]. Specifically, hyper-optimizers aim to learn the hyper-parameters such as the learning rate in an optimizer by treating them as learnable parameters through nested differentiation and optimization because manual tuning of those hyper-parameters can be time-consuming and can lead to sub-optimal performance. FTP stems from the same motivation as the manual specification of projection constraints can be computationally infeasible [10].

To understand the algorithmic similarity better, let's use SGD as an example. Suppose at iteration t -1 , we have updated the model parameters W t -2 through SGD with a learning rate α t -1 .

At the current step t , hyper-optimizers first calculate the gradient w.r.t to the learning rate α t -1 and update it using another SGD optimizer with a new learning rate parameter κ .

Finally, using the updated α t , hyper-optimizers update the main model parameters.

It's not hard to spot the algorithmic similarity between the FTP algorithm and hyper-optimizers. Both algorithms first update the hyper-parameters (projection constraints γ t in Eq. 4 vs. the learning rate α t in Eq. 7) using the cached information from the previous iteration and the gradient from the current iteration (known as hyper-gradients). Then, they apply the updated hyper-parameters to calculate the current model. Finally, hyper-optimizers make the same assumption of smoothness in the update of the hyper-parameters such that the update of the hyper-parameters and the model parameters can be performed consecutively in a single forward pass. In this regard, FTP can be seen as a special instance of hyper-optimizer for fine-tuning.

3.4 Dual Perspective: Fine-tuning Robustness in Feature Space and Weight Space

It is not immediately clear why FTP's and other methods' projections in the weight space maintain the robustness of the pre-trained model in the feature space besides the intuition that the closer to the pre-train model the more likely the fine-tuned model will behave like it. To fully understand this, we study the mathematical connection between projection and robustness. Let x ∈ R m denote an input vector and h ( x ) : R m → R n a function mapping it to a feature space. Given two input vectors x , x ′ ∈ R m , we denote the distance between them in the original space by their vector norms ∥ x -x ′ ∥ x and in the feature space by ∥ h ( x ) -h ( x ′ ) ∥ h . Let h f ( · ) and h 0 ( · ) denote a fine-tuned model and its pre-trained initialization, and ∆ h ( · ) ≡ h f ( · ) -h 0 ( · ) denotes the difference function .

To capture the robustness of a fine-tuned model, we apply the notion of Lipschitz continuity on the difference function because its Lipschitz constant captures the maximum rate of change of differences between the fine-tuned model and the pre-trained model in the feature space. Formally,

where L d ≥ 0 is the Lipschitz constant of the difference function ∆ h ( x ) . If the inequality is satisfied, in this paper, we call h f ( · ) L d -Lipschitz-robust w.r.t. the pre-trained initialization h 0 ( · ) .

The definition has a natural intuition stemming from Lipschitz continuity, a measure of robustness [19, 20, 21]. A Lipschitz function is limited by how fast it can change, governed by the Lipschitz constant. Traditionally, a small Lipschitz constant is associated with better robustness, because a small constant means less sensitivity to changes in the input. We provide the following lemma (proof in Appendix 8.1) to illustrate the connection between the difference function and the robustness of the fine-tuned model.

Lemma 1. If a fine-tuned model h f ( · ) is L d -Lipschitz-robust with respect to its L 0 -Lipschitz pretrained initialization h 0 ( · ) , i.e., ∀ ( x , x ′ ) ∈ R m ,

then, h ( · ) is ( L + L )

Feature Space. From Lemma 1, we can see that minimizing L d can improve the robustness of the fine-tuned model, defined by its Lipschitz constant ( L d + L 0 ) , which equals L 0 when L d = 0 . Therefore, the fine-tuned model can achieve a similar level of robustness as the pre-trained model if L d is minimized. Colloquially, given two inputs ( x , x ′ ) , where x ′ is a perturbed version of x , h f ( · ) will be just as sensitive/robust to the perturbation as h 0 ( · ) is if L d is small.

Weight Space. The definition of fine-tuning robustness (Eq. 9) not only leads to an interpretation of robustness in the feature space (Lemma 1) but also conveniently a projection operation in the weight space. Specifically, we investigate a single linear layer in a neural network and show that enforcing the inequality in Eq. 9 leads to a projection operation by virtue of linear operators and matrix norms. We illustrate this in the following lemma with a full discussion and proof in Appendix 8.2.

Lemma 2. Assuming linear models h ( x ) = Wx + b , W ∈ R n × m , b ∈ R n , and both the input space vector norm ∥·∥ x and the feature space vector norm ∥·∥ h are defined by l ∞ norm. w i p satisfies the inequality in Eq. 9 if

where w i denotes the i -th row of the matrix W and w i p is the new projected fine-tuned model.

This is an equation of projection between W f and W 0 defined by the MARS norm in the weight space for a single linear layer and is the projection operation used in FTP and prior works [9, 10] (Eq. 2). It indicates that we can choose an arbitrarily small L d and enforce it through Eq. 10, potentially trading off fitting the downstream task and preserving robustness. In summary, this section demonstrates the connection between robustness and projection. Specifically, we have shown that to achieve good fine-tuning robustness, we can enforce a small Lipschitz constant L d on the difference function ∆ h ( x ) in the feature space (Lemma 1), which can be physically enforced through the projection of the fine-tuned model towards the pre-trained model in the weight space (Lemma. 2).

4 Experiments

Overview. To validate the effectiveness of FTP in fine-tuning pre-trained models, we benchmark FTP on both image classification (Sec. 4.1) and dense vision tasks (Sec. 4.2) with different network architectures and pre-trained models. For each benchmark, we report both in-distribution (ID) performance as well as out-of-distribution (OOD) performance. We show that FTP not only achieves competitive ID performance and superior OOD performance but is also much more computationally efficient than prior works. We further test FTP's regularization capability on a continual learning benchmark and show state of art performance against recent SOTA methods (Sec. 4.3).

4.1 Image Classification Experiments

4.1.1 DomainNet

For the DomainNet experiment (image classification), which consists of five domains, Real, Sketch, Painting, Infographics, and Clipart, we follow the setup of the prior work [10] and use its released code to train FTP. Specifically, we use two pre-trained models, an ImageNet pre-trained MOCO-V3 ResNet50 [3] and a CLIP pre-trained ResNet50 [4]. For FTP, we only tuned the learning rate while keeping the other hyper-parameters fixed as in the prior work. We use the Real domain as the ID training dataset and the rest as OOD testing datasets. Please refer to Appendix 8.4 for more details.

FTP achieves the best OOD accuracy and is much more efficient. In Tab. 1 and Tab. 2, we show results training on 100 % DomainNet-Real data using CLIP and MOCO-V3 pre-trained initialization respectively. Compared to the previous SOTA methods TPGM [10], FTP achieves competitive ID accuracy and better OOD generalization performance. More importantly, in addition to favorable results, FTP is 36% faster on average on both benchmarks compared to TPGM . Following TPGM [10], we also report results training only on 10% DomainNet-Real data in Appendix Tab. 6.

4.1.2 ImageNet

Recently, zero-shot language-vision pre-trained models such as CLIP [4] have demonstrated strong generalization capability to other tasks. Notably, WISE [8] showed that linear interpolation between a fine-tuned model and its initialization achieves significant improvement in OOD generalization.

45

47

49

51

53

55

57

65

67

69

71

73

75

77

79

81

83

85

OOD performance

ID performance

WISE-FT

WISE-FTP

However, there are two limitations: 1) not all pre-trained models have this property of linear connectivity and 2) a zero-shot classifier head is needed to initialize the linear classifier head. Our contribution is orthogonal to WISE because FTP is a general optimization algorithm whereas WISE is a post-training algorithm for specific zero-shot models. Therefore, we first compare FTP to vanilla fine-tuning and then apply WISE to both models. We follow the public code base of DEIT [39] to train our CLIP pre-trained ViT-Base. Specifically, we use weight-decay ( 0 . 1 ), drop-path ( 0 . 2 ) [40], label-smoothing ( 0 . 1 ) [41], Mixup ( 0 . 8 ) [42] and Cutmix ( 1 . 0 ) [43]. We train our model on ImageNet and report OOD performance on ImageNet-V2 [44], ImageNet-A [45], ImageNet-R [46], and ImageNet-S [47]. Please refer to Appendix 8.4 for more details on implementation.

FTP outperforms vanilla fine-tuning and improves WISE performance. In Tab. 3, we report performance for competing methods. Even with various regularizations and augmentations in place, FTP can further improve ID performance on ImageNet. Furthermore, FTP brings better OOD performance on all four OOD datasets. This shows that FTP successfully maintains the robustness of the pre-trained CLIP model while existing regularization such as weight decay and drop-path do not. We also report the interpolation results using WISE [8] for the vanilla fine-tuned and FTP fine-tuned models. We sweep a range of interpolation ratios ∈ { 0 . 1 , 0 . 2 , ..., 0 . 9 } and show the trajectory of ID vs. OOD performance plot in Fig. 3. The models with the best average performance are reported in the lower portion of Tab. 3. As expected, WISE interpolation significantly improves OOD generalization for both methods. However, WISE-FTP has significantly better ID performance while still having better OOD performance. This shows that improvement to the base fine-tuning strategy can further benefit pose-training methods such as WISE.

4.2 PASCAL Dense Vision Task Experiments

To further demonstrate the effectiveness of FTP in more diverse scenarios, we test it on PASCALContext [49]. Specifically, following the prior work [50], we use the PASCAL-Context datasets [49], which consist of labels for semantic segmentation, human parts segmentation, and surface normal estimation. For OOD performance, following the popular natural robustness literature [51], we report results on various degradations including fog, defocus blur, Gaussian noise, and brightness corruption, with 5 severity each. We use a combination of Swin ViT-Tiny [52] (pre-trained on ImageNet-22K) and Segformer [53]. In this architecture, Swin Transformer serves as the feature extraction backbone and Segformer is the task-specific decoder. While the feature backbone is initialized with pre-trained weights, a significant part of the entire model (the Segformer decoder) is randomly initialized; In contrast, in simple classification (Sec. 4.1.1), only the last linear classification layer is randomly initialized. Please refer to Appendix 8.5 for details.

FTP achieves the best ID performance and OOD generalization. We report results for semantic segmentation, human parts segmentation, and surface normal estimation in Tab. 4, Appendix Tab. 7, and Appendix Tab. 8 respectively. We additionally add Layer-Wise Learning Rate decay [48] ( LLRD) as a strong baseline. Notably, in all three tasks, FTP outperforms vanilla fine-tuning on ID performance by 11 . 71% , 4 . 48% , and 18 . 30% respectively. This demonstrates the effectiveness of projection as a regularization technique for transfer learning. More importantly, the OOD performance improves as large as 33 . 02% in semantic segmentation. This shows that 1) FTP can effectively maintain the robustness of the original pre-trained model; 2) even though the entire decoder component is randomly initialized, it is worthwhile to put regularizations on the pre-trained feature backbone.

0

0.02

0.04

0.06

0.08

0.1

0.12

0.14

Task1

Task3

Task4

Task5

Task6

Task7

Task8

Task9

Task10

Magnitude

4.3 Continual Learning (CL) Experiments

Recently, pre-trained models have been shown to greatly improve the performance of CL algorithms [59]. We follow the settings in this work [59] to partition ImageNet-R (200 classes) into 10 sequential tasks with 20 non-overlapping classes in each task. A model is trained on each task only once sequentially. To use FTP for CL tasks, unlike supervised vision tasks (Sec. 4.1, 4.2), we re-initialize FTP after each task and use the current model as the 'pre-trained model' for the next task. Moreover, inspired by the prior work [59], we use FTP to only fine-tune the attention blocks. We report both the final task accuracy across all tasks A 1: N ↑ and the global forgetting F N ↓ in Tab. 5 to analyze plasticity and forgetting . Please refer to Appendix 8.6 for more on the metrics and experimental setup. In Table 5, we benchmark against the popular and recent rehearsal-free continual learning methods. FTP alone achieves state of art accuracy against all methods and relatively good forgetting compared to vanilla FT, a sign of superior plasticity and balanced forgetting . We visualize the learned constraints for each task in Fig. 4. We observe that while each task is independent and FTP is re-initialized each time, FTP learns stronger regularization for later tasks. This contributes to lower forgetting compared to FT. We found that FTP combined with a simple continual learning method, EWC [58], achieves state-of-the-art in this setting. Compared to the prompting methods L2P, DualPrompt, and the recent CODA-Prompt, FTP has clear and significant improvements. Our intuition is that the combination of the superior plasticity of FTP and the low forgetting of EWC is the key to the improvement.

5 Limitations

Like any regularization method, FTP has a hyper-parameter to adjust its regularization strength. In this case, the positive gradient annealing factor 0 ≤ κ ≤ 1 (default 1 ) (Sec. 3.2) controls the strength of projection with smaller values indicating weaker regularization. Note that κ = 0 means that the projection constraints are non-decreasing during training. In this case, FTP still provides regularization. For example, we found that a κ = 0 is necessary to obtain the best performance for some dense vision tasks in Appendix 8.5. Generally, we recommend starting with the default κ and only tuning it if underfitting is observed.

6 Conclusion

In this paper, we proposed Fast Trainable Projection, a fine-tuning algorithm to maintain the robustness and the generalization capability of the pre-trained model. FTP learns projection constraints for each layer in a neural network efficiently by carefully re-using past information to save computation. To understand the connection between robustness and projection, we provided a holistic discussion of fine-tuning robustness from its feature space definition to the weight space dual. The new perspective lends a mathematical foundation to the idea of using projection in fine-tuning. Across four vision tasks with different pre-trained models, FTP demonstrated superior ID and OOD generalization capability and significantly better computation efficiency. Furthermore, the continual learning experiments demonstrated FTP's potential in other deep learning paradigms beyond simple fine-tuning. Combined with its compatibility with popular optimization algorithms, we believe FTP can be broadly beneficial in improving the performance of learning tasks using pre-trained initialization.

References

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 16000-16009, 2022.

Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision , pages 9650-9660, 2021.

Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 9640-9649, 2021.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning , pages 8748-8763. PMLR, 2021.

Kaiming He, Ross Girshick, and Piotr Dollár. Rethinking imagenet pre-training. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 4918-4927, 2019.

Dan Hendrycks, Kimin Lee, and Mantas Mazeika. Using pre-training can improve model robustness and uncertainty. In International Conference on Machine Learning , pages 27122721. PMLR, 2019.

Ananya Kumar et al. Fine-tuning can distort pretrained features and underperform out-ofdistribution. ICLR , 2022.

Mitchell Wortsman, Gabriel Ilharco, Jong Wook Kim, Mike Li, Simon Kornblith, Rebecca Roelofs, Raphael Gontijo Lopes, Hannaneh Hajishirzi, Ali Farhadi, Hongseok Namkoong, et al. Robust fine-tuning of zero-shot models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 7959-7971, 2022.

Henry Gouk, Timothy M Hospedales, and Massimiliano Pontil. Distance-based regularisation of deep networks for fine-tuning. ICLR , 2021.

Junjiao Tian, Xiaoliang Dai, Chih-Yao Ma, Zecheng He, Yen-Cheng Liu, and Zsolt Kira. Trainable projected gradient method for robust fine-tuning. arXiv preprint arXiv:2303.10720 , 2023.

Kartik Chandra, Audrey Xie, Jonathan Ragan-Kelley, and Erik Meijer. Gradient descent: The ultimate optimizer. Advances in Neural Information Processing Systems , 35:8214-8225, 2022.

Xiang Wang, Shuai Yuan, Chenwei Wu, and Rong Ge. Guarantees for tuning the step size using a learning-to-learn approach. In International Conference on Machine Learning , pages 10981-10990. PMLR, 2021.

Matthias Feurer and Frank Hutter. Hyperparameter optimization. Automated machine learning: Methods, systems, challenges , pages 3-33, 2019.

Atilim Gunes Baydin, Robert Cornish, David Martinez Rubio, Mark Schmidt, and Frank Wood. Online learning rate adaptation with hypergradient descent. arXiv preprint arXiv:1703.04782 , 2017.

Fabian Pedregosa. Hyperparameter optimization with approximate gradient. In International conference on machine learning , pages 737-746. PMLR, 2016.

Justin Domke. Generic methods for optimization-based modeling. In Artificial Intelligence and Statistics , pages 318-326. PMLR, 2012.

Luís B Almeida, Thibault Langlois, José D Amaral, and Alexander Plakhov. Parameter adaptation in stochastic optimization. In On-line learning in neural networks , pages 111-134. 1999.

Yoshua Bengio. Gradient-based optimization of hyperparameters. Neural computation , 12(8):1889-1900, 2000.

Patricia Pauli, Anne Koch, Julian Berberich, Paul Kohler, and Frank Allgöwer. Training robust neural networks using lipschitz bounds. IEEE Control Systems Letters , 6:121-126, 2021.

Yujia Huang, Huan Zhang, Yuanyuan Shi, J Zico Kolter, and Anima Anandkumar. Training certifiably robust neural networks with efficient local lipschitz bounds. Advances in Neural Information Processing Systems , 34:22745-22757, 2021.

Henry Gouk, Eibe Frank, Bernhard Pfahringer, and Michael J Cree. Regularisation of neural networks by enforcing lipschitz continuity. Machine Learning , 110:393-416, 2021.

Yunhui Guo, Honghui Shi, Abhishek Kumar, Kristen Grauman, Tajana Rosing, and Rogerio Feris. Spottune: transfer learning through adaptive fine-tuning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 4805-4814, 2019.

Yoonho Lee, Annie S Chen, Fahim Tajwar, Ananya Kumar, Huaxiu Yao, Percy Liang, and Chelsea Finn. Surgical fine-tuning improves adaptation to distribution shifts. arXiv preprint arXiv:2210.11466 , 2022.

Neil Houlsby, Andrei Giurgiu, Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International Conference on Machine Learning , pages 2790-2799. PMLR, 2019.

Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2101.00190 , 2021.

Sang Michael Xie, Tengyu Ma, and Percy Liang. Composed fine-tuning: Freezing pre-trained denoising autoencoders for improved generalization. In International Conference on Machine Learning , pages 11424-11435. PMLR, 2021.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104.08691 , 2021.

Prasetya Ajie Utama, Nafise Sadat Moosavi, Victor Sanh, and Iryna Gurevych. Avoiding inference heuristics in few-shot prompt-based finetuning. arXiv preprint arXiv:2109.04144 , 2021.

Kaiyang Zhou, Jingkang Yang, Chen Change Loy, and Ziwei Liu. Learning to prompt for vision-language models. International Journal of Computer Vision , 130(9):2337-2348, 2022.

Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models. arXiv preprint arXiv:2106.10199 , 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Fahdi Kanavati and Masayuki Tsuneki. Partial transfusion: on the expressive influence of trainable batch norm parameters for transfer learning. In Medical Imaging with Deep Learning , pages 338-353. PMLR, 2021.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Xingjian Li, Haoyi Xiong, Hanchao Wang, Yuxuan Rao, Liping Liu, Zeyu Chen, and Jun Huan. Delta: Deep learning transfer using feature map with attention for convolutional networks. arXiv preprint arXiv:1901.09229 , 2019.

LI Xuhong, Yves Grandvalet, and Franck Davoine. Explicit inductive bias for transfer learning with convolutional networks. In International Conference on Machine Learning , pages 28252834. PMLR, 2018.

Sachin Goyal, Ananya Kumar, Sankalp Garg, Zico Kolter, and Aditi Raghunathan. Finetune like you pretrain: Improved finetuning of zero-shot vision models. arXiv preprint arXiv:2212.00638 , 2022.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 , 2017.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers amp; distillation through attention. In International Conference on Machine Learning , volume 139, pages 10347-10357, July 2021.

Gustav Larsson, Michael Maire, and Gregory Shakhnarovich. Fractalnet: Ultra-deep neural networks without residuals. arXiv preprint arXiv:1605.07648 , 2016.

Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2818-2826, 2016.

Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710.09412 , 2017.

Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVF international conference on computer vision , pages 6023-6032, 2019.

Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In International Conference on Machine Learning , pages 5389-5400. PMLR, 2019.

Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 15262-15271, 2021.

Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, et al. The many faces of robustness: A critical analysis of out-of-distribution generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 8340-8349, 2021.

Haohan Wang, Songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global representations by penalizing local predictive power. Advances in Neural Information Processing Systems , 32, 2019.

Tianyi Zhang, Felix Wu, Arzoo Katiyar, Kilian Q Weinberger, and Yoav Artzi. Revisiting few-sample bert fine-tuning. arXiv preprint arXiv:2006.05987 , 2020.

Mark Everingham, Luc Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman. The pascal visual object classes (voc) challenge. IJCV , 2010.

Yen-Cheng Liu, Chih-Yao Ma, Junjiao Tian, Zijian He, and Zsolt Kira. Polyhistor: Parameterefficient multi-task adaptation for dense vision tasks. arXiv preprint arXiv:2210.03265 , 2022.

Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. arXiv preprint arXiv:1903.12261 , 2019.

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022, 2021.

Enze Xie, Wenhai Wang, Zhiding Yu, Anima Anandkumar, Jose M Alvarez, and Ping Luo. Segformer: Simple and efficient design for semantic segmentation with transformers. Advances in Neural Information Processing Systems , 34:12077-12090, 2021.

James Seale Smith, Leonid Karlinsky, Vyshnavi Gutta, Paola Cascante-Bonilla, Donghyun Kim, Assaf Arbelle, Rameswar Panda, Rogerio Feris, and Zsolt Kira. Coda-prompt: Continual decomposed attention-based prompting for rehearsal-free continual learning. arXiv preprint arXiv:2211.13218 , 2022.

Zhizhong Li and Derek Hoiem. Learning without forgetting. IEEE transactions on pattern analysis and machine intelligence , 40(12):2935-2947, 2017.

Zifeng Wang, Zizhao Zhang, Chen-Yu Lee, Han Zhang, Ruoxi Sun, Xiaoqi Ren, Guolong Su, Vincent Perot, Jennifer Dy, and Tomas Pfister. Learning to prompt for continual learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 139-149, 2022.

Zifeng Wang, Zizhao Zhang, Sayna Ebrahimi, Ruoxi Sun, Han Zhang, Chen-Yu Lee, Xiaoqi Ren, Guolong Su, Vincent Perot, Jennifer Dy, et al. Dualprompt: Complementary prompting for rehearsal-free continual learning. arXiv preprint arXiv:2204.04799 , 2022.

James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al. Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences , 2017.

James Seale Smith, Junjiao Tian, Yen-Chang Hsu, and Zsolt Kira. A closer look at rehearsal-free continual learning. arXiv preprint arXiv:2203.17269 , 2022.

Sungyoon Lee, Jaewook Lee, and Saerom Park. Lipschitz-certifiable training with a tight outer bound. Advances in Neural Information Processing Systems , 33:16891-16902, 2020.

Kibok Lee, Kimin Lee, Jinwoo Shin, and Honglak Lee. Overcoming catastrophic forgetting with unlabeled data in the wild. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 312-321, 2019.

8 Appendix

8.1 Proof of Lemma 1

Prior works [9, 10] have used the high-level notion that staying 'close' to the pre-trained model can help maintain its robustness capability to justify using projection for fine-tuning. However, there is more than one way to encourage this, for example, regularization [35], a small learning rate [7], and projection [10]. It is not immediately clear why projection is a principled approach. To understand FTP's capability to maintain the pre-trained mode's robustness, we first propose to establish a connection between Lipschitz continuity, a commonly used measure of robustness [19, 20, 21], and fine-tuning through a new definition of difference function in the Lemma 1.

Proof. We first expand the difference functions in Eq. 9, i.e. plugging in ∆ h ( · ) = h f ( · ) -h 0 ( · )

,

Then we apply the reverse triangular inequality to the left-hand side of Eq. 11.

Therefore, we have,

Assuming that the pre-trained model h 0 is L 0 -Lipschitz, we know that ∥ h 0 ( x ) -h 0 ( x ′ ) ∥ h ≤ L 0 ∥ x -x ′ ∥ x , ∀ ( x , x ′ ) ∈ R m . Plug this into Eq. 12,

8.2 Proof of Lemma 2

In the previous section, we established a connection between the robustness of a fine-tuned model h f ( · ) and its difference function ∆ h ( · ) . Naturally, if we can limit the Lipschitz constant L d of the difference function, we can maintain the robustness of the pre-trained model. In this section, we show projection as an effective method to enforce the L d -Lipschitz condition in Eq 9.

Proof. Linear Operators. A neural network is composed of linear operators with connecting non-linear activations. Following prior works [9, 10], we analyze the linear operators 3 : h ( x ) = Wx + b , W ∈ R n × m , b ∈ R n . Let's define h f ( x ) = W f x + b f and h 0 ( x ) = W 0 x + b 0 , and plug them in Eq. 9.

Rearranging the above equation gives us an upper bound on L d ,

̸

Matrix Norms. Eq. 14 matches the definition of a matrix norm for a matrix W ∈ R n × m : ∥ W ∥ h , x = sup { ∥ W x ∥ h ∥ x ∥ x , ∀ x ∈ R n with x = 0 } . Therefore, to minimize L d in Eq. 9, we just need to minimize the matrix norm ∥ W f -W 0 ∥ h , x . Note that different vector norm combinations ( ∥ · ∥ h and ∥ · ∥ x ) will lead to a different matrix norm ∥ · ∥ h , x . Certain vector norm combinations have a closed-form matrix norm while the majority do not. Following prior works [9, 10], we use Maximum Absolute Row Sum (MARS) matrix norm, which is characterized by l ∞ vector norms in both domains.

Specifically, given a desired constraint L d , we want ∥ W f -W 0 ∥ ∞ , ∞ ≤ L d . Per the definition of the MARS matrix norm, which is the largest l 1 norm of each row of a matrix, the inequality can be equivalently enforced for each row independently, i.e.,

where w i denotes the i -th row of the matrix W .

Projection. To ensure the inequality in Eq. 15, we can project W f towards W 0 using the following projection equation. For each row w i in a matrix W , the projected weight ˜ w i p is calculated by

It is easy to check that w i p satisfies Eq. 15, i.e., ∥ w i p -w i 0 ∥ 1 ≤ L d if 0 ≤ γ ≤ L d .

Lipschitz Bound. Since a neural network is a composition of linear operators and non-linear activations, by the composition rule of the Lipschitz functions, an upper bound of the entire network is just the product of the Lipschitz constant for each linear operator and non-linear activations, where most non-linear activations are 1-Lipschitz [21]. However, the Lipschitz bound obtained by using the composition rule is not a tight bound on the entire network. While it is an active research area to find tighter bounds for neural networks without relying on the layer-wise composition rule [60, 20], the layer-wise approach is particularly suitable for connecting the fine-tuning process and Lipschitz continuity because it leads to layer-wise regularization techniques as we demonstrated above.

8.3 FTP: Additional Discussion

In the main paper Sec. 3.2, we described the algorithmic difference between TPGM and FTP. However, there is an implicit assumption made as a result of the difference. We now discuss the implications of it. After obtaining the updated constraints γ t in Eq. 5, if the algorithm were to follow TPGM, the next step would be applying the updated constraints to re-calculate the previous model W t -1 . However, instead of rolling back, FTP applies the updated constraints directly to the current unconstrained model ˜ W t . This step assumes smoothness in the update of γ t , i.e., the γ t does not change drastically in consecutive steps. The assumption is valid since γ t is updated by AdampUpdate (Alg. 2 below) which uses a moving average update with a momentum of 0 . 9 . So the change of γ t is very smooth because of the high discount factor of 0 . 9 . Importantly, we have re-used the same gradient g t available for computing the current unconstrained model ˜ W t . This is the key to saving computation because calculating the forward and backward pass through the model is the main computation bottleneck in TPGM because it requires a separate training loop as a result of 'rolling back'.

8.4 Image Classification Experiments Details and Additional Results

In Sec. 4.1.1, we presented image classification results on DomainNet100% data (111,307 images). Now we further present results using only 10% (11,031 images) of the training data in Tab. 6. In this case, projection-based methods, TPGM and FTP achieved the best performance, demonstrating their regularization capability under low-label conditions. Similar to findings in the main paper, FTP is

conv1.weight

layer1.0.bn2.bias

layer1.1.bn3.weight

layer2.0.conv2.weight

layer2.1.bn2.bias

layer2.3.bn1.weight

layer3.0.conv3.weight

layer3.1.bn3.bias

layer3.3.bn2.weight

layer3.5.conv1.weight

layer4.0.bn2.bias

layer4.1.bn3.weight

attnpool.q_proj.weight

0

10

20

30

40

50

60

1

7

13

19

25

31

37

43

49

55

61

67

73

79

85

91

97

103

109

115

121

127

133

139

145

Layers

Magnitude

Epoch

FTP Constraints for Each Layer in a ResNet50

up to 37% faster than TPGM during training. Next, we describe the hyper-parameters for all image classification experiments in Sec. 4.1 and above.

DomainNet. We use the released code from the prior work, TPGM [10] to train our FTP model. Therefore, we directly use the reported results from TPGM for competing methods. For FTP, we apply constraints to all trainable layers except for the last linear classification layers. For all experiments, we use SGD as the base optimizer with a weight decay of 5 e -4 . For DomainNet100% and DomainNet10% experiments, we train models for 50 and 150 epochs respectively with a batch size of 256. We sweep a range of learning rates and use the validation split to determine the best learning rate for FTP for each experiment. Here is the list of best-validated learning rates for all DomainNet experiments. We also provide a visualization of the learned constraints in Fig. 1b.

DomainNet100% MOCO-V3 ResNet50 (Tab. 1): 1 e -2

DomainNet100% CLIP ResNet50 (Tab. 2): 1 e -2

DomainNet10% CLIP ResNet50 (Tab. 6): 1 e -1

Note that we use the default κ = 1 for all these experiments. Every DomainNet experiment was conducted using 4 RTX 2080 GPUs.

ImageNet. For ImageNet experiments (Tab. 3, Fig. 3), we use a CLIP pre-trained ViT-Base [4]. Unlike the DomainNet experiments, we also initialize the last linear layer with zero-shot weights extracted from a CLIP text encoder, following the prior work WISE [8]. Therefore, FTP is applied to all trainable layers including the last linear layer. Training Transformers have been well-studied with abundant regularization and augmentation techniques. To obtain the best fine-tuning performance, we follow the public code base of DEIT [39] to fine-tune all methods. Specifically, we use weight-decay ( 0 . 1 ), drop-path ( 0 . 2 ) [40], label-smoothing ( 0 . 1 ) [41], Mixup ( 0 . 8 ) [42] and Cutmix ( 1 . 0 ) [43]. One exception is Linear-Probing (LP), where we do not use any of the above augmentations because they have been shown to degrade linear probing performance [3, 1]. We train all methods using AdamW [37] as the base optimizer with a weight decay of 0 . 1 , cosine learning rate schedule, and a batch size of 256 for 30 epochs. We also sweep relevant hyper-parameters for each method and document them below.

FT: learning rate 2 e -5

LP: learning rate 5 e -3

LP-FT: learning rate 2 e -5 . We take the best LP model (trained for 30 epochs) and then fine-tune it for another 15 epochs with the learning rate specified above.

L2-SP: learning rate 2 e -5 , regularization hyper-parameter 1 e -5 .

FTP: learning rate 3 e -5 , regularization hyper-parameter default κ = 1 .

Every ImageNet classification experiment was conducted on 2 A40 GPUs.

8.5 PASCAL Dense Vision Task Experiments Details and Additional Results

In Sec. 4.2, we presented results on semantic segmentation. In this section, we provide the additional results on semantic segmentation and surface normal estimation in Tab, 7 and Tab. 8. FTP achieves the best ID and OOD performance with significantly improved computation efficiency over TPGM [10]. Next, we will give more details on implementation.

Following prior works [50], we use a combination of Swin-Tiny Transformer [52] encoder and Segformer [53] decoder. The decoder is customized to allow different output formats. Only the Swin

segmentation,fog

segmentation,defocus

segmentation,brightness

70

Adapter

Vanilla FT

70

Adapter

Vanilla FT

75

70

60

L2-SP

BitFit

MARS-SP

LLRD

60

MARS-SP

Bitit

L2-SP

LLRD

70

TPGM

TPGM

65

65

50

FTP

FTP

mloU

40

Vanilla FT

30

60

Vanilla FT

60

Adapter

Adapter

BitFit

BitFit

L2-SP

MARS-SP

30

55

dS-21

MARS-SP

55

LLRD

TPGM

FTP

20

10

0

50

LLRD

FTP

TPGM

0

1

2

3

4

0

1

2

E

4

1

2

3

4

0

1

2

3

4

human parts,fog

human parts,defocus

human parts,gaussian

human parts,brightness

66

Vanilla FT

65

Vanilla FT

Vanilla FT

66

Vanilla FT

Adapter

Adapter

Adapter

Adapter

64

BitFit

BitFit

60

Bitit

64

BitFit

L2-SP

L2-SP

L2-SP

L2-SP

62

MARS-SP

LLRD

55

MARS-SP

LLRD

50

MARS-SP

LLRD

62

MARS-SP

LLRD

TPGM

TPGM

FTP

50

FTP

FTP

FTP

no

40

60

45

58

40

30

58

561

56

35

30

20

54

54

25

10

52

0

1

2

3

4

5

0

1

2

3

4

5

0

1

2

3

4

0

1

2

E

4

normals,fog

normals,defocus

normals,gaussian

normals,brightness

26

35.0

Vanilla FT

Vanilla FT

23

32.5

Adapter

BitFit

22

Bitit

Adapter

22

24

30.0

L2-SP

MARS-SP

21

L2-SP

MARS-SP

21

LLRD

LLRD

no

20

19

18

BitFit

Adapter

Vanilla FT

22

20

BitFit

Vanilla FT

Adapter

27.5

25.0

22.5

TPGM

FTP

20

19

18

FTP

TPGM

17

MARS-SP

L2-SP

18

L2-SP

MARS-SP

20.0

17

16

TPGM

FTP

LLRD

16

FTP

TPGM

LLRD

17.5

15.0

16

0

1

2

3

4

5

0

1

2

3

4

5

0

3

4

0

2

3

encoder is initialized with pre-trained weights (pre-trained on ImageNet-22k). Therefore, we only apply FTP to the encoder. For all methods, we use Adam as the base optimizer with a weight decay of 1 e -4 and a learning rate of 1 e -4 for 60 epochs. For methods with regularization hyper-parameters, we sweep a range of values and report the best one. We provide Tab. 9 for reference.

To test OOD robustness on the PASCAL-Context benchmark, we apply natural corruptions to the original clean images. Specifically, we select four types of corruptions from the popular benchmark [51], each of which is sampled from a main category: noise, blur, weather, and digital. Each corruption has five levels of severity. We report the average values over the five severity in our paper. Here, we also provide a detailed breakdown for each level of corruption in Fig. 6. Every PASCAL experiment was conducted on a single RTX 2080 GPU.

8.6 Continual Learning Experiments Details and Additional Results

In this section, we provide a brief overview of the settings in continual learning (CL). In CL, a model θ is trained on a sequence of task n ∈ { 1 , ..., N } . Each task has a non-overlapping set of class labels T n , and we denote the number of classes as | T n | . For ImageNet-R, we split the 200 classes into 10 tasks with 20 labels each, i.e., N = 10 , | T n | = 20 . Our experiments belong to the class-incremental category in CL. With each new task, the final linear classifier layer is expanded

with randomly initialized weights. We denote θ i, 1: n as the model that has been trained on i tasks and the classifier has all classes up to and including the n -th task ( i ≥ n ).

To measure global performance, we first define the global task accuracy A 1: N as,

where D test is the test dataset which has data from all N tasks and ˆ y ( x, θ ) denotes the predicted class from the model with weights θ . Then we define the global forgetting F N [61] as,

Following the prior work [59], all experiments in Tab. 5 use a ViT-Base pre-trained on ImageNet. We tune FTP with the code provided by the authors and directly compare it to the results from the prior work. Specifically, all methods use Adam as the base optimizer with no weight decay and a batch size of 128. All results are averaged over 3 random seed trials where the class allocation to each task is shuffled. For FTP, we train the model for 25 epochs with an initial learning rate of 5 e -4 and a cosine learning rate schedule. For all methods, we freeze the majority of the backbones and only fine-tune the QKV attention layers in the ViT. Please refer to the prior work for a more detailed description of the compared methods. Every CL experiment was conducted on 4 RTX2080 GPUs.

8.7 Pytorch Code Example of FTP

Here is an example of using SGDP (SGD+FTP) in Pytorch format. SGDP requires the common arguments for initializing an SGD optimizer class in Pytorch with two additional inputs: k and exclude_set. k is the hyper-parameter for positive gradient annealing (Sec. 3.2) and exclude_set contains the set of the names of parameters to be excluded from the projection operation. A complete demonstration of image classification is provided in the supplementary. You should be able to reproduce FTP results in Tab. 1 and Tab. 2.

where,

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a robust fine-tuning technique that has the decent scalability and efficiency to achieve higher ID and OOD performance when transferring a pre-trained model to downstream tasks. A new projection-based fine-tuning algorithm, Fast Trainable Projection (FTP) is proposed for computationally efficient learning of per-layer projection constraints with the theoretically analyses in the lens of Lipschitz continuity. Extensive experiments show its effectiveness.\n\nSTRENGTHS:\nThe paper is well-written and easy to read.\n\nThe proposed FTP is reasonable and easy to implement with the provided pseudo code. \n\nThe theoretical analysis is also appreciated.\n\nExtensive experiments on several downstream tasks and datasets are conducted to show the strong performances of the proposed method\n\nWEAKNESSES:\nWhile the method is simple with strong performances, the novelty is somewhat limited. It seems like a direct extension of TPGM. Using the previous model and cached gradients to update the constraints is trivial to me and makes limited technical contributions. Gradient annealing is also like a trick.\n\nCan the theoretical analyses prove that the proposed FTP is more powerful than TPGM? Since the experimental results show the superior performance of FTP.\n\nThe ablation study is insufficient. The ablation should be done on ProjUpdate and Gradient Annealing to show the effectiveness of each component.\n\nQUESTIONS:\nI understand FTP is more efficient than TPGM. But why FTP achieves better performance? Could the author give an intuitive explanation?\n\nFor the continual learning experiment, the comparison is somewhat not fair, since the proposed method needs to know the task id. (Page 9 line 288 we re-initialize FTP after each task and use the current model as the “pre-trained model” for the next task.)\n\nLIMITATIONS:\nWell addressed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: SUMMARY:
This paper introduces Fast Trainable Projection (FTP), a new algorithm aimed at improving the robustness and efficiency of fine-tuning pre-trained models. FTP achieves visible improvement in terms of computational speed and adaptability, demonstrated across various vision tasks and models, contributing to an average 35% speedup on benchmarks compared to prior works. The authors also provide a theoretical explanation for FTP's ability to maintain the robustness of pre-trained models through the lens of Lipschitz continuity.

STRENGTHS:
1. papers are well-written, figures are well-made and very easy to read
2. problem being studied is important -- the efficient finetuning of models are particularly relevant given the popularity of foundation models today.

WEAKNESSES:
1. Limited to Vision Tasks: The experiments are largely focused on vision tasks. Additional experiments in other application areas such as natural language processing could further demonstrate the algorithm's effectiveness.
2. Does the method work for foundation-level models? For example, in vision, there's stable diffusion, segment anything etc. There are already very mature finetuning pipeline and benchmarks for these models that can be experimented with. If the proposed method works for these models, the impact will be significantly increased.

QUESTIONS:
1. Does the proposed algorithm work on low-rank adaptation (LoRA)? Would be interesting to see related experiments for that.

LIMITATIONS:
The authors have addressed the limitations
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors propose an efficient fine-tuning method for deep learning models that learns the constraint for each layer more efficiently.\nDuring neural network training, the proposed method optimizes the layer-wise constraint using the last batch of data, unlike a prior method (TPGM) that updates layer-wise constraints on a \"validation\" dataset. This new method is less computationally burdensome since the gradient of the constraint can be computed using the chain rule, which reuses the gradient stored for the last batch.\nThe authors also propose a notation to describe fine-tuning robustness and provide theoretical justification for the proposed method under this notation. To demonstrate the effectiveness of the new method, the authors conduct a comprehensive list of experiments and show that the proposed method is not only more efficient but also has better performance.\n\nSTRENGTHS:\n-The writing in this paper is of high quality, as the authors have presented the methodology in a clear manner.\n\n-The related work section is well-discussed. The authors have categorized the existing fine-tuning studies into three categories: when, where, and how much to fine-tune. This is an insightful way to approach the current studies.\n\n-The authors have provided theoretical justification for the proposed method.\n\n-The experimental results are comprehensive, as the authors have included a list of different datasets and model architectures. The proposed method has been shown to outperform all baseline methods.\n\nWEAKNESSES:\n-My first concern is that the proposed method's novelty might be overshadowed due to its similarity to a previous method called TPGM. The major difference between the two seems to be the selection of batch data for updating the constraint parameter. The previous work used validation data, which may not be a common choice in the supervised learning setting, while this work uses the previous batch data. Although using the stored gradient indeed improves computational efficiency, the authors did not clearly state why this method has better performance than TPGM. The validation data and the other training data should have the same distribution.\n\n-The proposed difference function's meaning requires in-depth justification. While (Lipschitz) smoothness is usually associated with better performance and (adversarial) robustness for deep neural networks, it is not directly implied that the proposed difference function will improve the downstream performance of a fine-tuned model. Additionally, this notation only describes the difference in feature space; how does it describe the linear-probing fine-tuning method, which only changes the linear last layer? Can it be shown that this proposed notation indicates fine-tuning downstream performance?\n\n-The robustness notation here seems to be related to the OOD generalization capability. Are they the same thing here? Usually, the robustness is used to describe the model’s performance under some unwanted perturbation (adversarial samples or corruption)\n\n-The Low-rank fine-tuning methods have been really popular recently. Is it possible to compare the proposed method with them?\n\nMinor typo:\n\n Line 52, Liptschitz → Lipschitz\n\nLine 127, “Then FTP calculates the gradients”, is here FTP meant to be TPGM?\n\nQUESTIONS:\nWhat’s the pattern of learned constraint for each layer, with regard to training epochs? Do they converge during the training?\n\nI am really curious about the statement of discrepancies here. The author mentioned multiple times that the discrepancy between the validation data and training data or between the different batches of data enables the learning of meaningful projection constraints but without further elaborations. Typically, the training and validation data should be subject to the same distribution.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a new algorithm for robust fine-tuning of pre-trained models, specifically focusing on maintaining out-of-distribution (OOD) robustness while achieving competitive in-distribution (ID) performance. The algorithm, dubbed Fast Trainable Projection (FTP), is designed to overcome the scalability and efficiency limitations of current methods, offering a 35% speedup compared to prior work (to be specific, TPGM). FTP achieves this by efficiently learning per-layer projection constraints during the fine-tuning process.\n\nThe claimed contributions of this paper are:\n\n1. The introduction of the FTP algorithm that significantly improves computational efficiency while learning projection constraints and fine-tuning the model simultaneously. This algorithm can be integrated with existing optimizers like SGD and AdamW and can be adopted as a new drop-in fine-tuning optimizer.\n\n2. Empirical validation of the FTP algorithm's robustness on OOD datasets. They tested the algorithm across four vision tasks with five different pre-trained models, demonstrating robustness especially in scenarios with domain shifts and natural corruptions. The FTP algorithm also achieves state-of-the-art performance on a continual learning benchmark.\n\n3. Theoretical explanation of FTP's robustness maintenance capabilities. The authors provide a mathematical perspective that explains why FTP is effective at preserving the robustness of pre-trained models. They explore this through the lens of Lipschitz continuity, taking into account both the feature space and weight space of a model.\n\nSTRENGTHS:\nThe introduction of the Fast Trainable Projection (FTP) algorithm is an interesting contribution. It addresses the real-world limitations of scalability and efficiency associated with TPGM, making it a practical solution for various tasks.\n\nThe robustness of the FTP algorithm has been validated with experiments across four vision tasks and five pre-trained models. Superior results on OOD datasets and state-of-the-art performance on a continual learning benchmark provide empirical support for the authors' claims.\n\nWEAKNESSES:\n1. **Limited Scope:** the proposed method focuses on a specific Trainable projected gradient method (TPGM) and aims to improve its efficiency. However, the extent of its practical application could be questioned as it's not entirely clear how well it would generalize to other types of models or algorithms. Its significance might be seen as limited if it only optimizes a specific kind of model, especially when there are many alternative and possibly more efficient methods available.\n\n2. **Lack of Efficiency:** Despite the fact that the FTP improves on the TPGM's efficiency, it remains slower than many other methods, which raises questions about its practicality in real-world applications where computational resources and processing times are often crucial factors.\n\n3. **Theoretical Contributions:** The theoretical contribution of this paper may be seen as somewhat trivial, as the authors basically just formalize an intuitive understanding of Lipschitz continuity in relation to model robustness. The assumption that robustness equates to Lipschitz continuity may oversimplify the complex nature of robustness in deep learning models. Real-world robustness often depends on various other factors such as data variance, architecture, and loss landscape, which aren't addressed in this paper. Furthermore, they fail to delve deeper into the specifics of how the generalization relates to pre-trained weights, which might have offered more novel insights.\n\nQUESTIONS:\nWhat is the memory footprint and runtime overhead of the proposed method?\n\nHow does the proposed method affect generalization of neural networks?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents Fast Trainable Projection (FTP), an efficient fine-tuning algorithm. FTP learns a projection radius based on the current training batch. Through both analysis and experiments, the paper shows that FTP improves out-of-distribution robustness while maintaining in-distribution performance. Experiments also show that FTP can accelerate learning by up to 35% compared to previous methods.\n\nSTRENGTHS:\n- The idea seems to be a natural improvement following the lines of MARS-SP (project after unconstrained GD step) and TPGM (learn a different radius for each layer).\n- Section 3.3 was clear and interesting; it establishes a connection between weight-space projection and \"robustness\" in terms of the Lipschitz constant.\n- The experimental results are strong, across DomainNet, ImageNet, and PASCAL fine-tuning experiments.\n\nWEAKNESSES:\n- A main claimed benefit of FTP (in the abstract and intro) is its computational efficiency. To what extent is computational cost a bottleneck in the experimental settings you consider? I think fine-tuning is generally considered to be pretty computationally light, especially compared to the training of the foundation model itself. Given the experimental results, maybe the more relevant benefits are (1) better OOD acc (2) no need for held-out val data like TPGM.\n- While section 3.3 was interesting, I'm not sure if it contributes the to the main point of the paper (benefits of FTP over e.g TPGM), since this simplified analysis really applies to all projection-based methods. Does this analysis motivate FTP over other ways of using projection for fine-tuning?\nminor typo: Fig 2 caption PorjUpdate -> ProjUpdate\n\nQUESTIONS:\n- There seem to be similarities between FTP and the general hyperparameter tuning method proposed in https://arxiv.org/abs/1909.13371, in that both methods delay updates with the current batch to optimize hyperparameters in an online fashion (gamma here corresponds to alpha there). Could the authors comment on how these two methods relate?\n- see other questions in \"Weaknesses\" above\n\nLIMITATIONS:\nYes.\nNo particular negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}