OmniCosmos：跨宇宙传输粒子物理学知识

《PRX IntelligenceNEW》：OmniCosmos: Transferring Particle Physics Knowledge Across the Cosmos

　　基础模型能够构建数据的有效表示，并应用于各种下游任务。先前的研究表明，为对撞机物理学设计的OmniLearned基础模型能够显著提升对撞机实验的发现潜力。本文表明，在对撞机数据上训练的基础模型可以迁移至宇宙学。具体而言，我们改进了CosmoBench数据集中不同数据集的宇宙学参数以及晕和星系速度的预测。据我们所知，这是首次展示对撞机物理学模型能够跨学科领域泛化。

一、引言

当前和近期的宇宙学巡天正在彻底改变我们对暗物质和暗能量的理解。迄今为止，从巡天数据中获取宇宙学信息的方法通常是先压缩数据，再将低维摘要与理论预测进行比较。现代机器学习推动了超越这些摘要、开展场级分析的运动，以利用所有可用信息[1]。虽然这些方法有潜力显著提高巡天数据的精度，但这一方向的发展需要两个目前尚不完整的要素：（1）能准确描述所有相关长度尺度和时间尺度的模拟；（2）能够适应宇宙学数据和模拟结构的机器学习方法。尽管机器学习可能有助于解决（1），特别是超越暗物质N体演化的部分，但我们关注的是第二个挑战。具体而言，宇宙学数据由空间中可变长度的物体集合（如星系）组成（在机器学习中称为点云），而模拟整个宇宙的计算量巨大，因此推断必须仅依赖于相对较少的合成宇宙。

行业驱动的神经网络架构主要聚焦于规则网格（用于图像）和序列（用于文本）。我们可以将宇宙体素化到规则网格或序列上，以直接使用这些方法。相反，我们采用了一种专为点云设计的架构——它能够处理可变长度的输入，并且对输入的排列具有不变性（集合与序列的区别）。此类方法的最先进技术基于变换器[2]。虽然变换器非常有效，但需要大量数据才能充分发挥大型模型的容量。这对宇宙学推断而言是一个重大挑战，因为无法获得大规模的合成宇宙数据集。

我们提出通过构建一个专为粒子对撞机数据设计的点云基础模型来应对这些挑战。与宇宙学数据一样，对撞机数据也是可变长度的集合。由于量子力学效应，每次碰撞产生随机数量的出射粒子，每个粒子由三维动量及其他属性定义。完全模拟粒子碰撞的计算成本很高，但远低于模拟整个宇宙。此外还存在多个层次的模拟保真度。我们的假设是：一个足够具有表达能力的对撞机物理学基础模型，在相对较少的宇宙学模拟上进行微调后，将优于经典方法以及在宇宙学模拟上从头训练的机器学习模型。尽管宇宙学最相关的物理定律与粒子碰撞最相关的物理定律并不相同，但这些数据集具有足够的共性，足以证明这种迁移学习目标的合理性。

为验证这一假设，我们使用OmniLearned [3–5]基础模型和CosmoBench [6]数据集。OmniLearned是一个点边变换器（PET），通过监督和无监督学习在来自各种真实和合成粒子碰撞的十亿个标记和未标记强子喷流上训练。OmniLearned已被证明能够有效地在低保真和高保真模拟之间、模拟与真实数据之间以及不同碰撞系统之间迁移。本文中，我们证明同一模型可以用于与粒子喷流无关的基础物理数据。CosmoBench数据集是一组模拟点云，用于基准测试推断宇宙学标准模型（ΛCDM）参数的工具。具体而言，使用三种不同的模拟器来构建粗粒化[7]的宇宙模拟。点代表（有效）暗物质晕或星系。其中一个数据集中存在并合历史，但我们关注当前时刻的任务[8]。尽管引力是确定性的，但由于宇宙的初始条件，宇宙学模拟与对撞机模拟一样具有随机性。除了数据之外，CosmoBench还包括我们用作参考的经典和机器学习基准算法。

大多数与粒子物理学基础模型相关的研究使用专门的标记化来描述数据，并通过自监督学习进行训练[9–15]。虽然这些模型也是研究粒子物理学与宇宙学之间知识迁移的有前景的候选者，但它们引入了额外挑战。因为学习的标记化依赖于用于描述粒子相互作用的特征集，所以必须为宇宙学数据定义新的标记器。这与OmniLearned形成对比，后者只需替换单个层即可适应不同的输入表示。此外，大多数标记器需要显式的输入排序。由于输入晕的排序不应影响回归性能，因此设计为排列不变的点云表示更符合数据的结构。

许多相关工作涉及我们研究的某些方面。例如，多项研究探索了应用于宇宙学的点云/图机器学习模型[16–35]。CosmoBench包含一个此类模型，我们还将我们的基础模型与相同架构但从头训练的模型进行比较。这使我们能够区分优化架构和预训练带来的益处。OmniLearned的理念是通过预训练将大部分物理信息隐式编码在模型中。已提出了用于宇宙学推断的显式等变模型[36–51]，未来探索预训练是否优于这些模型将是很有趣的，正如在对撞机领域似乎表现的那样[5]。目前有许多积极努力创建专门针对宇宙学的基础模型[52–58]，但我们不知道有任何模型经过训练以点云形式摄取整个宇宙。CosmoBench中的不确定性量化方法相对简单，未来探索先进不确定性量化（UQ）（例如，参见文献[59,60]及其参考文献）与基础模型之间的相互作用将是很有趣的。

本文结构如下。第二节描述OmniLearned模型架构以及为适应宇宙学问题所做的修改。第三节介绍用于训练OmniCosmos的数据集。第四节描述不同任务和数据集的结果。最后，第五节给出结论和展望。

二、知识迁移

在粒子物理学中，碰撞后产生的重建粒子集合用于研究基本粒子如何相互作用。在大型强子对撞机（LHC）等高能对撞机中，通过粒子相互作用产生的一种普遍存在的物体是喷流，即粒子的准直喷射流。喷流由色粒子通过强力相互作用产生，是基础研究对象，因为它们提供了关于量子色动力学性质以及粒子碰撞所涌现的底层物理过程的信息。由于喷流的丰富性和复杂性，粒子物理学中的许多机器学习应用最初是针对喷流提出的[9,61,62]，包括卷积神经网络（CNN）[63]、递归神经网络（RNN）[64]、图神经网络（GNN）[65,66]、生成对抗网络（GAN）[67]、深度集合[68]以及注意力[69]/变换器[70,71]。当前最先进的技术[5,71–73]将喷流表示为点云，即度量空间中无序的对象集合。这些点云使用变换器模型处理，该模型能够考虑粒子之间的长程相关性，并且可以有效地扩展到在更多数据上训练的更大模型。

OmniLearned模型[3–5]使用聚集为喷流的粒子作为输入的一部分，并采用不同任务来创建丰富的点云表示，从而能够高效地适应不同数据集和任务。底层神经网络模型PET-v2由局部和全局变换器层的混合组成，旨在学习粒子之间的局部和全局相关性。至关重要的是，通用变换器层逻辑值通过物理启发的相互作用偏置项进行增强：

z_i = z_i + b_i, (1)

即使在存在大型训练数据集的情况下，这已被证明能提高性能。该模型随后使用包含真实粒子碰撞混合以及不同模拟保真度的模拟粒子相互作用的十亿个喷流进行训练。此预训练阶段的结果构成了我们称之为OmniLearned的核心基础模型。为了适应特定应用，我们将OmniLearned微调到特定数据集和待解决任务。第一个挑战是确定一种高效策略，将OmniLearned适应到完全不同的数据集。OmniLearned的输入粒子由其四动量描述，而本工作中使用的暗物质晕模拟由其三维空间坐标描述。此外，OmniLearned预训练阶段的相关任务是点云分类和生成，而在本案例中，我们关注不同宇宙学参数的回归以及每个输入对象的晕速度预测。最后，平均一个喷流包含约30–100个粒子，而我们考虑CosmoBench数据集中最多5000个最质量大的晕。我们通过引入OmniCosmos来应对所有这些挑战，其总体适应策略如图1所示。

图1. OmniLearned模型（在粒子碰撞上训练）到宇宙学任务（命名为OmniCosmos）的知识适应策略概要。

首先，我们处理点云之间对象表示的差异。在PET-v2模型中，点云中每个点的初始输入特征通过特征嵌入转换为大的潜在空间。具有非线性激活函数的全连接层仅接收初始特征表示作为输入，对点云中的所有点应用相同的变换，以确保输出的排列等变性。

使用每个输入点的k个最近邻信息并行确定包含局部粒子信息的局部嵌入。对于每个对象x_i，我们使用输入空间中的欧几里得距离识别k个最近邻x_{i_j}。接下来，我们通过包含在点对之间计算的相关函数f来增强原始输入特征。在此步骤中，我们能够纳入针对物理问题和所用数据集定制的函数。虽然在OmniLearned中我们包含了几种与粒子物理学相关的特征，例如两个粒子组合的洛伦兹不变质量，但我们确定了三个几何特征，除了差值Δx之外还要包含它们。我们包括绝对距离d，余弦距离cos θ，以及初始晕x_i与差值Δx之间的余弦距离，由cos θ_i描述。由此得到的成对特征为：

f(x_i, x_j) = (Δx, d, cos θ, cos θ_i) (2)

这些成对特征用作变换器模块的输入，其中注意力机制仅关注粒子x_i及其k个邻居。最后，我们取k个邻居的平均响应来定义每个点的新表示。

点云中每个点的最终潜在表示是前述两个操作的和：一个对所有粒子应用相同的可学习函数，另一个利用局部邻域。

为了通过OmniLearned预训练创建OmniCosmos，我们首先在新网络架构中加载所有OmniLearned的预训练权重。当权重因输入形状差异而不兼容时，使用随机初始化权重。这些仅占模型总权重的0.09%。由于新层需要更多训练步骤才能收敛，我们将新初始化层的学习率设置为高于预训练层的学习率。这也能有效保留预训练中学到的强大几何先验，同时防止“灾难性遗忘”。

第二个挑战是处理新的训练任务。虽然OmniLearned专注于分类和粒子生成，但我们感兴趣的是回归宇宙学参数以及预测点云中每个对象的三维晕速度。虽然这些任务在范围上显著不同，但它们在输出上相似。在分类任务中，我们基于所有可用类别预测每个喷流的类别分数。对于回归任务，唯一需要的修改是将输出大小调整为待预测参数的数量。对于速度预测，我们可以重新利用粒子生成器网络，该网络原本为每个粒子分配一个输出预测，现在改为预测晕的速度。在这种情况下，我们同样只需替换最后一层，以修正输出形状的差异。这一层是新建初始化的，因此分配的学习率也更高。最后一个挑战是处理原始喷流数据与当前数据集之间点数差异的问题。由于OmniLearned不强制固定粒子多重数，因此无需额外修改。然而，粒子数量的十倍增加限制了可使用的变换器块数量，因为内存需求较高。因此，我们只研究小型OmniLearned模型的适应，该模型有八个变换器块，总共约两百万个可训练参数。OmniLearned模型最初在粒子碰撞上训练，其知识适应得到的结果是OmniCosmos模型，用于处理宇宙学中的多个任务。

III. CosmoBench数据集
CosmoBench [6] 数据集利用多个物理模拟器引入了暗物质晕和星系的大规模宇宙学模拟，总计耗费超过4100万核小时。在本研究中，我们考虑CosmoBench中的两个数据集：Quijote [74] 和 CAMELS-SAM [75] 模拟套件。在这两个数据集中，我们感兴趣的是估计宇宙学参数Ωm和σ8的值，分别代表宇宙总能量密度中物质（包括暗物质和正常物质）所占的比例以及物质密度涨落的振幅。从随机初始条件出发，不同的程序使用N体模拟将系统演化至当前时刻。CAMELS-SAM模拟由暗物质-only N体模拟组成，并经过半解析星系形成模型后处理，盒尺寸为100 Mpc/h，暗物质粒子质量为1.5×10^8 M⊙/h。Quijote模拟使用1000 Mpc/h的盒尺寸，并考虑质量为8×10^11 M⊙/h的暗物质粒子。由于分辨率较低，Quijote模拟比CAMELS-SAM更易生成，后者更高的质量分辨率允许其解析更小的、高度非线性的尺度，但代价是模拟体积更小。在两个数据集中，我们使用模拟中最多5000个最大质量的晕。我们对每个数据集使用官方数据划分，CAMELS-SAM数据集包含600/204/196个样本用于训练、验证和测试，Quijote模拟则为19651/6550/6551。

IV. 结果
我们评估OmniCosmos在使用不同数量的可用训练样本时，对CAMELS-SAM模拟中宇宙学参数回归的性能。我们使用决定系数R2来量化性能，参数y定义为
R2 = 1 - (Σ(y_i - ?_i)2) / (Σ(y_i - ?)2) (3)
其中模型预测?_i与真实值y_i比较，?为测试数据集中每个条目y_i的平均值。我们通过自助法（bootstrap）对测试数据重采样1000次来确定R2的不确定性，报告均值作为最终预测，自助法样本的标准差作为不确定性。我们直接将结果与参考文献[6]中提供的结果进行比较，无需重新训练。输入在训练前进行标准化。我们测试OmniCosmos和从零开始训练的模型（即所有可训练权重从随机权重初始化）的不同超参数。由于网络参数来自原始OmniLearned模型固定，我们考虑不同的学习率、批量大小以及OmniCosmos微调中较低学习率层与较高学习率层之间的乘性因子。对应最佳性能模型的设置最终结果中报告。为确定每个点的局部邻域大小，我们使用k=20的值。对Ωm和σ8的预测结果见表I和图2，作为训练集中样本数的函数。除了比较从零开始训练的模型和OmniCosmos，我们还考虑两个额外的消融结果：一个是原始OmniLearned模型的变换器层被冻结，即在适应阶段不更新；另一个是同一架构首先在合成点云上预训练。第一个消融用于研究OmniLearned模型的学习表示在适应新数据集时有多大用处；第二个用于理解用于预训练模型的粒子物理表示有多大用处。合成点云包含十个等数量类别，每个类别由四个独立高斯分布特征表示。对于类别c和特征f，对应的高斯均值和标准差定义为
μ_cf = (c-1) * (5-1) / (10-1) + f * 0.1 (4)
σ_cf = 0.1 + (c-1) * (1-0.1) / (10-1) (5)
图2. 使用CAMELS-SAM模拟获得的不同宇宙学参数预测结果。
表I. 不同算法在CAMELS-SAM数据集上报告的性能比较。粗体结果表示性能最高的算法。
2PCF LLS GNN GNN (w/o edgeMP) Scratch OmniCosmos
Ωm 0.01 0.01
σ8 0.01 0.01
OmniCosmos (frozen)
Ωm 0.01 0.01
σ8 0.01 0.01
Gaussian pretrain
Ωm 0.04 0.03
σ8 0.04 0.03
特征均值和标准差都随类别标签单调增加，产生与用于预训练OmniLearned模型的特征范围相当的特征。与原始OmniLearned模型一样，预训练联合目标两个任务：对输入点云分类和生成点云。合成预训练数据集包含一千万个点云，每个点云由150个点组成，选择这些以匹配用于预训练OmniLearned模型的典型点云大小。在适应阶段，该预训练模型的变换器块也保持冻结，以便与OmniLearned模型直接比较。我们观察到，当使用完整的可用训练数据集时，微调后的OmniCosmos模型和从零开始训练的模型的性能均超过之前基准。有趣的是，仅使用100个可用模拟样本，OmniCosmos已经能够匹配或超越之前基准的性能。此外，与从零开始训练相比，微调策略显示出强大优势，在所有数据集大小下，OmniCosmos模型都能获得更好的参数预测，即使在极小的100个模拟数据集下也显示出有希望的结果。我们还观察到，完整微调策略和变换器块冻结的策略在不确定性范围内产生相似结果。这表明OmniLearned模型预训练期间学习到的表示对宇宙学也有用。这一结论进一步得到高斯预训练样本结果的支持，在小的数据集大小下与从零开始训练的模型匹配，随着数据增加仅略有提升，因为架构中最具表达力的部分保持冻结。

接下来，我们使用仅输入位置来评估晕速度预测。我们再次使用R2度量评估结果，但对所有晕和速度坐标求和，得到每个模拟的单一度量值。结果显示在图3中。与回归任务类似，我们观察到OmniCosmos和从零开始训练的模型性能良好，在使用所有可用训练事件时超过之前基准。这项任务也突出了使用更复杂深度学习模型相对于更简单摘要的好处，参考文献[6]中引入的GNN基准也表现出强大性能。与从零开始训练的模型相比，OmniCosmos观察到的改进有所减少，尽管在所有数据集大小下仍然存在。在小的数据集大小下，具有完整微调或冻结骨干的OmniCosmos模型显著优于从零开始训练的模型，后者表现与高斯预训练模型相似。这表明在小的数据集大小下，OmniCosmos学习到的表示有益于提高预测质量，而高斯预训练没有产生显著好处。这些结果令人鼓舞，表明在大型对撞机物理数据集上训练的模型编码了相关的点云结构，可以转移到其他科学领域。
图3. 使用CAMELS-SAM模拟的晕速度预测结果。

接下来，我们研究模型使用QUIJOTE模拟确定相同宇宙学参数集的能力。我们将考虑的邻居数量增加到k=32，因为我们观察到使用更大邻域性能略有提升，并扫描从零开始训练模型以及不同预训练和适应策略的超参数设置。最佳性能模型的结果显示在表II和图4中，用于不同数据集大小下Ωm和σ8的回归。鉴于更大的模拟样本，我们可以评估模型在不同数量级模拟尺寸下的扩展。与使用CAMELS-SAM模拟观察到的结果类似，OmniCosmos的性能要么匹配最佳性能基准，要么超越它们。特别地，OmniCosmos能够使用不到可用模拟的0.5%匹配Ωm的先前基准。有趣的是，对于σ8，OmniCosmos匹配两点相关函数（2PCF）的预测，但大大改善了先前的深度学习模型（GNN）。我们再次观察到冻结变换器层的模型与完整微调OmniCosmos表现同样好，与CAMELS-SAM模拟结果一致。然而，在合成数据集上预训练的微调模型表现显著更差。特别地，当仅使用100个示例时，适应模型无法预测σ8，而是无论输入数据如何，返回训练样本中存在的平均值。这导致R2值接近零，因此在图4中未显示。这种行为在本文所有超参数扫描中保持一致。这些结果表明，合成数据集不仅无法产生有用的表示，而且损害泛化，导致比随机权重初始化的模型更差的性能。最后，我们使用QUIJOTE模拟评估速度预测，结果显示在图5中。在之前基准中，一个简单的线性最小二乘（LLS）模型获得晕速度的最佳预测。OmniCosmos能够进一步改善先前结果，使用不到可用数据集的2%超越LLS结果。此外，我们再次观察到知识转移的好处。与从零开始训练的模型相比，OmniCosmos在多达数千个模拟样本中表现出卓越性能，而在使用更多数据时收敛到相同结果。与CAMELS-SAM结果类似，在合成高斯数据集上预训练的模型表现更接近从零开始训练的模型。相反，OmniCosmos在完整微调和冻结变换器层两种情况下都实现了强大性能。两个数据集研究中速度预测的结果列于表III。
图4. 使用QUIJOTE模拟获得的不同宇宙学参数预测结果。
图5. 使用QUIJOTE模拟的晕速度预测结果。
表II. 不同算法在QUIJOTE数据集上报告的性能比较。粗体结果表示性能最高的算法。
2PCF 0.004
LLS
GNN
GNN (w/o edgeMP)
Scratch
Ωm 0.003 0.003
σ8 0.003 0.003
OmniCosmos
Ωm 0.003 0.003
σ8 0.003 0.003
OmniCosmos (frozen)
Ωm 0.004 0.003
σ8 0.004 0.003
Gaussian pretrain
Ωm 0.004 0.003
σ8 0.004 0.003
表III. 不同算法在CAMELS-SAM和QUIJOTE数据集上晕速度预测报告的性能比较。粗体结果表示性能最高的算法。
CAMELS-SAM QUIJOTE
Linear oracle 0.2372 0.3769
LLS 0.2107 0.4347
GNN 0.2865 0.4100
Scratch 0.004 0.003
OmniCosmos 0.004 0.003
OmniCosmos (frozen) 0.004 0.003
Gaussian pretrain 0.004 0.003

V. 结论与展望
在这项工作中，我们介绍了OmniCosmos，一个利用在大规模粒子物理数据集上预训练的网络来执行宇宙学不同任务的机器学习模型。我们展示了粒子相互作用学习到的表示如何通过快速适应输入表示以及在模型架构中纳入几何洞察力来转移到新数据集。尽管原始OmniLearned模型在点云上训练，点数少一个数量级且任务不同，OmniCosmos模型通常能够获得比从零开始训练的相同模型更优越的性能。我们通过对比另外两种场景的结果，展示了预训练策略的优势：一种是在适应过程中保持模型的Transformer模块冻结，另一种是同一架构在合成的Gaussian数据集上进行预训练。我们观察到，冻结主干网络的模型与完整的OmniCosmos模型表现相当，表明从对撞机物理中学习到的点云中各个点之间的关系可以迁移到宇宙学。使用合成数据集得到的结果进一步支持了这一观察，因为这些结果通常不如OmniCosmos。我们使用CosmoBench提供的模拟数据对OmniCosmos进行评估。特别是，对于像CAMELS-SAM这样计算昂贵的星系晕模拟，OmniCosmos能够利用一半的可用数据集改善所有先前的基准。在QUIJOTE模拟中，OmniCosmos能够匹配或超越所有其他基准，通常使用的可用数据不到其十分之一。由于大规模模拟计算昂贵，OmniCosmos所表现出的性能有可能大幅减少为实现宇宙学精度所需的模拟数量。适应步骤只需替换不到0.1%的原始模型权重为随机初始化的权重。此外，虽然在对撞机物理数据上进行预训练需要大量计算资源，但适应OmniLearned模型的成本要低得多。例如，将OmniLearned适应到完整的CAMELS-SAM数据集，每个任务大约需要8个GPU小时。虽然较慢，但这与文献[6]中报告的GNN的1-3个GPU小时相当。当仅使用100个模拟样本时，适应大约需要4个GPU小时，同时在参数回归上匹配或超越所有基线。对于QUIJOTE模拟，适应完整数据集大约需要32个GPU小时，与文献[6]中报告的1天训练时间相当。虽然我们在数百个样本中观察到性能提升，但当有数千个训练样本可用时，从头训练的模型与OmniCosmos之间的差异会缩小。在这项工作中，我们探索了一种简单的适应策略，即修改或冻结模型不同部分的学习率。对于未来的工作，我们计划研究其他适应步骤的策略，使用机器学习社区中引入的其他技术来适应基础模型以解决特定问题，如文献[76,77]所述。最后，除了宇宙学参数的回归和晕速度预测之外，直接利用OmniLearned在OmniCosmos适应步骤中的生成能力进行点云生成将是有趣的，因为使用预训练模型可能潜在地减少实现高保真生成所需的模拟数据集数量，从而为宇宙学提供快速准确的替代模型。