GEM：用于最优LLM数据策展的几何熵混合

# GEM: Geometric Entropy Mixing for Optimal LLM Data Curation
来源: https://arxiv.org/html/2605.26121
###### 摘要

LLM 预训练的有效性日益取决于数据组成而非单纯的数据量。然而，最优混合受到分类缺陷的阻碍：人工分类法存在本体论不对齐，而欧几里得聚类无法解决嵌入各向异性。我们引入 GEM (Geometric Entropy Mixing)，这是一个将数据整理重新表述为超球面上的变分问题，并辅以混合平衡正则化器的框架。通过解耦生成先验并使用可证明的 MM (Minorize-Maximize) 算法优化目标，GEM 有效对抗聚类塌缩，从而发现欧几里得启发式方法无法察觉的平衡语义结构。我们采用师生蒸馏将这种几何保真度扩展到网络规模的语料库，并引入几何影响分数 (Geometric Influence Score, GIS) 用于可解释的分类法生成。使用 11 亿参数模型的实验表明，GEM 在集成到 DoReMi 和 RegMix 等混合策略时，建立了新的最先进水平，将平均下游准确率提升高达 1.2%，并为可预测的数据混合提供了稳健的坐标系统。

预训练, 数据混合, 数据整理

## 1 引言

数据整理已成为大型语言模型 (LLM) 性能的决定性因素 (Hoffmann et al., 2022 (https://arxiv.org/html/2605.26121#bib.bib25); Gunasekar et al., 2023 (https://arxiv.org/html/2605.26121#bib.bib26); Penedo et al., 2023a (https://arxiv.org/html/2605.26121#bib.bib27))，将研究前沿从单纯的参数缩放转向异构数据源的战略性“混合”。随着缩放定律 (Kaplan et al., 2020 (https://arxiv.org/html/2605.26121#bib.bib24)) 的演进，核心挑战在于将大规模、非结构化的语料库划分为语义上不同且平衡的聚类，这是任何有原则的数据混合策略的先决条件 (Ye et al., 2024 (https://arxiv.org/html/2605.26121#bib.bib32))。然而，当代数据分类方法通常分为两类，两者都面临基本的理论和实践瓶颈。

参考图注图 1：语义聚类中几何失配的示意图。虽然 (a) 基于分类法的方法受限于刚性不对齐和高成本，且 (b) 欧几里得聚类无法处理嵌入各向异性导致聚类塌缩，但 (c) 我们提出的 GEM 框架利用基于 MM 的推理在超球面上生成平衡、语义不同的分区，且效率更高。第一类方法，基于分类法，依赖于人工定义的分类层级 (Brown et al., 2020 (https://arxiv.org/html/2605.26121#bib.bib15); Touvron et al., 2023b (https://arxiv.org/html/2605.26121#bib.bib29))。这些方法通常利用高容量 LLM 或集成模型为文档分配标签。然而，如图 1 (https://arxiv.org/html/2605.26121#S1.F1) (a) 所示，这种范式存在关键的本体论不对齐：以人为中心的类别通常不能反映自监督学习所需的潜在语义粒度。经验证据表明，即使在分类细微网络数据时，最先进的模型也展现出较低的人工标注一致性，这表明人工分类法未能捕捉模型相关知识真正潜在分布 (Maini et al., 2024 (https://arxiv.org/html/2605.26121#bib.bib31); Abbas et al., 2023 (https://arxiv.org/html/2605.26121#bib.bib30))。此外，标注成本使得这种方法不可持续，尤其是在模型开发动态变化、数据不断更新的情况下，持续重新标注语料库在操作上不可行。

另外，诸如 K-Means (MacQueen, 1967 (https://arxiv.org/html/2605.26121#bib.bib22)) 的无监督方法提供了可扩展的选项，但基于欧几里得几何。这与现代神经嵌入（例如 BGE (Xiao et al., 2024 (https://arxiv.org/html/2605.26121#bib.bib35))、RoBERTa (Liu et al., 2019 (https://arxiv.org/html/2605.26121#bib.bib39))）存在根本性不匹配，后者本质上位于为余弦相似度优化的高维超球面流形上。这种几何差异因各向异性而加剧，即所谓的“锥形效应”(Li et al., 2020 (https://arxiv.org/html/2605.26121#bib.bib37))，其中表示集中在狭窄、非均匀的子区域。因此，将欧几里得聚类应用于此黎曼空间会引发“聚类塌缩”，如图 1 (https://arxiv.org/html/2605.26121#S1.F1) (b) 所示，其中主导聚类吞噬了语义长尾，严重限制了模型泛化所需的多样性 (Ethayarajh, 2019 (https://arxiv.org/html/2605.26121#bib.bib36); Gao et al., 2021 (https://arxiv.org/html/2605.26121#bib.bib38))。

为了弥合这一差距，我们引入了 GEM (Geometric Entropy Mixing)，它将语义划分与神经表示的内在黎曼几何对齐。如图 1 (https://arxiv.org/html/2605.26121#S1.F1) (c) 所示，GEM 脱离了欧几里得启发式方法，将聚类任务表述为在单位超球面上使用混合平衡正则化器增强的熵正则化变分目标。通过显式解耦生成先验并将关于经验质量的平衡正则化器集成到 von Mises-Fisher 混合模型 (vMFMM) 中，我们的方法有效缓解了嵌入各向异性并防止了聚类塌缩。这使得 GEM 能够发现传统基于距离的方法无法察觉的细粒度语义结构和长尾分布，为数据混合提供了更具表达力的语义基础。从系统角度来看，GEM 通过师生蒸馏管道架构设计用于网络规模部署，实现了与语料库大小成线性时间复杂度的推理。此外，为了弥合几何聚类与以人为中心的数据整理之间的差距，我们引入了一种基于几何影响分数 (Geometric Influence Score, GIS) 的采样方法，用于为每个语义类别生成可解释的细粒度分类法及描述。大量实验表明，通过 GEM 导出的数据混合一贯产生更优的缩放定律，表现为相比竞争基线，验证困惑度更低，且在各种下游基准测试中性能增强。

我们的主要贡献总结如下：

- **具有平衡正则化器的几何公式。** 我们提出了一种超球面变分框架，并引入新颖的混合平衡正则化器，以有效防止嵌入各向异性下的聚类塌缩。
- **可证明的基于 MM 的推理算法。** 我们推导了一个可证明的 MM (Minorize-Maximize) 算法，保证单调上升，确保正则化目标的稳定收敛。
- **具有可解释性的可扩展部署。** 我们通过师生蒸馏实现线性时间推理，并引入几何影响分数 (Geometric Influence Score, GIS) 用于可解释的分类法生成。
- **数据混合中的持续增益。** 使用 11 亿参数模型的实验表明，在各种基准测试中，相比强基线具有一致的性能提升。

参考图注图 2：GEM 框架示意图。管道包括两个阶段：(1) 几何优化（教师）：我们使用混合 von Mises-Fisher (vMF) 分布在超球面上执行熵正则化聚类。一种 Minorize-Maximize (MM) 算法在种子语料库 $\mathcal{X}_{seed}$ 上迭代更新黎曼参数 $(\mu, \kappa)$，以发现语义结构。(2) 可扩展蒸馏（学生）：收敛的几何分区用于在 GIS 分数指导下对大规模语料库 $\mathcal{X}$ 进行伪标签化。这些标签随后被蒸馏到一个轻量级 FastText 分类器中，实现大规模高效推理。
## 2 相关工作

**用于 LLM 的数据选择与混合。** 数据混合策略对于优化 LLM 训练稳定性和泛化能力至关重要 (Brown et al., 2020 (https://arxiv.org/html/2605.26121#bib.bib15); Touvron et al., 2023a (https://arxiv.org/html/2605.26121#bib.bib10); Team, 2023 (https://arxiv.org/html/2605.26121#bib.bib16))。最近的研究引入了自适应重加权框架。例如，DoReMi (Xie et al., 2023 (https://arxiv.org/html/2605.26121#bib.bib14))、DoGE (Fan et al., 2023 (https://arxiv.org/html/2605.26121#bib.bib49))、Aioli (Chen et al., 2024 (https://arxiv.org/html/2605.26121#bib.bib50)) 和 RegMix (Liu et al., 2024 (https://arxiv.org/html/2605.26121#bib.bib13)) 利用梯度对齐、超额损失或性能回归等训练信号动态调整领域权重。为了进一步细化粒度，SampleMix (Xi et al., 2025 (https://arxiv.org/html/2605.26121#bib.bib28)) 对每个样本进行评估，而 TikMix (Wang et al., 2025 (https://arxiv.org/html/2605.26121#bib.bib19)) 基于数据影响动态校准混合权重。此外，QuadMix (Liu et al., 2025 (https://arxiv.org/html/2605.26121#bib.bib18)) 引入了一个统一目标来评估数据质量和多样性。尽管如此，这些方法通常将底层分类视为外生常数。它们的有效性根本上受限于初始划分的质量。如果分类法在语义上不对齐或存在噪声，即使是这些复杂的混合算法也难以隔离高效用数据特征。为了解决这一限制，我们认为精细的结构粒度是有效混合的先决条件。我们提出了一种几何感知分类方案，从潜在空间中诱导出语义连贯的划分，从而实现对高熵网络数据的稳健混合。

**预训练数据分类。** 最近有几项工作明确解决了大规模预训练数据的分类问题，大致可分为基于分类法和无监督方法。**基于分类法的方法**依赖于预定义的标签系统，并使用监督分类器或 LLM 分配类别。诸如 WebOrganizer 和 TnT-LLM 等系统采用基于 LLM 的管道，将网络文档注释到手动设计的分类法中 (Wan et al., 2024 (https://arxiv.org/html/2605.26121#bib.bib20); Wettig et al., 2025 (https://arxiv.org/html/2605.26121#bib.bib12))。虽然这些方法产生人类可解释的标签，但它们存在两个限制：(i) 强加的分类法反映的是人类定义的本体论，而非模型学习的潜在语义结构，导致潜在的本体论不对齐；(ii) 使用 LLM 进行大规模推理会带来巨大的计算成本，限制了可扩展性。**无监督分类方法**避免手动标签，而是对预训练编码器产生的表示进行聚类。典型技术包括 K-Means 或基于密度的聚类算法，如 HDBSCAN (MacQueen, 1967 (https://arxiv.org/html/2605.26121#bib.bib22); McInnes et al., 2017 (https://arxiv.org/html/2605.26121#bib.bib21))，如 NVIDIA Climb 等系统所采用 (Diao et al., 2025 (https://arxiv.org/html/2605.26121#bib.bib23))。虽然可扩展且无需标签，但这些方法通常在欧几里得空间中运行，并依赖于基于距离的目标。然而，在高维嵌入空间中，距离趋于集中，使得欧几里得邻近性成为语义相似性的一个弱代理。

## 3 方法

我们引入 GEM (Geometric Entropy Mixing)，这是一个球形混合建模框架，用于对网络规模的文本嵌入进行无监督语义划分。GEM 建立在单位超球面 $\mathcal{S}^{d-1}$ 上的方向统计基础上，并优化一个熵正则化变分目标，该目标显式地加入混合平衡正则化器，以减轻由嵌入各向异性引起的聚类塌缩。图 2 (https://arxiv.org/html/2605.26121#S1.F2) 提供了一个概览。下面，我们描述 (i) 几何问题设置，(ii) 带有平衡正则化器的熵感知 vMF 混合公式，(iii) 一个可扩展的基于 MM (minorize-maximize) 的推理方案，具有可证明的单调上升性，(iv) 利用几何影响分数 (Geometric Influence Scores, GIS) 的可解释分类法生成管道，以及 (v) 用于在万亿令牌语料库上高效部署的师生蒸馏框架。

### 3.1 问题重新表述

我们考虑对大规模归一化文本嵌入语料库 $\mathcal{X} = \{x_i\}_{i=1}^N \subset \mathbb{R}^d$ 进行无监督语义划分，其中每个 $x_i$ 经过 $\ell_2$ 归一化，因此位于单位超球面上：$x_i \in \mathcal{S}^{d-1} \coloneqq \{x \in \mathbb{R}^d : \|x\|_2 = 1\}$。我们的目标是学习一个划分 $\mathcal{C} = \{C_1, \dots, C_K\}$，使得聚类通过“语义方向性”可区分，从而为 LLM 预训练中的下游数据混合提供稳健的语义基础。

**动机：高维球面上的集中现象。** 一个经典的集中现象表明，在高维空间中欧几里得邻近性变得信息量更少；在 $\mathcal{S}^{d-1}$ 上，随机方向之间的角度集中在 $\pi/2$ 附近。

###### 引理 3.1 (超球面上的集中现象 (Ledoux, 2001 (https://arxiv.org/html/2605.26121#bib.bib51)))。

设 $x \sim \mathrm{Unif}(\mathcal{S}^{d-1})$。对于任意固定的 $p \in \mathcal{S}^{d-1}$ 和任意 $\epsilon > 0$，

$$\mathbb{P}\!\left( \left\| \langle x, p \rangle \right\| \leq \epsilon \right) ~ \geq ~ 1 - 2 \exp\!\left( -\frac{d \epsilon^2}{2} \right).$$ (1)

**备注。** 引理 3.1 (https://arxiv.org/html/2605.26121#S3.Thmtheorem1) 提供了一个直观理解：当 $d \gg 1$ 时，随机方向几乎正交。虽然真实的神经嵌入在 $\mathcal{S}^{d-1}$ 上并非均匀分布，但它们通常表现出强各向异性和“中心性”，使得纯欧几里得聚类不稳定。这激发了使用充分统计量为余弦相似度的球形分布来建模方向一致性。

**变分学习目标。** 我们寻找方向参数 $\Theta$ 和软分配 $\Gamma = \{\gamma_{ik}\}$，以拟合一个球形混合模型，同时显式鼓励“平衡”的聚类质量，以改善诱导数据混合的多样性。设经验（软）聚类质量为

$$\pi_k(\Gamma) \coloneqq \frac{1}{N} \sum_{i=1}^N \gamma_{ik}, \quad \boldsymbol{\pi}(\Gamma) \in \Delta^{K-1}, \quad \text{其中 } \mathbf{u} \coloneqq \tfrac{1}{K} \mathbf{1}.$$ (2)

我们优化一个熵正则化变分下界 (ELBO)，并加入一个混合平衡正则化器：

$$
\begin{split}
\max_{\Theta, \Gamma} \; & \underbrace{\sum_{i=1}^N \sum_{k=1}^K \gamma_{ik} \log\!\big( \alpha_k f_{ik}(\Theta) \big) + \sum_{i=1}^N H(\gamma_i)}_{\text{几何保真度 (ELBO)}} \\
& - \frac{\lambda}{2} \underbrace{\big\| \boldsymbol{\pi}(\Gamma) - \mathbf{u} \big\|_2^2}_{\text{混合平衡}}, \qquad \lambda > 0.
\end{split}
$$ (3)