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

TL;DR

LiNeS addresses catastrophic forgetting and cross-task interference that arise when fine-tuning and merging pre-trained models. It introduces a post-training editing method that linearly scales residual updates by layer depth, with scaling factors $\lambda^{(\ell)} = \alpha + \beta \frac{\ell-1}{L-1}$, so shallow layers retain general pre-trained features while deeper layers preserve task-specific representations. The approach yields strong results across vision and NLP in single-task fine-tuning, multi-task model merging, model soups, and rewarded soups, including maintaining nearly all fine-tuned performance while restoring zero-shot generalization and improving robustness to distribution shifts. LiNeS is simple, computationally efficient, and complementary to existing techniques, enabling practical improvements in generalization and task merging, with open-source code available.

Abstract

Fine-tuning pre-trained models has become the standard approach to endow them with specialized knowledge, but it poses fundamental challenges. In particular, \textit{(i)} fine-tuning often leads to catastrophic forgetting, where improvements on a target domain degrade generalization on other tasks, and \textit{(ii)} merging fine-tuned checkpoints from disparate tasks can lead to significant performance loss. To address these challenges, 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. In multi-task model merging scenarios, 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. Our method is simple to implement, computationally efficient and complementary to many existing techniques. Our source code is available at https://github.com/wang-kee/LiNeS

Figures (16)

• 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.
• Figure 2: Our linear scaling (LiNeS) retains performance on both control and fine-tuned target tasks. Each dot represents a different model.
• Figure 3: Application of LiNeS to WiSE-FT wortsman2022robust improves performance on ImageNet and five different distribution shifts, resulting in a dominating Pareto Front over WiSE-FT.
• Figure 4: Applying LiNeS to Rewarded Soups rame2023model improves merging of LLM policies RL fine-tuned on different rewards with a dominating Pareto Front.
• Figure 5: Comparison of the scalings obtained by different methods on 8-task merging benchmark with CLIP ViT-B/32.