Chameleon: A Flexible Data-mixing Framework for Language Model Pretraining and Finetuning

📄 arXiv: 2505.24844v1 📥 PDF

作者: Wanyun Xie, Francesco Tonin, Volkan Cevher

分类: cs.LG, cs.CL

发布日期: 2025-05-30

备注: ICML 2025

🔗 代码/项目: GITHUB

💡 一句话要点

Chameleon：一种灵活的数据混合框架，用于语言模型预训练和微调。

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

关键词: 数据混合 领域重加权 杠杆分数 语言模型预训练 语言模型微调

📋 核心要点

  1. 现有领域重加权方法计算成本高昂，且引入新数据时需重新训练。
  2. Chameleon利用杠杆分数量化领域重要性，并构建领域亲和力矩阵。
  3. 实验表明，Chameleon在预训练、小样本推理和微调中均有提升。

📝 摘要（中文）

训练数据混合对大型语言模型的泛化性能有很大影响。现有的领域重加权方法通常依赖于代价高昂的权重计算，并且在引入新数据时需要重新训练。为此，我们引入了一种灵活高效的数据混合框架Chameleon，该框架采用杠杆分数来量化学习嵌入空间内的领域重要性。我们首先构建一个基于领域嵌入的领域亲和力矩阵。由此产生的杠杆分数确定了一种混合方式，该混合方式对在嵌入空间中共享共同表示的领域进行加权。这种公式允许通过计算新的领域嵌入直接转移到新数据。在实验中，我们展示了在三个关键场景中的改进：（i）我们计算的权重以现有方法一小部分的计算量提高了预训练领域的性能；（ii）Chameleon可以适应数据变化而无需代理重新训练，从而在转移到新数据时提高小样本推理精度；（iii）我们的方法能够在微调中实现有效的领域重加权，始终提高所有微调领域上的一致测试困惑度。

🔬 方法详解

问题定义：现有的大语言模型训练中，数据混合策略对模型性能至关重要。然而，现有的领域重加权方法存在两个主要痛点：一是计算领域权重需要大量的计算资源；二是当有新的数据加入时，需要重新进行权重计算和模型训练，缺乏灵活性和可扩展性。

核心思路：Chameleon的核心思路是利用领域嵌入空间中的杠杆分数来量化领域的重要性，并以此为依据进行数据混合。通过构建领域亲和力矩阵，使得在嵌入空间中具有相似表示的领域能够被赋予更高的权重。这种方法避免了直接计算复杂的权重，并且能够方便地适应新数据的加入。

技术框架：Chameleon框架主要包含以下几个步骤：1. 领域嵌入：首先，将每个领域的数据映射到一个嵌入向量。可以使用预训练的语言模型或其他方法来获得这些嵌入。2. 领域亲和力矩阵：基于领域嵌入，构建一个领域亲和力矩阵，该矩阵表示了不同领域之间的相似程度。3. 杠杆分数计算：利用领域亲和力矩阵计算每个领域的杠杆分数，杠杆分数越高，表示该领域越重要。4. 数据混合：根据杠杆分数对不同领域的数据进行加权混合，用于模型的预训练或微调。

关键创新：Chameleon的关键创新在于使用杠杆分数来量化领域重要性，并将其应用于数据混合。与传统的领域重加权方法相比，Chameleon具有更高的计算效率和更好的可扩展性，能够适应新数据的加入而无需重新训练。此外，Chameleon通过领域亲和力矩阵，考虑了领域之间的关系，从而能够更准确地评估领域的重要性。

关键设计：领域嵌入可以使用预训练语言模型的CLS token的表示，或者对领域数据进行平均池化得到。领域亲和力矩阵可以使用高斯核函数计算，其中带宽参数需要根据具体数据集进行调整。杠杆分数可以通过对领域亲和力矩阵进行特征分解得到。数据混合时，可以使用杠杆分数作为权重，对不同领域的数据进行采样。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Chameleon在预训练阶段能够以更少的计算资源达到与现有方法相当甚至更好的性能。在小样本推理任务中，Chameleon能够适应新数据，显著提升推理精度。在微调阶段，Chameleon在所有微调领域上均优于均匀混合策略，降低了测试困惑度。具体提升幅度未知，但原文强调了“consistently improving”。

🎯 应用场景

Chameleon框架可广泛应用于各种语言模型的预训练和微调任务中，尤其是在数据分布不均匀或存在新数据加入的场景下。例如，可以用于提升特定领域语言模型的性能，或者用于构建能够适应不断变化的数据环境的自适应语言模型。该方法在教育、医疗、金融等领域具有潜在的应用价值。

📄 摘要（原文）

Training data mixtures greatly impact the generalization performance of large language models. Existing domain reweighting methods often rely on costly weight computations and require retraining when new data is introduced. To this end, we introduce a flexible and efficient data mixing framework, Chameleon, that employs leverage scores to quantify domain importance within a learned embedding space. We first construct a domain affinity matrix over domain embeddings. The induced leverage scores determine a mixture that upweights domains sharing common representations in embedding space. This formulation allows direct transfer to new data by computing the new domain embeddings. In experiments, we demonstrate improvements over three key scenarios: (i) our computed weights improve performance on pretraining domains with a fraction of the compute of existing methods; (ii) Chameleon can adapt to data changes without proxy retraining, boosting few-shot reasoning accuracies when transferred to new data; (iii) our method enables efficient domain reweighting in finetuning, consistently improving test perplexity on all finetuning domains over uniform mixture. Our code is available at https://github.com/LIONS-EPFL/Chameleon.