CABS: Conflict-Aware and Balanced Sparsification for Enhancing Model Merging

zongzhen yang, Binhang Qi, Hailong Sun, Wenrui Long, Ruobing Zhao, Xiang Gao

TL;DR: We propose CABS, a method to address parameter overlap and weight imbalance in model merging, achieving superior results across various tasks and model sizes.

Abstract: Model merging based on task vectors, i.e., the parameter differences between fine-tuned models and a shared base model, provides an efficient way to integrate multiple task-specific models into a multitask model without retraining. Recent works have endeavored to address the conflicts between task vectors, one of the significant challenges faced by model merging, through sparsification; however, two issues significantly limit their performance: *high parameter overlap* and *unbalanced weight distribution*. To address these issues, we propose a simple yet effective framework called **CABS** (Conflict-Aware and Balanced Sparsification), consisting of **C**onflict-**A**ware Sparsification (CA) and **B**alanced **S**parsification (BS). CA reduces parameter overlap by applying masks during sequential pruning, ensuring that each task vector retains distinct, non-overlapping parameters. BS leverages $n$:$m$ pruning to preserve critical weights while maintaining an even distribution across layers. Our comprehensive experiments demonstrate that CABS outperforms state-of-the-art methods across diverse tasks and model sizes.

Lay Summary:

(1) Problem: Imagine you have several AI models, each an expert in a different task (like one for translating French, another for German). Combining them into a single, multi-talented AI is desirable but tricky. A common method, "model merging," often fails because the experts' internal "knowledge" clashes. Current attempts to simplify this knowledge to reduce clashes often leave too much overlap between experts, or the remaining important knowledge is unevenly distributed, leading to a poorly performing merged AI.

(2) Solution: We developed CABS (Conflict-Aware and Balanced Sparsification). First, its "Conflict-Aware" part carefully simplifies each expert's knowledge one by one, ensuring their unique contributions don't get muddled by using separate, non-overlapping internal "brain regions." Second, its "Balanced Sparsification" part ensures that the most critical remaining knowledge is evenly spread throughout the model, not all clumped in one area, preventing performance bottlenecks.

(3) Impact: CABS creates merged AI models that perform significantly better across a wide variety of tasks and model sizes. In our experiments, models merged with CABS were more versatile and often even outperformed a hypothetical "ideal" scenario where one would just pick the single best expert for each individual task. This research helps build more efficient and powerful multi-talented AI systems.

Link To Code: https://github.com/zongzhenyang/CABS

Primary Area: General Machine Learning->Transfer, Multitask and Meta-learning

Keywords: Model Merging, Pruning Technique, Task Vectors, Large Language Models, Conflict-Aware Sparsity (CA), Balanced Sparsity (BS)

Submission Number: 5789