Disperse-Then-Merge: Pushing the Limits of Instruction Tuning via Alignment Tax Reduction (2024)

Abstract: Supervised fine-tuning (SFT) on instruction-following corpus is a crucial approach toward the alignment of large language models (LLMs). However, the performance of LLMs on standard knowledge and reasoning benchmarks tends to suffer from deterioration at the latter stage of the SFT process, echoing the phenomenon of alignment tax. Through our pilot study, we put a hypothesis that the data biases are probably one cause behind the phenomenon. To address the issue, we introduce a simple disperse-then-merge framework. To be concrete, we disperse the instruction-following data into portions and then train multiple sub-models using different data portions. Lastly, we merge multiple models into a single one via model merging techniques. Despite its simplicity, our framework outperforms various sophisticated methods such as data curation and training regularization on a series of standard knowledge and reasoning benchmarks.

## introduction trained on trillions of tokens from webpages @xcite @xcite , large language models (llms) have demonstrated impressive capacity on obtaining general-purpose representations for various downstream nlp tasks. however, pre-trained language models may not follow human instructions @xcite and produce toxic, hallucinated, or biased content @xcite @xcite . to address the issue, supervised fine-tuning @xcite on instructionfollowing data has emerged as one of the de facto paradigms @xcite for aligning llms with human preferences. performance base llm llama-2-7b llama-2-13b (a) the performance on mmlu (5-shot, accuracy) vs. data size. 20% 40% 60% 80% 100% percentage of instruction data used 42 44 46 48 50 52 performance base llm llama-2-7b llama-2-13b (b) the performance on bbh (3-shot, exact match) vs. data size. however, with the size of instruction-following data rising, it has been observed that the performance of llm on standard knowledge and reasoning benchmarks does not always improve but exhibits degradation @xcite , i.e., the alignment tax @xcite , as is shown in figure 1 . in other words, simply scaling up the instruction-following data leads to a quick bump into the upper bound where the marginal return of increasing data size approaches zero or even minus. it is therefore non-trivial to unleash the full potential of large-scale instruction-following data. prior studies tend to attribute the alignment tax phenomenon to the low-quality samples within the instruction-following corpus @xcite , or the knowledge forgetting during the sft process @xcite . however, our pilot study in section 3 reveals that the quality issue and the catastrophic forgetting of pre-training knowledge are probably not the main cause of the alignment tax since the decline can be observed across corpora with varied sizes and quality. by analyzing the trend of loss descent during the sft process, we alternatively posit that the data biases fitted on the instruction data are probably one of the major causes behind it. specifically, during the tuning process, llms fit on dataset biases while acquiring instruction-following ability. in the beginning, the acquisition of generalizable ability predominates so the performance on knowledge and reasoning benchmarks improves. however, during the tuning process, the learning of generalization quickly stagnates and the model tends to acquire more data biases instead, which harms the parametric knowledge of llm and leads to a decline in related benchmarks. we propose a frustratingly simple dtm (disperse-then-merge) framework composed of three steps: to empirically verify the efficacy of the dtm framework, we conduct extensive experiments and evaluations across 9 benchmarks involving math reasoning, world knowledge, and code generation. the experiment results exhibit that dtm outperforms both (1) data selection methods that filter out low-quality samples @xcite ; and ## pilot study the existence of alignment tax indicates an upper bound of performance when directly increasing the data size at supervised fine-tuning. it is thus necessary to analyze the underlying cause for the alignment tax to unleash the full potential of instruction-following data. specifically, we first examine the intuitions that data quality and knowledge forgetting are responsible for the decline in conventional knowledge and reasoning benchmarks (section 3.1), and then posit our hypothesis that the biases during fitting the instruction-following data is probably one of the major causes (section 3.2). ## methodology as analyzed above, vanilla sft on the full volume of instruction-following data suffers from the assimilation of dataset biases, leading to inefficiency in exploiting large-scale instruction-following corpus. previously, @xcite discovered that when two bert-based classification models are merged together, the unshared knowledge within each model is mostly forgotten while the common knowledge is enhanced. getting inspiration from this, to unleash the full potential of large-scale instruction-following data, we propose a dtm framework, as shown in figure 7 . in a nutshell, we disperse the instruction-following data into multiple portions to obtain a series of submodels with different data biases. then through the fusion of multiple sub-models, we can aggregate their instruction-following capacities and eliminate their dataset biases at the same time. ## analysis the sft experiments on tülu-v2-mix have proven the efficacy of the proposed approach. to gain more in-depth insights, further exploration and analysis are detailed below. question1: how does dtm help? answer1: measuring the data biases or shortcuts on instruction following is challenging since we are agnostic to the specific form of the bias. therefore, we choose to quantify the data biases of llm through their error sets on mmlu @xcite . we plot the venn diagram for the error set of @xmath0. it can be observed that their error sets share a large portion (48.7%) but every sub-model has its own error cases (accounting for 4.2%, 4.6%, 4.8% and 3.3% of the entire error cases respectively), attributing to their unique shortcut. next, we bucket the test case of mmlu into different bins according to their count of appearance n in the error sets. for example, @xmath1. specifically, @xmath2. then we plot the accuracy of the fused model on each bin in figure 9 . from the figure, the accuracy on the first bin (@xmath3, suggesting that the common knowledge is retained. on the other hand, the high accuracy on the second bin means that the unique error cases of four sub-models are likely to be correctly solved by the fused model, which is evidence for the forgetting of unique data biases in four sub-models. question2: does dtm yield effective results across instruction-following data of varying sizes and domains? 5 , from which we can conclude that dtm is not constrained by the domain of the instruction-following data, but its superiority is influenced by the data size. question3: how is model fusion in comparison to model ensemble? answer: different from model fusion which aggregates the parameter of multiple models in the weight space, model ensemble aggregates multiple models by manipulating their logits. to draw a comparison of their effects, we substitute the model merging procedure in uniform soup @xcite and our approach with model ensemble, and the evaluation results on mmlu are shown in ## conclusion in this study, we target the alignment tax during the sft. through a series of pilot studies, we hypothesize that data biases are the root cause for the decline in standard benchmarks after an llm goes through the sft process. to deal with the problem, we propose a simple three-step framework to disperse the biases apart and employ model merging techniques to mitigate the effect of data biases. extensive experiments are conducted to empirically verify the efficacy of our approach and we hope our research will inspire more future work exploring the essence and mechanism of llm alignment together with its effects on the capacity of llm.