Spokes：优化预训练数据选择的多样性策略

背景

在大规模语言模型的预训练阶段，数据选择（Data Selection）的质量直接决定了模型的最终性能。随着预训练语料库规模的爆炸式增长，从海量数据中筛选出高质量、高价值的样本变得至关重要。

多样性（Diversity）在数据选择中扮演着核心角色。通过减少数据中的冗余和重复，多样性能够显著提升在固定数据预算下的模型表现。然而，优化多样性是一个极具挑战性的任务，因为它是集合级属性（set-level property），取决于数据点之间的相互作用，而非单个样本的独立特征。

现有的主流方法通常依赖于代理指标（proxies）或近似算法来估算多样性。这些方法往往无法确保生成的子集具有足够高的多样性，导致模型可能仍然接触到大量重复或高度相似的内容，从而限制了泛化能力的提升。

核心内容

本文提出了一种名为 SPOKES 的新方法，旨在直接优化数据选择的多样性。该方法基于概率化的多样化框架，并引入了 G-Vendi 分数 作为核心优化目标，通过**指数梯度下降（exponentiated gradient descent）**算法进行求解。

1. 方法论：直接优化多样性

SPOKES 摒弃了传统的近似手段，转而直接对多样性进行优化。其核心在于使用 G-Vendi 分数来衡量数据集的多样性。G-Vendi 分数基于特征矩阵的特征值分布，能够更准确地捕捉数据子集内部的多样性程度。通过指数梯度下降算法，SPOKES 能够高效地找到使 G-Vendi 分数最大化的数据子集。

2. 实验评估

研究者在两个主流大规模预训练数据集上评估了 SPOKES 的性能：FineWeb 和 DCLM (Data Compilations for Language Models)。

• 多样性提升显著：在 50 万样本的子集上，SPOKES 生成的子集比随机采样（random sampling）的 G-Vendi 分数高出 +489，证明了其在减少冗余方面的巨大优势。
• 下游任务表现：
  • 仅优化多样性（SPOKES diversity-only）：在 DCLM 和 FineWeb 上，相比随机采样，平均下游性能分别提升了 +0.4 和 +0.5 个点。
  • 联合优化质量与多样性：当同时优化数据质量和多样性时，SPOKES 取得了最佳结果。在 DCLM 和 FineWeb 上，分别实现了 +1.5 和 +1.4 个点的性能增益。
• 对比基线：SPOKES 的表现优于所有基线方法，包括语义去重（semantic deduplication）和质量过滤（quality filtering）等现有主流技术。

关键要点

• 直接优化而非近似：SPOKES 通过基于 G-Vendi 分数的概率化框架，直接优化集合级的多样性，解决了传统代理指标无法保证足够多样性的问题。
• 算法效率：采用指数梯度下降（exponentiated gradient descent）作为优化器，使得在大规模数据集上优化多样性变得可行且高效。
• 显著的性能增益：
  • 在多样性指标上，SPOKES 比随机采样高出 +489 (G-Vendi score)。
  • 在联合优化质量和多样性的设置下，SPOKES 在 DCLM 和 FineWeb 数据集上分别带来 +1.5 和 +1.4 的性能提升，超越了语义去重和质量过滤等基线。
• 通用性验证：该方法在 FineWeb 和 DCLM 两个不同的预训练数据集中均表现出一致的优势，证明了其鲁棒性和通用性。

意义与影响

SPOKES 的提出为预训练数据选择领域提供了一个新的视角：多样性不应仅仅是事后筛选的辅助标准，而应作为核心优化目标之一。

1. 突破冗余瓶颈：通过直接优化集合级多样性，SPOKES 能够有效打破数据中的冗余循环，让模型接触到更多样化的语言模式和知识领域，这对于提升模型的泛化能力和鲁棒性至关重要。
2. 提升数据利用效率：在计算资源有限的情况下，SPOKES 证明了通过更智能的数据选择，可以在不增加数据量的前提下显著提升模型性能。这对于降低大模型训练的碳足迹和计算成本具有实际意义。
3. 推动数据工程精细化：SPOKES 展示了将复杂的数学指标（如 G-Vendi 分数）与高效的优化算法（如指数梯度下降）结合在数据工程中的可行性，为后续研究如何更精细地控制预训练数据的分布提供了方法论参考。

总之，SPOKES 不仅在理论上解决了多样性优化的难题，更在实践中证明了“高质量+高多样性”数据策略相对于传统“去重+质量过滤”策略的显著优势，为大模型预训练的数据构建提供了强有力的新工具。