Towards Scalable Data Diversification for Language Model Pretraining via Leverage Score Sampling¶
会议: NeurIPS 2026
arXiv: 2609.32484
OpenReview: https://openreview.net/forum?id=eNBJGhDB8u
代码: https://github.com/mazailin/LevDataDiversity
领域: 预训练
关键词: 预训练, 数据多样性, 杠杆分数采样, 维度塌陷, 行列式体积扩张
一句话总结¶
针对大语言模型预训练数据过滤中质量筛选引发的“维度塌陷”与现有几何多样性方法算力代价过高的问题,本文提出基于杠杆分数采样(Lev)的可扩展数据多样化方法,利用单次共享 Gram 矩阵与 Cholesky 分解最大化表示子空间体积扩张,实现高达 72 倍加速并显著提升常识与代码下游表现。
研究背景与动机¶
在大语言模型预训练数据清洗流程中,基于分类器的模型质量过滤(如 DCLM-fastText、FineWeb-Edu 等)长期占据主导地位。这类方法通常依赖由高质量问答对或合成指令构成的狭窄正样本集来训练轻量级分类器,从而在海量网络语料(如 CommonCrawl)中筛选高分样本。然而,这种追求表面文本质量的做法不可避免地导致了严重的“维度塌陷”(dimensional collapse):入选文本在嵌入向量空间中高度局限于极窄的低维子空间,系统性排除了大量有价值的非典型领域与长尾知识(如代码片段、技术文档),致使预训练模型在特定下游任务上的表征能力显著退化。
为了缓解这种塌陷,研究者逐渐转向多样性驱动的数据选择。早期的领域平衡或基于属性的粗粒度配比只能间接缓解不均衡;而基于局部密度或聚类覆盖的采样虽然促进了样本在空间上的均匀分布,却无法在几何层面保证表征向量张成足够丰满的高维子空间——均匀的散点依然可能退化在平坦的低维流形上。近期直接优化协方差矩阵谱均匀性的方法(如 DiSF)虽然在几何层面上优化了集合多样性,但其贪婪选择过程要求在每轮迭代中为候选池里的每一个样本重新计算协方差矩阵和 Frobenius 范数,导致计算复杂度随批次大小呈平方级爆炸,面对百亿至万亿 token 规模的海量网络语料时完全无法扩展。
因此,该领域的核心矛盾在于:如何在直接最大化特征空间几何体积的同时,摆脱对候选样本逐一进行协方差矩阵重算的计算泥潭。本文从核回归与矩阵近似中的杠杆分数(Leverage Scores)获得启发,在理论上证明了候选样本的杠杆分数严格等价于向当前选定集合引入该样本时所带来的正则化 Gram 矩阵行列式体积的相对增量。核心 idea:将预训练数据几何多样性最大化转化为基于杠杆分数的逐批贪婪采样(Lev),通过共享正则化 Gram 矩阵的 Cholesky 分解将全体候选样本的多样性边际贡献评估转化为高效的下三角线性方程组求解,以 \(O(DBd^2)\) 的低复杂度实现高达 72 倍的算力加速与跨域表征无损保留。
方法详解¶
整体框架¶
杠杆分数采样(Lev)以整个语料库的文档嵌入向量集合 \(\mathcal{Z} = \{z_1, \dots, z_D\} \subset \mathbb{R}^d\)(单位范数归一化,\(\|z_i\|_2 = 1\))为输入,在给定的目标保留比例 \(\alpha\) 下,按批次(batch size \(B\))与更新步长(step size \(b\))迭代筛选兼具高几何多样性的子集 \(\mathcal{S}\)。整个流水线分为“分批初始化”、“共享 Gram 矩阵维护与 Cholesky 分解”、“批内全体候选杠杆分数并行求解”、“Top-\(b\) 贪婪吸纳与协方差增量更新”四个阶段,最终输出高多样性索引集合。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["语料嵌入输入<br/>Qwen3-Embedding 提取并归一化"] --> B["分批处理与初始子集构建<br/>批次大小 B 内随机初始化 b 个样本"]
B --> C["行列式体积扩张与杠杆分数映射<br/>计算正定 Gram 矩阵 G 并执行 Cholesky 分解 L"]
C --> D["共享 Gram 矩阵与 Cholesky 分解求解<br/>前代求解 L w = z 计算所有候选杠杆分数"]
D --> E["奇异值逆加权机制<br/>自适应放大低频奇异值分量对应的长尾样本"]
E --> F["贪婪选入 Top-b 样本与增量更新<br/>外积更新矩阵 K 迭代直至达到选择比例"]
F --> G["高质量多样化预训练子集"]
关键设计¶
1. 行列式体积扩张与杠杆分数映射:建立几何多样性测度
传统方法在评估新增样本对几何多样性的贡献时,往往缺乏直观且计算闭环的测度。Lev 基于多面体几何理论,将集合在表征空间中的多样性定义为其样本张成的多面体体积平方,即正则化 Gram 矩阵 \(G = \frac{1}{n} X^\top X + \lambda I_d\) 的行列式 \(\det(G)\),其中 \(X = [x_1, \dots, x_n]^\top \in \mathbb{R}^{n \times d}\) 为已选样本矩阵,\(\lambda > 0\) 为保证数值稳定性的微小扰动常数(实证设为 \(\lambda = (100 \times d)^{-1}\))。根据矩阵行列式引理(Matrix Determinant Lemma),向当前池中加入候选样本 \(z \in \mathbb{R}^d\) 后,新 Gram 矩阵的行列式满足: $\(\det(G + z z^\top) = \det(G) \cdot (1 + z^\top G^{-1} z)\)$ 由此推导出候选样本 \(z\) 的岭杠杆分数(Ridge Leverage Score)与体积扩张率之间精确的代数对应关系(定理 3.1): $\(\ell^\lambda(z; X) := z^\top \left(\frac{1}{n} X^\top X + \lambda I_d\right)^{-1} z = \frac{\det(G + z z^\top)}{\det(G)} - 1\)$ 该设计直接证明了杠杆分数就是候选样本对当前集合子空间几何体积的边际相对贡献率。选择杠杆分数最大的样本,等价于贪婪地最大化表征多面体的几何体积扩张,从而在数学上确立了高效挑选几何互补样本的理论准则。
2. 共享 Gram 矩阵与 Cholesky 分解求解:消除候选逐个协方差重算瓶颈
DiSF 等前序方法之所以极为缓慢,根本原因在于其目标函数要求在每次评估候选样本时,都必须将该候选样本并入集合并重新计算归一化的特征协方差矩阵及其 Frobenius 范数,单轮迭代评估 \(B\) 个候选需要 \(O(B^2 d^2)\) 复杂度。而 Lev 发现,在当前已选集合 \(X\) 给定后,所有未入选候选样本 \(z_i\) 的杠杆分数计算共享同一个 Gram 矩阵 \(G\)。
为此,Lev 采用 Cholesky 分解对半正定矩阵 \(G\) 进行因式分解:\(G = L L^\top\),其中 \(L \in \mathbb{R}^{d \times d}\) 为下三角矩阵。对于批次内剩余的所有候选样本构成的矩阵 \(Z \in \mathbb{R}^{(B-n) \times d}\),求解下三角线性方程组 \(L W^\top = Z^\top\)(即对每个候选求解 \(L w_i = z_i\)),随后候选样本的杠杆分数即可直接通过下三角解向量的二范数平方求得: $\(\ell^\lambda(z_i; X) = \|w_i\|_2^2\)$ 这一设计将昂贵的矩阵求逆转化为数值极其稳定的前代三角消元。对于批次中的所有候选,Cholesky 分解仅需执行一次(耗时 \(O(d^3)\)),而全体候选的得分评估总复杂度仅为 \(O(B d^2)\)。当批次大小 \(B > d\) 时,全流程复杂度从 DiSF 的 \(O(D B^2 d^2)\) 彻底降至 \(O(D B d^2)\),抹去了一个 \(B\) 的维度阶次,使大规模并行运行成为可能。
3. 奇异值逆加权机制:自适应放大低频长尾语义信号
为什么杠杆分数能够精准捕捉被质量过滤丢弃的长尾语义与跨域知识?定理 3.2 给出了统计层面的解析。对已选样本特征矩阵进行奇异值分解(SVD):\(\frac{1}{\sqrt{n}} X = U \Sigma V^\top\),其中 \(V = [v_1, \dots, v_d]\) 的每一列代表当前数据池的主语义正交基向量,特征值 \(\sigma_i^2\) 对应各主语义的出现频率,且满足归一化条件 \(\sum_{i=1}^d \sigma_i^2 = 1\)。将候选样本 \(z\) 在第 \(i\) 个主语义基上的投影记为 \(\tilde{z}_i = v_i^\top z\),则杠杆分数可以精确分解为: $\(\ell^\lambda(z; X) = \sum_{i=1}^d \frac{\tilde{z}_i^2}{\sigma_i^2 + \lambda}\)$ 该公式揭示了核心机制:各语义方向的贡献权重与该语义在已有集合中的出现频率 \(\sigma_i^2\) 成反比。当某一语义方向 \(v_i\) 在已选数据中几乎未出现时,\(\sigma_i^2 \approx 0\),其权重直接激增至 \(\frac{1}{\lambda}\);反之,高频成熟语义的权重被严重压缩。相比之下,DiSF 的方向导数准则等价于线性加权得分 \(d(z) = -\sum_{i=1}^d \sigma_i^2 \tilde{z}_i^2\)。当一个候选样本同时包含稀有语义和高频语义时,DiSF 的线性惩罚会完全掩盖稀有语义的贡献导致漏选;而 Lev 的倒数加权 \(\frac{1}{\sigma_i^2 + \lambda}\) 具有无限放大的灵敏度,能无条件将包含未见语义的样本推向采样前列。
损失函数 / 训练策略¶
数据筛选完成后,模型在筛选出的子集上采用标准的自回归因果语言建模损失进行预训练: $\(\mathcal{L}_{\text{pretrain}}(\theta) = -\sum_{t=1}^T \log P_\theta(u_t \mid u_{<t})\)$ 在实验中,统一采用 Llama 3.1 架构配置的 1.1B 模型(20 层 Transformer、隐层维度 2048、中间层 4096、注意力头数 32、SwiGLU 激活函数、RoPE 位置编码、上下文长度 2048),在 Megatron-LM 分布式框架下使用 AdamW 优化器(\(\beta_1=0.9, \beta_2=0.95\)、权重衰减 0.1、学习率 \(4\times 10^{-4}\) 配合余弦退火与 1000 步 warmup),总计训练 60B tokens(28,610 步迭代,全局 batch size 1024)。
实验关键数据¶
主实验¶
实验采用 Qwen3-Embedding-0.6B 作为文档嵌入模型,对 CommonCrawl(CC)无质量过滤的清洗语料(200B tokens)及 StarCoderData 六种主流编程语言(Python, Java, JavaScript, Go, PHP, Ruby,80B tokens)进行筛选。下游常识推理采用 OLMES 评测套件涵盖的 7 个自然语言任务(ARC-e, ARC-c, WinoGrande, HellaSwag, PIQA, Social IQa, CommonSenseQA),报告选择比例 50% 下的准确率;代码能力采用 MultiPL-E MBPP 与 HumanEval,报告 Bits-Per-Byte(BPB \(\times 10^{-2}\),越低越好)。
| 筛选方法 | 挑选比例 | ARC-e (%) | ARC-c (%) | HellaSwag (%) | PIQA (%) | CSQA (%) | Base Mean (%) ↑ | 代码 Mean (BPB) ↓ |
|---|---|---|---|---|---|---|---|---|
| Random | 50% | 48.50 | 27.82 | 46.50 | 69.40 | 43.00 | 47.66 | 65.96 |
| DSIR (目标重采样) | 50% | 53.50 | 27.73 | 42.00 | 65.20 | 42.26 | 46.86 | - |
| Density (局部密度) | 50% | 51.70 | 27.99 | 48.20 | 66.40 | 42.51 | 47.51 | 63.44 |
| DiSF (协方差谱优化) | 50% | 52.90 | 28.16 | 46.40 | 67.80 | 42.59 | 47.83 | 65.24 |
| Lev (本文方法) | 50% | 52.70 | 29.35 | 47.10 | 68.30 | 42.34 | 48.17 | 62.70 |
在 40% 筛选比例下,Lev 的 Base Mean 进一步达到 48.91%(高于 DiSF 的 48.35% 与 Random 的 47.66%)。DSIR 虽然在特定匹配领域(如 ARC-e 53.50%)表现突出,但在 HellaSwag 上暴跌至 42.00%,印证了单一领域重采样带来的严重维度塌陷;而 Lev 在各项常识任务上表现高度平稳,平均准确率稳居第一。
消融实验与多维多样性评估¶
在集合多样性客观指标(40% 筛选比例,Table 5)上,对聚类惯性(CluIne)、谱熵(Entropy)、局部近邻距离(KNN)、综合新颖度(NovelSum)及表征有效维度数(Vendi score)进行对比,验证了几何体积最大化的优越性;同时对比了在 8 张 RTX 5090 GPU 上不同方法处理 CommonCrawl 数据的端到端耗时(Table 3)。
| 评估配置 / 方法 | Vendi 分数 ↑ | 谱熵 (Entropy) ↑ | LogDet 体积 ↑ | KNN 独特性 ↑ | CC 筛选耗时 (\(b=1\)) | 相对加速比 |
|---|---|---|---|---|---|---|
| Random | 168.06 | 5.124 | 1.601 | 1.052 | - | - |
| DSIR | 167.73 | 5.122 | 1.599 | 1.053 | - | - |
| Density | 177.80 | 5.181 | 1.612 | 1.056 | - | - |
| DCLM-fastText | 178.85 | 5.187 | 1.611 | 1.055 | - | - |
| DiSF | 186.63 | 5.229 | 1.623 | 1.071 | 217.28 小时 | 基准 (1×) |
| Lev (本文方法) | 203.81 | 5.317 | 1.638 | 1.076 | 3.03 小时 | 72× |
在步长超参消融中,随着更新步长 \(b\) 从 1 增大至 8,Lev 的筛选耗时进一步由 3.03 小时急剧缩短至 35 分钟(代码数据从 1.03 小时缩短至 11 分钟),而 Vendi score 仍保持在 200 以上,下游 Base Mean 稳定波动,展现出极强的超参鲁棒性。
关键发现¶
- 质量过滤诱发跨域维度塌陷:当单独使用 DCLM-fastText 筛选 CC 语料并预训练模型时,其代码测试集的 Code Mean BPB 高达 266.77,而 Lev 筛选的 CC 仅为 256.43。启发式分类器统计显示,在挑选出的子集中,DCLM-fastText 召回的代码相关文档数量比 Lev 减少了 17.83%,Dolma3-fastText 与 Qurater 甚至减少超过 50%,证明主流质量打分器对非 QA 样式的代码等关键模态存在严重的系统性误杀。
- 几何多样性是无监督代码筛选的有效替代:在代码这种难以直接定义通用质量启发式规则的领域,单纯引入 Lev 进行多样性筛选即可使 StarCoderData 的下游 BPB 下降 3.08%~4.94%,且与质量过滤的网页数据混合训练时达到了最优的综合表现(Base Mean 51.53%,Code Mean 63.66)。
- 计算效率的代际跃迁:Lev 依靠单次 Cholesky 分解加前代三角消元,理论上省去了一个 \(B\) 的复杂度阶次,在真实硬件吞吐上实现了高达 72 倍的速度飞跃,彻底扫清了集合级几何多样性算法无法应用于工业级大规模预训练的工程障碍。
亮点与洞察¶
- 严格将杠杆分数与空间多面体体积建立映射:巧妙利用矩阵行列式引理,证明了岭回归中的统计杠杆分数在几何上严谨对应于样本对 Gram 矩阵行列式体积的相对扩张率,为抽象的多样性启发式规则赋予了坚实的代数根基。
- 单次 Cholesky 三角方程求解替代密集协方差更新:在算法实现上避开了矩阵直接求逆与候选遍历的组合爆炸,将所有未选候选的评分转化为共享三角矩阵下的线性求解,设计极其优雅且工程实现极其友好。
- 揭示了“质量过滤偏见”对通用能力的潜在危害:通过扎实的跨领域证据表明,追求单一参考分布对齐的质量过滤器会造成不可逆的跨域能力损失,为下一代预训练数据工程指明了“质量与多样性协同过滤”的必要范式。
局限与展望¶
- 依赖预训练嵌入模型的表征空间质量:Lev 的所有多样性计算均建立在冻结的文本嵌入模型(如 Qwen3-Embedding)之上,若嵌入模型本身对某些小众语种或特殊符号表征不佳,该偏差将被线性继承至采样结果中。
- 纯多样性目标缺乏质量甄别底线:作为纯粹的无监督几何扩张算法,Lev 无法自动剔除嵌入空间中位于稀疏离群点上的乱码、低俗文本或爬虫噪声,实际落地时仍需前置极其严格的启发式清洗与反垃圾流水线。
- 动态联合采样的演进空间:目前将 Lev 与质量过滤结合的方式仍是分步执行或分桶配比,如何将质量分数与杠杆分数统一内化至单一的可微或贪婪闭环准则中,是极具价值的未来探索方向。
相关工作与启发¶
- vs DiSF (Fan et al., 2025):DiSF 通过优化协方差矩阵的 Frobenius 范数来追求特征值均匀分布,但每选一个样本都需要对全体候选重新跟踪协方差并计算范数,复杂度高达 \(O(DB^2d^2)\);Lev 改用行列式体积准则与杠杆分数,复杂度降为 \(O(DBd^2)\),不仅提速 72 倍,且由于奇异值倒数加权的非线性放大机制,对长尾稀有语义的捕捉比 DiSF 更加灵敏。
- vs Density Sampling (Sachdeva et al., 2026):Density 依赖符号随机投影(SRP)哈希估计局部样本浓度并做逆密度加权采样,属于覆盖驱动(coverage-based)方法;然而局部均匀覆盖并不能保证全局几何子空间维度的充分张开(点可能均匀聚集在低维超平面上),而 Lev 直接优化 Gram 矩阵体积,实验证明其光谱熵与下游表征能力全面占优。
- vs DCLM-fastText (Li et al., 2024):DCLM-fastText 依赖 QA 和指令正样本对文本进行二分类质量打分,容易产生严重的单模态收敛;本文揭示了其对代码等关键模态的过滤缺陷,确立了 Lev 与质量过滤器互为补充、协同提质的预训练组合范式。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [严谨证明了杠杆分数与表示空间行列式体积扩张的等价性,从根本上重塑了预训练几何多样性采样的理论基础。]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 200B 网页与 80B 代码数据,涵盖 7 项语言理解与 7 项代码基准,兼具理论复杂度分析、跨域塌陷透视与真实 GPU 吞吐实测。]
- 写作质量: ⭐⭐⭐⭐⭐ [数学推导流畅自洽,算法结构与几何图景阐释清晰透彻,问题定义与实证结论逻辑闭环严密。]
- 价值: ⭐⭐⭐⭐⭐ [高达 72 倍的算力加速打破了多样性采样的工业应用瓶颈,揭示的质量过滤跨域偏见对前沿大模型数据清洗实践具有关键指导意义。]