LiNeS: Post-training Layer Scaling Prevents Forgetting and Enhances Model Merging

## Abstract

Abstract Large pre-trained models exhibit impressive zero-shot performance across diverse tasks, but fine-tuning often leads to catastrophic forgetting, where improvements on a target domain degrade generalization on other tasks. To address this challenge, we introduce LiNeS, Layer-increasing Network Scaling, a post-training editing technique designed to preserve pre-trained generalization while enhancing fine-tuned task performance. LiNeS scales parameter updates linearly based on their layer depth within the network, maintaining shallow layers close to their pre-trained values to preserve general features while allowing deeper layers to retain task-specific representations.
We further extend this approach to multi-task model merging scenarios, where layer-wise scaling of merged parameters reduces negative task interference. LiNeS demonstrates significant improvements in both single-task and multi-task settings across various benchmarks in vision and natural language processing.
It mitigates forgetting, enhances out-of-distribution generalization, integrates seamlessly with existing multi-task model merging baselines improving their performance across benchmarks and model sizes, and can boost generalization when merging LLM policies aligned with different rewards via RLHF. Importantly, our method is simple to implement and complementary to many existing techniques.

### 1 Introduction

Pre-trained models have become the backbone of modern machine learning pipelines (Bommasani et al., ; Touvron et al., ). Their introduction has shifted the paradigm from end-to-end training to fine-tuning (Zhuang et al., ), leading to the proliferation of thousands of fine-tuned checkpoints derived from a few foundation models (Rombach et al., ; Team et al., ). To improve downstream performance across multiple tasks or align with multiple preferences (Singh & Jaggi, ; Matena & Raffel, ; Ilharco et al., ; Yadav et al., ; Rame et al., ), model merging techniques combine available checkpoints, avoiding the costly process of joint fine-tuning (Ilharco et al., ; Yadav et al., ).
However, specializing models introduces trade-offs, such as the forgetting of previously acquired knowledge (Aghajanyan et al., ) – a phenomenon known as catastrophic forgetting (McCloskey & Cohen, ). Furthermore, merging checkpoints fine-tuned on different tasks can lead to significant performance degradation due to task interference (Yadav et al., ; Wang et al., ).

To mitigate catastrophic forgetting, many works propose regularizing the fine-tuning process (Aghajanyan et al., ; Kumar et al., ; Gouk et al., ; Razdaibiedina et al., ). Leveraging the insight that shallow layers capture generalizable representations (Yosinski et al., ; Neyshabur et al., ), Howard & Ruder (); Dong et al. () apply lower learning rates to the shallow layers to retain general features.
However, modifying the fine-tuning process can be complex and computationally expensive. This motivated the development of post-training model editing and model merging methods that directly edit the checkpoints in the weight space. For instance, Wortsman et al. (); Rame et al. () mitigate catastrophic forgetting by interpolating weights between pre-trained and fine-tuned models. In multi-task settings, Yadav et al. (); Wang et al. () proposed methods to reduce interference among tasks when merging multiple checkpoints. Yet, significant performance degradation persists when merging multiple models, leaving this as an open challenge.

Most model merging methods, however, treat all layers equally, overlooking the earlier insight that shallow layers should remain close to their pre-trained weights to avoid losing the general representations they encode. In this paper, we explore whether this insight can be leveraged post-training.
We find that reducing the magnitude of shallow-layer updates after fine-tuning can retain single-task performance gains while significantly mitigating forgetting.

We propose LiNeS, Layer-increasing Network Scaling, a post-training, plug-and-play method that directly edits the residual, i.e., the difference between the fine-tuned and pre-trained checkpoint, by applying a scaling coefficient that linearly increases with layer depth.
This scaling effectively preserves the general features captured in the shallow layers of the pre-trained model while retaining task-specific features in the deep layers of the fine-tuned model.
Moreover, we extend LiNeS to the multi-task model merging setting, where contributions from one task distort the general features also required by other tasks.
By preserving the general features in the shallow layers, LiNeS mitigates task interference and improves multi-task performance.

LiNeS demonstrates remarkable performance on diverse test scenarios and is orthogonal to existing post-training merging algorithms. It modifies the fine-tuned checkpoint to consistently retrieve nearly full performance on the fine-tuned task while significantly restoring generalization on other tasks. Furthermore, it can be seamlessly integrated with existing weight interpolation methods for improving out-of-distribution generalization (Wortsman et al., ). When merging multiple models, LiNeS improves baseline methods for merging checkpoints fine-tuned on multiple tasks in both computer vision and NLP benchmarks (Ilharco et al., ; Yadav et al., ; Wang et al., ) and also enhances performance when merging checkpoints fine-tuned on the same task (Wortsman et al., ) and merging LLM policies aligned with different rewards (Rame et al., ) via Reinforcement Learning with Human Feedback (RLHF) (Christiano et al., ).

Our contributions are as follows:

- •
We propose LiNeS, a post-training editing technique that preserves the zero-shot generalization of pre-trained models while retaining fine-tuned knowledge by applying layer-wise scaling on parameter updates. For example, in image-classification tasks with CLIP ViT-B/32 checkpoints,
LiNeS maintains on average 99.8% of performance on the fine-tuned task while preserving 97.9% performance of the pre-trained model on other control tasks, effectively mitigating catastrophic forgetting.
- •
We demonstrate that LiNeS significantly enhances multi-task model merging baselines, consistently improving performance across benchmarks and architectures in both vision and NLP domains. For instance, we observe a 3.1% and 4.0% improvement over Task Arithmetic (Ilharco et al., ) and Ties-merging (Yadav et al., ) respectively, for a 20-task computer vision benchmark with ViT-L/14.
- •
We show that LiNeS can be applied to enhance existing weight interpolation methods across various scenarios, improving out-of-distribution generalization, merging multiple checkpoints fine-tuned on the same task with different hyper-parameter configurations, and merging LLM policies aligned with different rewards.

Our proposed method is simple to implement(^1^11PyTorch pseudo-code in Appendix A; source code at [https://github.com/wang-kee/LiNeS](https://github.com/wang-kee/LiNeS).), orthogonal to many existing approaches, and improves performance in a wide variety of settings.

### 2 Related Works

###### Representation collapse and regularized fine-tuning

Pre-trained models such as CLIP exhibit strong zero-shot performance across diverse data distributions due to the robust and transferable feature representations learned during pre-training (Radford et al., ; Jia et al., ). However, fine-tuning on specific tasks often harms the zero-shot generalization performance on distributions different from the fine-tuning domain (Wortsman et al., ; Goyal et al., ; Aghajanyan et al., ). This degradation arises from the distortion of pre-trained features during fine-tuning (Kumar et al., ), a phenomenon referred to as representation collapse by Aghajanyan et al. (). To mitigate representation collapse, many works have proposed to regularize the fine-tuning process to preserve the general pre-trained features (Aghajanyan et al., ; Kumar et al., ; Goyal et al., ; Gouk et al., ; Zhang et al., ; Razdaibiedina et al., ; Shen et al., ; Lee et al., ). Some of these approaches take into account that different layers of a model learn distinct features, with the shallower layers capturing more general features and deeper layers specializing in task-specific representations (Neyshabur et al., ; Yosinski et al., ; Raghu et al., ; ). Specifically, they apply layer-wise learning rate decay, preserving more of the pre-trained features in the shallow layers while allowing deeper layers to specialize for the target domain (Clark, ; Bao et al., ; Dong et al., ; Howard & Ruder, ; Zhang et al., ). However, modifying the fine-tuning process is orders of magnitude more computationally expensive compared to post-training merging methods.

###### Weight interpolation and model merging

Garipov et al. (); Draxler et al. () showed that two solutions derived from separate training runs can be connected by nonlinear paths of low loss, while linear mode connectivity (Frankle et al., ) extended the paths to the linear case.
These insights enabled the transfer of the benefits regarding robustness of (traditional) output ensembles (Hansen & Salamon, ; Lakshminarayanan et al., ) to weight ensembles, reconciling the bias-variance trade-off (Belkin et al., ) while eliminating the computational cost of multiple inferences (Fort et al., ). These findings can be leveraged to improve performance on single-task (Izmailov et al., ; Wortsman et al., ; Rame et al., ; Wortsman et al., ), out-of-distribution (Wortsman et al., ; Ramé et al., ), multi-task (Ilharco et al., ; Dimitriadis et al., ) and multi-objective alignment (Zhong et al., ; Ramé et al., ) settings. Furthermore, model merging can also applied as a scalable approach to unify multiple task-specific models into a single model with multi-task capabilities (Ilharco et al., ; Yadav et al., ), despite performance loss compared to individual models.
Several methods have tried to improve multi-task model merging by preserving the important parameters defined via the Fisher Information Matrix (Matena & Raffel, ; Tam et al., ), using heuristics (Davari & Belilovsky, ; Luo et al., ; Jin et al., ),
randomly dropping and rescaling the task vector parameters (Yu et al., ) or by focusing on resolving weight interference caused by sign disagreements and redundant parameters (Yadav et al., ; Wang et al., ). Recent works use gradient descent to learn the layer-specific merging coefficients per task, e.g., Ada-merging (Yang et al., ) minimizes entropy in unlabeled test data while aTLAS (Zhang et al., ) optimizes using cross-entropy loss on validation data. Compared to LiNeS, these methods do not incorporate any prior knowledge on early vs. deep layers and require training, resulting in significant computational overheads.

### 3 Post-training Layer-wise Scaling Mitigates Forgetting

In this section, we present the key insight of our work: Scaling
down the updates of shallow layers after fine-tuning can mitigate catastrophic forgetting and restore zero-shot generalization while preserving performance on the target task.

###### Notation

We consider a pre-trained model $\bm{\theta}_{0}\in\mathbb{R}^{N}$ with $N$ parameters. Fine-tuning on a specific task $t$ results in the fine-tuned weights $\bm{\theta}_{t}$. The difference between these two sets of weights, $\bm{\tau}_{t}=\bm{\theta}_{t}-\bm{\theta}_{0}$, is referred to as the task vector or residual for task $t$ (Ilharco et al., ) and represents the updates made during fine-tuning.

###### Fine-tuning leads to catastrophic forgetting

We quantitatively demonstrate the phenomenon of catastrophic forgetting with the following experiments. Consider the 8-task image classification benchmark studied in Ilharco et al. (). We fine-tune a CLIP ViT-B/32 model on each task, measuring performance on the fine-tuned task – referred to as the target task – and the remaining 7 tasks – the control tasks. The results, shown in Table 1, demonstrate that while fine-tuning significantly improves accuracy on the target task, it drastically reduces accuracy on the control tasks, underscoring the loss of the model’s zero-shot generalization abilities.

**Table 1: Fine-tuning harms generalization on control tasks. Our proposed post-training edition leads to a superior trade-off between performance on target and control tasks.**
| Model / Accuracy | Target | Control |
| --- | --- | --- |
| Pre-trained | 48.3 | 48.3 |
| Fine-tuned | 90.5 | 38.0 |
| Fine-tuned+LiNeS (ours) | 90.3 | 48.0 |

###### Shallow-layer updates impact minimally on target task accuracy

Most parameter updates during the fine-tuning process are redundant, as similar performance is achievable without updating most pre-trained parameters (Yadav et al., ; Wang et al., ; He et al., ).
Moreover, prior work shows that task-specific features are often concentrated in deeper layers of the network (Neyshabur et al., ; Yosinski et al., ; Raghu et al., ). Based on these observations, we hypothesize that updates to the shallow layers contribute minimally to target tasks. To test this, we progressively downscale the updates to shallow layers after fine-tuning.
Specifically, we apply a scaling factor to the updates to the $\ell$-th layer $\bm{\tau}^{(\ell)}$, defined as:
$\lambda^{(\ell)}=\gamma+(1-\gamma)\frac{\ell-1}{L-1},\ \forall\ell\in[L],$
This linearly scales the updates from a factor of $\gamma$ for the first layer to $1$ for the last one.
As a result, fine-tuning updates to the shallow layers are scaled down more aggressively, with later layers experiencing progressively smaller reductions. We then reintroduce the scaled task vector into the pre-trained model and measure its performance on the fine-tuned task.
Figure 1 shows the results of this experiment for the CLIP ViT-B/32 checkpoint fine-tuned across the 8 tasks, where $\gamma$ is progressively decreased to strengthen the downscaling effect. We observe that, even with strong downscaling of shallow layers, the target task accuracy remains nearly unaffected. In contrast, when we downscale the deeper layers, target task accuracy drops significantly.
These results support our hypothesis that shallow-layer updates are largely unnecessary for maintaining accuracy on the target task.

Figure: Figure 1: Downscaling the shallow layers maintains the fine-tuned performance on target tasks (orange line, left), while restoring zero-shot performance from pre-trained model on control tasks (orange line, right). The performance for downscaling deep layers instead is presented in blue lines, which underperforms downscaling shallow layers in both cases. $\gamma$ represents the minimum scaling factor applied to the layers, where a smaller $\gamma$ leads to stronger downscaling strength, with $\gamma=1$ restoring the original fine-tuned model.
Refer to caption: x1.png

###### Shallow-layer updates undermine zero-shot generalization

While shallow-layer updates have minimal impact on target-task accuracy, they distort the general features learned during pre-training, which reside primarily in the shallow layers (Neyshabur et al., ; Yosinski et al., ; Raghu et al., ). We hypothesize that the degradation of performance on control tasks is largely due to these distortions in the shallow layers.
Using the same experimental setup, we now evaluate the zero-shot performance on the control tasks, i.e., the other 7 unseen tasks. As shown in Figure 1 (right), as the strength of the shallow-layer downscaling increases, the accuracy on control tasks approaches the original pre-trained model’s performance. This shows that by reducing the shallow-layer updates, we can restore much of the zero-shot performance that is lost during fine-tuning.

###### Improved trade-off between target and control performance

To optimize the trade-off between target and control task performance, we select a scaling coefficient $\gamma$ for each model

Figure: Figure 2: Our linear scaling (LiNeS) retains performance on both control and fine-tuned target tasks. Each scatter dot represents a different model.
Refer to caption: x2.png

that maximizes a weighted balance between these two objectives (detailed in Appendix C.1).
After selecting the optimal scaling coefficient, the test results are shown in the final row of Table 1.
Our post-training method preserves target task accuracy with a minimal 0.2% difference while improving control task performance by 10%, compared to the fine-tuned model.

We further apply the same method to a 20-task computer vision benchmark (Wang et al., ). For evaluation, we report both the target task normalized accuracy and the control task normalized accuracy on the 19 tasks, where accuracy is normalized by the performance of the fine-tuned model for the target task and the zero-shot accuracy of the pre-trained model for the control tasks.
We compare to fine-tuned models on each task and the pre-trained model as baselines. Figure 2 shows that fine-tuning degrades zero-shot generalization, as indicated by the performance drop on control tasks. In contrast, our post-training scaling method significantly improves generalization while maintaining near-full target task accuracy. On average, our method achieves a target task normalized accuracy of 99.8% and a control task normalized accuracy 97.9%. This demonstrates its effectiveness in preserving both task-specific knowledge from fine-tuned checkpoints and the generalization capabilities of the pre-trained model. The full breakdown of results by task is available in Figure 6 in Appendix.

### 4 Method

Motivated by the results of the previous section for mitigating forgetting, we propose LiNeS for Layer-increasing Network Scaling, a simple post-training technique that linearly rescales the updates of different layers in the task vector based on their depth in the network. LiNeS is designed to retain general features in the shallow layers while preserving the task-specific adaptations in the deeper layers.

Given a task vector $\bm{\tau}$ with $L$ layer blocks
we apply the layer-wise linear scaling to adjust the contributions of shallow and deep layers using the following formulation:

$$ $\bm{\tau}_{\textrm{LiNeS}}=\texttt{concat}\left(\lambda^{(1)}\bm{\tau}_{ \textrm{}}^{(1)},\dots,\lambda^{(L)}\bm{\tau}_{\textrm{}}^{(L)}\right),\quad \text{where}\ \ \lambda^{(\ell)}=\alpha+\beta\frac{\ell-1}{L-1},\quad\forall \ell\in[L].$ (1) $$

As a result, the layers in $\bm{\tau}$ are progressively scaled with a factor between $\alpha$ for the first layer and $\alpha+\beta$ for the last layer, with intermediate layers scaled with a linearly increasing schedule depending on their depth. The final model $\bm{\theta}$ is then obtained by summing the pre-trained model weights and the edited task vector, i.e., $\bm{\theta}=\bm{\theta}_{0}+\bm{\tau}_{\text{LiNeS}}$. Notice that, in Equation 1, $\bm{\tau}$ can correspond to either a single-task residual or, in the context of model merging, a multi-task vector obtained by merging the residuals of multiple checkpoints fine-tuned starting from a common initialization. Additional details on this process are provided in the next section.
Setting $\alpha=\beta=0$ corresponds to the pre-trained model, while $\alpha=1,\beta=0$ is the fine-tuned model in the case that $\bm{\tau}$ is a single-task vector.

In practice, we find that tuning just one hyper-parameter (either $\alpha$ or $\beta$) is often sufficient to achieve a good balance between target task performance and generalization. Specific details on hyper-parameter tuning for different applications are provided in the experimental sections.
The linear scaling method introduced in Section 3 corresponds to LiNeS by setting $\alpha=\gamma$ and $\beta=1-\gamma$. This formulation generalizes our previous approach, offering a flexible way to adjust the contributions of different layers based on the task requirements.

### 5 Model Merging Experiments

We empirically verify the effectiveness of applying LiNeS across diverse application domains.
Section 5.1 presents results for improving robust fine-tuning (Wortsman et al., ) for OOD generalization;
Section 5.2 focuses on improving existing multi-task merging methods (Ilharco et al., ; Yadav et al., ; Wang et al., ) in both vision and NLP benchmarks. In Section 5.3, we apply LiNeS and improve the merging of single-task fine-tuned models within the setting of Model Soups (Wortsman et al., ), and finally, we enhance merging foundation models fine-tuned on different rewards (Rame et al., ) in Section 5.4.

#### 5.1 Improving Robust Fine-tuning for OOD Generalization

We first consider the setting of robust fine-tuning or WiSE-FT (Wortsman et al., ), where linearly interpolating between the pre-trained and the fine-tuned weights improves model performance on OOD datasets. The interpolation is equivalent to scaling the residual $\bm{\tau}$: $(1-\gamma)\bm{\theta}_{0}+\gamma\bm{\theta}=\bm{\theta}_{0}+\gamma\bm{\tau}$, for $\gamma\in[0,1]$. We apply LiNeS to the residual $\bm{\tau}$.
Following Wortsman et al. (), we evaluate CLIP models fine-tuned on ImageNet (Deng et al., ), considering 5 OOD datasets, namely ImageNetSketch (Wang et al., ),
ImageNet-A (Hendrycks et al., ),
ImageNet-R (Hendrycks et al., ), ObjectNet (Barbu et al., ), ImageNet-V2 (Recht et al., ).

We apply this LiNeS to each of the 70 fine-tuned checkpoints(^2^22The checkpoints are CLIP ViT-B/32 models fine-tuned on ImageNet with different hyper-parameters.) provided by Wortsman et al. () setting $\alpha=\beta=0.5$. We present the average results in Figure 3, comparing the performance of WiSE-FT with and without applying LiNeS on the 5 OOD datasets.
Without applying WiSE-FT, LiNeS already enhances both the ID and OOD performance of the fine-tuned models by a noticeable margin. Starting from this edited model and applying the WiSE-FT interpolation with the pre-trained weights leads to a Pareto Front (Caruana, ) that consistently dominates the one by WiSE-FT across all distribution shifts, illustrating the applicability of the proposed method across various distribution shifts.. A granular result for applying LiNeS to each of the 70 checkpoints is provided in Appendix C.6.2, further highlighting its universal effectiveness across models.
We also report similar findings in Figure 8 in Appendix for a CLIP ViT-B/16 checkpoint fine-tuned on ImageNet, using the same hyper-parameters as Wortsman et al. ().

Figure: Figure 3: Application of LiNeS to WiSE-FT (Wortsman et al., ) improves performance on ImageNet and five different distribution shifts, resulting in a dominating Pareto Front over WiSE-FT.
Refer to caption: x3.png

#### 5.2 Improving Multi-task Model Merging

In this section, we extend LiNeS to improve multi-task merging algorithms, aiming to combine multiple models fine-tuned independently on different tasks into a single model (Matena & Raffel, ; Ilharco et al., ; Ortiz-Jimenez et al., ; Yadav et al., ). Task arithmetic (Ilharco et al., ) proposed to decouple the contributions of the pre-trained model and individual task vectors, first generating a multi-task vector $\bm{\tau}_{\textrm{MTL}}=g(\bm{\tau}_{1},\dots,\bm{\tau}_{T})$ with a merging function $g:\mathbb{R}^{N}\times\cdots\times\mathbb{R}^{N}\mapsto\mathbb{R}^{N}$, and then adding back to the pre-trained checkpoint with a scaling factor to construct a multi-task model $\bm{\theta}=\bm{\theta}_{0}+\lambda\cdot\bm{\tau}_{\textrm{MTL}}$. The scalar coefficient $\lambda$ is tuned using a held-out validation set. Recent works (Yadav et al., ; Wang et al., ) follow the same protocol while improving the merging function $g$ for retaining more task information. We refer to Appendix B.1 for a more detailed explanation of these methods.

However, significant performance loss occurs between the merged multi-task model and the original fine-tuned checkpoints. This performance decrease partially stems from interference (Yadav et al., ; Wang et al., ) among task vectors, where the contribution of one task negatively impacts performance on others, leading to overall degradation.
Task interference is linked to catastrophic forgetting, as the individual task vectors lose a significant amount of generalization ability to other tasks after fine-tuning and merging them leads to interference among each other. Therefore, we can edit each task vector with LiNeS before merging to restore the generalization to other tasks, or for simplicity, edit directly the merged multi-task vector to preserve the shallow and general features that are beneficial across tasks.

We enhance the merging methods by applying LiNeS on the merged multi-task vector $\bm{\tau}_{\textrm{MTL}}$.
For the linear scaling schedule, we tune only $\beta$ and set $\alpha$ using a heuristic that adjusts based on both the number of merged models and the merging method. Specifically, for task arithmetic which aggregates the individual task vectors through a simple summation operation: $\bm{\tau}_{\text{sum}}=\sum_{i=1}^{N_{\textrm{models}}}\bm{\tau}_{i}$, we set $\alpha=1/N_{\textrm{models}}$. For other merging strategies which result in $\bm{\tau}_{\textrm{MTL}}$ with different magnitudes of norms, e.g., aggregation with summation leads to a norm $\times N_{\textrm{models}}$ larger compared to averaging, we further multiply by a scaling term $\left\|\bm{\tau}_{\text{sum}}\right\|/\|\bm{\tau}_{\textrm{MTL}}\|$ to normalize their norm to simple summation. Overall, we set the intercept $\alpha$ to:

$$ $\alpha=\frac{1}{N_{\textrm{models}}}\frac{\left\|\bm{\tau}_{\text{sum}}\right \|}{\|\bm{\tau}_{\textrm{MTL}}\|},\textrm{ where }\tau_{\text{sum}}=\sum_{i=1} ^{N_{\textrm{models}}}{\bm{\tau}_{i}}$ (2) $$

Therefore, we only tune $\beta$ and search over the same range as the constant scaling $\lambda$ used by the aforementioned merging techniques. As a result, LiNeS shares the same computational requirements as the baseline merging methods; we provide more details for the hyper-parameters in Appendix B.2.1.
Specifically, we consider various multi-task model merging baselines. namely Task Arithmetic (Ilharco et al., ), Ties-merging (Yadav et al., ), Consensus Merging (Wang et al., ), enhancing them with LiNeS and evaluate on both computer vision and NLP benchmarks.

**Table 2: Results for multi-task model merging in vision classification benchmarks of 8 tasks (Ilharco et al., ), 14 tasks, and 20 tasks (Wang et al., ) for different vision transformer architectures. Applying LiNeS improves baseline performance for all benchmark/architecture combinations.**
| Method | with<br>LiNeS |  |  | ViT-B/32 |  | ViT-L/14 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  | 8 tasks | 14 tasks | 20 tasks |  | 8 tasks | 14 tasks | 20 tasks |  |  |
| Zero-shot |  |  |  | 48.3 | 57.3 | 56.1 |  | 64.8 | 68.3 | 65.3 |
| Fine-tuned |  |  |  | 90.5 | 89.5 | 90.4 |  | 94.0 | 93.3 | 94.0 |
| Task Arithmetic | ✗ |  |  | 69.7 | 65.0 | 60.3 |  | 84.0 | 79.2 | 74.0 |
| ✓ |  |  | $\text{74.2}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}4.5})$ | $\text{69.1}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}4.1})$ | $\text{63.4}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.1})$ |  | $\text{86.5}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.5})$ | $\text{82.2}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.0})$ | $\text{77.1}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.1})$ |  |
| Ties-Merging | ✗ |  |  | 73.6 | 67.6 | 63.1 |  | 85.6 | 79.3 | 75.6 |
| ✓ |  |  | $\text{77.2}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.6})$ | $\text{72.1}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}4.5})$ | $\text{67.2}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}4.1})$ |  | $\text{88.0}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.4})$ | $\text{82.5}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.2})$ | $\text{79.6}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}4.0})$ |  |
| Consensus Merging | ✗ |  |  | 74.5 | 70.1 | 65.3 |  | 85.2 | 81.9 | 78.7 |
| ✓ |  |  | $\text{77.6}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.1})$ | $\text{73.6}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.5})$ | $\text{68.6}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.3})$ |  | $\text{87.3}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.1})$ | $\text{84.0}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.1})$ | $\text{81.0}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.3})$ |  |

##### 5.2.1 Computer vision

We experiment with the 8-task image classification benchmark proposed by Ilharco et al. (), as well as the more challenging 14-task and 20-task benchmarks from Wang et al. (). Detailed descriptions of task composition appear in Appendix B.2.2. We also examine the efficacy of LiNeS across the model scale axis, studying three vision transformer (Dosovitskiy et al., ), namely ViT-B/32, ViT-B/16 and ViT-L/14, as CLIP visual encoders (Radford et al., ) .

Table 2 presents the results for ViT-B/32 and ViT-L/14, while Appendix C.4 contains the ViT-B/16 experiments. We observe that LiNeS provides a significant improvement to all baseline merging methods across all tested scenarios, regardless of model sizes and total number of tasks. For example, for the 8-task benchmark with ViT-B/32, LiNeS improves task arithmetic by 4.5%, Ties-merging by 3.6% and consensus merging by 3.1%. For the challenging 20-task benchmark with ViT-L/14, LiNeS leads to consistent and significant improvements, improving task arithmetic by 3.1%, Ties-merging by 4.0% and consensus merging by 2.3%. The detailed performance on individual tasks for each tested scenario is presented in Appendix C.7.

##### 5.2.2 Natural Language Processing

We also evaluate the effectiveness of LiNeS in NLP domain, including a 7-task NLP benchmark (Yadav et al., ), an 8-task Question-Answering benchmark (Zhou et al., ), and their combined 11-task benchmark (Wang et al., ). Appendix B details the experimental settings. Following Tam et al. (), we adopt a variant of T5-large model (Raffel et al., ), namely T5-large-LM-Adapt (Lester et al., ), and use their provided checkpoints. While T5-large contains both encoder and decoder networks, we apply LiNeS only to the decoder, as our findings in Appendix C.5 indicate that applying the edition to the decoder leads to similar observations to vision.

The performance of applying LiNeS to baseline methods with T5-large across various NLP tasks is summarized in Table 3. LiNeS consistently improves multi-task performance across baseline merging methods and benchmarks with a notable margin. For example, on the 7 NLP tasks benchmark, LiNeS improves task arithmetic by 4.5 points, and consensus merging by 1.9 points. Meanwhile, LiNeS outperforms Ties-merging by 3.0% and 2.4% for the 8-QA benchmark and 11-NLP benchmark, respectively.

**Table 3: Results for multi-task model merging methods in three NLP benchmarks with T5-large model. LiNeS improves baseline performance across merging methods and benchmarks.**
| Method | with<br>LiNeS | T5-large (Lester et al., ) |  |  |
| --- | --- | --- | --- | --- |
| 7 NLP tasks (Yadav et al., ) | 8 QA tasks (Zhou et al., ) | 11 NLP tasks (Wang et al., ) |  |  |
| Zero-shot |  | 44.9 | 33.1 | 36.9 |
| Fine-tuned |  | 85.9 | 80.7 | 78.7 |
| Task Arithmetic | ✗ | 71.9 | 63.8 | 63.6 |
| ✓ | $\text{76.4}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}4.5})$ | $\text{67.6}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.8})$ | $\text{66.2}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.6})$ |  |
| Ties-Merging | ✗ | 71.6 | 63.0 | 64.0 |
| ✓ | $\text{72.0}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}0.4})$ | $\text{66.0}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.0})$ | $\text{66.4}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.4})$ |  |
| Consensus Merging | ✗ | 73.5 | 68.6 | 67.5 |
| ✓ | $\text{75.4}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}1.9})$ | $\text{69.3}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}0.7})$ | $\text{67.5}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}0.0})$ |  |

#### 5.3 Improving Model Soups for Merging single-task models

Averaging in weight space multiple models fine-tuned on the same task derived from the same pre-trained model has been shown to increase target performance (Wortsman et al., ; Rame et al., ). In this section, we investigate whether LiNeS can enhance the test performance when merging single-task models.

We follow the setting in Model Soups (Wortsman et al., ) and merge 70 CLIP ViT-B/32 checkpoints fine-tuned on ImageNet (Deng et al., ) using different hyper-parameters, plus the pre-trained checkpoint. We consider both variants introduced in Wortsman et al. (), namely uniform and greedy soup. We refer to Appendix B.3 for details regarding these methods and experimental settings.
For both cases, the weight-averaging process can be decomposed as follows:

$$ $\vspace{-2pt}\bm{\theta}_{\text{soup}}=\bm{\theta}_{0}+\bm{\tau}_{\text{soup}} ,\ \text{where}\ \bm{\tau}_{\text{soup}}=\frac{1}{N_{\textrm{models}}}\sum_{i= 1}^{N_{\textrm{models}}}{(\bm{\theta}_{i}-\bm{\theta}_{0})}$ (3) $$

We apply LiNeS to $\bm{\tau}_{\text{soup}}$ fixing $\alpha=1$ and searching over $\beta$. As a baseline, we also consider task arithmetic, where we search for a constant scaling factor on $\bm{\tau}_{\text{soup}}$. Note that both settings introduce one hyper-parameter to vanilla model soups, and refer to Appendix B.3 for a detailed description of the modifications. Table 4 summarizes the results and shows that LiNeS improves over vanilla soups and task arithmetic for both uniform and greedy soup by $0.48\%$ and $0.15\%$ on ImageNet, respectively.
We report the best-performing model and the average performance as baselines.
Finally, our proposed method compounds the gains from the greedy soup and leads to the best-performing model.

#### 5.4 Improving Rewarded Soups

In this section, we explore the effectiveness of LiNeS for merging foundation models fine-tuned on different rewards. We consider the Rewarded Soups setting (Ramé et al., ), which interpolates the weights $\bm{\theta}_{1}$ and $\bm{\theta}_{2}$ of two LLM policies, each optimized for a distinct reward $R_{1}$ and $R_{2}$, respectively.

Starting with an LLM parameterized by weights $\bm{\theta}_{0}$, we first fine-tune it using supervised fine-tuning (SFT) on labeled demonstrations. From the resulting weights $\bm{\theta}_{\rm SFT}$, we then apply Reinforcement Learning from Human Feedback (RLHF) (Christiano et al., ; Ouyang et al., ), training two independent policies via Proximal Policy Optimization (PPO) (Schulman et al., ) to maximize the rewards $R_{1}$ and $R_{2}$ respectively. To merge these policies, we linearly interpolate the residuals $\bm{\tau}_{1}=\bm{\theta}_{1}-\bm{\theta}_{\rm SFT}$ and $\bm{\tau}_{2}=\bm{\theta}_{2}-\bm{\theta}_{\rm SFT}$. This interpolation defines a continuous set of rewarded policies:

$$ $\bm{\theta}_{RS}=\bm{\theta}_{\rm SFT}+\lambda\bm{\tau}_{1}+(1-\lambda)\bm{ \tau}_{2},\quad\lambda\in[0,1],$ (4) $$

where the coefficient $\lambda$ models the user’s preferences.
We apply LiNeS to the weighted-sum residual: $\lambda\bm{\tau}_{1}+(1-\lambda)\bm{\tau}_{2}$, fixing $\alpha=\beta=1$ for computational reasons.

**Table 4: LiNeS improves performance over Model Soups (Wortsman et al., ), for both uniform and greedy soup in merging multiple checkpoints fine-tuned on ImageNet with different hyper-parameter configurations.**
| Method | Enhancements | ImageNet Acc. |
| --- | --- | --- |
| Averaged accuracy | / | 77.98 |
| Best individual model | / | 80.36 |
| Uniform soup | / | 79.99 |
| Task Arithmetic | 80.17 |  |
| LiNeS | $\text{80.47}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}0.48})$ |  |
| Greedy soup | / | 81.01 |
| Task Arithmetic | 81.01 |  |
| LiNeS | $\text{81.16}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}0.15})$ |  |

In our experiment, we use LLaMA-2 7B (Touvron et al., ) and the Reddit Summary task (Stiennon et al., ), which consists of 14.9k post-summary pairs. We employ two reward models: GPT2-reward-summarization – which scores summaries based on human preferences – and BART-faithful-summary-detector (Chen et al., ) – which evaluates the faithfulness of the generated summary to the source post. To evaluate the models, we use a subset of 1k samples from the test set, generate the responses, and compute the average score for each reward dimension.

In Figure 4, we present the empirical Pareto Fronts for both Rewarded Soups and Rewarded Soups+LiNeS. LiNeS consistently outperforms the vanilla Rewarded Soups across the full preference space, Pareto dominating the baseline. This result highlights the generality of LiNeS.

### 6 Discussion

Figure: Figure 5: Comparison of the scalings obtained by different methods on 8-task merging benchmark with CLIP ViT-B/32.
Refer to caption: x5.png

We compare LiNeS with prior work that optimizes the scaling coefficients via backpropagation. Specifically, Ada-merging (Yang et al., ) minimizes the entropy loss of the predictions on the test set, while aTLAS (Zhang et al., ) minimizes a cross entropy loss on validation samples. Both methods operate on a more fine-grained level and introduce coefficients per layer and per task, requiring all $T+1$ checkpoints, for task vectors and pre-trained model respectively, to be stored in memory during their fine-tuning process.

We consider the 8-task computer vision benchmark with the ViT-B/32 visual encoder, and present the per-layer scalings in Figure 5. For aTLAS and Ada-merging, we report the average optimized scaling coefficients for attention and linear layers in each block across tasks. Without requiring training, LiNeS leverages the inductive bias of neural networks to achieve scaling very close to Ada-merging or aTLAS , but with much less computational cost. Apart from the excessive memory overhead, both aTLAS and Ada-merging require multiple training epochs, making it challenging to scale for large models. As we demonstrate in Section 5.4, LiNeS efficiently scales to large models like LLaMA (Touvron et al., ).

### 7 Conclusion

In this work, we presented LiNeS, a novel method designed to mitigate catastrophic forgetting after fine-tuning process. By reducing the magnitude for parameter updates in the shallower layers, LiNeS improves the generalization performance of the edited model on control tasks while almost fully preserving performance on the fine-tuned tasks. Furthermore, we demonstrated the versatility of LiNeS in addressing task interference in multi-task model merging, where it consistently improves the baseline model merging methods across vision and NLP benchmarks. Our experiments confirmed the broad applicability of LiNeS across various scenarios, from improving OOD generalization to enhancing multi-task and single-task model merging strategies, as well as improving merging LLM policies aligned with different rewards. Given its simplicity and ease of integration with existing methods, LiNeS offers a practical and inexpensive solution for boosting the generalization and robustness of fine-tuned models in diverse application domains.

##### Acknowledgments

The authors thank Adam Hazimeh for constructive discussions and comments.

## Appendix

### Appendix A LiNeS Pseudocode

We provide here a python pseudocode for the scaling the task vectors.

[bgcolor=white]python
def line_scaling(task_vector, alpha=0.0, beta=1.0, num_blocks=12):
”””
Progressively scales the task vector based on layer depth.

Parameters:
———–
task_vector : dict
A dictionary representing the residual between the fine-tuned checkpoint
and the pre-trained checkpoint.
alpha : float
The minimum scaling factor for the blocks.
beta : float
The maximum scaling coefficient difference between the last and first block.
num_blocks : int
The total number of layer blocks in the model.
Returns:
——–
scaled_task_vector : dict
A copy of ‘task_vector‘ where each key is scaled based on the layer depth.
”””

import copy

# Deep copy the task vector to avoid modifying the original
scaled_task_vector = copy.deepcopy(task_vector)

# Generate the key blocks corresponding to the layers of the model
key_blocks = [f”.layeri.” for i in range(num_blocks)]

# Create a scaling dictionary to store the scaling factor for each key
scaling_dic =
for k in task_vector.keys():
# Find the layer block in the key and assign scaling factor based on layer depth
for layer, block in enumerate(key_blocks):
if block in k:
scaling_dic[k] = alpha + beta * (layer / (num_blocks - 1))
break

# Scale the task vector based on the scaling dictionary
scaled_task_vector.vector =
# Use alpha if layer is outside residual blocks
k: task_vector.vector[k] * scaling_dic.get(k, alpha)
for k in task_vector.keys()

return scaled_task_vector

# example: scale single-task fine-tuned residual
task_vector = k:theta_t[k] - theta_0[k] for k in theta_0.keys()
scaled_task_vector = line_scaling(task_vector, alpha=gamma, beta=1.0-gamma, num_blocks=12)

# example: Scale the multi-task vectors
mtv = k: sum(theta_ft[k] - theta_0[k] for theta_ft in ft_models) for k in theta_0.keys()
scaled_mtv = line_scaling(mtv, alpha=1/len(ft_models), beta=beta, num_blocks=12)

### Appendix B Experimental Details

#### B.1 Descriptions of baseline model merging methods

- •
Task Arithmetic (Ilharco et al., ) generates a multi-task vector by summing the individual task vectors for each task. This multi-task vector is then added to the pre-trained checkpoint, with a scaling factor chosen based on validation set performance.
- •
Ties-Merging (Yadav et al., ) resolves parameter conflicts during model merging by first pruning parameters with lower magnitudes from the individual task vectors, followed by addressing sign mismatches, and finally merging parameters with consistent signs with averaging operation. The resulting multi-task vector is then added to the pre-trained checkpoint using a scaling factor determined from the validation set.
- •
Consensus Merging (Wang et al., ) enhances existing model merging techniques by eliminating redundant weights in the multi-task vector. It first identifies the relevant subset of parameters for each task, then filters out weights that are relevant to either none or only one task. After removing these redundant weights, the refined multi-task vector is added to the pre-trained checkpoint with a scaling factor selected from the validation set.

While consensus merging can be applied to various merging methods, in all our experiments, we evaluate only its application to task arithmetic.

#### B.2 Experimental details for multi-task model merging

##### B.2.1 Hyper-parameters tuning

We list here the hyper-parameter search space for each model merging method in Table B.2.1, while we suggest the authors to the original papers for a detailed description of these hyper-parameters. We highlight that applying LiNeS does not introduce extra computational cost in hyper-parameter search for the baseline merging methods.

| Method | With<br>LiNeS | Hyper-parameter search space |
| --- | --- | --- |
| Task Arithmetic | ✗ | constant scaling term for multi-task vector: [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0] |
| ✓ | scaling term $\beta$ in Eq. 1 for the multi-task vector: [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0] |  |
| Ties-Merging | ✗ | constant scaling term for multi-task vector: [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0,1.1,1.2,1.3,1.4,1.5] |
| ✓ | scaling term $\beta$ in Eq. 1 for the multi-task vector: [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0,1.1,1.2,1.3,1.4,1.5] |  |
| Consensus Merging | ✗ | constant scaling term for multi-task vector: [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0];<br>weight-pruning threshold: [1, 2] |
| ✓ | scaling term $\beta$ in Eq. 1 for the multi-task vector: [0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1.0];<br>weight-pruning threshold: [1, 2] |  |

##### B.2.2 Benchmarks

###### Image classification

For the benchmarks used in image classification, we utilized the 8-task benchmark initially proposed by Ilharco et al. (), as well as the 14-task and 20-task benchmarks expanded by Wang et al. ().

- •
The 8-task benchmark comprises the following tasks: Cars (Krause et al., ), DTD (Cimpoi et al., ), EuroSAT (Helber et al., ), GTSRB (Stallkamp et al., ), MNIST (LeCun, ), RESISC45 (Cheng et al., ), SUN397 (Xiao et al., ), and SVHN (Netzer et al., ).
- •
The 14-task benchmark includes the original eight tasks plus additional ones: CIFAR100 (Krizhevsky et al., ), STL10 (Coates et al., ), Flowers102 (Nilsback & Zisserman, ), OxfordIIITPet (Parkhi et al., ), PCAM (Veeling et al., ), and FER2013 (Goodfellow et al., ).
- •
The 20-task benchmark builds on the 14-task benchmark with the addition of: EMNIST (Cohen et al., ), CIFAR10 (Krizhevsky et al., ), Food101 (Bossard et al., ), FashionMNIST (Xiao et al., ), RenderedSST2 (Socher et al., ; Radford et al., ), and KMNIST (Clanuwat et al., ).

###### Natural Language Processing

For our NLP experiments, we utilized benchmarks established by Yadav et al. (), Tam et al. (), and Wang et al. ().

- •
The 7 NLP Tasks benchmark, as explored in Yadav et al. (), includes the following datasets: QASC (Khot et al., ), QuaRTz (Tafjord et al., ), PAWS (Zhang et al., ), Story Cloze (Sharma et al., ), WikiQA (Yang et al., ), Winogrande (Sakaguchi et al., ), and WSC (Levesque et al., ).
- •
The 8 QA Tasks (Tam et al., ) comprises the following datasets: CosmosQA Huang et al. (), QASC (Khot et al., ), QuAIL Rogers et al. (), QuaRTz (Tafjord et al., ), PAWS (Zhang et al., ), ROPES Lin et al. (), SocialIQA Sap et al. (), and WikiQA (Yang et al., ).
- •
The 11 NLP Tasks benchmark is a union of these two benchmarks, as studied in Wang et al. (). It contains the following tasks: QASC (Khot et al., ), QuaRTz (Tafjord et al., ), PAWS (Zhang et al., ), Story Cloze (Sharma et al., ), WikiQA (Yang et al., ), Winogrande (Sakaguchi et al., ), WSC (Levesque et al., ), CosmosQA Huang et al. (), QuAIL Rogers et al. (), ROPES Lin et al. (), and SocialIQA Sap et al. ().

#### B.3 Experimental details for single-task model merging

##### B.3.1 Description of Model Soups

Model soups (Wortsman et al., ) is a model merging method which averages the weights of multiple fine-tuned models with different hyper-parameter configurations, improving accuracy of the merged model without increasing inference or memory costs. The authors of model soups proposed two methods:

- •
Uniform soup: Averages the weights of all fine-tuned checkpoints, providing a simple and efficient way to improve performance.
- •
Greedy soup: Starting with the best-performing checkpoint, greedily and iteratively adds the next best-performing checkpoint to the soup, keeping those that improve accuracy of current collection of model checkpoints.

##### B.3.2 Experimental details for modifications to model soups

We describe in detail the modifications to model soups, namely task arithmetic and our proposed LiNeS. For reference, model soups merges the checkpoints by averaging the weights of the individual checkpoints:

$$ $\bm{\theta}_{\text{soup}}^{\text{vanilla}}=\bm{\theta}_{0}+\bm{\tau}_{\text{ soup}}$ (5) $$

###### Enhancing Model Soups with Task Arithmetic

We enhance model soups with task arithmetic, by introducing a scaling factor $\lambda_{\text{ta}}$ to $\bm{\tau}_{\text{soup}}$ in Equation 5. We search for this hyper-parameter within the range of $[1.0,1.1,1.2,1.3,1.4,1.5,1.6,1.7,1.8,1.9,2.0]$. Note that $\lambda_{\text{ta}}=1.0$ yields the vanilla model soups.

$$ $\bm{\theta}_{\text{soup}}^{\text{ta}}=\bm{\theta}_{0}+\lambda_{\text{ta}}\cdot \bm{\tau}_{\text{soup}}$ (6) $$

###### Enhancing Model Soups with LiNeS

We enhance model soups with LiNeS, by applying LiNeS to $\bm{\tau}_{\text{soup}}$ in Equation 5. For the scaling, we apply directly the scaling introduced in Equation 1 to create a scaled task vector $\bm{\tau}_{\text{soup}}^{\text{LiNeS}}$, fixing $\alpha$ to $1$ while searching the value for $\beta$ within the range of [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]. Note that $\beta=0.0$ yields the vanilla model soups.

$$ $\bm{\theta}_{\text{soup}}^{\text{LiNeS}}=\bm{\theta}_{0}+\bm{\tau}_{\text{soup }}^{\text{LiNeS}}$ (7) $$

We further note that, both Task Arithmetic and our proposed method introduce only one hyper-parameter to model soups, while the computational cost for hyper-parameter search is the same. When applying the modifications to greedy soup, we only apply them directly to the selected subset of checkpoints after the greedy selection process.

We search for the hyper-parameter within the validation set and report the performance on test set with the best hyper-parameter based on validation performance.

### Appendix C Additional Results

#### C.1 Different trade-offs for retention of task and control task performance

In Section 3, we need to balance two competing objectives: maximizing accuracy on the target task while preserving performance on the control task. To account for different user preferences, we scalarize these objectives by assigning varying weights to the target task accuracy. This weighting scheme can be adjusted depending on the scenario to reflect different priorities. Let $w_{\text{target}}$ represent the weight assigned to the target task accuracy, and $\text{M}_{\text{target}}$ and $\text{M}_{\text{control}}$ denote the normalized accuracies for the target and control tasks, respectively. The optimal value of $\gamma$ is selected to maximize the following weighted trade-off on the validation set:

$$ $w_{\text{target}}\text{M}_{\text{target}}+\text{M}_{\text{control}}$ $$

To account for the high variance in control task performance and to emphasize the target task, we assign it a weight of 2, signifying that its accuracy is prioritized twice as much as the control task’s accuracy.

**Table 5: Validation results on the target vs control performance benchmark, presented in Section 3, averaged over the 8 tasks. We balance two competing objectives with various scalarization weights $w_{\text{target}}.$ In the main text, we use $w_{\text{target}}=2$.**
| $w_{\text{target}}$ | Averaged normalized accuracy (%) |  |
| --- | --- | --- |
| Target task | Control tasks |  |
| 1 | 99.8 | 101.9 |
| 2 | 100.0 | 101.5 |
| 5 | 100.2 | 100.8 |

#### C.2 Detailed labels for Figure 2

We provide in Figure 6 the detailed labels corresponding to each scatter dot. Each scatter dot corresponds to applying a specific model (FT for fine-tuned model; PT for pre-trained model; LS for fine-tuned model edited with LiNeS) on different tasks.

Figure: Figure 6: Figure 2 with detailed label information. Each scatter dot corresponds to applying a specific model (FT for fine-tuned model; PT for pre-trained model; LS for fine-tuned model edited with LiNeS) on different task.
Refer to caption: x6.png

#### C.3 Ablations of different choices of scaling function

We provide in this section an ablation study on applying different scaling functions for LiNeS. In LiNeS we used directly $\lambda^{(\ell)}=\alpha+\beta\cdot\frac{\ell-1}{L-1}$ to scale different layers. Here we test the performance on multi-task model merging in vision benchmarks with the following choices for scaling functions $f(\cdot)$: linear scaling, quadratic scaling and square root scaling:

- •
linear scaling: $\lambda^{(\ell)}=\alpha+\beta\cdot\frac{\ell-1}{L-1}$
- •
square root scaling: $\lambda^{(\ell)}=\alpha+\beta\cdot{\left(\frac{\ell-1}{L-1}\right)}^{\frac{1}{
2}}$
- •
quadratic scaling: $\lambda^{(\ell)}=\alpha+\beta\cdot{\left(\frac{\ell-1}{L-1}\right)}^{2}$

We provide in Table  6 the performance of different choices of scaling on vision benchmarks with ViT-B/32. While using quadratic scaling sometimes outperforms using identify function, especially with a larger number of tasks during merging, the improvement is not substantial. Therefore, we choose the linear scaling to keep the method simple and general.

**Table 6: Ablation study for applying different scalings for LiNeS on vision benchmarks with ViT-B/32.**
| Method | Scaling function |  |  | ViT-B/32 |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  |  | 8 tasks | 14 tasks | 20 tasks |  |  |
| Task Arithmetic | linear |  |  | 74.2 | 69.1 | 63.4 |
| square root |  |  | 73.9 | 67.5 | 62.4 |  |
| quadratic |  |  | 73.8 | 69.2 | 64.6 |  |
| Ties-Merging | linear |  |  | 77.2 | 72.1 | 67.2 |
| square root |  |  | 76.9 | 70.4 | 65.6 |  |
| quadratic |  |  | 76.1 | 71.6 | 67.4 |  |
| Consensus Merging | linear |  |  | 77.6 | 73.6 | 68.6 |
| square root |  |  | 77.1 | 72.5 | 67.3 |  |
| quadratic |  |  | 77.1 | 73.9 | 69.0 |  |

#### C.4 Results for ViT-B/16 for multi-task merging

We provide in Table 7 the results complementary to Table 2 for using ViT-B/16 as the image encoder, where we observe similar performance gains and observations by using LiNeS as in Table 2.

**Table 7: Complementary to Table 2, for results obtained with ViT-B/16 as image encoder.**
| Method | with<br>LiNeS |  |  | ViT-B/16 |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  |  | 8 tasks | 14 tasks | 20 tasks |  |  |
| Zero-shot |  |  |  | 55.5 | 61.4 | 59.8 |
| Fine-tuned |  |  |  | 92.6 | 91.6 | 92.3 |
| Task Arithmetic | ✗ |  |  | 74.6 | 70.4 | 65.7 |
| ✓ |  |  | $\text{77.6}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.0})$ | $\text{72.7}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.3})$ | $\text{67.7}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.0})$ |  |
| Ties-Merging | ✗ |  |  | 79.1 | 73 | 68.1 |
| ✓ |  |  | $\text{79.9}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}0.8})$ | $\text{75.2}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}2.2})$ | $\text{71.2}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}3.1})$ |  |
| Consensus Merging | ✗ |  |  | 78.9 | 73.9 | 70.2 |
| ✓ |  |  | $\text{79.5}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}0.6})$ | $\text{75.8}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}1.9})$ | $\text{72.0}\ (+{\color[rgb]{0.0,0.5,0.0}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.0,0.5,0.0}1.8})$ |  |

#### C.5 Results for editing T5 with LiNeS

We repeat here similar experiments we performed on ViT-B/32 model in Section 3 with T5-large (Raffel et al., ). T5-large contains both encoder and decoder structure, with sequential residual blocks in both structures. We investigate separately how the shallow-layer updates in the encoder and decoder infect the target and control task accuracy.

We consider the 8-question-answering benchmark (Zhou et al., ), and plot in Figure 7 the averaged target and control task accuracy after applying LiNeS to

- 1.
only the decoder part (left),
- 2.
only the encoder part (middle),
- 3.
both the encoder and decoder part (right).

Figure: Figure 7: The impact of downscaling the shallower-layer parameter updates on T5-large model within the 8-question-answering benchmark (Raffel et al., ). Downscaling only on the decoder (left) architecture preserves full target performance, while slightly improving the control tasks performance. Downscaling on the encoder leads to performance degradation on target tasks.
Refer to caption: x7.png

We observe that, only downscaling on the decoder architecture fully preserves full target performance, while downscaling on the encoder, or on both encoder and decoder, leads to performance drop on the target tasks. On the other hand, downscaling on the decoder part slight improves control generalization of the fine-tuned model, which we do not observe from downscaling on the encoder, or simultaneously on the encoder and decoder. We also note that, unlike the case in vision, the fine-tuned checkpoints on this NLP benchmark actually improve over the zero-shot performance of the pre-trained model on control tasks.

These results motivate us to apply LiNeS to only the decoder part of T5-large when merging multiple checkpoints, which preserves full target task accuracy while slightly improving control task performance, leading to similar observation in applying LiNeS to the ViT-B/32 architecture in vision,

#### C.6 Additional results for improving WiSE-FT with LiNeS

##### C.6.1 Results for using ViT-B/16 as visual encoder

We provide in Figure 8 the results for applying LiNeS for improving WiSE-FT, using ViT-B/16 as visual encoder. The ViT-B/16 checkpoint obtained through fine-tuning the CLIP checkpoint on ImageNet with the same hyper-parameter configurations in Wortsman et al. (). From Figure 8 we observe that LiNeS improves over WiSE-FT for both ID and OOD accuracies, leading to similar observations as the results obtained with ViT-B/32.

Figure: Figure 8: Results for improving WiSE-FT with LiNeS on with ViT-B/16 model fine-tuned on ImageNet.
Refer to caption: x8.png

##### C.6.2 Individual results for 70 checkpoints

We provide in Figure 9 individual results separately for the experiments on the 70 individual model checkpoints. Note that here y-axis represents the averaged accuracy over 5 OOD datasets. From the figure, we observe that LiNeS consistently improves WiSE-FT in terms of both ID and OOD accuracies for most of the individual checkpoints.

Figure: Figure 9: Performance of applying LiNeS to WiSE-FT to each ViT-B/32 checkpoint fine-tuned on ImageNet.
Refer to caption: x9.png

#### C.7 Detailed performance on individual tasks for multi-task model merging

###### Image Classification

We provide the detailed performance on each individual task for multi-task model merging in image classification benchmarks, complementary to the results in Table 2 and Table 7 where the accuracies are averaged on the individual tasks.

The single-task performance is presented in Figure 10 for ViT-B/32, Figure 11 for ViT-B/16, and Figure 12 for ViT-L/14. From the results we observe that our method demonstrates a noticeable improvement over baseline merging techniques across individual tasks in all test scenarios.

Figure: Figure 10: Single-task accuracies for multi-task merging on image classification benchmarks for ViT-B/32.
Refer to caption: x10.png

Figure: Figure 11: Single-task accuracies for multi-task merging on image classification benchmarks for ViT-B/16.
Refer to caption: x11.png

Figure: Figure 12: Single-task accuracies for multi-task merging on image classification benchmarks for ViT-L/14.
Refer to caption: x12.png

###### Natural Language Processing

We provide in Figure 13 the detailed single-task performance for the three NLP benchmarks using T5-large, complementary to the results in Table 3. Similar to the observation in vision, our method provides a consistent improvement over baselines across individual tasks.

Figure: Figure 13: Single-task accuracies for multi-task merging on NLP benchmarks for T5-large.
Refer to caption: x13.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. The basic assumption of this work is that the degradation of performance on control tasks is largely due to these distortions in the shallow layers. However, it lacks of formal theory.

2. From Fig. 2 and Fig. 7, fine-tuned models have obtained high performance for both target task and control tasks in some cases. Can catastrophic forgettinh happen when the result is high?  Please discuss or analyze these specific cases where fine-tuning seems to improve performance on both target and control tasks.

3. The result of "Model 2" in Fig. 10 is different from the other 69 curves, what is the reason?
soundness: 3
strengths: 1. The paper is easy-to-follow.

2. The topic is essential yet the idea is moderate. Both regularized fine-tuning and model merging are important techniques for the community and the paper addresses two formulations simultaneously.

3. A new method that seems well-motivated and performs well on a lot of benchmarks.
summary: This paper introduces a post-training editing technique, LiNeS, designed to address catastrophic forgetting and facilitate model merging after fine-tuning. LiNeS scales parameter updates linearly with the depth of layers within the network. The technique has shown significant improvements across various domains, including OOD generalization, single-task and multi-task model merging, demonstrating its effectiveness.
weaknesses: 1. I like the fact that the authors simply scale the task vectors. But I was not clear why directly edits the difference between the fine-tuned and pre-trained checkpoint via a linearly scaling coefficient ($\alpha+\beta\frac{l-1}{L-1}$). Also, compare to existing methods, like `Ties-Merging` and `Consensus Merging`, what are the core strengths of this work? More efficient or simpler?

2. The point above also points to the limitation of the current framework, which is a lack of formal theory. It works very well experimentally, but some aspects are not very clear. Providing concrete theoretical analyses or proofs for linearly layer-wise scaling would be better.

3. There are numerous typos and bad grammar throughout the paper. The authors should do a very careful proofreading and fix all the errors.

4. Minor comments:
- Some of the content in the appendix can be put into the body to ensure that the body is a full 10 pages.
- It must be noted that the consistency in verb tense is not maintained throughout the document. For example, the second paragraph in the Introduction section.
- line 112 "task=specific"
-  Two identical sentences appear from line 223 to line 226.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: - What could be an explanation for why the benefit of LiNeS appear to be smaller for deeper models (e.g. ViT-L/14 in Table 2) and models with smaller patch sizes (e.g. ViT-B/16 in Table 5)?
soundness: 3
strengths: - The paper is well-motivated and easy to follow.
- The experimental results seem impressive.
summary: This paper introduces LiNeS, a post-training editing technique designed to enhance the generalization ability of fine-tuned models. LiNeS applies a layer-wise linear scaling function to modulate parameter updates, reducing them more in shallow layers than in deep layers. The method is evaluated across several settings to demonstrate its effectiveness.
weaknesses: - The validation set has a large influence on the OOD performance of the selected models [1]. I suggest the authors discuss the construction of the validation set and whether the same empirical trend can be observed with different validation sets, as suggested by [1].
- The evaluation is limited to the transformer architecture. Including other model architectures, such as ResNet-based CLIP models, would validate whether the benefits of layer-wise scaling extend beyond transformer architecture.
- There is a lack of sensitivity analysis for the hyper-parameter $\alpha$ and $\beta$ across different evaluation settings.
- Typos should be checked. For example,
    - enchancing → enhancing (Line 23)
    - bcenchmarks → benchmarks (Line 24)
    - task=specific → task-specific (Line 112)
    - Figure 4 → Table 4

[1] In search of lost domain generalization. In ICLR, 2021

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: - Are there additional experimental results for LLMs, including different models, datasets, and tasks?
soundness: 3
strengths: - The LiNeS method introduces a novel post-training editing technique that prevents catastrophic forgetting through hierarchical scaling of parameter updates.
- Extensive experiments validate the effectiveness of LiNeS across various scenarios, including single-task fine-tuning, multi-task model merging, and improved out-of-distribution (OOD) generalization.
- The LiNeS method is applicable not only to visual tasks but also to natural language processing tasks, demonstrating its cross-domain versatility.
- LiNeS can be integrated with existing multi-task model merging baselines.
summary: This paper presents a novel post-training editing technique called LiNeS, designed to address the problem of catastrophic forgetting in large pre-trained models during fine-tuning. Catastrophic forgetting refers to the phenomenon where a model loses its generalization ability on other tasks when fine-tuned for a specific task. LiNeS balances the trade-off between retaining the generalization capability of the pre-trained model and enhancing performance on specific tasks through hierarchical scaling of parameter updates.
weaknesses: - The article does not discuss the performance of LiNeS in long-term maintenance and updating of models, particularly when faced with evolving data distributions.
- There is a lack of experimental results for LLM tasks.
- Comparisons with the performance of other fine-tuning methods (such as LoRA) are missing.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1) Why a shallow layer must be updated less compared to a deeper layer? Theoretical analysis should be provided to why shallow layers need to be updated less than deeper layers and how does this method affects the model's representation space?
2) Clarify exactly which tasks were used as target and control in your results, and to explain their task selection criteria. Add more baselines that aims at protecting both generalization and the fine-tuning to a downstream task. This would clarify the performance improvements?
3) Conduct an ablation study comparing the linear scaling approach to a learned scaling approach, and analyze the resulting patterns in the learned scaling factors. 
4) How does the proposed method differ from applying different learning rates per layer (low LR for shallow layers and higher LR for deep layers) during fine-tuning? Include a comparison between the proposed post-training method and the layer-wise learning rate scheduling during fine-tuning. This would help clarify the unique benefits of the proposed approach.
5) Provide an experiment to compare your proposed approach to freezing different shallow layers and updating the deeper layers using fine-tuning. This includes, but not limited to, freezing the feature extractor and finetuning the classifier.
soundness: 3
strengths: 1) Overall the idea of the paper is pretty simple and should be easily reproducible. Further, the paper is clearly structured and mostly well-written.
2) The paper thoroughly evaluates the proposed method across multiple tasks. 
3) The results clearly demonstrates strong trade-off performance on both target and control tasks.
summary: The paper introduces a post-training technique designed to mitigate catastrophic forgetting; more specifically preserving pre-trained generalization while enhancing fine-tuned task performance. This is done by reducing the magnitude of parameter updates for the shallow-layers compared to the deeper-layers. The method boils down to a layer-wise linear interpolation between the fine-tuned and the pretrained model.
weaknesses: 1) Theoretical analysis have not been provided to support the proposed idea. 
2) Some results lack details. For example, Table 1 what is the selected target and control tasks? Were all other combinations tried out? Not enough baselines have been added. Comparing to pretraining and fine-tuning performances is not enough.
3) How does learning the scaling factors perform compared to linearly scaling the parameter updates? Moreover, does the learned parameters have a consistent increasing trend from shallow-to-deep layers?

Meta-review: Summary: LiNeS is a post-training method that applies layer-wise scaling of weight distances, preserving general features while enhancing task-specific deeper features in single- or multi-task settings. It can mitigate catastrophic forgetting and improve model merging.

Strengths: simple yet novel post-training technique; extensive experiments across vision and NLP; compatible with existing model merging baselines.

Weaknesses: only on transformer architectures; more baselines needed; lacking theoretical analysis (not required)

Reasons for decision: LiNeS' ability to effectively preserve pre-trained generalization while enhancing task performance through a simple, versatile post-training method outweighs its minor weaknesses above.

Decision: Accept (Poster)