Awesome World Models / OpenReview

HEP-JEPA: A foundation model for collider physics

作者：作者信息待提取 · 单位：Jai Bardhan, Radhikesh Agarwal , Abhiram Potula , Cyrin Neeraj, Subhadip Mitra, Center for Computational Natural Sciences and Bioinformatics, International Institute of Information Technology, Hyderabad · 来源：World Models for Science · 发布日期：2025 · 分类：Social Science

科学建模基础世界模型

精读笔记

Problem Setting

这篇论文并不是在提出新的 top tagger 或 quark-gluon tagger，而是在问：collider physics 是否可以用类似 vision foundation model 的自监督预训练方式，得到一个可迁移的 jet representation backbone。实际问题是，在 HEP 中标签往往来自 simulation，强监督模型很容易学习到任务标签和 generator / detector simulation 的耦合偏差；但完全无监督/自监督目标如果设计不好，又会被低层粒子重建细节牵着走。

真正困难点在于 jet 数据同时具有点云、集合、局部几何、QCD shower 层级结构和 detector response 的混合属性。一个好的预训练目标需要捕捉局部 substructure 与全局 topology 的统计关系，同时避免把容量浪费在逐粒子 feature reconstruction 上。以前 masked particle modeling / generative modeling 的瓶颈正是在这里：它们可以学到数据分布，但不一定学到对 tagging 最有用的抽象判别结构；监督式 foundation model 又没有摆脱标签和模拟先验。

Motivation

作者的核心观察是：jet 内部结构天然适合“部分观察预测未观察部分”的自监督学习，但预测对象不应是原始粒子本身，而应是已经抽象化的局部表示。因为原始粒子层面的可重建细节太多，其中大量信息对下游判别是冗余甚至有害的；而 latent-space prediction 可以迫使模型学习更稳定、更语义化的相关结构。

这篇论文真正填的缺口是：HEP 中还缺少一个系统验证 JEPA paradigm 的大规模 jet pretraining 实例。它不是发明新的物理 observable，也不是提出新的 SOTA tagging architecture，而是把“decoder-free latent prediction”作为 HEP foundation representation 的候选训练范式。

Core Idea

核心思想可以概括为：把 jet constituents 组织成局部 patch tokens，然后用可见 patch 的 embedding 去预测被遮挡 patch 的 teacher embedding。这个改变的本质是把建模目标从“还原粒子”改成“对齐局部结构的潜在表示”。因此模型不需要学习一个高保真的 jet generator，也不需要像 contrastive learning 那样构造负样本；它只需要学习 jet 内部哪些局部结构在表示空间中互相约束。

这个方向理论上可能有效，是因为 jet substructure 不是任意点云：boosted object decay、QCD radiation pattern、heavy flavor、lepton content 都会在局部区域之间产生强相关。只给模型一部分 constituents，如果它能预测另一部分的 representation，就说明它捕捉了某种 event topology / shower pattern 的统计规律。与 prior 的本质区别在于，HEP-JEPA 不做 full reconstruction，也不靠 label supervision，而是在 latent space 中学习 conditional structure。

Method

方法的必要机制主要有三层。

第一，particle group tokenization 解决的是 jet 输入可变长且局部密集的问题。直接对所有 constituents 做 Transformer 既昂贵又容易过拟合低层细节；将粒子按 η-ϕ 几何分组并用 permutation-invariant encoder 得到 patch token，相当于先构造局部 receptive field，使模型关注局部 shower / decay fragment 的组合模式。

第二，context-target masking 解决的是自监督信号如何构造的问题。目标不是随机预测某个粒子 feature，而是给定一部分 patch，预测另一部分 patch 的表示。contiguous masking 比完全随机 masking 更接近“补全局部区域”的结构学习，但文中 ablation 显示该设计的增益并不稳定，说明 masking strategy 还不是完全搞清楚的核心。

第三，EMA target encoder 和 latent loss 解决 collapse 与重建负担问题。teacher encoder 给出稳定 target embedding，student/predictor 只需在 representation space 对齐。这个机制的核心变化是去掉 decoder，把训练压力从像素/粒子级 fidelity 转到抽象表征一致性。physics bias 和 registers 是合理的 inductive bias / Transformer 修补项，但更像辅助工程选择。

Key Insight / Why It Works

这篇论文最重要的 insight 是：对 jet foundation model 来说，预训练目标不一定要生成或重建 jet；预测局部结构的 latent representation 可能更贴近下游判别任务。HEP 中很多 tagging 信号本来就是 distributional / relational 的，例如 prong structure、radiation pattern、energy sharing、leptonic vs hadronic composition。latent prediction 让模型学习这些统计依赖，而不是纠结某个 constituent 的精确四动量。

我认为真正有效的部分大概率是“JetClass 大规模覆盖 + patch-level latent prediction + fine-tuning initialization”的组合，而不是某个单独模块。few-shot 增益说明预训练确实提供了可复用特征；但当标签充足时优势消失，说明 representation 并没有显著超越监督学习可从同分布数据中学到的结构。换言之，它目前更像一个有效的 representation initialization，而不是已经展现出强 emergent physical world model。

physics bias 和 registers 的 ablation 有正向效果，但幅度有限，且实验设置是缩减版 linear/SVM proxy，增益来源不清。augmentations 没有明显收益也值得注意：这说明模型当前可能主要依赖数据分布覆盖，而不是显式物理对称性约束。若要声称对真实实验数据更 scalable，需要证明它在 generator shift / detector shift 下的稳定性；文中未充分说明。

Relation To Prior Work

HEP-JEPA 最接近三条路线：Masked Particle Modeling、OmniJet-α 这类生成/重建式自监督；SimCLR-style contrastive HEP representation learning；以及 OmniLearn / ParT 这类监督式大模型或强 jet tagger。它的真正差异在于 objective：不重建输入、不生成 token、不使用负样本、不依赖标签，而是预测 teacher embedding。

看似新的地方有不少是已有思想重组：PointNet-style local grouping、ViT/Point-JEPA masking、EMA teacher、Particle Transformer 的 pairwise physics bias、register tokens 都来自已有技术谱系。实质创新在于把这些组合成一个适配 jet 数据的 JEPA pretraining framework，并在 JetClass 规模上验证其可行性。

相对 MPM，它少了 decoder reconstruction，因此更可能学抽象表示；相对 OmniJet-α，它不是 generative model，不能直接建模 jet likelihood 或生成事件；相对 supervised ParT，它弱化标签依赖，但当前性能也没有超过 task-specific supervised SOTA。它属于“decoder-free self-supervised representation learning for scientific point sets”这条技术谱系，而不是严格意义上的物理 world model。

Dataset / Evaluation

评估覆盖了 JetClass few-shot、top tagging、quark-gluon tagging，基本覆盖 jet classification 的典型场景，但仍然主要是同类模拟数据之间的迁移。JetClass 本身巨大且包含多类 boosted/object jets，预训练数据覆盖很强；下游 TQTR 和 quark-gluon benchmark 与预训练域并非完全相同，但仍是公开模拟 benchmark，不足以验证真实实验数据上的 foundation claim。

实验最有说服力的是低标签 regime：HEP-JEPA 在少量标签下收敛更快、效果更好，说明预训练 representation 有实际价值。但在 full-label 设置下优势基本消失，且下游任务未超过 ParT/ParticleNet，说明论文验证的是“自监督预训练有用”，不是“HEP-JEPA 已形成统一 collider world model”。

另一个问题是 evaluation 的归因不够干净。对比包括 from-scratch 同架构、supervised pretrained、SOTA，但 JEPA objective、数据规模、patch tokenizer、physics bias、registers、fine-tuning 策略之间的贡献仍没有完全拆开。benchmark 是否存在隐性 overlap 或 generator-level shortcut，文中没有深入分析。

Limitation

方法成立依赖几个强前提。第一，jet 局部结构之间必须存在足够强的可预测关系，否则 latent prediction 只能学到分布均值或 trivial correlation。第二，pretraining distribution 必须覆盖下游任务相关结构；如果下游涉及新的 detector condition、pileup regime、rare topology 或 generator mismatch，当前证据不足以说明表示仍可靠。第三，teacher embedding 本身来自同一架构和同一数据分布，模型可能是在对齐模拟数据内部的统计编码，而不是真正学习物理不变量。

scalability 的上限也不清楚。Jet-level patch prediction 可以扩展到更大数据，但 event-level collider analysis 涉及多 jet、多 lepton、MET、trigger、detector geometry、systematics，context-target 定义会复杂得多。当前模型仍停留在单 jet representation；把它称为 foundation model 可以接受，把它称为 HEP world model 则偏早。

此外，增益可能主要来自 scaling / data coverage。JetClass 100M jets 的预训练本身就是很强的先验；few-shot 提升未必证明 JEPA objective 最优，只证明大规模自监督初始化有效。文中没有充分比较 MAE/MPM/contrastive 在同架构、同预算、同数据下的表现，因此“JEPA 优于其他 SSL paradigm”的判断还不能成立。

Takeaway

  • 1. 最值得记住的是目标函数的转向：对科学点云/jet 数据，latent-space prediction 可能比 raw reconstruction 更适合作为 foundation representation 的预训练目标。
  • 2. HEP-JEPA 的价值主要在少标签和初始化效率，而不是当前绝对性能；它推动的是训练范式，而非 SOTA tagger。
  • 3. patch-level tokenization + physics-aware attention bias 是一个可迁移设计：先把低层粒子集合压成局部物理对象，再在对象间学习关系，这比直接 tokenizing every particle 更接近可扩展 foundation backbone。
  • 4. 未来真正值得做的不是继续在同类 benchmark 上刷 accuracy，而是检验 generator shift、detector shift、real data adaptation、event-level multi-object tasks，以及与 unfolding/anomaly detection 等非分类任务的关系。

一句话总结

HEP-JEPA 是把 JEPA 式 decoder-free latent prediction 系统迁移到 jet physics 的一次有效范式验证，真正贡献在于证明大规模自监督 jet 表征预训练能改善低标签迁移，但它目前更像 representation initialization 方法，而不是已经成熟的 collider world model。