SuperMerge: An Approach For Gradient-Based Model Merging

TL;DR

SuperMerge presents a gradient-based approach to merge multiple fine-tuned models into a single multitask model by learning per-layer merging weights with a small validation-driven objective. The key idea is to represent the merged weights as ${m{ heta}}^m(j) = {m{ heta}}^p(j) + \sum_{i=1}^k \tanh(w(i, j)) ({m{ heta}}^f(i,j) - {m{ heta}}^p(i,j))$, where ${m{ au}}(i,j) = {m{ heta}}^f(i,j) - {m{ heta}}^p(i,j)$ and $W = [w(i,j)]$ is trained on a small validation set. The method uses a tanh activation to bound contributions, enabling negative and positive adjustments with a tiny number of parameters, and introduces hierarchical merging to reduce peak memory when combining many tasks. Experiments across NLP and CV tasks show that SuperMerge generally outperforms existing merging methods in-domain and generalizes better to unseen tasks, while hierarchical merging substantially lowers memory requirements with little to no loss in accuracy. These results imply a practical, scalable pathway to extend task coverage for large pretrained models without full re-fine-tuning.

Abstract

Large language models, such as ChatGPT, Claude, or LLaMA, are gigantic, monolithic, and possess the superpower to simultaneously support thousands of tasks. However, high-throughput applications often prefer smaller task-specific models because of their lower latency and cost. One challenge of using task-specific models is the incremental need for solving newer tasks after the model is already deployed for existing tasks. A straightforward solution requires fine-tuning the model again for both existing and new tasks, which is computationally expensive and time-consuming. To address this issue, we propose a model merging based approach called SUPERMERGE. SUPERMERGE is a gradient-based method to systematically merge several fine-tuned models trained on existing and new tasks. SUPERMERGE is designed to be lightweight and fast, and the merged model achieves similar performance to fully fine-tuned models on all tasks. Furthermore, we proposed a hierarchical model merging strategy to reduce the peak space requirement without sacrificing the performance of the merged model. We experimentally demonstrate that SUPERMERGE outperforms existing model merging methods on common natural language processing and computer vision tasks.

Figures (6)

• Figure 1: Overview of our proposed SuperMerge: The figure visualizes the merging steps of SuperMerge. The blue and yellow blocks represent the layer-wise task vectors and parameters in the large language model. The red blocks represent the layer-wise trainable parameters introduced by SuperMerge. For clarity, we only show the computation for the last layer of the merged model, as we presented in \ref{['merge_eqn']}
• Figure 2: Average performance of different model merging methods on 11 NLP tasks, using different hyperparameter $\lambda$.
• Figure 3: Variation of task vectors in PEFT fine-tuned model T5-IA$^3$. Each violin plot along the $x$-axis visualizes the distribution of the task vector per weight matrix (per layer). The $y$-axis shows the magnitude of task vector. The first half of the figure (blue and green) plots the layers in the encoder and the second half (blue, green, and red) plots the layers in the decoder. Observe that different layers in the same models have different magnitudes, and there is a clear cut-off between encoder and decoder. Similar pattern is observed for same layer in different models. For example, the first few layers of encoder (left most ones) show smaller magnitude than the latter layers. We provide the full size figures in the appendix for better readability (\ref{['fig:distribution_of_tv_copa_full']} and \ref{['fig:distribution_of_tv_rte_full']}).
• Figure 4: Illustration of hierarchical merging. Models from similar tasks (represented with a box of the same color) are merged to from intermediate models (represented with grey boxes).
• Figure 5: Visualization of merging weights.