DEM : Distribution Edited Model for Training with Mixed Data Distributions (2024)

Abstract: Training with mixed data distributions is a common and important part of creating multi-task and instruction-following models. The diversity of the data distributions and cost of joint training makes the optimization procedure extremely challenging. Data mixing methods partially address this problem, albeit having a sub-optimal performance across data sources and require multiple expensive training runs. In this paper, we propose a simple and efficient alternative for better optimization of the data sources by combining models individually trained on each data source with the base model using basic element-wise vector operations. The resulting model, namely Distribution Edited Model (DEM), is cheaper than standard data mixing and outperforms strong baselines on a variety of benchmarks, yielding upto 6.2% improvement on MMLU, 11.5% on BBH, 16.1% on DROP, 6% MathQA and 9.3% on HELM with models of size 3B to 13B. Notably, DEM does not require full re-training when modifying a single data-source, thus making it very flexible and scalable for training with diverse data sources. The code is available at https://github.com/amazon-science/dem-distribution-edited-model.

## introduction large language models (llm) go through an extensive pretraining on billions or trillions of tokens @xcite @xcite @xcite , but they typically require supervised finetuning on diverse instruction-following datasets for properly following human instructions @xcite @xcite . supervised training is crucial for ensuring that generated outputs meet user expectations and perform well on downstream tasks @xcite . the datasets for supervised training are often of different sizes and follow different distributions. recent state-of-the-art fine-tuning approaches @xcite demonstrate that training on multiple data distributions requires careful tuning of the mixing weights for each data source to capture the combined distribution and improve downstream task performance. tuning these weights in a data-mixing approach is a computationally expensive process. although, there are techniques to speed-up the search @xcite , the process remains time-consuming. moreover, when one or more new datasets are introduced, the weights for each dataset need to be re-tuned. this requirement makes the data-mixing approach inflexible and hard to maintain in a production environment. to address these challenges when fine-tuning an llm on a set of diverse data distributions, we propose a simple and efficient approach that combines individually trained versions of the base model using element-wise vector operations. our method focuses on the challenging setting of combining diverse data distributions that correspond to multiple tasks from different domains such as math, reasoning, conversations and coding. in particular, our goal is to better capture a diverse data distributions as opposed to editing the model on a single downstream task @xcite . hence, we call resulting model distribution edited model (dem, shown in figure 1 ). our experiments on a variety of downstream tasks show that dem is an effective, highly capable and low cost alternative to the models trained using data mixing methods. the primary benefit of the proposed approach is its ability to efficiently identify the optimal combination of data sources for training a model. instead of exhaustively training and validating on all possible combinations of data sources, which can be computationally expensive, we take a more streamlined approach. first, we finetune the original model on each individual data source independently with early stopping to obtain the optimal model. second, we extract source distribution vectors by subtracting the original model from the finetuned ones. lastly, we create the final model by adding a weighted combination of these distribution vectors to the base model, allowing it to capture the joint distribution of different data sources in a controlled manner while enabling incremental updates with new datasets. our contributions can be summarized as follows: hyperparameter tuning of data-mixing weights is a compute intensive process, and methods such as doremi @xcite and online data mixing @xcite have been proposed to speed-up the process for pretrained data-mixing either through a proxy-model training or through a multi-armed bandit approach respectively. @xcite used a submodular function to assign importance scores to tasks which are then used to determine the mixture weights. @xcite built a framework to find multiple diverse solutions in the pareto front of many objectives. in this work, we propose an alternative strategy for training with multiple data sources by using vector arithmetic to combine models fine-tuned on individual datasets, rather than mixing training data in specific proportions. ## background: data mixing let us consider a pretrained language model with parameters θ, and d 1 , d 2 , ..., d n denote n different supervised fine tuning datasets. each dataset can consist of a single or multiple tasks. the exact tasks may have an overlap between these datasets, however, the corresponding samples are unique to each dataset. standard data mixing @xcite methods create training batches by performing a weighted sampling from each of the training datasets d i . the goal is to learn a joint data distribution that can span all training datasets. ## proposed approach: distribution edited model (dem) in contrast to standard data mixing, we propose to learn each data distribution separately and combine them post training. in the following subsections, we present two variants of that lead to a distribution edited model that achieves this goal. ## conclusions and future work we proposed a simple and efficient approach for training on diverse data distributions that trains checkpoints individually on each data source and then combines them with basic element-wise vector operations. dem siginficantly outperforms the standard weighted data mixing in terms of downstream performance and overall compute cost. our experiments demonstrate that dem works with both single-task (e.g. math qa) and multi-task data distributions (e.g. sni, p3), and that they can be incrementally added to the pretrained model, resulting in improved downstream performance. we further performed extensive analysis to better understand the properties of the learned distribution vectors, finding that dem is better aligned with the individual models than baseline while remaining close to the original model. in future, it is important to evaluate the proposed approach using other model architectures e.g. encoder-decoder or mixture of experts model to better understand its effectiveness with other model designs. additionally, dem can be further improved by using more sophisticated methods for combining the individual checkpoints that can reduce the negative effects of interfering data distributions.