DEM: Distribution Edited Model for Training with Mixed Data Distributions

📄 arXiv: 2406.15570v2 📥 PDF

作者: Dhananjay Ram, Aditya Rawal, Momchil Hardalov, Nikolaos Pappas, Sheng Zha

分类: cs.CL, cs.LG

发布日期: 2024-06-21 (更新: 2024-11-05)

备注: Accepted to EMNLP 2024 (Main Conference)

💡 一句话要点

提出分布编辑模型(DEM)，高效解决混合数据分布下的模型训练难题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合数据训练 多任务学习 指令跟随模型 模型组合 分布编辑模型

📋 核心要点

  1. 多任务和指令跟随模型训练面临混合数据分布的挑战，现有数据混合方法性能欠佳且训练成本高昂。
  2. 论文提出分布编辑模型（DEM），通过组合独立训练的模型，利用简单的向量运算实现高效优化。
  3. 实验表明，DEM在多个基准测试中显著优于现有方法，且训练成本更低，并具备良好的灵活性和可扩展性。

📝 摘要（中文）

在使用混合数据分布进行训练是创建多任务和指令跟随模型的常见且重要的组成部分。数据分布的多样性和联合训练的成本使得优化过程极具挑战性。数据混合方法在一定程度上解决了这个问题，但它在不同数据源上的表现并非最优，并且需要多次昂贵的训练运行。本文提出了一种简单而有效的替代方案，通过使用基本的逐元素向量运算将分别在每个数据源上训练的模型与基础模型相结合，从而更好地优化数据源。由此产生的模型，即分布编辑模型（DEM），比标准数据混合便宜11倍，并且在各种基准测试中优于强大的基线，在使用3B到13B大小的模型时，在MMLU上提高了高达6.2%，在BBH上提高了11.5%，在DROP上提高了16.1%，在MathQA上提高了6%，在HELM上提高了9.3%。值得注意的是，当修改单个数据源时，DEM不需要完全重新训练，因此对于使用多样化数据源进行训练来说，它非常灵活且可扩展。

🔬 方法详解

问题定义：在多任务学习和指令跟随模型训练中，如何有效地利用混合数据分布是一个关键问题。现有的数据混合方法虽然能够缓解这个问题，但往往在不同数据源上表现不佳，并且需要多次昂贵的训练过程，效率较低。因此，需要一种更高效、更灵活的方法来优化混合数据分布下的模型训练。

核心思路：论文的核心思路是将分别在不同数据源上训练的模型与一个基础模型进行组合，通过简单的逐元素向量运算来实现知识的融合和迁移。这种方法避免了从头开始训练整个模型，从而大大降低了训练成本，并提高了训练效率。同时，由于每个数据源的模型是独立训练的，因此可以灵活地调整和更新模型，而无需重新训练整个模型。

技术框架：DEM的整体框架包括以下几个步骤：1) 首先，在一个基础数据集上训练一个基础模型。2) 然后，在每个特定的数据源上独立训练一个模型。3) 接下来，使用逐元素向量运算将每个数据源的模型与基础模型进行组合，得到最终的DEM模型。这个过程可以表示为：DEM = Base Model + Σ (α_i * Source Model_i)，其中α_i是每个数据源模型的权重。

关键创新：DEM的关键创新在于它提出了一种简单而有效的模型组合方法，避免了昂贵的联合训练过程。与传统的数据混合方法相比，DEM能够更好地利用每个数据源的知识，并且具有更高的训练效率和灵活性。此外，DEM还具有良好的可扩展性，可以方便地添加或删除数据源，而无需重新训练整个模型。

关键设计：DEM的关键设计包括以下几个方面：1) 权重α_i的设置：可以通过验证集上的性能来调整每个数据源模型的权重，以达到最佳的性能。2) 逐元素向量运算的选择：可以使用加法、乘法等不同的逐元素向量运算来进行模型组合，具体选择取决于具体的任务和数据分布。3) 基础模型的选择：可以选择一个在通用数据集上训练的模型作为基础模型，以提供通用的知识。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DEM在多个基准测试中取得了显著的性能提升。例如，在MMLU上提高了高达6.2%，在BBH上提高了11.5%，在DROP上提高了16.1%，在MathQA上提高了6%，在HELM上提高了9.3%。更重要的是，DEM的训练成本比标准数据混合方法降低了11倍，这使得它成为一种更具吸引力的选择。

🎯 应用场景

DEM可广泛应用于多任务学习、指令跟随模型训练等领域，尤其适用于数据来源多样、数据量大的场景。该方法能够有效降低训练成本，提高模型性能，并具备良好的可扩展性，有助于加速人工智能应用在各个领域的落地，例如智能客服、自动驾驶、医疗诊断等。

📄 摘要（原文）

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.