高能物理机器学习模型应如何在预算轴上扩展与比较?

稳健的高能物理扩展律:在每个预算下重调学习率与批大小,再在算力最优前沿上比较设计。

直接答案

把训练动力学与模型、数据集一起扩展。这篇论文表明,对于在约 110 亿喷注的 ATLAS JetSet2 数据集上训练的多任务喷注标记 Transformer,最优学习率和批大小会随训练预算系统性移动,因此在任何架构、目标或输入比较变得有意义之前,必须在每个规模上重新调优并迁移超参数。做到这一点后,算力最优前沿就成为架构的属性,而不是超参数选择的属性。

1篇文献引用

本文由 WisPaper 驱动的搜索和论文分析生成。

固定超参数会把真实扩展与移动的最优点混在一起

扩展律把可达到的损失与模型规模 N、数据集规模 D 和算力 C 联系起来。这篇论文的核心方法论观点是:只有当每次运行都在适当调优的动力学下训练时,这一关系才成立。在固定模型规模和形状下,最优学习率和批大小会随训练预算系统性变化:随着 token 预算增长,固定批大小下的最优学习率下降,而固定学习率下的最优批大小上升,联合最优点向更大的批大小和学习率迁移。

后果很具体。如果把超参数固定在某个参考预算上,测得的损失曲线就会把真实的扩展行为与训练最优点的移动混在一起。论文的教师—学生玩具实验让这种污染的程度变得可见:在小规模端调优一次然后保持固定的学习率,并不会相对逐格调优的基线保持恒定偏移。它会偏离调优后的幂律,并在大规模端向上翘起,因此这样训练的网格会偏置拟合出的指数,甚至可能完全掩盖扩展行为。

同样的效应也出现在压缩表示系统中,那里的一层隐藏层 MLP 探针在冻结的集合 Transformer 自编码器潜变量上训练。相对逐格调优的基线,在单个格上调优一次然后保持固定的学习率会留下一个损失差距,并朝网格的大 D 端和大宽度端同时增大。在固定学习率网格上重新拟合 L(N, D) 曲面会得到完全不同的定律,其地板估计高于最佳实测格,因此误调优的网格会同时误报指数和扩展结束的位置。

联合最优点可预测并能跨规模迁移

这套配方建立在关于联合(学习率、批大小、token 预算)最优点的理论描述,以及让超参数可跨宽度和深度迁移的参数化之上。作者拟合联合最优学习率、批大小和模型形状如何随 token 预算移动,然后检验这些预测能否在拟合从未见过的配置上成立。

验证是直接的。在批大小、token 预算、宽度和深度的留出组合上,预测的最优学习率落在每条损失—学习率曲线的底部或附近,因此它能跨批大小、token 预算和模型形状迁移。联合最优点(学习率、批大小)也能迁移,只可能存在一个随模型规模变化的最优批大小偏移(在最优学习率下)。对于喷注分类,这一偏移实际上不存在,因为损失在最优学习率下对批大小是平坦的。对于多任务目标,最优批大小按 token 预算的平方根增长,使这一偏移成为模型规模偏离参考值时唯一可能的残余自由度。

使用预测的批大小而不是扫描中损失最低的批大小,损失代价至多为几个 10^-3,即千分位量级的惩罚,与运行间噪声相当。作者指出,这种迁移可能得益于损失在批大小最优点附近的平坦性。实际解读是:联合最优点可以跨规模携带,而无需再增加模型规模依赖,这正是使算力最优前沿在不同配置间可比较的原因。

算力最优扩展把预算近乎均分给 N 和 D

应用于 JetSet2 上的 GN 系列 Transformer 标记器时,提取出的扩展律显示模型规模和数据集规模对 C 的依赖近乎相等的平方根关系。换句话说,随着算力增长,模型参数和唯一训练喷注以相近速率增长,而不是某一方主导分配。算力最优模型在最优模型规模下对数据做单遍训练,扩展律投影也假设同样如此。

这种分配在不同设置间并不一致。论文报告了五种设置的算力最优分配:两种训练目标(全辅助多任务和仅喷注分类单任务)以及三种缩减输入集。仅喷注分类设置是明显的离群者,在每个预算下都偏好更多数据和更小的模型,而输入层级似乎不改变分配。辅助目标把算力最优分配推向更大的模型和更小的数据集。

这一区别对读者如何比较设计很重要。单任务标记器和多任务标记器不在同一条分配线上,因此在某个固定模型规模或固定数据集规模上比较它们,会把目标与预算分配混在一起。论文的框架改为在比较前给每个配置各自的算力最优分配。

辅助任务在同等算力下有益,更丰富输入降低损失

两条设计轴在各自的算力最优前沿上被比较。第一条是训练目标。辅助目标在同等算力预算下降低主喷注分类损失,并且把算力最优分配推向更大的模型和更小的数据集。论文指出,辅助目标在固定训练样本数下降低喷注分类损失,但也增加训练的计算成本,因此它们的用处取决于分类增益是否超过这一额外成本。算力预算比较正是让这一权衡在此处倾向于它们的原因。

第二条轴是输入表示。把输入扩展到更低层级的数据会系统性降低损失,同时让扩展指数几乎不变。研究的缩减输入集包括 GN2 输入集、无 flow 变体和仅四动量模型。因此更低层级输入改善了损失曲线的水平,而不改变其斜率,这与辅助目标带来的增益是不同类型。

把这两种效应分开正是关键。目标变化同时移动损失水平和最优分配;输入变化移动损失水平,而分配和指数大致固定。读者在大规模比较标记器设计时,需要知道自己看的是哪一种。

数据受限情形下四个量被联合优化

论文把该流程扩展到数据受限情形,此时运行在早停下训练,而不是单遍训练。在这里,作者据其所知首次预测了早停下联合最优的模型规模、训练时长、学习率和批大小。这与算力受限情形是不同的优化问题:训练时长本身是自由变量,而不是由单遍假设固定。

两种情形的区别并非表面文章。在算力最优设置中,模型在最优模型规模下对数据做单遍训练,投影也假设同样如此。在数据受限设置中,早停决定训练何时结束,因此最优时长必须与模型规模和两个训练动力学超参数一起预测。

论文还报告了在批大小、token 预算、宽度和深度的留出组合上对联合最优点迁移的验证,其中预测的最优学习率落在每个最小值处或附近。残余自由度是在最优学习率下可能随模型规模变化的批大小偏移,而使用预测批大小而非扫描最优批大小所测得的惩罚保持在千分位量级。这正是联合最优点在实践中可被视为可迁移的原因。

幂律区间只有在看到足够数据后才开始

幂律区间的起点本身由规模决定。在数据集规模的某个阈值以下,损失几乎不携带关于高算力扩展的信息。因此,基于较小数据集的拟合可能返回有偏的扩展指数,即使得到的曲线看起来很好地符合幂律。这是对表面拟合的警告,而不只是对噪声拟合的警告:一条曲线可以看起来像干净的幂律,却仍然给出错误的指数。

作者把这一点直接联系到数据基础设施。像 JetSet2 这样的大型全模拟数据集是确立高能物理模型在大规模下如何表现的前提,而 LHC 实验产生并保留此类样本的能力被描述为高能物理机器学习研究的一项主要资产。约 110 亿喷注的 JetSet2 数据集通过组合 ttbar 和 Z' 到 qqbar 样本与 VH 到 bbbar、ccbar、tau+tau-、Z' 到 ttbar、多喷注以及 gamma* 到 tau+tau- 样本,并重新加权到跨味类的平均 pT 和 eta 谱而构建,正是使这项研究成为可能的资源。

投入的规模也有报告。整项研究约需 97k A100 GPU 小时,其中跨所有变体的算力最优网格占 92k 小时,涉及 1219 次运行和 1.5 x 10^12 个喷注。作者将其描述为按高能物理标准相当可观,但与工业级模型相比仍属适中。

损失投影转化为排除率,而非标记器排名

论文在 ttbar 样本上把损失转化为背景排除率,固定全局 b 喷注效率为 70%,条件为 20 < pT < 250 GeV 和 |eta| < 2.5,使用 GN3 的判别分数 (fc, ftau) = (0.2, 0.1)。它报告了 GN3、TN25-86M 模型以及在完整输入和任务集上训练的算力最优模型投影的轻喷注和 c 喷注排除率,投影由三种拟合方法得到。

这一比较需要仔细阅读。GN3、TN25-86M 和算力最优模型共享同一架构,因此属于同一模型类;它们主要在模型规模、唯一训练数据量和算力上不同。论文明确指出这不是标记器设计的比较。GN3 约有 12M 参数,在 350M 唯一喷注上训练数遍;TN25-86M 有 86M 参数,在约 7.7B 唯一喷注上训练数遍,两者都训练到验证损失不再改善。相比之下,算力最优模型在最优模型规模下做单遍训练。

作者还指出,部署的标记器是一个冻结的预测器,必须在少数工作点上针对数据校准,而最优背景分数一般是 (pT, eta) 的函数,并因分析而异,因此没有单一固定选择是普遍最优的。排除率投影来自排除率与损失之间的拟合关系,三种拟合方法之间的离散度也一并展示。

一项扩展结论在被信任前必须展示什么

论文自身对进展路径的总结是一份可用的清单。一种方法应当证明:随着模型规模、数据集规模和算力增长,其优势持续存在;在获得实验可用的完整输入信息时仍具竞争力;并最终在真实碰撞数据上成立。每一条都对应研究记录的一种具体失败模式:一旦训练动力学被扩展就消失的优势、依赖扣留输入的优势,以及只在模拟上展示过的优势。

论文中最具决策相关性的阈值是数据集规模阈值,低于它损失几乎不携带关于高算力扩展的信息。该阈值决定一项扩展研究能否对大预算说出任何东西,也是作者把大型全模拟数据集视为前提而非便利的原因。读者评估一项扩展结论时,应先问拟合前看到了多少唯一数据,以及超参数是在每个预算下重新调优,还是从参考运行固定下来。

前瞻性建议是从每个应用单独训练模型,转向在大规模上预训练单个模型,再通过微调把得到的表示适配到不同实验、探测器配置和物理任务。观察到的扩展行为被描述为暗示可以有效地投入大得多的预算,尤其是在这一基础模型方向下。论文的框架被提供为此类未来扩展研究的实用蓝图。

本文引用的文献

1

How to scale your HEP ML models: A recipe for robust architecture comparisons at scale

可靠的高能物理扩展律要求把训练动力学与模型和数据集一起扩展;借助可预测的超参数迁移,在约 110 亿喷注的 JetSet2 数据集上的多任务喷注标记 Transformer 显示模型规模和数据集规模对 C 的依赖近乎相等的平方根关系,辅助目标在同等算力下降低主喷注分类损失,更丰富的输入降低损失而不改变指数,且幂律区间只在数据集规模阈值之上才开始。