本文介绍了一项关于大型语言模型（Large Language Model, LLM）高效训练方法的研究。该研究由来自普林斯顿大学（Princeton University）、加州大学伯克利分校（UC Berkeley）和弗吉尼亚理工大学（Virginia Tech）的研究团队合作完成，论文题为《GREATS: Online Selection of High-Quality Data for LLM Training in Every Iteration》。该研究已发表于人工智能领域的顶级会议——第38届神经信息处理系统大会（NeurIPS 2024）。
这项研究的核心背景是解决当前LLM训练面临的两大挑战：漫长的训练周期和训练数据中普遍存在的低质量、冗余或偏差信息。静态数据选择方法在训练前一次性筛选数据，无法适应模型训练过程中的动态学习状态，可能导致所选数据在整个训练周期中并非最优。相比之下，在线批量选择（Online Batch Selection）方法在每个训练迭代中动态地从采样批次中选择数据，更具适应性。然而，现有的在线批量选择方法存在显著局限：要么依赖额外的、计算成本高昂的参考模型，要么基于简单的启发式规则（如高损失值或大梯度范数），这些规则可能无法准确衡量数据样本的真实信息价值或相关性，导致性能不佳，甚至不如简单的随机选择。
针对上述问题，本研究提出了一个名为GREATS（Greedy Approximation Taylor Selection）的新型在线批量选择算法。其研究目标是开发一个原则性强、计算高效、无需参考模型且能捕捉数据真实信息价值的方法，以显著加速LLM训练的收敛速度并提升其泛化性能。
本研究的具体工作流程主要包括五个环节，涵盖了算法原理推导、高效实现技术开发以及全面的实验验证。
第一环节：提出原则性强的在线批量选择问题定义。 研究人员将在线批量选择问题形式化为一个集合效用函数（Utility Function）优化任务。其核心思想是，在每个训练步骤t，给定一个来自目标领域的小型验证数据集（可视为任务指引），最优的训练数据子集应当是那些能在更新模型后，最大程度降低验证集损失（Loss）的数据。具体而言，效用函数定义为： u(t)(s; z(val)) := ℓ(wt, z(val)) - ℓ(w̃t+1(s), z(val)) 其中，wt是当前模型参数，s是从候选批次bt中选出的子集，w̃t+1(s) = wt - ηt Σ_{z∈s} ∇ℓ(wt, z) 是用子集s进行一次梯度下降更新后的“虚拟模型”。直接优化这个效用函数是组合爆炸问题，因为需要评估所有可能的子集。
第二环节：设计基于贪婪算法和泰勒展开的高效近似求解方案（GREATS算法核心）。 为了高效地近似求解上述集合优化问题，研究团队采用了贪婪算法框架。然而，直接应用贪婪算法需要频繁评估效用函数，这涉及到计算聚合梯度、更新模型和计算验证损失，计算开销巨大。研究的关键创新在于使用一阶和二阶泰勒展开（Taylor Expansion） 来近似计算在贪婪选择过程中，添加一个数据点所带来的边际效用增益。推导过程如下： 1. 边际增益近似：将添加一个数据点zi到已选集合b̂t的边际增益近似为：u(t)(zi|b̂t) ≈ ηt g(zi) · ∇ℓ(w̃t+1(b̂t), z(val))，其中g(zi)是训练样本zi在当前模型wt下的梯度。 2. 虚拟模型梯度近似：进一步使用泰勒展开近似虚拟模型w̃t+1(b̂t)下验证样本的梯度：∇ℓ(w̃t+1(b̂t), z(val)) ≈ g(z(val)) - ηt H(z(val)) Σ{z∈b̂t} g(z)，其中g(z(val))是验证样本梯度，H(z(val))是其Hessian矩阵。 3. 最终近似公式：将两步近似合并，得到边际增益的最终高效近似表达式： u(t)(zi|b̂t) ≈ ηt g(zi) · g(z(val)) - ηt² g(zi) H(z(val)) Σ{z∈b̂t} g(z) 这个公式包含两部分： * 初始重要性分数：ηt g(zi) · g(z(val))，即训练样本梯度与验证样本梯度的内积。这衡量了在更新前状态下，用zi更新模型对降低验证损失的直接贡献，与TracIn评分思想类似。 * 多样性校正项：- ηt² g(zi) H(z(val)) Σ_{z∈b̂t} g(z)，这是对已选集合b̂t的校正。它惩罚与已选样本梯度相似（通过Hessian加权衡量）的训练样本，从而鼓励选择能提供互补信息、多样化的数据点。
基于此近似公式，研究人员设计了GREATS算法。算法开始时，为候选批次bt中的每个数据点计算初始重要性分数ϕz = ηt g(z) · g(z(val))。然后迭代执行以下步骤直至选出k个样本：1) 选择当前分数最高的数据点z加入已选集合b̂t；2) 对剩余每个数据点的分数进行更新：ϕz ← ϕz - ηt² g(z) H(z(val)) g(z)。这个过程巧妙地模拟了贪婪算法的行为，但完全避免了显式的效用函数评估和模型更新，仅需梯度与Hessian信息。
第三环节：开发“Ghost Inner-Product”技术以实现可扩展的高效计算。 虽然GREATS算法避免了显式的模型更新，但直接实现仍面临挑战：计算初始重要性分数需要每个训练样本与验证样本的梯度内积，而计算校正项则需要成对的梯度内积或梯度-海森-梯度乘积。如果为每个样本单独计算并存储完整模型大小的梯度向量，计算和内存开销将不可承受。 为此，研究团队创新性地提出了“Ghost Inner-Product”技术。该技术的核心洞察是，对于线性层等常见网络层，可以分解梯度计算，从而在不实例化任何完整梯度向量的情况下，仅通过一次反向传播就高效地计算出所有训练样本与验证样本之间、以及训练样本两两之间的梯度内积。其基本思想是在反向传播时将验证数据与训练数据一起送入批次，并利用链式法则，将梯度内积的计算转换为对中间激活（input）和输出梯度（output gradient）的张量运算。例如，对于线性层s = Aw，个体损失ℓ(i)对权重W的梯度可以表示为输入向量a(i)与输出梯度向量b(i)的外积。而两个梯度g(z1)与g(z2)的内积可以转化为(a(1)与a(2)的点积)乘以(b(1)与b(2)的点积)，从而避免了构造大梯度矩阵。这一技术也适用于LoRA（Low-Rank Adaptation）层和序列数据。通过这项技术，GREATS的整个选择和梯度更新过程可以在一次反向传播中完成，运行时间与常规训练相当。
第四环节：进行全面的实验评估。 为了验证GREATS的有效性、高效性和鲁棒性，研究团队设计了广泛的实验，涵盖了不同的模型、训练数据集、评估任务以及设置。 * 实验设置：使用了多种模型（LLaMA-2-7B, Mistral-7B, LLaMA-3-8B, GPT-Small）和数据集组合，包括指令微调（使用LESS、Alpaca数据集在MMLU、TyDiQA、SamSum任务上评估）和预训练（使用OpenWebText数据集）。验证集数据量被限制在很小（≤16）以模拟实际场景。比较的基线方法包括：常规训练（使用全部批次数据）、基于启发式的在线选择方法（MaxLoss：选择损失最高的样本；GradNorm：选择梯度范数最大的样本）、基于参考模型的方法（使用LLaMA-3.1-8B-Instruct作为参考模型的RhoLoss）以及基于语义相似度的静态选择方法（SBERT）。 * 实验流程与结果分析： 1. 收敛速度与泛化性能：通过绘制训练过程中验证集和测试集困惑度（Perplexity）的变化曲线来评估。实验结果表明，GREATS在所有实验设置下都显著加速了训练收敛，并且通常能达到更优的最终测试性能。相比之下，MaxLoss和GradNorm等方法因可能选择到低质量或不可学习的噪声数据，其效果不稳定，有时甚至不如常规训练。 2. 下游任务准确率：在MMLU（9个子科目）和TyDiQA测试集上评估准确率/F1分数。GREATS持续优于或至少与所有基线方法持平。在MMLU的平均准确率上，GREATS相比最佳基线至少提升了3.4个百分点，证明了其提升模型实际任务能力的有效性。 3. 对验证数据量的鲁棒性：即使仅使用2个验证样本，GREATS在MMLU任务上的测试困惑度也持续低于常规训练，表明其对极少量验证数据具有鲁棒性。 4. 在预训练场景下的效果：即使在预训练场景下（验证数据与训练数据同分布，均来自OpenWebText），GREATS也能带来测试性能的提升，尽管提升幅度小于在分布差异更大的微调场景中。这证明了其在不同学习场景中的适用性。 5. 运行效率验证：通过比较不同实现的吞吐量（每秒处理的训练数据点数）来评估效率。结果显示，采用“Ghost Inner-Product”技术的GREATS实现，其运行时间与常规训练高度接近，吞吐量损失很小。而直接计算逐样本梯度的GREATS和GradNorm实现则要慢10-20倍，证明了所提高效实现技术的巨大价值。
第五环节：得出研究结论并讨论局限性。 本研究成功提出并验证了GREATS这一新颖的在线批量选择算法。其科学价值在于：1) 提供了一个原则性强、数学推导清晰的在线批量选择问题形式化框架；2) 创新性地结合了贪婪算法、泰勒展开和高效的“Ghost Inner-Product”计算技术，解决了以往方法在效率与效果上的矛盾；3) 通过大量实验，系统性地证明了该方法在加速LLM训练收敛、提升下游任务性能方面的有效性，并验证了其计算效率。其应用价值在于为大规模语言模型的训练提供了一种实用、高效的数据选择工具，有助于缩短研发周期、降低计算资源消耗。
研究的亮点包括：1) 方法论的创新：将集合优化、泰勒近似和高效计算技术巧妙结合，形成了完整且可实施的算法。2) 显著的性能提升：在多个基准测试中稳定地超越了现有方法。3) 卓越的计算效率：提出的“Ghost Inner-Product”技术使算法达到了接近常规训练的运行时开销，具备了大规模应用的可行性。4) 广泛的适用性：在微调和预训练、不同模型和数据集上都证明了有效性。
研究人员也坦诚地讨论了局限性和未来方向：1) 需要验证数据：GREATS需要一个小型验证集，这在某些场景（如纯无监督预训练）可能无法获得，未来可探索无验证数据的变体。2) 当前基于SGD优化器：“Ghost Inner-Product”技术目前主要针对SGD推导，对于Adam等带有自适应学习率的优化器需要进一步扩展，尽管实验中用SGD作为代理已被证明有效。3) 大批次下的内存约束：若GPU内存不足以将验证数据与训练批次一起进行反向传播，可通过梯度累积等技术解决，但可能增加运行时间。4) 优化目标为困惑度：目前直接优化验证损失（困惑度），虽然下游任务表现通常也提升，但未来可探索直接针对特定下游任务指标进行优化的方法。
GREATS研究为解决LLM训练中的数据选择难题提供了一条新颖且高效的路径，其提出的算法框架和实现技术对机器学习社区，特别是在高效训练和优化领域，具有重要的参考价值和启发意义。