LM-Cocktail: Resilient Tuning of Language Models via Model Merging

Abstract: The pre-trained LLMs are continually fine-tuned to better support downstream applications. However, this operation may result in significant performance degeneration on general tasks beyond the targeted domain. To overcome this problem, we propose LM-Cocktail which enables the fine-tuned model to stay resilient in general perspectives. Our method is conducted in the form of model merging, where the fine-tuned LLM is merged with the pre-trained base model or the peer models from other domains through weighted average. Despite simplicity, LM-Cocktail is surprisingly effective: the resulted model is able to achieve a strong empirical performance in the whole scope of general tasks while preserving a superior capacity in its targeted domain. We conduct comprehensive experiments with LLama and BGE model on popular benchmarks, including FLAN, MMLU, MTEB, whose results validate the efficacy of our proposed method. The code and checkpoints are available at https://github.com/FlagOpen/FlagEmbedding/tree/master/LM_Cocktail.

• The paper introduces LM-Cocktail, a method that merges fine-tuned and pre-trained models to mitigate catastrophic forgetting.
• It uses weighted parameter averaging based on few-shot examples to balance performance across general and target tasks.
• Experiments on both encoder (BGE) and decoder (LLaMA) models demonstrate consistent improvements across multiple benchmarks.

"LM-Cocktail: Resilient Tuning of LLMs via Model Merging"

Abstract and Introduction

The paper "LM-Cocktail: Resilient Tuning of LLMs via Model Merging" (2311.13534) addresses the challenge of fine-tuning pre-trained LLMs for specific tasks without compromising their performance on general tasks. Traditional fine-tuning often results in catastrophic forgetting, where a model loses its ability to perform well on tasks outside its fine-tuning domain. This research proposes LM-Cocktail, a model merging method that preserves general capabilities while enhancing performance on targeted tasks.

The proposed method combines fine-tuned models with pre-trained base models and models from other domains using weighted averages. The approach is both simple and effective, applying as a post-processing step after fine-tuning, thereby maintaining compatibility with existing training pipelines. Empirical evaluations on LLaMA and BGE models demonstrate the efficacy of LM-Cocktail across various benchmarks, proving its potential for universal application.

Figure 1: The illustration of LM-Cocktail demonstrating improved accuracy on new target tasks while maintaining performance on other tasks.

LM-Cocktail: Framework and Variations

General Paradigm

The LM-Cocktail approach involves merging models by averaging their parameters based on performance on few-shot examples from the target domain. Given a pre-trained base model and a set of domain-specific fine-tuned models, LM-Cocktail constructs a resilient-tuned model that integrates strengths from multiple models. The merging formula is expressed as:

$$\mathcal{M}_r \leftarrow \alpha \mathcal{M}_t + (1-\alpha) \sum (w_i * \mathcal{M}_i)$$

Where $\mathcal{M}_r$ is the resilient-tuned model, $\alpha$ is a hyper-parameter, and $w_i$ are weights computed from prediction losses on target domain examples.

Variations

LM-Cocktail includes adaptations for scenarios where fine-tuned models in general domains are not available:

• Mono-Specialist: Merges only the base and targeted fine-tuned model.
• Without Fine-tuning: Uses peer models from other general domains when targeted data is insufficient.

Experimental Setup

Decoder-based and Encoder-based Models

Experiments are performed using decoder-based (LLaMA) and encoder-based (BGE) models, across benchmarks like FLAN, MMLU, and MTEB. Fine-tuning involves various tasks and datasets, with evaluations conducted on unseen tasks using LM-Cocktail variants.

Results

LM-Cocktail consistently enhances performance compared to standard fine-tuning:

Figure 2: Performance with different alpha values demonstrates fine-tuning efficacy across tasks.

• Decoder Models: Retain performance in the target domain and exhibit enhanced general task capabilities. LM-Cocktail$_{2}$ (merging base and target models) and LM-Cocktail$_{10}$ (including general domain specialists) outperform fine-tuned models on unrelated tasks.
• Encoder Models: Provide similar improvements, validating LM-Cocktail's universality across model types.

LM-Cocktail without Fine-tuning

Investigating scenarios where fine-tuning is infeasible, LM-Cocktail leverages models from other domains to enhance performance on new tasks. Results from additional tasks in MMLU show LM-Cocktail surpasses traditional methods, proving effective even with minimal data.

Analysis of Merging Approach

Impact of Weight Alpha

Adjustments to the merging weight $\alpha$ reveal improvements on general tasks with slight trade-offs in target task performance, emphasizing the method's flexibility to tune model characteristics dynamically.

Figure 3: Performance of encoder-based LMs with different merging weights.

Example Number Effectiveness

Testing various example sizes confirms that LM-Cocktail achieves stability with few samples, highlighting efficiency in model merging processes.

Conclusion

LM-Cocktail presents a practical solution for the problems of catastrophic forgetting in LLMs. Its ability to merge models efficiently broadens deployment possibilities without significant computational overhead. Future work may explore further optimization strategies for merging weights and extend the approach to other model types.

Overall, LM-Cocktail is a versatile tool in the domain of AI model tuning, offering resilience without sacrificing performance. The provided empirical evidence establishes a foundation for continued exploration and application in diverse AI systems.