DoGE: Domain Reweighting with Generalization Estimation

📄 arXiv: 2310.15393v2 📥 PDF

作者: Simin Fan, Matteo Pagliardini, Martin Jaggi

分类: cs.LG, cs.AI, cs.CL

发布日期: 2023-10-23 (更新: 2024-02-05)

💡 一句话要点

提出DoGE以优化大语言模型的领域重加权问题

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

关键词: 领域重加权 泛化能力 双层优化 大语言模型 机器学习

📋 核心要点

  1. 现有方法在调整数据领域影响时依赖启发式和试错，缺乏系统性和原则性。
  2. DoGE通过双层优化算法获取领域权重，并根据这些权重对基础模型进行训练，优化了领域采样过程。
  3. 在SlimPajama数据集上，DoGE显著提升了基础模型的困惑度和少样本推理准确性，尤其在未见领域上表现优异。

📝 摘要（中文）

预训练数据的覆盖和组成对大型语言模型（LLMs）的泛化能力有显著影响。尽管这一点至关重要，现有的LLMs仍依赖启发式方法和试错来调整数据领域的影响。我们提出了领域重加权与泛化估计（DoGE），以原则性的方法优化从每个领域采样的概率（领域权重）。该方法包括两个阶段：首先，训练一个代理模型，通过双层优化算法获取领域权重；其次，根据学习到的领域权重对更大基础模型进行训练。实验表明，DoGE在SlimPajama数据集上显著提升了基础模型在六个任务上的困惑度和少样本推理准确性，并有效识别领域间依赖关系，在未见的目标领域上实现了更好的测试困惑度。

🔬 方法详解

问题定义：本论文旨在解决大型语言模型在预训练阶段对不同数据领域的重加权问题。现有方法往往依赖于启发式调整，缺乏有效的系统性策略，导致模型泛化能力不足。

核心思路：论文提出的DoGE方法通过双层优化算法，系统性地学习每个领域的权重，从而优化模型在不同领域的采样策略。这样的设计使得模型能够更好地适应目标任务的数据分布。

技术框架：DoGE的整体架构分为两个主要阶段：第一阶段是训练一个代理模型，通过双层优化获取领域权重；第二阶段是使用这些权重对更大基础模型进行训练，确保模型在各领域的泛化能力。

关键创新：DoGE的核心创新在于其系统性地优化领域权重的能力，区别于传统方法的启发式调整，能够更有效地识别和利用领域间的依赖关系。

关键设计：在实现过程中，采用了双层优化算法来获取领域权重，并在训练过程中动态调整领域采样比例，以确保模型在不同领域的训练效果最优。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在SlimPajama数据集上，DoGE显著提升了基础模型的性能，具体表现为在六个任务中困惑度和少样本推理准确性均优于基线方法，尤其在未见领域的测试困惑度上表现出色，展示了其有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等。通过优化领域重加权，DoGE能够提升模型在多样化任务中的泛化能力，具有重要的实际价值和广泛的应用前景，尤其是在处理未见领域数据时的表现。

📄 摘要（原文）

The coverage and composition of the pretraining data significantly impacts the generalization ability of Large Language Models (LLMs). Despite its importance, recent LLMs still rely on heuristics and trial and error to increase or reduce the influence of data-domains. We propose DOmain reweighting with Generalization Estimation (DoGE), which optimizes the probability of sampling from each domain (domain weights) in a principled way. Our approach is a two-stage process consisting of (i) training a proxy model to obtain domain weights using a bi-level optimization algorithm; (ii) training a larger base model by sampling training domains according to the learned domain weights. In our experiments, we extensively show how DoGE improves the generalization of the base model to any target data mixture. On the SlimPajama dataset, our base model gets better perplexity and few-shot reasoning accuracies across $6$ tasks compared to baseline methods. Moreover, aiming to generalize to out-of-domain target tasks, which is unseen in the pretraining corpus (OOD domain), DoGE can effectively identify inter-domain dependencies, and consistently achieves better test perplexity on the target domain.