Rethinking Weight Decay for Robust Fine-Tuning of Foundation Models

Authors: junjiao  tian
chengyue  huang
zsolt  kira
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 18:48:43

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

Main paper: 
 introduction: 
 Modern optimizers, such as Adam [1] , LARS [2] , and LAMB [3] usually include momentum and adaptive learning rates. They help optimizers avoid local minima and accelerate learning [4, 5] to explore wider parameter spaces. However, we hypothesize that this behavior is not always beneficial for fine-tuning from a well pre-trained foundation model, especially when fine-tuning a few layers is already sufficient for fitting the target data [6, 7, 8, 9] . Several prior works have found that unnecessary exploration will lead to large deviation from the initialization and worse robustness [10, 11] , and constraining the deviation can improve a model's generalization on in-distribution (ID) data and robustness to out-of-distribution (OOD) data [12, 13, 14] 1 . For example, L2-SP [13] imposes a regularization term on the distance between the current and pre-trained models. More recently, TPGM [10] and FTP [11] propose to learn different hard constraints for each layer. These new works have demonstrated impressive results on benchmarks. However, they are either difficult to tune, specialized to specific settings, or require significant computation and storage overhead. This motivates us to ask whether a simple few-liner solution exists for this fundamental problem.We propose re-examining the existing methods and summarizing their findings to find this solution. Starting from the simplest: L2-SP [13] . Specifically, L2-SP adds an L2 regularization term to the original objective function. Formally,L(θ) = L(θ) + λ 2 ∥θ -θ 0 ∥ 2 2 (1)where θ denotes the model parameters, θ 0 the initialization, L(θ) the original objective function, and λ the hyper-parameter for regularization strength. When θ 0 = 0, L2-SP reduces to an ordinary weight decay. This simple method should be effective enough to constrain the model, as our experiments show it can reduce the deviation between the fine-tuned and pre-trained models (Sec. 4.1). However, it is held back by an important design choice: the penalty is always applied to all model parameters. Our empirical results identify that a large λ prevents every layer from deviating too much and leads to poor fitting, while a small λ cannot provide enough regularization. This significantly limits the otherwise effective design (Sec. 4.1) . So, what is missing in this algorithm?Recent works in robust fine-tuning and parameter-efficient fine-tuning (PEFT) have shown that customizing constraints for each layer and selectively choosing layers for fine-tunning can improve robustness [10, 11, 6] . Inspired by these findings, we hypothesize that selectively imposing the regularization to different layers is the key. Therefore, we propose a simple selective version of L2-SP weight decay: selective projection decay (SPD). This new algorithm innovates in two aspects: a selection condition and a regularization strength ratio. The former determines when to apply regularization to a layer, and the latter determines the strength of regularization for intuitive hyperparameter tuning. Specifically, we derive the selection condition from hyper-optimization [15, 16, 17] by treating the condition as an optimizable parameter (Sec. 3.3), and the regularization strength ratio by re-writing L2-SP as a projection operation (Sec. 3.4) . Intuitively, when the condition is met, the algorithm imposes large regularizations on selected layers. This allows the algorithm to avoid unnecessary deviation and simultaneously fit into the fine-tuning data. We test SPD on large-scale computer vision, and NLP benchmarks with popular foundation models and test ID and OOD performance on various distribution and domain shifts. SPD achieves SOTA performance while being much simpler than other competing methods. Our contributions are:• We propose a selective projection decay, a selective variant of the popular L2-SP/weight decay regularization methods, for robust fine-tuning of large foundation models. We show that selectivity is important to make regularization effective. • We conduct a detailed study of ID and OOD performance on image classification and semantic segmentation with natural distribution and domain shifts. SPD improves ID and OOD performance on these benchmarks. • We show that SPD consistently improves the performance of PEFT methods (e.g. LoRA [7] and adapters [9] ) on 8 common sense reasoning language tasks with LLaMA-7B (-13B).Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes]Justification:Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related works: 
 Robust Fine-Tuning with Distance Regularization. Constraining the distance or deviation between the fine-tuned and pre-trained models has been studied in several prior works. L2-SP [13] explicitly adds an L2 norm penalty on the deviation and shows improved ID generalization for fine-tuning. MARS-SP [12] studies different forms of norms as the penalty. It shows that the Matrix Row Sum (MARS) norm can be a superior alternative to the L2 norm. These two methods impose "soft" penalties and can be less effective [18] . Instead, LCC [18] proposes constraining the deviation through direct projection on the parameters, which also enforces a hard constraint on the Lipschitz continuity of the fine-tuned model. However, LCC is hard to tune because the projection radius is not an intuitive hyper-parameter. Furthermore, using a single projection constraint for all layers is not an ideal strategy [10] . More recently, TPGM [10] proposes to automatically learn the constraints in LCC during fine-tuning, customizing a different projection radius for each layer through a bilevel optimization scheme. FTP [11] further improves the computation efficiency of TPGM by adopting hyper-optimization [15, 16, 17] in its computation. Nevertheless, FTP is still difficult to control because hyper-optimization requires a secondary optimizer with additional optimization hyper-parameters, and the learned regularization can be too strong with no intuitive way to adjust. In contrast, SPD is a much simpler and more intuitive method, which can be implemented with just a few lines of code. The superior controllability makes SPD potentially applicable to more applications.Parameter Efficient Fine-Tuning (PEFT). PEFT methods such as adapters [9, 8] and LoRA [7] have been proposed to reduce training memory usage and computation complexity. Recent works have found that PEFT methods also provide good robustness because they modify fewer parameters and retain more knowledge of the pre-trained models [11] . Surgical fine-tuning [6] concludes that fine-tuning a selective few layers can improve ID generalization. These new works motivate us to re-evaluate L2-SP and weight decay, often uniformly applied to all layers. We identify that the inferior performance of the simple methods is because of this uniformity, which exhibits a strong trade-off between fitting and regularization. Other robust fine-tuning methods, such as LP-FT [19] and FLYP [20] , focus on feature distortion. We will review them in the Appendix 8.1.Other Robust Fine-Tuning Methods. WiSE-FT [14] discovers that linearly interpolating between the fine-tuned and pre-trained models after fine-tuning can improve out-of-distribution robustness. This demonstrates that a closer distance to the pre-trained model can improve robustness. However, it only applies to models with zero-shot capabilities. Another orthogonal line of research for robust finetuning focuses on feature distortion. LP-FT [19] shows that fine-tuning with a randomly initialized head layer distorts learned features. It proposes a simple two-stage method to train the head layer first and then fine-tune the entire model. FLYP [20] shows that fine-tuning a foundation model using the same objective as pre-training can better preserve the learned features. Our contribution is an optimization method to penalize the derivation between the fine-tuned and pre-trained models explicitly during fine-tuning, which is orthogonal to them.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification:Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methods: 
 In this section, we first provide an overview of the Selective Projection Decay (SPD) method and then describe the intuition behind SPD with a numerical example. Finally, we provide a concrete mathematical motivation for our method's algorithmic design.Formulation. SPD is a regularization technique that penalizes significant deviation from the pretrained model. We motivate the formulation from an existing method: L2-SP [13] (Eq. 1). L2-SP adds a distance penalty on the deviation between the fine-tuned and pre-trained models. The penalty is applied to all model parameters at all times. A large λ prevents every layer from deviating too much and empirically leads to poor fitting, while a small λ cannot provide enough regularization. This significantly limits the otherwise effective design. We propose a selective version of this simple technique: selective projection decay (SPD). We will examine L2-SP and SPD in Alg. 1 and Alg. 2.Notations. We follow the notations in prior works [1, 21] . Let m t , v t denote the moving average of the gradient and squared gradient, β 1 , β 2 their hyper-parameters, and α the learning rate.Alg. 1 shows the Adam optimizer with the L2-SP regularization in Eq. 1. The effects of the regularization are highlighted in blue, also shown in Eq. 2. Intuitively, the regularization leads to an interpolation-like equation 2 . If the product λα = 1, then θ t ← θ 0 and if λα = 0, then θ t ← θt , where θ 0 and θt denote the initialization and the updated model without regularization.EQUATIONAlg. 2 shows the proposed SPD. There are two changes compared to Alg. 1.• a condition,EQUATION• a new interpolation-like equation with a multiplier, r t , replacing the learning rate α,EQUATIONCompared to L2-SP, SPD only imposes a penalty when the condition is met (c t < 0), and the strength of the penalty is controlled by a hyper-parameter λ and an analytical quantity deviation ratio r t , which we will introduce later.Algorithm 1: Adam with L2-Regularization Initialize m 0 ← 0, v 0 ← 0, t ← 0 While θ t not converged t ← t + 1 g t ← ∇ θ L(θ t-1 ) m t ← β 1 m t-1 + (1 -β 1 )g t v t ← β 2 v t-1 + (1 -β 2 )g 2 t Bias Correction m t ← mt 1-β t 1 , v t ← vt 1-β t 2 Update θt ← θ t-1 -α mt √ vt+ϵ θ t ← θt -λα( θt -θ 0 ) Algorithm 2: Adam with Selective L2-Reg. Initialize m 0 ← 0, v 0 ← 0, t ← 0,c 0 ← 0 While θ t not converged t ← t + 1 g t ← ∇ θ L(θ t-1 ) m t ← β 1 m t-1 + (1 -β 1 )g t v t ← β 2 v t-1 + (1 -β 2 )g 2 t Bias Correction m t ← mt 1-β t 1 , v t ← vt 1-β t 2 Update θt ← θ t-1 -α mt √ vt+ϵ c t = -g ⊺ t (θ t-1 -θ 0 ) If c t < 0 : θ t ← θt -λr t ( θt -θ 0 )3.2 Intuition Behind SPD SPD prioritizes layers with consistent improvement. SPD adds regularization on layers that meet the condition c t < 0 to slow their growth. The condition is determined by the sign of the inner product between two vectors. One vector is the negative gradient direction (-g t ), i.e., the descent direction, and the other is the current progress direction (θ t-1 -θ 0 ). The inner product between them measures the alignment between the vanilla3 update direction and the progress so far. When the inner product is positive, the current progress direction generally points to a low loss region, and following it will lead to consistent loss reduction. Conversely, if the inner product is negative, the current progress direction will likely lead to a higher loss region, indicating inconsistent improvement. In this case, SPD will impose a penalty to slow down updates for those layers. Recall that modern optimizers use momentum to escape local minima and explore wider regions. Without this penalty, the model will likely head towards the higher loss region to overcome it. SPD chooses to slow down these layers and prioritizes layers with more consistent loss reduction. We will motivate this strategy in a principled manner and validate it in our experiments.Previously, we explained the intuition behind SPD. Specifically, we interpreted the condition c t as a measure of alignment and a test of update consistency. Nevertheless, there is a more profound reason why the quantity c t is a natural choice for selective regularization. In this section, we motivate SPD from a more mathematical perspective.Hyper-Optimization Setup. Hyper-optimization is a technique to optimize hyper-parameters inside an optimizer [15, 16, 17] . They treat the hyper-parameters as trainable parameters and optimize them using another gradient-based optimizer. Let's start from the vanilla Adam with L2-SP algorithm (Alg. 1) and treat the regularization strength hyper-parameter λ as a trainable parameter. To update λ, we need to obtain its gradient by taking a derivative w.r.t. λ after applying it.EQUATIONSelection Condition c t . Intuitively, if the quantity ∇λ is negative, applying the update in gradient descent will increase the value of λ, thus increasing the regularization strength of L2-SP. Conversely, a positive quantity will decrease the regularization strength. Therefore, the sign(-g ⊺ t+1 ( θt -θ 0 )) determines the change of regularization strength in the hyper-optimization of λ. Formally, we define the condition c t as,EQUATIONFor memory efficiency, we use (θ t -θ 0 ) instead of ( θt -θ 0 ) because both vectors point in the same direction and won't affect the sign of c t . This allows us to discard θt . Otherwise, we need to keep an additional copy in memory. In summary, when c t < 0, we apply a regularization for that layer as shown in Alg. 2. This calculation is done for each layer, and the regularization is selectively applied.Alternative Interpretation: We just interpreted the selection condition c t in SPD as a measure of consistency between the current heading direction and the gradient direction. This perspective is more valid when the algorithm has accumulated some updates, i.e., ∥θ t -θ 0 ∥ 2 ≫ 0, and less justified when a heading has not been established at the beginning of training. To analyze this, we discuss the behavior SPD from the perspective of stochastic optimization when ∥θ t -θ 0 ∥ 2 is small at the beginning of training in the Appendix 8.1.The selection condition c t determines when to apply regularization to which layers. However, one remaining question is the strength of regularization, which is not intuitive to tune. To overcome this, we introduced an analytical quantity, the deviation ratio r t , in Eq. 2 and Alg. 1. In this section, we will motivate it from the perspective of projection.L2-SP is projection. Projection onto a norm ball is common in constrained optimization. While L2-SP is not a constrained optimization problem, its operation bears similarity to projection. Suppose we project a model θt to an L 2 -norm ball with radius γ centered around its initialization θ 0 . The equation of projection is the following,EQUATIONEquivalently, we can rewrite the equation as,EQUATIONNow, we can equate this equation to the highlighted L2-SP equation in Eq. 2 and Alg. 1, we can see that if λα = 1 -γ max{γ,∥θt-θ0∥2} , the regularization is equivalent to projection with radius γ. Deviation Ratio r t . This equivalence inspires us to define a deviation ratio r t :EQUATIONwhere γ t := ∥ θt -θ 0 ∥ 2 and γ t := ∥θ t-1 -θ 0 ∥ 2 denote the current deviation (before regularization) and the previous deviation from the initialization θ 0 , respectively. We use r t in SPD (Alg. 2) to replace the learning rate α in L2-SP (Alg. 1) to make hyper-parameter (λ) tuning more intuitive. Specifically, suppose the hyper-parameter λ = 1, then the regularization in SPD is:EQUATIONIntuitively, with λ = 1, the regularization in SPD is equivalent to projection with a radius equal to the previous deviation if the current deviation is larger. In summary:• No regularization (λ = 0): the projection radius is ∥ θt -θ 0 ∥ 2 , meaning no projection.• Weak regularization (1 ≥ λ > 0): the projection radius lies between ∥ θt -θ 0 ∥ 2 and ∥θ t-1 -θ 0 ∥ 2 . Within this range, all layers will expand or remain unchanged.• Strong regularization (λ > 1): the projection radius lies between 0 and ∥θ t-1 -θ 0 ∥ 2 . In this range, it's possible that regularized layers can contract.We recommend starting with λ = 1 and adjusting the strength according to the specific needs. As shown in Alg. 2, SPD retains a copy of the pre-trained model in memory. This adds additional memory requirements to the overhead of vanilla optimizers. While this is practical for moderate-sized models, as fine-tuning focuses more and more on large models, additional memory requirements become undesirable. Fortunately, in extremely large models, the prevalent fine-tuning strategy is parameter-efficient fine-tuning (PEFT), such as LoRA [7] , series adapters [9] , and parallel adapters [8] . SPD is naturally compatible with these methods without the additional memory. Intuitively, SPD selectively projects the current model towards the pre-trained initialization. PEFT methods generally initialize new parameters to add to the original model weights. To recover the behavior of SPD, we can instead project the new parameters towards the origin, equivalent to a selective version of regular weight decay, i.e., replacing θ 0 with 0 in Alg. 2. Consequently, this does not require a memory copy of the pre-trained model. It consistently improves PEFT fine-tuning for large language models on common sense reasoning benchmarks in Sec. 4.4.For example, LoRA decomposes a linear layer h = W t x into two components, where h ∈ R m×1 , W t ∈ R m×n and x ∈ R n×1 are the output, weights, and input of this layer.EQUATIONwhere W 0 ∈ R m×n , W up ∈ R m×r and W down ∈ R r×n are the pre-trained model, up-projection and down-projection matrices. If r ≪ min{m, n}, (W up W down ) is a low-rank approximation of ∆W t .To regularize the overall deviation ∥W t -W 0 ∥ 2 , it suffices to regularize ∥W up W down ∥ 2 to be close to zero. In this case, SPD acts as selective weight decay on W up and W down individually.In summary, we propose selective projection decay (SPD) to impose strong regularization on layers during fine-tuning selectively. As shown in Fig. 1 , SPD regularizes the deviation of the fine-tuned model from the pre-trained model ∥W t -W 0 ∥ 2 for full fine-tuning and the deviation from the origin ∥∆W t ∥ 2 for PEFT fine-tuning.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 We test Selective Projection Decay on a diverse set of benchmarks, architectures, and tasks to demonstrate its effectiveness. We will test both ID generalization and OOD robustness across various domain and distribution shifts.Image Classification. We first analyze the behavior of SPD on conventional image classification datasets DomainNet [22] and ImageNet [23] . We use a CLIP ViT-Base model for both experiments as the pre-trained initialization [24] . Specifically, DomainNet consists of images from several domains with 345 classes. We fine-tune on one domain and test on all domains. ImageNet is a large-scale dataset with 1000 classes. We fine-tune on ImageNet and test on ImageNet and four variants, namely ImageNet-V2 [25] , ImageNet-A [26] , ImageNet-R [27] , and ImageNet-S [28] .Semantic Segmentation. We further test SPD on the PASCAL-Context semantic segmentation dataset [29] . Following prior works [30, 11] , we use a Swin ViT-Tiny [31] , pre-trained on ImageNet-22K, and Segformer [32] segmentation architecture. To construct the OOD datasets, we follow the popular natural robustness literature [33] and apply four representative image corruptions (fog, defocus blur, Gaussian noise, and brightness) with 5 severity each. We fine-tune on the clean segmentation data and test on clean and corrupted data.Common Sense Reasoning. Moreover, we show that SPD can benefit PEFT fine-tuning on large language models (LLMs). We use the Commonsense-170K dataset [34] , which consists of training data from eight common sense reasoning benchmarks. Following the prior work [34] , we fine-tune LLaMa-7B (-13B) [35] using LoRA [7] , series adapters [9] , and parallel adapters [8] .Visual Question Answering. Finally, we demonstrate SPD's superiority on multi-modal task. We use Google's recently released PaliGemma [36] pretrained on a broad mixture of large-scale visionlanguage tasks. We fine-tune on VQAv2 [37] and test on nine OOD datasets using LoRA [7] .For the near OODs, we evaluate on VQAv2's six variants, namely IV-VQA [38] , CV-VQA [38] , VQA-Rephrasings [39] , VQA-CP v2 [40] , VQA-CE [41] and AdVQA [42] , which cover uni-modal, multi-modal and adversarial distribution shifts from VQAv2. We also include TextVQA [43] , VizWiz [44] and OK-VQAv2 [45] , which are constructed from different sources than VQAv2, as the far OOD datasets. In this section, we utilize the DomainNet benchmark to test our claims. Specifically, we will show that Adam-SPD consistently outperforms AdamW in OOD robustness across multiple domains, and this is due to a much smaller deviation from the pre-trained model. Furthermore, by sweeping across a range of hyper-parameters, we show that uniform regularization, such as L2-SP fails to provide adequate constraints, while Adam-SPD shows robust performance.Small deviation correlates with better OOD performance. Earlier, we hypothesized that a significant deviation can lead to worse OOD performance. Theoretically, prior works [12, 11] have shown that large deviations from the initialization result in a large Liptschtz constant and, hence, worse robustness. In Tab. 1, we present a comprehensive study by fine-tuning a pre-trained CLIP model on different domains from DomainNet separately and reporting test results on all domains. Across all domains, Adam-SPD consistently achieves higher OOD performance and shows noticeably less deviation. This empirical result corroborates with prior works' findings and our hypothesis.Selective regularization exhibits stronger deviation-robustness correlation. In Tab. 2, we compare the behavior of L2-SP and SPD using DomainNet. Specifically, we fine-tune a CLIP model on the Clipart domain (ID domain) and report performance on Clipart and other domains (OOD domains). In Tab. 2a, we observe that while L2-SP successfully restrains the model's deviation from its initialization, it does not effectively improve OOD performance. With a very large regularization, the ID performance deteriorates as well. On the contrary, SPD effectively restrains the model's deviation and significantly improves OOD performance while matching the best ID performance. Under SPD, the correlation coefficient between OOD performance and deviation is -0.96, which indicates a strong negative correlation between the two quantities, i.e., smaller deviation and higher OOD accuracy. This experiment shows that selective regularization is superior to uniform regularization.Training Details. We use the vision transformer public repository for DEIT [46] to fine-tune all methods. We use λ = 1 for all Adam-SPD results in Tab. 1. More details are in Appendix 8.4. Training Details. For Adam-SPD, we fine-tune the model with a learning rate of 3e -5 and λ = 1.4.The regularization hyper-parameter is found through cross-validation, and the model with the best ID validation accuracy is taken. More details are in Appendix 8.4. SPD is compatible and consistently improves PEFT methods. Previous experiments have shown that SPD imposes effective regularization for full fine-tuning. Furthermore, SPD can also improve the performance of PEFT methods. We fine-tune LLaMa-7B (-13B) models on the Commonsense-170k dataset [34] . As shown in Tab. 5, SPD consistently improves regular fine-tuning with AdamW, which uses a uniform weight decay for all tested PEFT methods. This demonstrates that selective regularization benefits full fine-tuning and PEFT fine-tuning. Combined with its simplicity, SPD can potentially improve generalization and robustness for more tasks in deep learning.Training Details. We follow the training code released by a prior work [34] . We report the best performance from the original paper and compare them with Adam-SPD. More details are in Appendix 8.4. Training Details. For Adam-SPD, we fine-tune the model with a learning rate of 1e -3 and λ = 0.5.The regularization hyper-parameter is found through cross-validation, and the model with the best ID validation accuracy is taken. More details are in Appendix 8.4. 
limitations: 
 SPD is a selective regularization technique explicitly designed for fine-tuning. While it can be theoretically used for pre-training, it will likely lead to poor performance because it will hinder the training of some layers. For fine-tuning, it works well because the pre-trained foundation model is assumed to be a good initialization, and only small changes in a selected few layers can lead to a good local minimum. Furthermore, the level of performance gain depends on how well the foundation models are exposed to the fine-tuning and OOD data distributions during pre-training. For example, in the DomainNet experiment (Tab. 1), fine-tuning a CLIP ViT model on any other domain does not have reasonably good OOD robustness on the Quickdraw domain. One can deduce that Quickdraw is not well represented in the pre-training data of CLIP ViT. 
conclusion: 
 Fine-tuning differs from training from scratch because it starts from a good initialization. Therefore, effective regularization is critical to retaining the knowledge of the pre-trained foundation model while fitting a model to the target distribution. We identified that 1) regularization is necessary to keep the fine-tuned model close to its initialization and maintain robustness; 2) uniform regularization can hurt model fitting if regularization is too strong. In this paper, we proposed selective projection decay (SPD), a selective version of the popular weight decay/L2-SP regularization method. With an additional few lines of code, SPD can be integrated into existing optimizers and performs selective regularization. It demonstrates superior regularization performance on different tasks and modalities in our experiments. 
acknowledgement: 
 This work was supported by ONR grant N00014-18-1-2829.8 Appendix 
extended related works: 
 Other Robust Fine-Tuning Methods. WiSE-FT [14] discovers that linearly interpolating between the fine-tuned and pre-trained models after fine-tuning can improve out-of-distribution robustness. This demonstrates that a closer distance to the pre-trained model can improve robustness. However, it only applies to models with zero-shot capabilities. Another orthogonal line of research for robust finetuning focuses on feature distortion. LP-FT [19] shows that fine-tuning with a randomly initialized head layer distorts learned features. It proposes a simple two-stage method to train the head layer first and then fine-tune the entire model. FLYP [20] shows that fine-tuning a foundation model using the same objective as pre-training can better preserve the learned features. Our contribution is an optimization method to penalize the derivation between the fine-tuned and pre-trained models explicitly during fine-tuning, which is orthogonal to them.In previous sections, we interpreted the selection condition c t in SPD as a measure of consistency between the current heading direction and the gradient direction. This perspective is more valid when the algorithm has accumulated some updates, i.e., ∥θ t -θ 0 ∥ 2 ≫ 0, and less justified when a heading has not been established at the beginning of training. This section discusses SPD from the perspective of stochastic optimization when ∥θ t -θ 0 ∥ 2 is small at the beginning of training.Inner product of gradients captures gradient variance. Modern deep learning models are trained by stochastic optimization techniques, e.g., mini-batch SGD, leading to stochasticity due to sampling. We first show that the inner product of gradients captures the variance of a sampling process. We invoke a common assumption in the convergence analysis of stochastic gradient descent [1, 49, 21] .Assuming that the stochastic gradient g t is a stationary process G over a short period, with a small step size, successive gradients, e.g., g t , g t+1 , can be seen as samples drawn from the same distribution G. Given two successive draws g 1 and g 2 , we can approximate the first and second moment of G.EQUATION)Define the variation of gradients as V ar(g) := E ∥g -ḡ∥ 2 [50, 51] , where ḡ := E[g],we can show thatEQUATION≈ ∥ḡ∥ 2 -E ∥g∥ 2 -∥E [g] ∥ 2 = ∥ḡ∥ 2 -V ar(g)Remarks. Eq. 13 shows that the inner product of two consecutive stochastic gradients, under certain assumptions, can be seen as the estimator for the difference between the gradient norm and the variance of gradients. When the inner product is negative, this indicates that the variance outweighs the magnitude of the gradient.SPD prioritizes layers with higher expected gain. At the beginning of training, the heading direction (θ 1 -θ 0 ) is dominated by early gradients. For example, at t = 2 the direction of (θ 1 -θ 0 ) is the same as -g 1 in Adam. The sign of -g ⊺ 2 (θ 1 -θ 0 ) is the same as the sign of g ⊺ 2 g 1 . This shows that the condition c t captures the difference between gradient norm and gradient variance. With this interpretation, we show that c t reflects expected performance gain in stochastic optimization. To see it, we can invoke the descent lemma for SGD. For an L-smooth function f (W ) [50] , the descent lemma for SGD states that,Lemma 1. E[f (θ k+1 )] -f (θ k ) Expected Performance Gain ≤ -η k (1 - η k L 2 ) ≤0 ∥ḡ k ∥ 2 + η 2 k L 2 ≥0 V ar(g k ),where η k ≤ 2 L is the learning rate. Remarks. The term on the left hand side E[f (θ k+1 )] -f (θ k ) is the expected performance improvement for each step. Ideally, this should be a negative quantity. On the right-hand side, we observe that improvement depends on two quantities ∥ḡ k ∥ 2 and V ar(g k ). To lower the upper bound, we want a large ∥ḡ k ∥ 2 and a small V ar(g k ). According to the decoupling Eq. 13, the inner product between successive gradients approximates this proportionality. Consequently, a negative c t likely indicates a higher upper bound on the expected gain, meaning a smaller improvement. Therefore, SPD will prioritize layers with potentially larger expected gains.8.3 Approximation in L2-SP Algorithm 3: (Ours) Adam with L2- Regularization Initialize m 0 ← 0, v 0 ← 0, t ← 0 While θ t not converged t ← t + 1 g t ← ∇ θ L(θ t-1 ) m t ← β 1 m t-1 + (1 -β 1 )g t v t ← β 2 v t-1 + (1 -β 2 )g 2 t Bias Correction m t ← mt 1-β t 1 , v t ← vt 1-β t 2 Update θt ← θ t-1 -α mt √ vt+ϵ θ t ← θt -λα( θt -θ 0 ) Algorithm 4: (Original) Adam with L2- Regularization Initialize m 0 ← 0, v 0 ← 0, t ← 0 While θ t not converged t ← t + 1 g t ← ∇ θ L(θ t-1 ) m t ← β 1 m t-1 + (1 -β 1 )g t v t ← β 2 v t-1 + (1 -β 2 )g 2 t Bias Correction m t ← mt 1-β t 1 , v t ← vt 1-β t 2 Update θt ← θ t-1 -α mt √ vt+ϵ θ t ← θt -λα(θ t-1 -θ 0 )The Adam with L2-SP Regularization algorithm in the main paper is not the precise mathematical implementation of the original formulation written in Eq. 1. To see the difference, we compare ours and the accurate implementation here in Alg. 3 and Alg. 4. In our implementation, we replaced θ t-1 -θ 0 (Alg. 4) with θt -θ 0 (Alg. 3). This is done intentionally to improve memory efficiency when transitioning to the selective version (see Adam-SPD in Sec. 3.3). We can make the following substitution to see how this modification changes computation. Starting from our implementation,EQUATION= θt -λα(θ t-1 - α m t √ v t + ϵ -θ 0 ) = θt -λα(θ t-1 -θ 0 ) + λα 2 m t √ v t + ϵ .We can further combine the λα 2 mt √ vt+ϵ into the update of θt . The new θt isEQUATION= θ t-1 -(1 -λα) α m t √ v t + ϵTherefore, our implementation of L2-SP adds a minor additional dampening of the learning rate α by a factor of (1 -λα).What if we followed the original implementation of L2-SP as in Alg. 4? This would change the condition c t in the main paper (Eq. 3) toEQUATIONAt the current time step t, we only have access to the parameters θ t-1 from the previous step t -1. To calculate the c t in Eq. 16, we would have to store the weights from two steps back in memory. This increases memory consumption of the algorithm. As we have shown in Eq. 15, our implementation only differs from the original implementation slightly but reduces memory consumption. Therefore, we decided to make this approximation.NeurIPS Paper Checklist 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [No] Justification: Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In this paper, ID generalization and OOD robustness refer to test accuracy on the fine-tuning distribution and robust accuracy on other shifted distributions, respectively.Mathematically, this is not the precise formulation of L2-SP as written in Eq. 1. See the Appendix 8.3 for further discussion.the direction w/o momentum. 


Decision: The paper introduces a regularisation technique for finetuning large pretrained models by imposing penalties on specific layers to enhance both ID and OOD generalisation. 

The reviewer consensus is an acceptance. They appreciate the simplicity of the method, clear presentation and strong empirical performance. There were some remaining reservations:
- CVio: I still believe that increased analysis of the proposed method (from a more theoretical perspective rather than completely empirical) would be more meaningful given that the nature of the work.
- Jrpf: I agree with the other reviewer that the theoretical perspective is lacking in this work, though the empirical results are strong.

The merits greatly outweigh the concerns in this case. We recommend an acceptance.

Reviews: review: paper_summary: paper_summary: This work proposes a new regularization scheme called Selective Projection Decay (SPD), which selectively imposes penalties on layers where the current progress direction does not align with the vanilla update direction. This method is compatible with PEFT methods, including LoRA-type algorithms, making it practical for end users. Experiments show that ADAM equipped with SPD consistently outperforms baselines, including L2 regularization, in terms of in-distribution generalization and out-of-distribution robustness on multiple popular vision and language benchmarks.
strength_weakness: strength_weakness: No obvious weaknesses are observed. One minor issue is that the performance improvement, while consistently observed, is not significantly high. Additionally, I have noted some unclear justifications in the manuscript, as described in the questions section.
questions: questions: #1. How does the performance compare when using EWC regularization for fine-tuning, since EWC is generally better than L2 regularization?

#2. Comparing the performance of L2-SP and Adam-SPD in hyper-parameter sweeping experiments seems challenging since these are reported in tables. Perhaps heatmaps could be a clearer alternative to visualize and highlight the results.

#3. The authors argue that SPD selectively imposes penalties on large deviations in a layer-wise manner. However, in Algorithm 2, there is no layer-wise information. It appears that all parameters are updated based on the alignment score denoted by c_t. From my understanding of Algorithm 2, if the alignment is high (i.e., c_t is smaller than zero), the updated parameters are further adjusted by the interpolation-like equation. If c_t is larger than zero, do we use the vanilla ADAM update, since \theta_t has already been updated before calculating c_t? Am I correctly understanding Algorithm 2? In such cases, it seems that the parameters shouldn’t be updated, since SPD will impose a penalty to slow down updates for those layers with negative c_t.

#4. In the LLaMA PEFT experiments, there is no baseline with L2 regularization or its advanced versions (such as EWC). While I expect that ADAM with SPD is better than AdamW, I would like to see if there are also improvements compared to ADAM with L2 regularization in PEFT settings.
limitations: limitations: The authors state that the proposed approach only works well if the pretrained model is good enough to provide good initial parameters. No clear negative societal impacts have been observed.
paper_summary: paper_summary: This paper proposes a novel weight decay strategy called Selective Projection Decay (SPD).
SPD selectively imposes a stronger penalty on certain layers, and is designed to improve both in-distribution (ID) and out-of-distribution (OOD) performance.
The paper demonstrates the effectiveness of SPD through experiments on several benchmarks.
strength_weakness: strength_weakness: - Overall, I think this is a strong paper, and its strengths outweigh its weaknesses.
I think the paper could benefit from an ablation experiment, for example, trying Adam-SPD without the condition or benchmarking SPD's performance with optimizers other than Adam.
- [1] reports that ensembling with the initial weights (whether in model space or weight space) is a simple strategy that improves OOD robustness, and a few recent works [2, 3] report that ensembling continues to improve performance after their proposed fine-tuning strategies. Could the authors validate how the benefits from SPD extend to the ensembling setting?

[1] Robust fine-tuning of zero-shot models
[2] Finetune like you pretrain: Improved finetuning of zero-shot vision models
[3] AutoFT: Learning an Objective for Robust Fine-Tuning
questions: questions: Please see the weaknesses section.
limitations: limitations: Please see the weaknesses section.
paper_summary: paper_summary: The authors propose Selective Projection Decay, a modification on the AdamW optimizer where a regularization penalty is applied that controls for an overall drift from the prior weights such that as long as the overall shift is controlled for but individual layers can potentially adapt differently. The authors provide experiments on OOD generalization tasks and PEFT language modelling fine-tuning.
strength_weakness: strength_weakness: - The way the authors frame it, there appears to be limited novelty with the proposed method, with it appearing to simply involve an additional conditional check before applying a potential regularization penalty that has been explored quite a bit in existing literature. I do think that there is something of merit here, however it's hard to pinpoint it under the current framing of the work.
- The statement "We recommend starting with  $\lambda= 1$ and adjusting the strength according to the specific needs" is rather unsatisfying and doesn't provide the reader with a good idea of whether or not $\lambda$ is easy to tune. This can probably be resolved by adding additional analysis along this end compare do just Table 2.
- Experimentally, the results feel somewhat lacking and somewhat misleading. For example, in the Domain-Net experiments, the only comparison is with AdamW, which inherently should lack the exploration capabilities the authors are attempting to account for. There are quite a few baselines that are missing as a result in my opinion, such as sharpness aware minimization (SAM) and/or other exploration based optimizers.
- For language modelling, it would be more useful to show performance over a multitude of models and sizes, as the results appear to suggest for example that the gap between AdamW and Adam-SPD narrows quite significantly between sizes 7B and 13B, as well as when the PEFT method changes.
questions: questions: - I think a more interesting way to use the proposed method would be to have $\lambda$ to be a parameter or layer-specific hyper-parameter rather than a global hyper-parameter. This way, having individual $c_t$ for different components might make the method more robust and in fact help the model adapt relevant weights while keeping specific weights close to the pre-trained model. In many cases, different layers can begin to exhibit different functionalities and therefore there might be more use in having some layers change more than others based on the features that might be captured there.
- Although not a deal-breaker, I would appreciate a perhaps more in-depth analysis of the language modelling experiments, in particular without PEFT. While I understand the relationship with PEFT and the proposed optimizer, I believe that it would be quite valuable to show how AdamW-SPD works even under naive fine-tuning scenarios.
limitations: limitations: N/A
paper_summary: paper_summary: This paper proposes a new weight decay technique to adapt foundation models to target tasks, focusing on fitting the target data while maintaining the pre-trained knowledge. Specifically, the method, Selective Projection Decay (SPD), selectively imposes a strong penalty on certain layers while allowing others to change freely. Experimentally, the method consistently provides better in-distribution generalization and out-of-distribution robustness across multiple popular vision and language benchmarks.
strength_weakness: strength_weakness: There are no major weaknesses, please respond to my questions and issues.
questions: questions: - Why is L2-SP not compared in the experiments shown in Tables 1 and 5? The authors are requested to provide the results for this method as well.
- Tables 2a and 2b should be converted into figures to enhance readability.
- Minor Typo: On line 8, "retraining" should be corrected to "retaining".
- Several of the results are wrongly highlighted in Table 3.
- The appendix sections referred in the paper are missing.
limitations: limitations: The limitation of this work is clearly highlighted: the method's performance depends on the quality of the foundation model.
paper_summary: paper_summary: The paper proposes a method for robust fine-tuning of deep networks. The proposed method is a variant of L2-SP. The original formulation of L2-SP applied a uniform L2 penalty on the difference between the weights of the pre-trained model and the fine-tuned model. The proposed method (L2-SPD) aims to make this penalty non-uniform for different layers. This is achieved by only applying weight decay to layers where the direction of descent and weight decay are sufficiently similar. The efficacy of the method is demonstrated over various image classification, segmentation and reasoning tasks, showing better ID and OOD performance than baselines.
strength_weakness: strength_weakness: There are some minor issues with the presentation - 
* It would be nice to have a short preliminaries section to define all notations used in the paper, along with their dimensions. 
* It is not clear from Alg 2 how SPD is selective to different layers. It appears that $\theta_t$ represents the entire model’s parameters, and $c_t$ is a single scalar for the entire model. Hence, it seems like the “weight decay” only happens for the entire model if the total gradient is in the opposite direction of the weight decay. I might be incorrect here, but I am unable to see where the group structure comes in (possibly $c_t$ is separate for each layer?)

Apart from this, some of the experiments have the following weaknesses -  
* The Himmelbau example is not fully clear to me either. It would be interesting to compare the trajectories against L2-SP to actually understand the “selective penalization”, since regular L2-SP will also penalise deviations similarly. I tried running the code that the authors provided, and with some tuning of the hyper-parameters I was able to get similar trajectories with L2-SP as well. 
* Lack of ablation studies on the two components of the method. It would be good to understand the gains from selective projection and adaptive regularization separately as well.  
* I am unable to understand some of the trends in the results on ImageNet. It is not immediately clear to me how regularizing the model to minimize distance from the Zero-Shot model can outperform both the Zero-Shot model and the vanilla FT model on one domain (Im-Sketch), but is in between the two on another domain (Im-Rendition). I had the intuition that L2-SPD will pull the weights to be closer to the Zero-shot model, hence inheriting its robustness and biases, but I could be mistaken. 
* Minor - The caption for Table 3 states that "SPD... beats L2-SP by 8.8%", however, I cannot understand how this was computed.
* Minor - While the paper compares against L2-SP, it misses Elastic Weight Consolidation, a similar method using Fisher information, which also adaptively sets the regularization parameters for different layers.
questions: questions: * Is there a way to disentangle the gains from selective projection and adaptive regularization? Related to this, how often does the selective projection kick in in practice, and how does it change as training progresses (this is not crucial to the paper, but is a good-to-have analysis)? 
* I am unable to understand some of the trends in the results on ImageNet. It is not immediately clear to me how regularizing the model to minimize distance from the Zero-Shot model can outperform both the Zero-Shot model and the vanilla FT model on one domain (Im-Sketch), but is in between the two on another domain (Im-Rendition). I had the intuition that L2-SPD will pull the weights to be closer to the Zero-shot model, hence inheriting its robustness and biases, but I could be mistaken. 
* Can this technique of selective regularization be done at a more granular level beyond layers? 
* In section 4.1, why are all domainnet domains not used, omitting the infograph domain? The paper will be helped by a justification of this.
* For table 2a, the range of deviation is quite different for L2-SP and L2-SPD. While L2-SPD does dominate performance for similar deviations, I wonder if it is possible to use even larger $\lambda$ and make the deviation very low? Will the same correlation with OOD accuracy still hold?
* Hyper-parameter selection is not clear for ImageNet. Was k-fold cross validation used? Is the ID val used to select hyper-parameters?
* On quickdraw, will more aggressive regularization yield better OOD accuracy? Since the pretrained model is supposed to be closer to other domains? 
* For ImageNet, what is the linear layer for final classification?
* How does the method compare against a simpler baseline like Wise-FT?
limitations: limitations: The authors have addressed limitaions of their method, pointing out worse performance if the base model is not aligned with the OOD distribution.

Comments: title: Official Comment
comment: Thank you to the authors for addressing my questions. I agree with the other reviewer that the theoretical perspective is lacking in this work. However, I believe the empirical results highlight the significance of the study, so I will keep my score.
I thank the authors for their response. The response has addressed most of my concerns.

I believe that incorporating these into the paper will greatly improve the presentation, and hence I keep my rating of acceptance.
I appreciate the response provided by the authors. I believe the authors have addressed most of my concerns and I have adjusted my score accordingly. While I am now more convinced about the merits of the work, I still believe that increased analysis of the proposed method (from a more theoretical perspective rather than completely empirical) would be more meaningful given that the nature of the work.
We appreciate the reviewer's taking the time to examine the paper and the provided code closely. We clarified your concerns and added Infograph experiments as requested. If anything is not clear, we are happy to discuss it during the rebuttal period. 

**The Himmelbau example.**

We appreciate the reviewer's effort in running our code! In this simple example, you are correct that regular L2-SP would generate a similar trajectory. However, to see the difference between SPD's selective regularization and L2-SP, we suggest that the reviewer print out which direction is regularized during the optimization process.

Add the following printing statement to line 40 in the code before the return. 
* print(self.condition_buffer)
* return new_param_list

The condition buffer contains the selection condition $c_t$ for the X and Y directions. Whenever an entry is negative, regularization is applied. This demonstrates the underlying mechanism for SPD. 

**Is there a way to disentangle the gains..?**

Sorry for the confusion! SPD introduces only one component that affects performance, selective projection $c_t$ (Section 3.3),  which selectively applies regularization during fine-tuning. The selectivity is the only complement that provides the performance gain.

Section 3.4 introduces a re-parametrization of the regular L2 regularization as a projection for better hyper-parameter tuning (see a comment from the reviewer Ex5N). It's a different way of writing L2 regularization and does not introduce performance gain. 

Selectivity kicks in very often. However, it is hard to quantify, and it changes depending on the task, model, data, optimization hyper-parameters, etc. An analysis of the selectivity's behavior will be insightful. 

**Trends in the results on ImageNet.**

This is a very good question. Indeed, L2-SPD will pull the weights closer to the zero-shot model, as shown in Table 1 and Table 2 in the main paper, thus inheriting its robustness and biases. 

However, the extent of performance gain has a more subtle answer. It depends on the interplay between the pre-trained model, the fine-tuning (in-distribution) dataset, and the OOD datasets. Specifically, the ID dataset can provide useful information to the OOD datasets. In this benchmark, the useful information is that all datasets share the same label space. If fine-tuning leads to more gain, it's possible that after fine-tuning (assuming with proper regularization), the fine-tuned model can outperform both zero-shot and vanilla FT. 

**Can this technique of selective regularization be done at a more granular level beyond layers?**

Yes, theoretically, the method can be applied at a parameter level. However, the storage cost is significantly increased because we would need to store a different $c_t$ for each parameter in the model. 

**In section 4.1, why are all domainnet domains not used, omitting the infograph domain?.**

Infograph was omitted because none of the methods achieved good ID performance on this domain. We added the domain back to answer the reviewer's question (Table 1 in the rebuttal PDF). The superiority of SPD remains unchanged with the addition of this domain. 

**Larger regualrization and make the deviation very low? Will the same correlation with OOD accuracy still hold?**

We have tried larger hyper-parameters. At extremely large values, the model will eventually underfit and perform poorly on ID and OOD data like L2-SP. We intend to show that L2-SPD outperforms L2-SP for similar deviations as the reviewer suggested.  

**Hyper-parameter selection is not clear for ImageNet.**

We used the ID validation dataset following prior works. ImageNet-V2 can be seen as the clean test set for ImageNet. 

**On quickdraw, will more aggressive regularization yield better OOD accuracy?**

Yes, we conducted a similar experiment as in Table 2 in the main paper. More aggressive regularization yields better OOD on quickdraw. We swept hyper-parameters from 1.0 to 2.3 and calculated the correlation coefficient between the hyper-parameters and OOD performance. The correlation is 0.88, which means a robust positive correlation between large regularization and good OOD performance. 

**Linear layer for For ImageNet**

Because we used a CLIP ViT-Base model, we used a zero-shot linear head to initialize and fine-tune the linear layer. This practice follows from prior works such as WISE-FT. 

**Comparison with Wise-FT?**

Technically, WISE-FT is an orthogonal method. Therefore, SPD should further improve its performance. To validate this, we added experiments combining WISE-FT and SPD in our Imagenet experiments (Table 1 in the rebuttal PDF). Our results show that SPD further improves the performance of Wise-FT. It's also worth pointing out that SPD is more general than WISE-FT, which only applies to models capable of zero-shot. It does not apply to the segmentation experiment (Table 4 in the main paper), whereas SPD does not have this constraint.
We thank the reviewer for the positive comments and pointing our our typos. We clarified your concerns and added L2-SP experiments as requested. If anything is unclear, we are happy to discuss it during the rebuttal period. 


**Why is L2-SP not compared in the experiments shown in Tables 1 and 5? The authors are requested to provide the results for this method as well.**

* Table 1: Originally, we designed this table to present the improvement over a vanilla AdamW. For the rebuttal, we added L2-SP to Table 1 (Table 1 in the rebuttal PDF). SPD outperforms L2-SP in all domains.  

* Table 5: In the LLaMA PEFT experiments, we compared our method with Adam $+$ weight decay, as indicated in the table title. In this PEFT setting, L2-SP reduces to weight decay, as explained in section 3.5. In other words, the L2-SP baseline is already included.


**Tables 2a and 2b should be converted into figures to enhance readability.**

Thank you for the great suggestion. We will update the two tables to heat maps for better readability.
We really appreciate the positive comments and suggestions for ensemble methods. In the response, we clarified your concerns and included new experiments. If anything is not clear, we are happy to discuss it during the rebuttal period. 

**Overall, I think this is a strong paper, and its strengths outweigh its weaknesses. I think the paper could benefit from an ablation experiment, for example, trying Adam-SPD without the condition or benchmarking SPD's performance with optimizers other than Adam.**

Thank you for the positive comments on this work. The ablation study without the condition is shown in Table 2. Essentially, without the selecting condition (line 96-102), SPD reduces to a regular L2 regularization, which is shown ineffective in our comprehensive analysis (line 255). Furthermore, SPD is a general method and applicable to other optimizers. We hypothesize that it should have positive effects on optimizers using regular weight decay and L2 regularization. 


**[1] reports that ensembling with the initial weights (whether in model space or weight space) is a simple strategy that improves OOD robustness, and a few recent works [2, 3] report that ensembling continues to improve performance after their proposed fine-tuning strategies. Could the authors validate how the benefits from SPD extend to the ensembling setting?**

Thank you for mentioning these works, especially the ensemble technique such as Wise-FT. Technically, ensemble techniques are orthogonal methods. SPD should further improve the performance of them such as WISE-FT. To validate this, we added experiments combining WISE-FT and SPD in our Imagenet experiments (Table 1 in the rebuttal PDF). Our results show that SPD further improves the performance of Wise-FT. It's also worth pointing out that SPD is more general than WISE-FT, which only applies to models with linear connectivity and capable of zero-shot classification. It does not apply to the segmentation experiment (Table 4 in the main paper), whereas SPD does not have this constraint.
We thank the reviewer for the positive comments and for suggesting insightful new related works. In the response, we clarified all your concerns. If anything is not clear, we are happy to discuss it during the rebuttal period. 

**How does the performance compare when using EWC regularization for fine-tuning, since EWC is generally better than L2 regularization?**

Thank you for mentioning this work!  However, EWC does not apply to general fine-tuning and was explicitly developed for continual learning (CL). Under the CL setting, EWC requires the gradient information from the previous task, which is not available for fine-tuning a pre-trained model. Therefore, we cannot fairly compare to EWC in the fine-tuning setting. Nevertheless, EWC  points out that uniform regularization can be ineffective, which is also the main motivation of this paper for fine-tuning. 


**Comparing the performance of L2-SP and Adam-SPD in hyper-parameter sweeping experiments seems challenging since these are reported in tables. Perhaps heatmaps could be a clearer alternative to visualize and highlight the results.**

Yes, this is a very good suggestion. We will update the tables into heat maps in the main paper. 

**The authors argue that SPD selectively imposes penalties on large deviations in a layer-wise manner. However, in Algorithm 2, there is no layer-wise information. It appears that all parameters are updated based on the alignment score denoted by $c_t$. From my understanding of Algorithm 2, if the alignment is high (i.e., $c_t$ is smaller than zero), the updated parameters are further adjusted by the interpolation-like equation. If $c_t$ is larger than zero, do we use the vanilla ADAM update, since $\theta_t$ has already been updated before calculating $c_t$? Am I correctly understanding Algorithm 2? In such cases, it seems that the parameters shouldn’t be updated, since SPD will impose a penalty to slow down updates for those layers with negative $c_t$.**

SPD uses layer-wise information in the calculation of $c_t$. Specifically, each layer has a different $c_t$. Here, $\theta_t$ denotes the weight matrix of a layer, not an individual parameter or the entire network.  We will make the notation clear in the update. Yes, when $c_t>0$, the layer $\theta_t$ is updated using vanilla Adam. If this does not clarify your confusion, we are happy to discuss this more in the following rebuttal period. 


**In the LLaMA PEFT experiments, there is no baseline with L2 regularization or its advanced versions (such as EWC). While I expect that ADAM with SPD is better than AdamW, I would like to see if there are also improvements compared to ADAM with L2 regularization in PEFT settings.**

In the LLaMA PEFT experiments, we compared our method with Adam $+$ weight decay, as indicated in the table title. In this PEFT setting, L2 regularization reduces to weight decay as explained in section 3.5. In other words, the L2 regularization baseline is already included. We will make this clearer in the text; thanks for bringing this up.