DEM: Distribution Edited Model for Training with Mixed Data Distributions

TL;DR

The paper introduces DEM, a cost-effective method that combines individually trained models on different data sources to improve multi-task training, outperforming traditional data mixing approaches.

Contribution

DEM offers a simple, scalable, and more efficient alternative to data mixing by combining models with element-wise operations, reducing costs and enhancing performance.

Findings

DEM is 11x cheaper than standard data mixing.

DEM improves performance on multiple benchmarks, up to 16.1%.

DEM does not require full re-training when data sources change.

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 11x 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% on 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.