BLADE: Scalable Bi-Level Adaptive Data Selection for LLM Training¶

会议: NeurIPS 2026
arXiv: 2606.18650
代码: https://github.com/cxcscmu/MATES
领域: 模型压缩
关键词: 数据选择, 双层优化, 无Hessian矩阵, 超额损失, 在线批次课程

一句话总结¶

针对大语言模型海量预训练中样本选择面临的影响函数二阶逆Hessian计算爆炸及静态参考模型超额损失表征漂移难题，本文提出了基于拉格朗日罚函数单层重构的无Hessian双层自适应数据选择框架 BLADE，通过动态同步参考模型与无记忆随机块坐标 Frank-Wolfe 在线 token 级掩码选择，实现了严格的一阶收敛保证并显著超越了现有数据选择基线。

研究背景与动机¶

随着大语言模型（LLM）预训练语料迈向数万亿 token 级别，盲目在全量网页抓取语料上训练面临严重的边际收益递减。未经筛选的海量文本充斥着低质量噪声、重复内容以及非信息性 token，容易严重损害模型的深层推理与泛化能力。因此，如何根据模型当前学习状态自适应地挑选最具信息增益、最能提升下游验证能力的优质数据，成为提升预训练与持续预训练效率的核心前沿课题。

当前模型感知（model-aware）的主流数据选择方法主要分化为两大阵营，但均存在根本性缺陷：其一是基于影响函数（Influence-based）的方法，该类方法以目标验证集损失最小化为外层目标、以加权训练损失为内层约束，构建了严谨的双层优化（Bi-level Optimization）数学框架，但在计算超梯度时不可避免地涉及逆 Hessian 矩阵与向量乘积（inverse-HVP），在数十亿参数的 LLM 场景下计算复杂度与显存开销完全不可承受，迫使已有工作只能依赖粗糙的层级近似或引入脆弱的代理预测器；其二是基于超额损失（Excess-loss）的方法（如 RHO-1），通过计算待选样本在目标代理模型与预先训练好的参考模型之间的损失差值进行高效打分，计算轻量且易于扩展到 token 粒度，然而其参考模型在训练开始前即被冻结，随着代理模型在训练过程中的能力演进，静态参考模型迅速发生表征滞后与校准失效，导致后期筛选评分失去判别力。

本文立足于彻底打破这两大范式的割裂，寻找一种兼备双层优化严谨理论根基与超额损失高扩展性的通用数据选择范式。核心 idea：利用拉格朗日乘子法将双层数据选择问题等价转化为单层罚函数优化目标，不仅严格消除了逆 Hessian 矩阵求逆的高昂计算，而且在数学上自然推导出了超额损失打分结构，并用与代理模型动态周期同步的参考模型替代传统静态参考，构建了兼具一阶收敛保证与高扩展性的在线 token 级选择算法。

方法详解¶

整体框架¶

BLADE 框架将整个大模型数据选择体系抽象为协同演进的三元组：连续松弛的数据选择向量 $\boldsymbol{\alpha}$、正在学习的主训练代理模型 $\boldsymbol{u}$（Proxy Model），以及同时吸收训练分布与验证分布信号的动态参考模型 $\boldsymbol{w}$（Reference Model）。算法在执行过程中，首先利用拉格朗日罚函数将内层极小化约束转移至外层目标函数中；在每个迭代步中，直接利用当前参考模型与代理模型的 token 级损失差作为超梯度估计，通过无记忆随机块坐标 Frank-Wolfe（Randomized Block-Coordinate Frank-Wolfe）算法在线贪心选出损失差最大的前 $\gamma$ 比例的高价值 token，并以二值损失掩码的形式在保留因果上下文完整性的同时更新代理模型；每隔 $\tau$ 步，将代理模型的权重同步给参考模型并利用验证集与训练集联合梯度进行 $K$ 步微调更新，确保评估信号始终与当前训练阶段紧密对齐。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["候选文本 Batch 采样<br/>B ~ D_train (上下文长度 2048)"] --> B["双模型前向损失评估<br/>计算代理 u 与参考 w 的 token 损失"]
    B --> C["单层罚函数表征统一<br/>将双层优化消除逆 Hessian 映射至损失差"]
    C --> D["动态同步参考模型演进<br/>周期性 w_0 = u(e*tau) 消除表征滞后"]
    D --> E["随机块坐标 Frank-Wolfe 选择<br/>贪心选取 Top-gamma 损失差并施加损失掩码"]
    E --> F["代理模型在线梯度更新<br/>仅对保留的有效 token 反向传播更新 u"]
    F -->|每隔 tau 步循环| D
    F --> G["收敛的高质量语言模型<br/>理论满足 O(T^-1/4) 一阶稳定收敛"]

关键设计¶

1. 单层罚函数表征统一：彻底消除逆 Hessian 矩阵求逆的高昂开销

传统影响函数方法为了求解外层关于数据权重 $\boldsymbol{\alpha}$ 的验证损失极小化问题，通过隐函数定理（Implicit Function Theorem）推导出超梯度 $\nabla_{\boldsymbol{\alpha}_n} \mathcal{L}_{\mathrm{val}}(\boldsymbol{w}(\boldsymbol{\alpha})) \propto -\nabla_{\boldsymbol{w}} \mathcal{L}_{\mathrm{train}}^n(\boldsymbol{w})^T \boldsymbol{H}^{-1} \nabla_{\boldsymbol{w}} \mathcal{L}_{\mathrm{val}}(\boldsymbol{w})$，其中 $\boldsymbol{H} = \nabla_{\boldsymbol{w}\boldsymbol{w}}^2 \mathcal{L}_{\mathrm{train}}$ 是海量参数空间的二阶 Hessian 矩阵，求逆运算在大模型上根本无法实现。BLADE 从根源上将内层最优性约束转化为拉格朗日惩罚项，构造单层目标函数： $$ \min_{\boldsymbol{\alpha}\in\mathcal{A},\boldsymbol{w}} \mathcal{H}{\lambda}(\boldsymbol{\alpha},\boldsymbol{w}) := \mathcal{L}}}(\boldsymbol{w}) + \lambda \left( \mathcal{L{\mathrm{train}}(\boldsymbol{\alpha},\boldsymbol{w}) - \min) \right) $$ 在此单层表征下，关于数据权重向量 }} \mathcal{L}_{\mathrm{train}}(\boldsymbol{\alpha},\boldsymbol{u$\boldsymbol{\alpha}$ 的超梯度直接退化为内层加权损失在参考模型 $\boldsymbol{w}$ 与代理模型 $\boldsymbol{u}$ 之间的超额损失差 $\Delta_n = \mathcal{L}_{\mathrm{train}}^n(\boldsymbol{w}) - \mathcal{L}_{\mathrm{train}}^n(\boldsymbol{u})$。这一数学重构在完全保留双层优化对验证集目标对齐性质的同时，从解析层面彻底移除了任何二阶导数和逆 Hessian 逼近项，使算法能够无缝扩展至数十亿参数的稠密自回归语言模型中。

2. 动态同步参考模型演进：克服静态参考模型的表征滞后与判别力塌缩

现有的超额损失方法（如 RHO-1）依赖训练前单独预训练并全程冻结的静态参考模型 $\bar{\boldsymbol{w}}$，一旦代理模型在持续预训练中迅速掌握了领域基础知识，静态模型与代理模型之间的能力差距迅速抹平，其超额损失差分布会向零甚至负值塌缩（如图 3 所示），丧失对高质量复杂样本的判别精度。BLADE 将参考模型设计为动态演进机制：每隔固定步长 $\tau$（如 1000 步），开启一个新的演进周期 $e = \lfloor t/\tau \rfloor$，显式执行硬同步 $\boldsymbol{w}_0^{(e)} \leftarrow \boldsymbol{u}^{(e\tau)}$，将代理模型最新的参数状态作为参考模型的初值。随后，在验证损失与当前批次训练损失的联合惩罚目标下更新 $K$ 步： $$ \boldsymbol{w}{k+1}^{(e)} = \boldsymbol{w}k^{(e)} - \eta \left( \nabla}}} \mathcal{L{\mathrm{val}}(\boldsymbol{w}k^{(e)}) + \lambda \nabla}} \mathcal{L) \right) $$ 该周期性同步机制在数学上严格将参数漂移距离 }}(\boldsymbol{\alpha}^{(e\tau)}, \boldsymbol{w}_k^{(e)$\|\boldsymbol{u} - \boldsymbol{w}\|_2$ 控制在极小有界范围内（如图 4 所示），不仅避免了二者轨迹发散，而且使得实际工程中无需令理论惩罚系数 $\lambda \to \infty$，仅需设置较温和稳定的常数（如 $\lambda=1$）即可精确维持双层约束，彻底杜绝了梯度爆炸，使得超额损失分布始终保持显著的正右尾分布，持续捕获具有正边际效用的高质量样本。

3. 随机块坐标 Frank-Wolfe 选择：大尺度在线 token 级贪心课程更新

针对全量预训练语料规模庞大导致全局计算超梯度向量极具挑战的问题，BLADE 将全局约束集 $\mathcal{A}:=\{\boldsymbol{\alpha}\in[0,1]^N \mid \sum_{n=1}^N \boldsymbol{\alpha}_n = \gamma N\}$ 下的优化问题实例化为无记忆随机块坐标 Frank-Wolfe（Randomized Block-Coordinate Frank-Wolfe）算法。在每个训练迭代步 $t$，算法从数据流中均匀抽取候选批次 $\mathcal{B}_{\mathrm{cand}}^{(t)}$，这等价于在高维选择向量 $\boldsymbol{\alpha}$ 中随机选取一个局部坐标块。利用 Frank-Wolfe 的线性极小化子问题（Linear Minimization Oracle, LMO），当采用贪心步长 $\eta_{\boldsymbol{\alpha}}=1$ 时，连续松弛极小化解自然投影至极值点，导出闭式二值硬选择准则： $$ \boldsymbol{\alpha}n^{(t+1)} = \mathbb{I}}\text{-}\gamma |\mathcal{B}^{(t)}|} \left( \mathcal{L{\mathrm{train}}^n(\boldsymbol{u}^{(t)}) - \mathcal{L}) \right) $$ 为维持自回归语言模型自注意力机制的上下文因果完整性，被未选中的 token（}}^n(\tilde{\boldsymbol{w}}^{(e)$\boldsymbol{\alpha}_n=0$）并不在物理序列中切除，而是作为 token 级损失掩码（Loss Mask）参与代理模型梯度反向传播。在双时间尺度优化设定下，定理 1 证明了该随机块坐标 Frank-Wolfe 过程满足 Frank-Wolfe Gap 的全局一阶稳定收敛速度为 $\mathcal{O}(T^{-1/4})$。

实验关键数据¶

主实验¶

论文在两类主流预训练场景下进行了系统评测：其一是针对特定领域的数学持续预训练（Domain-Shift Pre-training，在 5B 规模的 OpenWebMath 候选池上采用 MetaMath 与 Mammoth 混合集作为验证集，选出 3B token 训练）；其二是通用能力持续预训练（General Pre-training，在 5B 混合候选池上以 OpenHermes-2.5 为验证集训练）。下游任务涵盖 9 个数学推理榜单及 8 个通用语言理解基准。

在数学领域持续预训练中，TinyLlama-1.1B 与 LLaMA-2-7B 上的少样本 CoT 推理评测结果如下表所示：

模型与方法	训练 Token	GSM8K	MATH	SVAMP	ASDiv	MAWPS	MMLU-STEM	平均准确率	耗时 (min)
TinyLlama-1.1B
Base (初始未训练)	-	2.9	3.2	11.0	18.1	20.4	16.1	13.4	-
Base-CT (全量训练)	5B	3.7	3.6	17.6	27.8	35.0	18.0	16.7	314.6
Random (随机采样)	3B	3.6	3.6	14.6	23.7	30.1	18.8	14.7	188.8
MATES (影响函数样本级)	3B	4.2	4.8	17.4	28.1	38.3	13.6	17.7	831.5
RHO-1 (静态参考超额损失)	3B	14.6	4.8	34.9	46.7	61.1	20.8	25.7	431.5
BLADE (本文)	3B	20.8	8.6	35.7	48.9	66.5	22.4	29.7	494.3
LLaMA-2-7B
Base-CT (全量训练)	5B	18.1	4.4	36.4	52.7	66.1	33.4	33.8	1670.3
Random (随机采样)	3B	17.3	5.2	36.9	51.9	64.2	32.9	32.3	1011.0
MATES (影响函数样本级)	3B	16.5	4.4	35.5	51.6	65.7	33.6	33.4	3274.7
RHO-1 (静态参考超额损失)	3B	33.0	7.6	57.9	60.6	78.9	30.1	41.1	2140.7
BLADE (本文)	3B	40.4	9.8	60.8	66.1	84.9	33.1	43.8	2375.1

消融实验¶

针对动态参考机制、跨尺度数据迁移性以及超参数敏感性，论文在 TinyLlama-1.1B 架构上进行了全面消融分析。

配置 / 变体	核心变动说明	平均准确率	核心指标表现 (GSM8K / MATH)	现象与结论
BLADE (完整配置)	动态参考模型 ($\tau=1000, \gamma=0.6$)，定期参数硬同步	29.7	20.8 / 8.6	完整模型建立强判别力，右尾损失分布持久保持
静态参考模型 (退化为 RHO-1)	移除动态更新与同步，全程使用离线冻结参考模型	25.7	14.6 / 4.8	训练中后期参考信号滞后，分值塌缩至 $\le 0$，掉点 4.0%
无参数同步 ($\boldsymbol{w}_0 \neq \boldsymbol{u}$)	参考模型与代理模型独立训练，不执行周期重置	训练震荡	- / -	参数漂移量 $\\|\boldsymbol{u}-\boldsymbol{w}\\|_2$ 迅速发散，导致数值不稳定与梯度爆炸
跨模型规模数据迁移 (BLADE-Tr)	使用 1.1B 模型选出的 3B token 直接训练 7B 模型	37.2 (7B模型)	23.3 / 5.8 (7B模型)	显著优于随机采样(32.3)与Base-CT(33.8)，但落后原生7B选择(43.8) 6.6%
保留比例 $\gamma=0.2$	极度保守的强过滤，仅保留 20% token	~24.5	- / -	训练 token 绝对体量过小，无法支撑充分参数更新
保留比例 $\gamma=0.8$	较为宽松的过滤，保留 80% 候选 token	~26.8	- / -	混入过多低质量网页冗余噪声，能力较 $\gamma=0.6$ 显著回退

关键发现¶

动态参考是打破性能天花板的决定性动力：RHO-1 在前 1000 步与 BLADE 趋势相当，但在 1000 步之后由于代理模型能力反超静态参考，RHO-1 的超额损失差距急剧塌缩至非正区间，丧失筛选能力；而 BLADE 在经历第 1000 步的参考模型同步与更新后，迅速恢复了明显的右尾正向分布（$\Delta > 0$），将学习焦点平滑过渡至高质量、结构严谨的推理 token，使 TinyLlama 上的综合表现超出 RHO-1 达 4.0 个百分点。
超参数更新周期 $\tau$ 展现出优异的鲁棒性：在 $\tau \in \{500, 1000, 1500, 2000\}$ 的宽广区间内，下游评测准确率几乎保持平稳无损。这意味着在实际生产落地中，工程师可以安全采用较大的更新间隔（如 $\tau=2000$），进一步削减参考模型更新带来的计算开销。
数据选择的容量依赖性（Capacity-dependence）：BLADE-Tr 实验揭示出极具启发性的现象——虽然小模型筛选出的数据具有明显的跨尺度通用性，但依然明显弱于大模型原生筛选（37.2 vs 43.8）。这是因为小容量模型往往会把“高难度但对大模型极具学习价值”的高阶样本错误当成不可拟合的噪声剔除，证明最优数据选择策略必然依赖于模型当前容量与认知边界。

亮点与洞察¶

拉格朗日对偶重构架起理论与效率的桥梁：将双层优化的内层条件融入外层罚函数，推导出等价于超额损失差的解析形式，一举省去了二阶导数与逆 Hessian 估算。不仅赋予了超额损失启发式方法严谨的双层优化理论背书，更赋予了其严格的一阶收敛性。
周期性参数重置化解深度学习中的极端罚参数矛盾：在经典优化理论中，罚函数严格逼近原问题往往需要罚参数 $\lambda \to \infty$，这在深度神经网络中必定导致梯度爆炸；BLADE 通过巧妙的周期性同步 $\boldsymbol{w}_0^{(e)} \leftarrow \boldsymbol{u}^{(e\tau)}$ 将两网络参数限制在微小邻域内，从而使极小的固定超参数 $\lambda=1$ 即可稳定逼近双层约束。
分阶段内存卸载机制维持单卡峰值显存恒定：为避免双模型驻留导致峰值显存翻倍，BLADE 在工程上采用周期性权重卸载（Offloading）策略，仅在每隔 $\tau$ 步更新参考模型时将其装载入显存，计算完成后卸载并固化分值，使得端到端训练的显存开销完全等同于单模型标准训练。

局限与展望¶

作者承认的局限：尽管相比 MATES 等基于影响函数的方法已大幅提速，但维持参考模型的周期性更新与前向双重打分依然相比朴素持续预训练增加了约 57% 的壁钟时间（494.3 min vs 314.6 min）。
外部观察的局限：虽然验证集在宏观上指导了参考模型的前进方向，但若下游目标任务十分多元或缺乏高质量、小规模的目标验证集作为参照锚点（例如纯无监督通用爬虫语料清洗），该方法的参考模型导向性可能会受到一定程度的干扰。
未来改进方向：探索利用更轻量级的紧凑孪生分支或低秩适配器（LoRA）充当参考模型更新的载体，进一步抹平与单模型训练之间的壁钟耗时差距；同时研究无显式验证集驱动下的自演进（Self-evolving）动态参考模型策略。

评分¶

新颖性: ⭐⭐⭐⭐⭐ [将影响函数的双层优化通过拉格朗日罚函数桥接到超额损失，并提出动态同步参考模型，理论与设计极为优美]
实验充分度: ⭐⭐⭐⭐⭐ [覆盖 1B 与 7B 模型、数学专业与通用两大多任务场景、详实的消融分析、可视化分布与跨尺度迁移探索]
写作质量: ⭐⭐⭐⭐⭐ [数学推导严谨自洽，论证逻辑环环相扣，从理论定理到工程落地的过渡极其自然流畅]
价值: ⭐⭐⭐⭐⭐ [为大规模自回归大模型的在线动态数据筛选和数据配比提供了兼备理论证明与工程可行的突破性范式]

配置 / 变体	核心变动说明	平均准确率	核心指标表现 (GSM8K / MATH)	现象与结论
BLADE (完整配置)	动态参考模型 (\(\tau=1000, \gamma=0.6\))，定期参数硬同步	29.7	20.8 / 8.6	完整模型建立强判别力，右尾损失分布持久保持
静态参考模型 (退化为 RHO-1)	移除动态更新与同步，全程使用离线冻结参考模型	25.7	14.6 / 4.8	训练中后期参考信号滞后，分值塌缩至 \(\le 0\)，掉点 4.0%
无参数同步 (\(\boldsymbol{w}_0 \neq \boldsymbol{u}\))	参考模型与代理模型独立训练，不执行周期重置	训练震荡	- / -	参数漂移量 \(\\|\boldsymbol{u}-\boldsymbol{w}\\|_2\) 迅速发散，导致数值不稳定与梯度爆炸
跨模型规模数据迁移 (BLADE-Tr)	使用 1.1B 模型选出的 3B token 直接训练 7B 模型	37.2 (7B模型)	23.3 / 5.8 (7B模型)	显著优于随机采样(32.3)与Base-CT(33.8)，但落后原生7B选择(43.8) 6.6%
保留比例 \(\gamma=0.2\)	极度保守的强过滤，仅保留 20% token	~24.5	- / -	训练 token 绝对体量过小，无法支撑充分参数更新
保留比例 \(\gamma=0.8\)	较为宽松的过滤，保留 80% 候选 token	~26.8	- / -	混入过多低质量网页冗余噪声，能力较 \(\gamma=0.6\) 显著回退