Fast Trainable Projection for Robust Fine-Tuning

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.

Paper

Similar papers

Peer review

Reviewer yjP35/10 · confidence 5/52023-07-05

Summary

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

Strengths

The paper is well-written and easy to read. The proposed FTP is reasonable and easy to implement with the provided pseudo code. The theoretical analysis is also appreciated. Extensive experiments on several downstream tasks and datasets are conducted to show the strong performances of the proposed method

Weaknesses

While 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. Can the theoretical analyses prove that the proposed FTP is more powerful than TPGM? Since the experimental results show the superior performance of FTP. The ablation study is insufficient. The ablation should be done on ProjUpdate and Gradient Annealing to show the effectiveness of each component.

Questions

I understand FTP is more efficient than TPGM. But why FTP achieves better performance? Could the author give an intuitive explanation? For 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.)

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Well addressed.

Reviewer wzxc5/10 · confidence 2/52023-07-07

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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have addressed the limitations

Reviewer 2kDG5/10 · confidence 4/52023-07-08

Summary

In this paper, the authors propose an efficient fine-tuning method for deep learning models that learns the constraint for each layer more efficiently. During 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. The 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.

Strengths

-The writing in this paper is of high quality, as the authors have presented the methodology in a clear manner. -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. -The authors have provided theoretical justification for the proposed method. -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.

Weaknesses

-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. -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? -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) -The Low-rank fine-tuning methods have been really popular recently. Is it possible to compare the proposed method with them? Minor typo: Line 52, Liptschitz → Lipschitz Line 127, “Then FTP calculates the gradients”, is here FTP meant to be TPGM?

Questions

What’s the pattern of learned constraint for each layer, with regard to training epochs? Do they converge during the training? I 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.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Reviewer HDd65/10 · confidence 5/52023-07-12

Summary

The 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. The claimed contributions of this paper are: 1. 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. 2. 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. 3. 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.

Strengths

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

Weaknesses

1. **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. 2. **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. 3. **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.

Questions

What is the memory footprint and runtime overhead of the proposed method? How does the proposed method affect generalization of neural networks?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes.

Reviewer FZ4t6/10 · confidence 4/52023-07-27

Summary

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

Strengths

- 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). - Section 3.3 was clear and interesting; it establishes a connection between weight-space projection and "robustness" in terms of the Lipschitz constant. - The experimental results are strong, across DomainNet, ImageNet, and PASCAL fine-tuning experiments.

Weaknesses

- 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. - 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? minor typo: Fig 2 caption PorjUpdate -> ProjUpdate

Questions

- 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? - see other questions in "Weaknesses" above

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes. No particular negative societal impact.

Reviewer HDd62023-08-15

I have raised my score, to appreciate the empirical results of this paper. The fact that it is just an addon of a previous paper that is not widely used, remains unchanged. And I insist that the "theoretical" analysis part should be removed if the paper is finally accepted. It indeed adds little value to the paper.

Reviewer yjP32023-08-18

The rebuttal address my concerns and I would like to maintain my rate.

Reviewer 2kDG2023-08-20

Response

I would like to thank the authors' response as well as the additional experimental results. While I still hold my concern regarding some of the limitations I raised here, (most importantly, the discrepancies between training/validation batches), most of my concerns have already been addressed and I will increase my score and lean to acceptance. Thanks!

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC