RankMean: Module-Level Importance Score for Merging Fine-tuned Large Language Models (2024)

Abstract: Traditionally, developing new language models (LMs) capable of addressing multiple tasks involves fine-tuning pre-trained LMs using a wide collection of datasets, a process that often incurs significant computational expenses. Model merging emerges as a cost-effective alternative, allowing the integration of existing models fine-tuned on different tasks into a single model that performs well across all tasks, eliminating the need for additional training. In this paper, we propose RankMean, an algorithm for merging fine-tuned LMs without requiring any downstream data. RankMean determines merging coefficients based on the relative rankings of weight change magnitudes and applies these coefficients for module-wise integration of various fine-tuned models. Our experimental results demonstrate that RankMean outperforms existing baseline methods on multiple benchmarks. The code is available at github.com/VITA-Group/RankMean.

heading: 1 Introduction
text: The LLM research community has invested significant effort in developing robust models capable of tackling diverse downstream tasks such as mathematical reasoning (Hendrycks et al., 2021; Cobbe et al., 2021), program completion (Chen et al., 2021), and general question answering (Welbl et al., 2017). One approach to craft these models is to fine-tune large pre-trained language models specifically for the targeted downstream tasks. However, such a fine-tuning process involving a large number of parameters often incurs significant training costs (Bartoldson et al., 2023), especially when constructing models intended to address multiple tasks simultaneously, where more samples from diverse datasets need to be used in training.
Instead of fine-tuning LMs for solving multiple downstream tasks, model merging seeks to amalgamate multiple already trained model weights into a single model without the need for additional training. This approach is advantageous, especially in
the context of LMs, as it leverages the community’s existing efforts in fine-tuning LMs across diverse tasks, while also preserving inference time. Following this streamline of research, different model merging methods have been proposed, such as Fisher Merging (Matena and Raffel, 2022), RegMean (Jin et al., 2022) and DARE (Yu et al., 2023). However, these methods either require access to training samples, which is not always cheap and feasible in practical scenarios, or suffer from performance degradation after merging compared to their fine-tuned counterpart.
In this work, we propose RankMean, an algorithm that aims to merge multiple LMs fine-tuned on downstream tasks and overcomes the aforementioned limitation. Our algorithm considers each module of LMs as a basic unit (e.g., attentions, feed-forward networks, or a whole transformer block), adopts a module-wise merging scheme, and derives the merging coefficient purely based on the relative rank of the magnitude of changes in module’s weights within each model. Such rankbased coefficients are straight-forward to obtain, and do not need access to any data sample. We conduct experiments on multiple architectures and datasets and evaluate the merged models on various benchmarks. Specifically, although still suffering from some level of performance degradation when compared to the fine-tuned models, RankMean outperforms the simple averaging algorithm by clear margins and the most competitive baseline method, DARE, by up to 1.82 when applied to merge models fine-tuned for coding and mathematical tasks.

heading: 2 Related Works
text: Model merging aims to combine multiple models into a single one that achieves higher accuracy or robustness on multiple tasks. One of the typical ways to perform the merging is to adopt a weights averaging approach. Numerous efforts have been made to study the effectiveness of weights averaging under different settings. For example, Na-
1776
garajan and Kolter (2019) discovered that models trained from the same initialization on MNIST can be averaged without losing much accuracy. Frankle et al. (2020) demonstrated that the two models can be averaged without accuracy drops if the two models share some common optimization trajectories. Neyshabur et al. (2020) showed that interpolating models fine-tuned from the same pretrained weights can attain accuracies. Wortsman et al. (2022) proposes to average the weights of multiple fine-tuned models on the same dataset with different hyperparameter configurations and often leads to improved accuracy. Ainsworth et al. (2022) studies the property of mode connectivity and proposes more effective algorithms.
Several recent works have considered the application of model merging in the context of language models (Matena and Raffel, 2022; Jin et al., 2022; Ilharco et al., 2022; Ramé et al., 2023; Yadav et al., 2023). For example, Fisher Merging (Matena and Raffel, 2022) merges models by approximating the Fisher matrices. RegMean (Jin et al., 2022) proposes to use the representation of data to calculate the merging ratio of models. DARE (Yu et al., 2023) proposes a simple algorithm that prunes and rescales the delta vector for model merging. In this work, we opt for the same philosophy and propose a model merging algorithm for fine-tuned LLMs that achieves better performance than prior art on various downstream tasks.

heading: 3 Methodology

heading: 3.1 Preliminaries and Problem Formulation
text: We denote a language model by f(·) and its weights by θ ∈ Rd where d indicates the hidden dimension. In this paper, our objective is to combine multiple weights that correspond to the same language model f , denoted by θ1, θ2, . . . ,θM into one θmerged. The series of weights, i.e. θi, can be derived from different training tasks, and the merged weights model is expected to have good performance on all these tasks.
Existing works on model merging aim to find a set of coefficients, denoted as c1, c2, . . . , cM , that is leveraged to construct the merged weights by θmerged = ∑M i=1 ciθi. A simple baseline would be setting ci = 1M , i = 1, 2, . . . ,M . More advanced methods assign different ci to different weights. For example, Fisher Merging (Matena and Raffel, 2022) assumes ci ∈ Rd and calculates it with the following formula: c(j)i = F (j) i /( ∑M i=1 F (j) i ), where Fi = Ex∼DiEy∼pθi (y|x)∇θi(log pθi(y|x)) 2,
and obtains the merged weights by computing θmerged = ∑M i=1 ci ⊙ θi., where ⊙ denotes the point-wise product. Instead of using the same coefficients for all the layers for a model, RegMean (Jin et al., 2022) calculates the coefficients ci in a layer-wise manner. More specifically, for every layer j the coefficients are calculated as cji = G j i/( ∑M i=1G j i ), where G j i is derived based on the Gram matrix of features of training data input to the j-th layer of the i-th model. The weights in each layer of the combined models are computed based on the coefficients. More recently, DARE (Yu et al., 2023) proposed a simple algorithm to merge different finetuned LLMs without access to training data. It prunes the “delta” vector (i.e., the difference of weights between the finetuned model and the original model), and conducts a simple average between the delta vectors. In our work, we follow a similar philosophy to merge language models that are fine-tuned on different downstream datasets without resorting to the utilization of training data, aiming to build a strong model that are capable of solving multiple downstream tasks.

heading: 3.2 RankMean: Module-Level Importance Score for Merging Models
text: Most existing works on model merging leverage the training data to calculate the coefficients. However, the process of iterating over samples from pre-training datasets incurs a computational cost that is prohibitively expensive. These datasets contain an extensive number of samples, with a scale reaching billions of tokens. This renders the act of processing all samples through the models exceedingly costly. Similarly, downstream task datasets usually comprise thousands of samples, making the complete traverse also highly burdensome in computation, particularly as the number of parameters in models increases. To overcome these issues, we propose RankMean, a framework for merging multiple fine-tuned LLMs without the access to original training data. RankMean consists of two important designs that are different from existing methods on LLMs: (1) it adopts module-wise merging coefficients to combine the weights from multiple finetuned models (derived from the same pre-trained LLM); and (2) it estimates the cross-model importance scores based on relative change within the same model through rank.
Module-level Merging. Existing literature has demonstrated that different layers vary in their contributions to a specific task (Kim et al., 2024). It
has also been discovered that the attention and the FFNs modules have diversified characteristics (Dai et al., 2022). Nonetheless, employing uniform merging coefficients across all layers neglects the consideration of the varying significance of each module. To rectify this oversight, we propose the adoption of distinct coefficients for each module of models, thereby acknowledging and accommodating the different importance inherent to their contribution to the overall model performance. We achieve this by first defining a function I that maps weight vectors into a series of coefficients. More specifically, we assume θi = {w1i + ∆w1i ,w2i + ∆w2i . . . ,wni + ∆wni } where wji indicates the pretrained weights of the j-th module, ∆wji indicates the weight changes caused by fine-tuning and n denotes the number of modules in the model. Note that a module does not necessarily mean a linear layer but can also be a Transformer block. We define I as
I(w11 +∆w 1 1, . . . ,w n 1 +∆w n 1 ; . . . ,w n N +∆w n N )
= c11, c 2 1, . . . , c n 1 ; . . . ; c 1 N , . . . , c n N .
Consequently, the merged weights can be written as {∑i c1i (w1i + ∆w1i ), ∑ i c 2 i (w 2 i +
∆w2i ), . . . , ∑ i c n i (w n i +∆w n i )}. Importance Score Estimation. Essentially, I estimates the importance of different components for model merging. A straight-forward method to estimate the importance is to use the mean magnitude of parameters, which has been proven to be effective in locating essential parameters for network pruning. However, since different models can be trained on different numbers of samples with different hyperparameters such as learning rate, they can have drastically different distributions of parameters (refer to Figure 1), which undermines the effectiveness of estimated importance and hence the final performance. Using the magnitude of changes in parameters suffer from the same problem. As an alternative, we seek to use the relative order of the Frobenius norm of weight changes to indicate the importance of components. The mathematical expression can be written as follows:
cji = ArgSort(∥∆wji∥F ; ∥∆w1i ∥F , . . . , ∥∆wni ∥F )∑N
k=1 ArgSort(∥∆w j k∥F ; ∥∆w1k∥F , . . . , ∥∆wnk∥F )
,
where ArgSort(∥∆wji∥F ; ∥∆w1i ∥F , . . . , ∥∆wni ∥F ) indicates the ranking (i.e., the index after sorting in ascending order where 1 means the smallest) of ∥∆wji∥F among ∥∆w1i ∥F , ∥∆w2i ∥F ,
. . . , ∥∆wni ∥F . Intuitively, if a weight change ∆w ranks high within a model, it has a higher chance to receive a higher coefficient in merging. In Section 4.3, we also compare with a “reversed” form of I where we observe worse performance.
In our experiments, we consider each transformer block to be a “module” as we discover this leads to the best performance (refer to Section 4.3).

heading: 4 Experiments

heading: 4.1 Implementation Details
text: Architectures and Baselines. We conduct experiments with Falcon (Almazrouei et al., 2023) and Llama-2 (Touvron et al., 2023). We compare with multiple baseline methods: (1) SFT: where we use directly the fine-tuned model for evaluation; (2) Simple Averaging, where we directly calculate the mean weights of multiple models; (3) DARE (Yu et al., 2023), which prunes and rescales the weight change vectors first, and applies the average of them to the pre-trained model.
Hyper-parameters. The models fine-tuned from Falcon-7B are trained with LoRA (Hu et al., 2021) for 3 epochs, using a learning rate of 5 × 10−5 with a cosine annealing learning rate scheduler and a gradient accumulation steps of 4. The modules adapted by LoRA in each layer are self_attention.query_key_value, mlp.dense_h_to_4h and mlp.dense_4h_to_h. The models fine-tuned from LlaMA-v2-13B are downloaded from the huggingface hub (see Appendix A.1). In all experiments that use DARE, we have used a masking rate of 0.9 as suggested.
Evaluation. We use different benchmarks to evaluate the model’s capabilities. For coding evaluation, we evaluate models on HumanEval (Chen et al., 2021) and we report the pass@1 metric. For mathematical reasoning, we evaluate on both GSM8K (5-shot) (Cobbe et al., 2021) and MATH (zero-shot) (Hendrycks et al., 2021) and report the accuracy. We also report the performance on AlpacaEval (Dubois et al., 2023).

heading: 4.2 Experimental Results
text: In the first set of experiments, we fine-tune Falcon-7B on three different datasets, MathInstruct (Yue et al., 2023), CodeAlpaca (Chaudhary, 2023), and SciQ (Welbl et al., 2017), covering a wide range of skills. Note that we randomly sample 20K samples from the MathInstruct dataset to maintain the number of samples at the same level across tasks. The evaluation results are presented in Table 1, where we observe that RankMean achieves
the best performance. Notably, when merging models fine-tuned on maths and coding data, RankMean achieves an improvement of 0.53 on GSM8K and 0.04 on HumanEval, when compared to the secondbest scores. When merging other models, we observe that RankMean achieves performance at a comparable level with DARE, the currently most competitive algorithm for LM merging.
We continue to conduct experiments on LlaMA-v2-13B using the setting introduced in DARE and present the performance in Table 2. It shows that RankMean out-performs substantially simple averaging on downstream benchmarks and achieves better overall performance compared to DARE when merging LMs fine-tuned on coding and mathematical data. For the other two combinations, RankMean remains a top performer, consistently achieving better performance than simple averaging and DARE especially on AlpacaEval and HumanEval.

heading: 4.3 Ablation Studies
text: We compare the performance of our method with different variants: (1) Finer-Grained: where we define the modules as weights of linear layers; (2)
Reversed, where we change the formula of cji to
1/(1 + n− ArgSort(∥∆wji∥F ; ∥∆w1i ∥F , . . . ))∑ k 1/(1 + n− ArgSort(∥∆w j k∥F ; ∥∆w1k∥F , . . . )) .
The results are shown in Table 3, which demonstrates that RankMean and its finer-grained counterpart achieve comparable performance, while the reversed version under-performs the other two.

heading: 4.4 Visualization
text: In Figure 1, we visualize the average magnitude of parameter changes in each module after finetuning Falcon-7B on different datasets, where we can observe that different models exhibit different distributions. Such a difference creates a bias when
using absolute values of weight changes for deriving merging coefficients, which we speculate to be the reason for inferior performance.

heading: 5 Conclusion
text: We present RankMean, an algorithm for merging language models that are fine-tuned on different downstream tasks. We propose to calculate the relative rank of the mean magnitude of weight changes in each module and use the obtained rank to derive the coefficients for model merging. Even though our algorithm still faces some level of performance degradation when compared to the fine-tuned models, we demonstrate its effectiveness and superiority in mitigating the issue on multiple benchmark datasets and architectures.

heading: Limitations
text: The experiments conducted throughout this work are performed using a limited variety of tasks. Investigating if RankMean can be used to combine certain skills of the models (such as readability, conciseness, completeness, ...) could be an interesting research path, and we will explore it as future works. Besides, our algorithm explores only the merging between models with the same architectures. It would be also interesting to explore if models with different architectures can be merged.

heading: Ethics Statement
text: Merging model methods allow users to cheaply combine features of different models. Even though most of the use cases are harmless, there are scenarios where these methods can facilitate the spread of malicious models. One example of such harmful use cases would be to fine-tune popular pre-trained models to be explicitly harmful and merge them with different popular task-specific fine-tuned models to create malicious versions of these models.

heading: Acknowledgements
text: This work was performed under the auspices of the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DE-AC52-07NA27344 and was supported by the Laboratory Directed Research and Development (LDRD) program under project tracking code 22- SI-007. This work is reviewed and released under: LLNL-CONF-860667. G. Jacob Perin acknowledges support from São Paulo Research Foundation (FAPESP), grants 2022/11645-1, 2023/15047- 4 and 2022/15304-4.

heading: A More Experimental Details

heading: A.1 Model Architectures
text: The models fine-tuned from LlaMA-v21 and used in the Table 2 are WizardLM2 (Xu et al., 2023), WizardMath3 (Luo et al., 2023) and llama-2-13b-code-alpaca4.

heading: A.2 Dataset Statistics
text: Most of the datasets we have used in our experiments are publicly available on Huggingface5. The usage of these datasets is consistent with their intended use. MathInstruct, MATH and GSM8K contain samples related to mathematical reasoning. CodeAlpaca and HumanEval contain samples related to programming. SciQ, Evol-Instruct and AlpacaEval are general samples.

heading: A.3 Computational Resources.
text: We conduct experiments on NVIDIA RTX A6000 and NVIDIA A100 80G.

heading: A.4 Additional results
text: We conduct an additional set of experiments comparing RankMean to Multi-task PEM merging (Zhang et al., 2023), which directly averages the LoRA modules without merging them to the pretrained models first. In this experiment, Multi-task PEM merging uses uniform coefficients.
1https://huggingface.co/meta-llama/ Llama-2-13b-hf
2https://huggingface.co/WizardLM/ WizardLM-13B-V1.2
3https://huggingface.co/WizardLM/ WizardMath-13B-V1.0
4https://huggingface.co/layoric/ llama-2-13b-code-alpaca
5https://huggingface.co/datasets/
From Table 5 we can see that the performance of our algorithm still matches and outperforms these additional baselines.

heading: Future Work
text: ) could be an interesting research path, and we will explore it as future works. Besides, our algorithm explores only the merging between models with the same architectures. It would be also interesting to explore if models with different architectures can be merged