Spokes: 优化多样化的预训练数据选择

# Spokes: 面向多样化预训练数据选择的优化
来源：https://arxiv.org/html/2606.15216
Clarence Lee DSO National Laboratories & Yejin Choi Stanford University & Luke Zettlemoyer University of Washington & Pang Wei Koh University of Washington & Hai Leong Chieu DSO National Laboratories

###### 摘要

多样性在数据选择中起着关键作用：在固定数据预算下，通过减少冗余和重复来提升性能。然而，直接优化多样性本身极具挑战性，因为它是一个集合级属性，取决于数据点之间的交互，而非单个样本。因此，现有方法通常依赖于代理指标或近似手段，这往往无法保证子集的充分多样性。本文通过引入一个基于 G-Vendi 评分的概率多样化框架，并采用指数梯度下降法进行优化，直接对多样性进行优化。我们的方法生成的子集在多样性上显著优于随机采样：在 50 万样本的子集上，G-Vendi 评分提升了 489 分。我们在 FineWeb 和 DCLM 上评估了该方法，结果一致优于现有方案。值得注意的是，Spokes（仅基于多样性）在 DCLM 和 FineWeb 上分别比随机采样平均提升了 0.4 和 0.5 个下游性能点。更重要的是，同时优化质量和多样性取得了最强结果：Spokes 在 DCLM 和 FineWeb 上分别取得了 1.5 和 1.4 个点的提升，超越了所有基线方法，包括语义去重和质量过滤。

## 1 引言

数据多样性是构建预训练语料库的关键因素。已有研究表明，显式地将多样性纳入数据混合设计与优化可以提升下游性能（Liu 等，2025 (https://arxiv.org/html/2606.15216#bib.bib4)；Fan 等，2025 (https://arxiv.org/html/2606.15216#bib.bib14)；Jung 等，2025 (https://arxiv.org/html/2606.15216#bib.bib1)）。然而，多样性本质上是一个集合级目标——它衡量的是集合的属性而非单个元素——直接优化它在计算上因子集选择的组合性质而变得困难，这一问题在预训练规模下尤为突出。为了解决这个问题，实际方法已被开发出来，并已应用于现代预训练流程中。这些方法通常依赖粗粒度的结构，例如将数据按主题（Wettig 等，2025 (https://arxiv.org/html/2606.15216#bib.bib18)）、技能（Chandramani 等，2026 (https://arxiv.org/html/2606.15216#bib.bib19)）或无监督聚类（Liu 等，2025 (https://arxiv.org/html/2606.15216#bib.bib4)；Diao 等，2025 (https://arxiv.org/html/2606.15216#bib.bib20)）进行聚类。虽然这些方法通过确保每个簇在最终混合中得到表征来有效近似多样性，但它们仍然依赖于低分辨率的数据划分，有效的多样性最多只能通过随机采样数据来近似。因此，长尾或代表性不足的知识可能无法被充分捕获。使用来自单个数据点的细粒度信号可以成为构建高多样性集合的更可靠代理，但如何在规模上可靠地实现这一点，文献中仍存在空白。

这一空白引出了我们的核心问题：如何在大规模预训练语料库中可靠地提取出多样化的子集，从而带来更好的下游结果？

就像多样化的归一化向量在二维空间中均匀分布在一个单位圆上——捕捉到轮辐（spokes）的意象——我们的方法 Spokes 是一种获取多样化数据集合的原则性方法。我们不依赖启发式规则，而是求解一个全局优化问题，分析每个数据点的贡献。然后我们利用这一点来提取对集合整体多样性有贡献的数据点。

我们的方法揭示出，可以从预训练语料库中提取出密集且高度多样化的子集。随机样本的多样性分数会迅速饱和，而使用 Spokes 优化后的权重选择数据点，在 50 万样本的子集上，G-Vendi（Jung 等，2025 (https://arxiv.org/html/2606.15216#bib.bib1)）分数显著提高了 489 分。这与现有尝试提高多样性的方法（如 SemDeDup（Abbas 等，2023 (https://arxiv.org/html/2606.15216#bib.bib7)），仅提升了 7 分）形成鲜明对比。

虽然现有方法将 G-Vendi 分数视为多样性的事后评估指标，但我們直接对其进行优化，并引入一套实用策略（参见第 4 节 (https://arxiv.org/html/2606.15216#S4)），使我们成功地从 Fineweb 和 DCLM 的大规模预训练语料库中提取出多样化集合。不仅所选子集的批量级别多样性分数有所提高，我们在两个数据集上也证明了其在改进下游性能方面的成功。这表明多样性的好处能够可靠地跨不同数据集（这些数据集包含不同程度的过滤和质量）传递。

虽然质量过滤传统上是数据选择的强基线（因为许多质量信号旨在与评估性能相关），但 Spokes 表明多样性也是一个需要的重要维度。通过平衡质量分数（基于模型分类器，2.95 vs 3.18）和更高的 G-Vendi 分数（425 vs 315），我们能够比仅基于质量的基线在评估分数上取得显著改进：在 DCLM 和 Fineweb 上分别提升了 1.0 和 1.9 分。

因此，我们研究了在预训练规模下直接优化多样性的数据选择方法。我们的贡献如下：(1) 我们引入了 Spokes，一种可扩展的多样性优化技术，能够高效地处理现代预训练数据集规模。(2) 我们展示了 Spokes 的有效性，能够从现有预训练语料库中提取出高度多样化的子集。(3) 我们同时优化质量和多样性，在预训练性能上取得了一致的增益。

## 2 背景

### 2.1 G-Vendi 作为多样性度量

*Vendi* 分数（Friedman and Dieng, 2022 (https://arxiv.org/html/2606.15216#bib.bib2)）被引入作为一个量化数据集多样性的原则性指标。具体来说，给定一组表示，Vendi 分数构建一个相似度矩阵，并考察其特征谱（特征值）熵的指数。当数据点高度相似时，大部分谱集中在少数方向上。当数据点多样化时，谱分布更均匀，表明存在许多独立的方向。

*G-Vendi 分数*（Jung 等，2025 (https://arxiv.org/html/2606.15216#bib.bib1)）通过测量*梯度空间*而非表示空间中的多样性，扩展了这一思想。G-Vendi 不直接比较输入样本，而是比较每个样本在训练过程中引起的梯度。形式上，令 ∇ℓ(x;θ) 表示损失函数关于模型参数 θ 对数据样本 x 的梯度，通过代理模型的反向传播计算得到。G-Vendi 分数的优点之一是它鼓励表示向量之间的正交性，从而促进学习集合的多样性。在梯度空间中，这减少了更新方向之间的冗余，导致更独立、信息更丰富的优化步骤。因此，参数更新之间的干扰减少，每次迭代的信息增益更高，从而提高了数据效率。

当使用余弦相似度时，两个样本之间的相似度由它们的 ℓ2 归一化表示的点积给出。令 g_i ∈ R^d 表示数据点 i 的逐样本梯度，并定义 X ∈ R^{n×d} 为矩阵，其第 i 行是 g_i^T，其中每个 g_i 已进行 ℓ2 归一化。所得相似度（核）矩阵为 K = X X^T ∈ R^{n×n}，因此每个元素 K_{ij} = g_i^T g_j 对应数据点 i 和 j 的梯度之间的余弦相似度。令 {σ_i} 表示 K 的特征值，并定义归一化特征值 λ_i = σ_i / Tr(K)，它们构成一个概率分布。G-Vendi 分数则定义为该谱的香农熵的指数：

G-Vendi(K) = exp( -∑_{i} λ_i log λ_i ). (1)

实际上，计算完整的核矩阵并非必要。我们利用 Vendi 分数核是半正定的这一事实，以及 Gram 矩阵与核矩阵共享相同的非零特征值。这使我们能够直接处理维度为 R^{d×d} 的 Gram 矩阵，这在计算上显著更高效且可扩展。

## 3 Spokes: 在梯度空间中优化数据多样性

### 3.1 Spokes: 实现高 G-Vendi 子集的可扩展优化

鉴于与 G-Vendi 相关的强大下游性能，我们引入了 Spokes，一种从大规模预训练语料库中提取高 G-Vendi 子集的方法。除了多样性之外，我们的公式还引入了每个样本的质量分数，这是受现代数据选择流程中质量过滤作用的启发。我们使用可调参数 α ∈ [0,1] 控制质量与多样性之间的权衡，其中 α=0 恢复为仅优化多样性。在我们的实验中，质量分数使用 FineWeb-Edu 分类器（Penedo 等，2024 (https://arxiv.org/html/2606.15216#bib.bib12)）获得，但该公式与任何逐样本质量度量兼容。

我们从一个离散的子集选择问题开始，目标是从 [n] 中选择一个固定大小 k 的子集 S，使其同时最大化质量和多样性。令每个样本 x_i 关联一个质量分数 q_i，并令 K_S 表示限制在所选子集上的相似度核。遵循 Nguyen 和 Dieng（2024 (https://arxiv.org/html/2606.15216#bib.bib3)）的对数形式质量加权 Vendi 目标，我们将优化问题定义为

max_{S⊆[n], |S|=k} α ln( (1/k) ∑_{i∈S} q_i ) + (1-α) ln Vendi(K_S). (2)

虽然这个公式直接捕获了质量与多样性之间的期望权衡，但由于子集选择的组合性质，优化它在计算上是棘手的。因此，我们采用下面的松弛优化。

为了构建 G-Vendi 使用的相似度表示，我们首先使用一个代理模型计算基于梯度的嵌入：

g_i = ∇_θ ℓ(x_i), i = 1,...,n. (3)

由于这些梯度是高维的，我们使用一个 Rademacher 矩阵 R ∈ {-1,+1}^{k×d} 应用随机投影：

z_i = (1/√k) R g_i. (4)

使用投影后的嵌入，我们通过余弦相似度核定义成对相似度。然后，我们通过为每个数据点添加连续的非负权重 w ∈ Δ^n 来松弛离散优化问题。在此松弛下，质量目标变为期望质量

Q(w) = ∑_{i=1}^{n} w_i q_i. (5)

为了将子集核 K_S 扩展到松弛设置，我们定义一个依赖于权重的核

K(w)_{ij} = √(w_i w_j) · (z_i^T z_j) / (‖z_i‖ ‖z_j‖), (6)

得到的松弛优化问题为

max_{w∈Δ^n} α ln Q(w) + (1-α) ln Vendi(K_w) (7)

我们使用指数梯度下降法优化此目标，同时在每次迭代中通过归一化强制执行单纯形约束，防止权重坍缩到少量高权重样本上。优化后，我们通过选择学习到的权重向量中前 k 个条目来恢复一个离散子集。我们在图 4 (https://arxiv.org/html/2606.15216#S6.F4) 中展示了该策略的有效性。Spokes 的完整公式总结在算法 1 (https://arxiv.org/html/2606.15216#alg1) 中。

算法 1 Spokes: 通过指数梯度下降进行概率 G-Vendi
1: 输入: 数据集 {x_1,...,x_n}，质量分数 {q_i}_{i=1}^n，学习率 η>0，迭代次数 T，初始分布 w^{(0)} ∈ Δ^n，权衡参数 α∈[0,1]
2: 步骤 1: 使用代理模型计算梯度嵌入: g_i = ∇_θ ℓ(x_i), i=1,...,n
3: 步骤 2: 采样一个 Rademacher 投影矩阵 R ∈ {-1,+1}^{k×d} 并投影嵌入: \~z_i = (1/√k) R g_i, ∀i
4: 步骤 3: 构建加权余弦相似度核: K(w)_{ij} = √(w_i w_j) · (z_i^T z_j) / (‖z_i‖ ‖z_j‖),
5: 步骤 4: 定义加权质量: Q(w) = ∑_{i=1}^{n} w_i q_i
6: 优化目标: 最大化 α ln Q(w) + (1-α) ln Vendi(K) 满足 w ∈ Δ^n
7: 步骤 5: 定义损失: L(w) = - (α ln Q(w) + (1-α) ln Vendi(K))
8: for t=0 到 T-1 do
9: 计算梯度: g^{(t)} = ∇_w L(w^{(t)})
10: 指数梯度更新: \~w_i^{(t+1)} = w_i^{(t)} exp(-η g_i^{(t)}), ∀i
11: 归一化: w^{(t+1)} = \~w^{(t+1)} / ∑_{j=1}^{n} \~w_j^{(t+1)}
12: end for
13: 输出: 子集 S = Top-k(w^{(T)})

### 3.2 优化在质量与多样性之间产生平滑的权衡

Spokes 在实践中非常有效，我们观察到质量与多样性之间存在平滑的权衡。图 2 (https://arxiv.org/html/2606.15216#S3.F2) 显示，增加 α 会提高平均质量，同时以一种可控的方式降低 G-Vendi 分数。这种优化在完整预训练集上并测量从该集中采样的批次的 G-Vendi 和质量分数时是一致的。

参见图注
图 1: 不同 α 下的质量-多样性权衡
参见图注
图 2: 仅质量、仅多样性和联合优化子集之间的重叠

为了选择最终的 α 值，我们测量联合优化子集与两个极端解（仅质量和仅多样性）之间的重叠。对于 DCLM，α=0.001 产生了平衡的权衡：与仅多样性子集的重叠为 64.8%，与仅质量子集的重叠为 68.8%。对于包含更高比例低质量文档的 FineWeb，