FM4NPP: A Scaling Foundation Model for Nuclear and Particle Physics
Authors: David Keetae Park
Shuhang Li
Yi Huang
Xihaier Luo
Haiwang Yu
Yeonju Go
Christopher Pinkenburg
Yuewei Lin
Shinjae Yoo
Joseph D. Osborn
Jin Huang
Yihui Ren
Venue: ICLR 2026 Poster (2026)
Keywords: Foundation Model
State Space Model
Neural Scaling
Particle Tracking
Nuclear Physics
Particle Physics
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We present a scalable foundation model for sparse nuclear and particle physics detector data, capable of adaptation to diverse downstream tasks after pretraining.

Abstract: Large language models have revolutionized artificial intelligence by enabling large, generalizable models trained through self-supervision. This paradigm has inspired the development of scientific foundation models (FMs). However, applying this capability to experimental particle physics is challenging due to the sparse, spatially distributed nature of detector data, which differs dramatically from natural language. This work addresses if an FM for particle physics can scale and generalize across diverse tasks. We introduce a new dataset with more than 11 million particle collision events and a suite of downstream tasks and labeled data for evaluation. We propose a novel self-supervised training method for detector data and demonstrate its neural scalability with models that feature up to 188 million parameters. With frozen weights and task-specific adapters, this FM consistently outperforms baseline models across all downstream tasks. The performance also exhibits robust data-efficient adaptation. Further analysis reveals that the representations extracted by the FM are task-agnostic but can be specialized via a single linear mapping for different downstream tasks.

Published as a conference paper at ICLR 2026
FM4NPP: A SCALING FOUNDATION MODEL FOR NUCLEAR AND PARTICLE PHYSICS
David Park", Shuhann Li"", Yi Huang" Xihaier Luo', Haiwang Yu", Yeonju Go",
Christopher Pinkenburg', Yuewei Lin', Shinjae Yoo', Joseph Osborn", Jin Huang'. Yihui Ren
1 Artificial Intelligence Department, Brookhaven National Laboratory, United States
Nuullar and Particle Physics Department, Brookhaven National Laboratory, United States "hese authors contributed equally. Corresponding author, yren@bnl... gov
ABSTRACT
Large language models have revolutionized artificial intelligence by enabling large, generalizable models trained through seff-supervision. This paradigm has inspired the development of scientific foundation models (FMs). However, apply- ing this capability experimental particle physics challenging due the sparse, spatially distributed nature of detector data, which differs dramatically from nat- ural language. This work addresses if an FM for particle physics can scale and generalize across diverse tasks. We introduce a new dataset with more than 10 mil- lion particle collision events and suite of downstream tasks and labeled data for evaluation. We propose a novel self-supervised training method for detector data and demonstrate its neural scalability with models that feature up to 188 million parameters. With frozen weights and task-specific adapters, this FM consistently outperforms baseline models across all downstream tasks. The performance also exhibits robust data-efficient adaptation. Further analysis reveals that the represen- tations extracted by the FM are task-agnostic but can be specialized via a single linear mapping for different downstream tasks.
INTRODUCTION
The emergence of large-scale language and vision models Wang et al. (2023) has marked a paradigm shift from specialized neural architectures tailored to individual tasks toward universal, scalable, and multitasking models. These large models, containing billions of parameters and trained through self- supervised learning massive unlabeled datasets, be efficiently adapted to diverse downstream tasks, ranging from language translation and code generation to general reasoning. Recognizing their transformative potential, the scientific community has termed these scalable, general-purpose models as foundation models (FMs) Bommasani et al. (202T). Among their underpinning features, FMs can employ self-supervised learning on extensive unlabeled datasets, allowing them develop generalized representations adaptable to various downstream tasks with minimal additional labeled training. However, scientific data often fundamentally differ from natural language or visual data. Hence, the design and implementation of FMs for scientific fields still faces challenges Li et al. 20244; Pyzerr-napp al. (2025)).
This work investigates developing FMs created for experimental nuclear and particle physics (NPP), using data from the Relativistic Heavy Ion Collider (RHIC) and the sPHENIX detector Brookhaven
National Laboratory (2025). NPP research uses particle colliders, such as RHIC or the Large Hadron Collider (LHC), to explore subatomic phenomena. Discovery of the Higgs boson exemplified the transformational significance of collider-based NPP Collaborraiio al. (2012). In particular, RHIC collides heavy ions and polarized protons, facilitating essential studies of quark-gluon plasma and the structure protons and nuclei Belmont et al. (2024). Commissioned in 2023 Moskowitz/ (003)) the sPHENIX detector features advanced tracking and calorimetry and generates extensive and com- plex data. The complexity of collider data and the breakthrough science it enables have motivated exploration of new data processing tools like FMs that employ self-supervised learning. In partic- ular, the high occupancy hadronic collisions at RHIC or the LHC is particularly challenging for traditional reconstruction algorithms, motivating exploration into methods.

---

Published as a conference paper at ICLR 2026
Stage 1: Pretrain FM4NPP on self-supervised autoregressive forecasting
pretrain prediction model
FMANPP
embedding
Stage 2: Adapt FMANPP for downstream tasks
Task 1 Track Finding
РMАNPP (frozen)
downstream adapter
embedding
Task 2: Partiele Identification
Task 3: Noise Tagging
downstream
downstrcam adapter
adapter
Pestio Othere
Fion
Racn
Figure 1:
Overview of pretrained foundation model that can be adapted to various downstream tasks. We answer two questions in this work: a) Whether the foundation model is scalable, i.e., can larger model and dataset sizes improve performance? and b) the foundation model adaptable to solve multiple downstream tasks.
However, developing an FM for NPP poses several challenges. The sparse, three-dimensional (3D)- spacepoint nature of collider data lacks an established framework for formulating self-supervised tasks. Additionally, optimal neural architectures and the scaling behavior of pretraining losses with respect to model and data size remain unknown. Crucially, it is uncertain if neural representations from frozen, pretrained FM can generalize effectively to various downstream tasks, thereby out- performing existing traditional solutions and specialized Al models.
Here, we take first step toward enabling the use of FMs for NPP by adopting a cost-effective two- stage paradigm: 1) pretrain a large FM using self- celf-supervised objective and pair the frozen FM with lightweight, task-specific adapters (Figure [ The core hypothesis is sufficiently trained FM encodes rich, task-agnostic representations that can be efficiently adapted to diverse downstream tasks with minimal additional training.
To this end, we construct large-scale dataset, exceeding 10 million simulated collision events and characterized by sparse, high-dimensional detector data. We also define three downstream tasks with corresponding labeled datasets to evaluate FM adaptability. We introduce a self-supervised pretraining strategy tailored to the sparsity and structure of detector data and demonstrate strong neural scaling behavior with models up to 188 million parameters. With frozen FMs and simple adapters, we achieve state-of-the-art performance across all downstream tasks. This analysis further reveals that FM representations are broadly task-agnostic and can be specialized using single linear transformation. In summary, our contributions are:
A large-scale, open benchmark dataset for FM training and evaluation in particle physics.
A self-supervised pretraining method designed for sparse detector data.
Empirical evidence of scaling behavior and data-efficient adaptation with frozen FMs.
Insight into the structure and adaptability of FM representations across diverse tasks.
2

---

Published as a conference paper at ICLR 2026
2 RELATED WORK
Scientific Foundation Models. Developing FMs for scientific domains is a promising yet formidable endeavor. Progress has been most evident in domains where data exhibit modality struc- tures similar to language or vision. For instance, Aurora Bodnar et al. (2025) is an atmospheric FM trained on continuous spatiotemporal climate data, and recent work has demonstrated FM-based disease detection from retinal images Zhou et al. (2023). In high energy physics, recently devel- oped FMs, such as OmniJet-a: Birk et al.] 2024) and Omnil high-level physics objects known as jets collimated sprays of partieles that can be represented as dense matrices. However, many scientific disciplines, including low-level detector data (e.g.. g. raw hits or clusters) in particle physics, materials science, and single-cell omics, present unique challenges like irreuularly structured and sparse data. Traditional approaches, such as graph neural networks (GNNs), are well suited for sparse data, but they face scalability issues due to certain phenomena, e.g., oversmoothing Rusch et al. (2023). Surveys in materials science Syyzer-Knapi et al. (2025b) and single-cell omics Ma et (2024) emphasize additional bottlenecks, including limited data availability and high computational costs. These challenges also apply to NPP, where it remains un- clear how best to model extremely sparse data, how much data are needed, and whether pretraining benefits can effectively transfer to downstream tasks. This work takes first step toward addressing these questions by developing scalable FM for sparse lower-level detector data in NPP, focused on efficient pretraining, architectural scalability, and downstream generalization.
[Bhijji al. (2025), focus on
Scalable Neural Architectures. Three neural network architectures are prominent for their scala- bility: Transformers, Mixture-of-Fii off-pprt- (MoE), and State Space Models (SSMs). The Transformer architecture Vaswani al. (2017) has revolutionized deep learning via self-attention, enabling ef- fective modeling of long-range dependencies. This has led to widespread adoption in both natural language processing (NIP) and computer vision Dosovitskiy et al. (202T). However, the quadratic time and space complexity of self-attention limits scalability on long sequences a critical bot- tleneck for scientific data. MoE architectures [edus et al. (2022) improve inference efficiency by activating only a subset of the model per input, although they face challenges, such training insta- bility and expert imbalance. The Mamba architecture Gu & Dao (2024). an SSM variant, achieves linear time complexity and shows competitive superior performance to Transformers. Given the relatively large number of spacepoints per collision event, which can result in especially long se- quences, we explore SSMs a backbone due to their favorable training efficiency and memory usage.
AI Model Tasks in NPP. In collider physics, high-energy particles collide to produce new par- ticles, whose trajectories, called tracks, are reconstructed from discrete spacepoints recorded by layered detector components. Track finding, or assigning spacepoints to different tracks, is one of the most important tasks. Traditional algorithms rely on combinatorial sceding followed by Kalman-filter-base refinement Kalman| (1960). These classical methods are computationally ex- pensive and difficult to parallelize on modern accelerators. GNN-based approaches have become popular in track finding. Exa. Trkx Ju et al. (2021) formulates the task as edge classification, whereas EggNet Calafiura et al. (2024) employs contrastive learning followed by clustering. Each predicted track then corresponds to a connected subgraph of spacepoints. Other recent work has introduced Transformer-
models, such as HEPT Miao et al. (2024), and SSM-based Diang & Qian| (2025) tracking models at (1M) parameter scale. Beyond track finding, another common task is particle identification (PID). Previous machine learning bMLs) PID approaches, such as MIPF Mokhtar et al. (005) and HGPF HGPF/Kakati et al.] (2025), rely on reconstructed tracks as key in- puts and typically incorporate calorimeter topoclusters. In contrast, employ a lightweight adapter model, designed as a probe FM generalizability, that performs PID directly at the low level using only time projection chamber (TPC) spacepoint data without requiring high-level track information or calorimeter inputs.
While these models are promising- - with some achieving robust results using fewer than one million parameters there is no systematic study of scaling behavior. Moreover, open datasets designed for scaling and evaluating FMs are limited. In this work, more than 10 million simulated collision events are generated, affording comprehensive scaling studies. We also develop FM with 188 million parameters, two orders of magnitude larger than prior models, and evaluate its performance in tracking and for other broader NPP tasks.
3

---

Published as a conference paper at ICLR 2026
(a) 3D grid partition and ordering
(b) Foundation model
last
Beng esquenee order
Mamba blooks
XL
Poeitional Enconding (NeRF)
start
frsst
Figure 2: (a) Hierarchical Raster Scan strategy to serialize the unordered spacepoints into one- dimensional sequence. (b) A Mamba FM backbone for k-Nex--Nearest-Neighbbo prediction.
3 PARTICLE DETECTOR DATASET
Dataset. As part of the sPHENIX detector's central tracking system, the high- -rah-granllarity TPC Klest (2020) records more than 85% of the total data volume. The TPC consists of 48 con- centric cylindrical readout layers, encompassing approximately 160,000 channels that each record 260 time samples, totaling 41.6 million voxels. Functioning as a three-dimensional (3D) camera, the TPC records the paths of particles emerging from collision events, delivering continuous 31) spacepoint information.
We use a publicly available dataset with over 10 million simulated p+p collision events at v 200 GeV [i al. (2025). The natural sparsity of events in p+p collisions makes them ideal testing ground for developing an FM for NPP applications. The simulation pipeline includes real detector geometry, electromagnetic fields, hadronic interactions, continuous energy loss, multiple scattering, decay processes, secondary particle production, and precise energy deposition. The raw detector hits subsequently are reconstructed to spacepoints that serve as inputs in this work. More concretely, a collision event E is represented as a set of spacepoints } where each spacepoint is expressed by its deposit energy and location (E, x. y: z). The number of spacepoints per event can vary from hundreds to thousands.
Downstream Tasks. We select three complementary downstream tasks to evaluate the generaliz- ability of an FM: Track Finding, PID, and Noise Tagging. Track Finding assigns each spacepoint to its corresponding predicted track as shown in Figure [ (Task 1, 2, and 3). Assume there are m tracks (T,) r where each track T, consists of its associated spacepoints {s; Tj}. }. The goal of track finding is to predict partition P over the set of spacepoints, where P = 1 if spacepoint s, is assigned track T,. The number of tracks can vary from event event. This task is analogous to instance segmentation in computer vision.
To evaluate performance, we employ both conventional physics-motivater metrics, tracking effi- ciency and purity Calafiura al. (2018), as well as the statistical metric Adjusted Rand Index (ARI) Hubert & Arabie (1955. As the exact definition of whether or not a predicted track matching a true track can differ among physics experiments, we adopt the "double-me majooiyy rule" from the TrackML challenge Amrouche et al. (2020). The rule enforces that a predicted track is successfully matched true track only when greater than 50% of the predicted track's spacepoints belong to that track and more than 50% of the true track's spacepoints are present in the predicted track. This stringent rule guarantees neither predicted tracks nor true tracks are matched more than once. Then, tracking efficiency (recall) is defined as the ratio between the true positive and total number of truth tracks, while tracking purity (precision) the ratio between the true positive and total number of predicted tracks.
PID aims to label each spacepoint to the particle species that produced it, i.e., pion, kaon, proton, electron, and others. This is comparable to a segmentation task in computer vision. Noise Tagging, the third downstream task, seeks to identify spacepoints associated to low-momentum secondary particles, primarily delta electrons as they typically are not associated with physics observables of interest. This also be considered segmentation task. For these two downstream tasks, report overall accuracy, macro-averaged precision, and recall. Additional information about the TPC detector, data generation pipeline, and statistical analysis are included in Appendix A.
4

---

Published as a conference paper at ICLR 2026
Spacepointa
brack instance
MLP
predictions
Transformer Decodle
krack queries
kay, valus
MLP
track embeddings
Instance
seg,
Flnding
spacepoint embedddings
prodlurt
mambba blocks
Lineaz
mamnh blocks
Adapter Model for Track Finding
Transformer Devodee
semantic predictions
semantic
nismbba blocks
Linewe
MLP
seg. loss
Poundacion Modell
*
Adapter Model for Particle Identification Noiee Tugging
Purtick leiutiffiation
Figure 3: The pretrained FM is kept frozen during training of the adapter models for downstream tasks. The adapter models for particle identification and noise lagging share the same architecture but are trained independently.
METHODOLOGY
This section introduces the scalable FM for NPP, including a novel serialization method for sparse spacepoints, adaptation to the Mamba architecture, self-supervised pretraining objective, and lightweight adapter models for downstream tasks. Additional information included in Appendix B.
4.1 SELF-SUPERVISED SCALING FOUNDATION MODEL.
Serialization of Spacepoints. key challenge in applying sequence- based models like Mamba to particle detector data is in serializing the unordered set of 3D spacepoints s; from an event E into meaningful one-dimensional (1D) sequence. The serialization strategy must balance two competing objectives: preserving the global structure of particle trajectories, which typically propagate outward from the collision point, and maintaining local continuity along individual tracks T3 to retain fine- grained geometric information. Naive serialization schemes struggle to achieve this balance. For example, space-filling curves (e.g., Hilbert or Z-order) prioritize spatial locality but often interleave points from different tracks, disrupting trajectory coherence. Conversely, sorting points by their radial distance preserves the outward particle flow but scatters spacepoints from the same track across distant positions in the sequence, breaking local continuity. An effective serialization must navigate this trade-off, allowing the model to learn both global and local physics from a sequential input.
We propose Hierarchical Raster Scan strategy to serialize the unordered spacepoints s; into ID sequence suitable for sequence models as shown in Figure aa) First, all spacepoints are transformed from Cartesian (x. y, 2) cylindrical-polat system (r, Φ, that better reflects the geometry and symmetries collider experiments, where r radial distance, & depicts the azimuthal angle, and represents the pseudorapidity (angle to the beam axis). The raster scan method operates els. The first is inter-box ordering, where spacepoints are initially partitioned into non-overlapping 3D spatial boxes. Then, these boxes are ordered based on the (r, Φ. 17) coordinates of their geometric centers, starting from the innermost region and progressing outward. This produces global ordering over the spatial domain. The second is intra-box ordering, where, within each box, spacepoints are sorted by their radial coordinate which generally aligns with the direction of particle propagation. By concatenating the intra-box sequences according to the inter-box order, we obtain globally se- rialized sequence that preserves both local spatial continuity and global physical progression. This hierarchical structure captures important geometric and physical priors while producing a format compatible with sequence models. Specifically,
lev-
partition the spatial domain into a 8 8 grid along the (r, 17, Φ) axes, respectively. The r bins are aligned with the physical boundaries of the TPC detector layers, while the yı and & bins determined using frequency-based binning to ensure balanced point distributions across the grid.
Mamba as the FM Model Backbone. Mamba is a selective SSM that efficiently processes long sequences, achieving linear time complexity Gu & Dao (003)). It features selection mechanism that makes its internal state matrices input-dependent, allowing the model to dynamically focus relevant information and filter out noise all while using a hardware-aware algorithm for fast computation. In this work, we employ Mamba-2|Dau & Gu Gul (2024), which further improves upon this foundation. Mamba-2 introduces structured State Space Duality (SSD), new theoretical framework
5

---

Published as a conference paper at ICLR 2026
that simplifies the architecture and enhances hardware utilization, leading to significant speedups in both training and inference.
We treat every spacepoint as an input "token" a sequence. To map an input tuple (E, r, Φ. v)) to the model width dmndel, we employ two-pathway process inspired by Neural Radiance Fields Milden- hall et al. (202T). First, the feature component & is projected into a feature embedding of dimension dmodel- Concurrently, the spatial coordinates (r, O, 11) are transformed with high-freqqenney po- sitional encoding function, 7(-), which uses sine and cosine transformations. Then, this encoded position is also projected into positional embedding of dimension dmodee- The final representation the element-wise sum of the feature and positional embeddings, yielding a single vector of size dmodel that holistically captures the event's properties and location.
Self-supervised Pretraining Objectives. To create a self-supervised pretraining task, the pre- diction objective must be decoupled from the sequence order naive "next-spacepoint predic- tion" would learn artifacts of the serialization itself. The target for any given spacepoint 8; must be defined by its geometric relationship to other spacepoints in 3D, not its 1D) sequence posi- tion. While predicting nearest neighbors a natural geometric objective, a standard k-Nearest- Neighbor task is unsuitable in autoregressive framework due information leakage from pre- viously seen spacepoints. We partially address this by introducing t-Next-Weerest-Neiieee diction (k-NNN) (Figure Z()),, which aligns the objective with particle propagation. For any query spacepoint st, the model's s task is to predict its k nearest spacepoints within its next neighborhood No((s)) = {sy E T3 > Tı}, i.c. those with larger radius. Let Y, (V.1..... ..........k y.,.} denote the predictions and Y, =
pre-
the ground truth neighbors, both ordered by increasing T -m=1 Pi Vi.ml Larger l: expands the geometric horizon
distance. The loss is then L: 1 and makes the task more difficult.
4.2 ADAPTIVE MODELS FOR DOWNSTREAM TASKS
Track Finding. Figure deppictss how our downstream adapter model for track finding, formulated instance segmentation task, is inspired by image panoptic segmentation models Cheng et al. 2202T 2022) and adapted to sparse spacepoints data.
Poiet features from the FM are first projected to spacepoint embeddings via single linear layer. This projection serves as task-alignment filter, compressing and reorienting the pretrained representation into a lower-dimensional space while also providing a probing point to assess the task relevance of the FM features. We initialize N learnable queries (track queries) Q {qx J and refine them over L transformer decoder layers. In each layer, cross-attention aggregates information from spacepoint embeddings, modulated by an additive attention mask computed from intermedi- ate assignment logits and followed by self-attention among the queries. The resulting refined track queries are passed through two separate multilayer perceptron (MLP) heads produce track em- bedding and a classification score Vn. Point-to- to-query assignment probability P.n. is computed as the sigmoid of the dot product between spacepoint embeddings and each track embedding.
Let E { be the set of true tracks of an event E. We match the refined track queries to E via the Hungarian algorithm, minimizing the combined cost of Dice loss Ldice, Focal loss Lrocal on the per-point assignments, and classification loss Lets for track versus no-object. For each matched pair (T) 9,), the loss is cli-)) = Adice ((im) + Atocal c(i.n) + Aets c(n). Unmatched track queries incur only c(m). We also apply auxiliary losses at each decoder layer. At inference time, each spacepoint i assigned the track n; = arg max, (pm (n.) and labeled accordingly.
Particle Identification and Noise Tagging. For both PID and noise tagging tasks, illustrated in Figure 3 our lightweight adapter first projects each d-dimensional point feature into de di-diimenll embedding via a linear layer then aggregates global context with a single self- attention layer. Finally, it feeds the result through an MLP classifier.
EXPERIMENTS AND RESULTS
Here, we begin by examining the scaling behavior of our FM with respect to model size, dataset size, and computational cost. We then benchmark the FM paired with lightweight adapters against strong baselines across three downstream tasks. Finally, we present additional analyses to better understand the FM's adaptation behavior.

---

Published as a conference paper at ICLR 2026
(a)
x
x_0=>
X
7.0
(M/1 31
12/21 10
63
MSE (1)
6.6
MSE (L)
131
209M
84
10
135
100)
100M
900
0,83
9.00
101
10
30-7
101
100
Model Pazameterss (M)
Specepoint Count (D)
FLOPS (C)
Figure 4: Neural Scaling Behaviors of FM4NPP. We evaluate neural scaling trends on increasing (a) model parameter size M, (b) training spacepoint count D, and (c) compute in FI.OPs. Cmin denotes the minimum for each compute.
5.1 NEURAL SCALING BEHAVIORS OF FM4NPP
Table 1: Model Sizes and Compute Resources
Model Sizes
Compute Resources
m2
m3
m4
m5
m5
NVIDIA GPU
H100 80GB
A100 80GB
Model Width
64
128
256
512
1024
1536
Num GPUs
8
24
64
Model Params
0.34M
1.3M
5.3M
21M
84M
188M
Train His
12
20
32
50
72
Table Performance on Track Finding. Results for the FM4NPP (m6) model are averaged over 10 random seeds. Uncertainties shown in parentheses indicate standard deviation of the mean in the last digit. The best- pest-performing results for each metric are highlighted in bold, while the second-best results are underlined.
Track Finding
#trnb para.
model
ARIT
efficiencyt
purityt
EggNet
0.16M
0.726
74.2%
75.1%
TrkX
3.86M
0.877
91.8%
66.4%
HEPT
0.31M
0.831
81.2%
78.0%
AdapterOnly
2.39M
0.724
78.0%
64.5%
FM4NPP (m6)
2.39M
0.945(3)
96.1(2)%
93.1(1)%
Table 3: Performance on Particle Identification and Noise Tagging. Results for the FM4NPP (m6) model are averaged over 10 random seeds. Uncertainties shown in parentheses indicate standard deviation of the mean in the last digit. The berf-peefo results for each metric are highlighted in bold, while the second-best results are underlined.
Particle Identification
Noise Tagging
#trnbl para.
model
acc.1
recallf
pre.
acc.1
recallf
pre.
GATConv
0.91M
0.692
0.397
0.637
0.910
0.673
0.806
GCNConv
0.91M
0.689
0.391
0.632
0.910
0.673
0.804
SAGEConV
0.91M
0.726
0.456
0.650
0.917
0.723
0.817
GraphConv
0.91M
0.7079
0.4176
0.6425
0.9190
0.7213
0.8252
OneFormer3D
44.95M
0.770
0.490
0.577
0.965
0.940
0.895
AdapterOnly
0.74M
0.663
0.339
0.611
0.911
0.622
0.836
FM4NPP (m6)
0.74M
0.904(1)
0.765(3)
0.878(3)
0.971(1)
0.937(1)
0.919(1)

---

Published as a conference paper at ICLR 2026
We evaluate our FM's scaling behavior across three axes: model size, dataset size, and compute budget. Results are summarized in Figures 4 a)
Model Scaling. We construct a series of FMs with varying capacities, denoted m1 through m6 in Table [ Figure 4a) a) shows the validation mean squared error (MSE) plotted against model size on a log-log scale, revealing a clear power-law relationship. As the number of parameters increases, validation loss consistently decreases, which aligns with neural scaling laws observed in language and other scientific domains Kaplan et al.J (2020); [Hffmaan et al. (2022): Nguyen et al. (2023); Bodnar et al. (2025). Notably, performance plateaus at m6, suggesting a possible saturation point, which we leave for future investigation.
Data Scaling. To isolate the effect of training dataset size, we train the m3 model on varying sub- sets (1%, 2.4%, 11.6%, 20%, 47.6%, 100%) of the full dataset. Figure 44b) shows how performance improves steadily with more data, again following a power-law trend. This suggests the FM can continue to benefit from the large-scale data routinely produced in collider experiments.
Compute Scaling. Finally, we study the relationship between compute and model performance. Figure c)) shows validation MSE against the total number of floating-point operations (FLOPs) used during training. Models up to m3 are trained at 25%, 50%, and 100% of their total iteration budget, while larger models (m4-m6) are trained to full completion only. The results show that smaller models are initially more compute-effiiint, but larger models outperform them when more resources are allocated. This highlights the importance of compute-optimal model scaling for de- ployment in high-throughput environments. All experiments have been conducted using A100 and H100 GPUs with corresponding hardware costs summarized Table
All models are trained with a batch size of 256. An optimal learning rate of x 10-4 is se- lected through hyperparameter tuning on the m3 model and reused across all variants using the y u-parameterrzaaion principle Vankadara et (2024). This approach ensures consistent gradient flow across model sizes and enables zero-shot hyperparameter transfer Yang et al. (2022). Smaller mod- els (m1, m2) are trained for 50,000 iterations, while larger models (m3-m6) are trained for 100,000 iterations. We apply cosine learning rate decay with 10,000-step linear warmup and use gradient clipping at a threshold of 0.1. All experiments employ the Adam optimizer Loshchilov & Hutter (217) with weight decay of 0.01. Additional training details are provided in Appendix B.
arinciple/Vankadaraa
5.2 PERFORMANCE ON DOWNSTREAM TASKS
Track Finding. Our baseline coverage prioritizes reproducible, end-to-end pipelines. Some cent SSM- and Transformer-based trackers cited in "Related Work" are therefore omitted as they either do not report full end-to-end results lack public, runnable releases. Thus, we focus on Exa. Trkx Ju et al. (202T) and EggNet Calafuura et al. (2024). both GNN-based methods de- signed specifically for tracking, as well as HEPT Miao et al. (2024). Transformer- based method that uses target different detector geometries, we adapt them to this dataset (see Appendix C). To verify em- beddings extracted using the pretrained FM provide richer information, we also train the lightweight adapter model alone.
hashing for efficient attention. Because their original implementations
Table 2 reports the track finding results of our FM with several baselines. All metrics are computed over the entire test set rather than averaged per event. For example, tracking efficiency is defined as the fraction of all true tracks the dataset that are successfully matched. Our model achieves higher performance on conventional clustering metrics, such as ARI, and outperforms other approaches in tracking efficiency (recall) and purity (precision).
We also compare this work against the official sPHENIX reconstruction pipeline, which employs Cellular Automaton seeding followed by a Kalman filter Osborn et al. (2021). As that algorithm is optimized for high transverse momentum (pT), long tracks within the TPC acceptance, we restrict this comparison to tracks that leave at least 20 spacepoints in the TPC and satisfy pr > 1 GeV and |y] < 1.1. Under these criteria, our model reaches a tracking efficiency of 99.6%, exceeding the sPHENIX pipeline's 's 94.6%.
Particle Identification and Noise Tagging. For the PID and noise tagging tasks, we experiment with four conventional GNN models and report on the best performing one, SAGEConV. The graph edge set is constructed by k-nearest neighbors with distance cap. We also adapt and train state-
8

---

Published as a conference paper at ICLR 2026
or- segmentation model for 3D point cloud data named OneFormer3D Kolodiazhnyi et al. 2024).
Table 3 reports the segmentation accuracy, as well as macro-averaged recall and precision. For PID task, our FM consistently outperforms all baselines, achieving the highest accuracy, recall, and precision. Meanwhile, for the noise tagging, the FM outperforms GNN-based baselines with similar performance compared to OneForme: armer33. It worth noting that On.eormerr3 has about 45 million trainable parameters, whereas our adapter head has 0.74 million. More details about comparative model implementations and sample outputs are provided in Appendix C.
5.3 INSIGHTS ABOUT FM ADAPTATION
We aim further understand FM adaptation behaviors by answering the following questions:
Q1: Does increasing the size of the FM lead to improved performance on downstream tasks?
Q2: Do larger FMs require fewer labeled examples achieve comparable performance (i.e., better data efficiency)?
Q3: Are the learned FM embeddings task-agnostic, and, if so, how much adaptation is needed to specialize them for specific tasks?
(a) Performance vs. FM model size
(b) Performance dataset size
ml
m3
Track Pinding
PHjmil + Adhater
(elon)) Achoper
Accuracy Accuracy ARI
0.000
0.9
PMEzZ - Adagler
P40mil Adagter
Adapter Only
0.8
PHinll
0.00
6,850
0.900
PID
0.85
0.800
0.5
0.00
0.970
TTgging
2.9x
0.960
0.950
0.2
10
10a
100
1000
10000
70000
Foundation Model Parameters
Number of lebeled events
Figure 5: The effect of FM model size (a) and dataset size (b) downstream task performance.
Figure 5 ()) shows the downstream performance of all three tasks plotted against a pretrained FM size. All tasks use the same frozen pretrained representation. Larger pretrained models consistently yield higher performance across every task, confirming that scaling up the pretrained model size im- proves various downstream performance even when only the lightweight decoder head is trained.
Figure b) depicts the ARI for the track finding task training on different numbers of labeled data, from 100 to 70,000. Larger FMs consistently outperform smaller ones across different levels of labeled data, indicating neural embeddings extracted from larger FMs contain richer information and can be generalized easier. This confirms common empirical observations that larger models can generalize better Novak et al. (2018). In addition, compared baseline adapter-only model trained solely on labeled data (dashed line in F pretraining self-supervised FM model on large amount of unlabeled data easy come by in NPP proves effective. The relative gain in ARI is greater in the fewer labeled data regions compared to those with abundant labeled data: 2.9x versus 1.3x.
5.4 ABLATION STUDIES
All ablations second-largest model m5 and identical training budgets. We report absolute changes in task metrics and the relative increase in the remaining gap to perfect performance in Table For bounded metries such as accuracy and ARI, this normalizes small absolute changes near saturation.
Neighborhood Size k. Our baseline uses 3= during pretraining for the k-NNN objective. Smaller neighhorhoods (k=1 k=5) reduce downstream performance the model is conditioned

---

Published as a conference paper at ICLR 2026
Table 4: Ablation study results on downstream tasks' performances. Numbers show absolute drops in metric with relative increases in the remaining gap to perfect performance (in parentheses). "Hilbert" represents Hilbert space-filling curve.
Ablation
Noise Tagging (Acc.)
PID (Acc.)
Track Finding (ARI)
Next-token prediction (vs. k-NNN)
-.,0100 (4.6%)
-0.0023 (2.5%)
-0.0009 (1.6%)
k == (vs. k = 30)
-0.0012 (5.7%)
-0.0049 (5.3%)
-0.0019 (3.3%)
k 5 (vs. k 30)
-0.0007 (3.6%)
-0.0016 (1.7%)
-0.0003 (0.5%)
Hilbert (vs. Hierarchical Raster Scan)
-0.0014 (7.0%)
-0.0075 (8.0%) -0.0051 (9.1%)
only very local geometry and misses longer-range structure. In contrast, moderately larger & al- lows the FM to capture richer global context while still preserving locality. Takeaway: incorporating broader geometric neighhorhoods during pretraining produces more transferable representations.
K-NNN versus Next-token Prediction. Autoregressive next-token prediction is widely used strategy in vision and language van den Oord al. (2016); Radford et al. (2018). Thus, we compare against it to isolate the value of geometry-ave aware neighborhoods. Takeaway: conditioning on local geometric neighborhoods (k-NNN) yields more transferable representations than generic next-token training.
Serialization Strategy. Space-filling curves are common for prresriyprrs serialization in imaging/point-cluv pipelines Chen et al. (2022). We test the popular Hilbert ordering against Hi- erarchical Raster Scan (Sec. 4.1) serialization. Takeaway: trajectory-consiste serialization beats purely spatial locality interleaving tracks harms downstream coherence.
CONCLUSION AND FUTURE WORK
With this work, we demonstrate that FMs can be effeettvelyy extended to experimental particle physics by introducing a scalable self-supervised training strategy tailored to sparse detector data. Our model, trained with frozen weights and lightweight adapters, consistently outperforming task-specific baselines. Its effective performance and data efficiency suggest the model learns rich, task-agnostic repre- sentations that are easily adapted using simple mappings. These findings reveal the potential for general-purposee scalable models in NPP.
than 10 million events, generalizes across diverse downstream tasks
Limitations and Future Work. While our work establishes a proof of principle for scaling foun- dation models on sparse detector data, we acknowledge limitations regarding the scope of our current evaluation. We demonstrate this approach using single collider experiment (sPHENIX). Trans- forming this into a universal FM, spanning diverse detector systems, multiple facilities (e.g.. LHC), and various collision systems, will require significant future research and community-wide data curation. Additionally, with TnC-only inputs, the number of well-defined downstream tasks is natu- rally limited. Although tested generalizability three downstream tasks, we recognize broader scope of multi-level and multi-modal tasks with multiple detector system is necessary to fully stress test the model's understanding. Most importantly, validating this approach on real experimental data remains essential to realizing the full potential of this technology for the NPP community.
ACKNOWLEDGMENTS
The authors would like to express their sincere gratitude to the SPHENIX Collaboration for sharing the simulation data and experimental knowledge, as well as Jubin Choi, Abhay Deshpande, Alexei Klimentov. Michael Begel, Torre Wenaus, Nicholas D' Imperio, James Dunlop, and John Hill from Brookhaven National Laboratory for their valuable support and feedback. This work was supported by the Laboratory Directed Research and Development (LDRD) Program at Brookhaven National Laboratory, LDRD 25-045, which is operated and managed for the U.S. Department of Energy Of- fice of Science (DOE-SC) by Brookhaven Science Associates under contract No. DE-SC0012704. Shuhang Li was partially supported by the DOE-SC through the Office of Nuclear Physics under Yihui Ren, Xihaier Luo, and Shinjae Yoo were partially sup-
Award No. ported by the DOE-SC through the Office of Advanced Scientific Computing Research and the Scientific Discovery through Advanced Computing (SciDAC) program. This research also utilized
10

---

Published as a conference paper at ICLR 2026
resources of the National Energy Research Scientific Computing Center (NERSC), DOE-SC User Facility, under NERSC Award EERB-EERAP004,,,,, and the authors are grateful to the NERSC staff for their support, particularly Shashank Subramanian and Wahid Bhimji.
REPRODUCIBILITY STATEMENT
We provide anonymized source code and run scripts the supplementary material to reproduce pretraining and all downstream experiments. Hyperparameters, training schedules, and model sizes (ml-m6) are specified in Sec. 5.1 and Table I; dataset provenance, preprocessing, and labeling rules are in Appx. A; architectural/scriiiiztoo
details and optimization settings (AdamW, batch size 256, learning rate 2 x 10 with cosine decay and 10k warmup, gradient clipping 0.1) are in Sec. 4 and Appx. B; evaluation metries and the TrackML double-majority matching protocol are described in Sec. 3; and baseline adaptations plus additional results are in Appx. C and Tables 2-3. We fix random seeds for all reported runs and include configuration files to replicate numbers. Due to storage limits, we do not include pretrained checkpoints, but the provided scripts reproduce them. Hardware used (A 100/HH100) reported in Table 1, and exact environment specs and seeds are listed in the supplementary run scripts.
REFERENCES
Sabrina Amrouche, Laurent Basara, Paolo Calafiura, Victor Estrade, Steven Farrell, Diogo R. Fer- reira, Liam Finmie, Nicole Finnie, Cécile Germain, Vladimir Vava Gligorov. Tobias Golling, Sergey Gorbunov, Heather Gray, Isabelle Guyon, Mikhail Hushchyn, Vincenzo Innocente, Moritz Kiehn, Edward Moyse, Jean-François Puget, Yuval Reina, David Rousseau, Andreas Salzburger, Andrey Ustyuzhanin, Jean-Roch Vlimant, Johan Sokrates Wind, Trian Xylouris, and Yetkin Yil- maz. The tracking machine learning challenge: Accuracy phase. In Sergio Escalera and Ralf Herbrich (eds.), The NeurlPS '18 Competition, pp. 231-264, Cham, 2020. Springer International Publishing.
Ron Belmont et al. Predictions for the sPHENIX physics program. Nucl. Phys. A, 1043: 122821, 2024. doi: 10.1016f.mcccphyys.2244.. 1222.1
Wahid Bhimji, Chris Harris, Vinicius Mikuni, and Benjamin Nachman. Omnilearned: foundation model framework for all tasks involving jet physics. arXiv preprint arXiv:2510.24066 2025.
Joschka Birk, Anna Hallin, and Gregor Kasieczka. Omnijet-a: the first cross-task foundation model for particle physics. Machine Learning: Science and Technology. 5(3):035031, 2024.
Cristian Bodnar, Wessel P Bruinsma, Ana Lucic, Megan Stanley, Anna Allen, Johannes Brandstetter, Patrick Garvan, Maik Riechert, Jonathan A Weyn, Haiyu Dong, et al. foundation model for the earth system. Nature, pp. 1-8, 2025.
Rishi Bommasani. Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg. Antoine Bosselut, Emma Brunskill, al. On the opportu- nities and risks of foundation models. arXiv preprint arXiv:2 2021.
Brookhaven National Laboratory. sPHENIX Detector at RHIC https:/wwww.bbl.gg//ttii//pppppp php 2025. Accessed: 2025-07-19.
Paolo Calafiura, Jay Chan, Loic Delabrouille, and Brandon Wang. Eggnet: An evolving graph-based graph attention network for particle track reconstruction, 2024. URL hntps://ariv..or/bb//44077
13925
Polo Calafiura, Steven Farrell, Heather Gray, Jean-Roch Vlimant, Vincenzo Innocente, Andreas Salzburger, Sabrina Amrouche, Tobias Golling, Moritz Kiehn, Victor Estrade, Cécile Germaint, Isabelle Guyon, Ed Moyse, David Rousseau, Yetkin Yilmaz, Vladimir Vava Gligorov, Mikhail Hushchyn, and Andrey Ustyuzhanin. Trackml: A high energy physics particle tracking challenge. In 2018 IEEE 14th International Conference on e-Science (e- (e-Science), pp. 344-344, 2018. doi: 10. 109/ 201.ieennee228..000 1.00088..
11

---

Published as a conference paper at ICLR 2026
Wanli Chen, Xinge Zhu, Guojin Chen, and Bei Yu. Efficient point cloud analysis using hilbert curve. In Shai Avidan, Gabriel Brostow, Moustapha Cissé, Giovanni Maria Farinella, and Tal Hassner (eds.), Computer Viston ECCV 2022, pp. 730-747, Cham, 2022. Springer Nature Switzerland. ISBN 978-3-031-20086-1.
Bowen Cheng, Alexander G. Schwing, and Alexander Kirillov. Per classification is not all you need for semantic segmentation, 2021. URL httpps//aaxxi..orggt////0......0
Bowen Cheng, Ishan Misra, Alexander G. Schwing, Alexander Kirillov, and Rohit Girdhar. Masked- attention mask transformer for universal image segmentation, 2022. URL. https://arxi..rr//bb//
2112.01527
ATLAS Collaboration, Georges Aad, et al. Observation of new particle in the search for the standard model higgs boson with the atlas detector at the lhc. Physics Letters B, 716(1):1-29, 2012. doi: 10.10166..hyyslct2221...88
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learn- ing. 2024. URL http://s///ppeeeviiieeeet//rrrro http:///forommmmdd//////CFCCC Id=d
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference Learming Representations, 2021. URL https:
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-33) 2022.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0077 2023.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024. URL
VY2
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203.15556 2022.
Lawrence Hubert and Phipps Arabie. Comparing partitions. of Classification, Classficcctiionn 2(1):193-218; 1985.
Cheng Jiang and Sitian Qian. Application of structured state space models to high energy physics with locality sensitive hashing. In International Conference Artificial Intelligence and Statis- tics, pp. 3961-3969. PMLR, 2025.
Xiangyang Ju, Daniel Murnane, Paolo Calafiura, Nicholas Choma, Sean Conlon, Steven Farrell, Yaoyuan Xu, Maria Spiropulu, Jean-Roch Vlimant, Adam Aurisano, Jeremy Hewes, Giuseppe Cerati, Lindsey Gray, Thomas Klijnsma, Jim Kowalkowski, Markus Atkinson, Mark Neubauer, Gage DeZoort, Savannah Thais, Aditi Chauhan, Alex Schuy, Shih-Chieh Hsu, Alex Ballow, and Alina Lazar. Performance of geometric deep learning pipeline for hl-lhe particle track- ing. The European Physical Journal C, 81(10), October 2021. ISSN 1434-6052, doi: URL http:/dxx.ooi.org/ 10. 1-tpepppcss10002222-0000666666
10.1 1e1e114/ep///ssl 10052-021-096677777
Nilotpal Kakati, Etienne Dreyer, Anna Ivina, Francesco Armando Di Bello, Lukas Heinrich, Marumi Kado, and Eilam Gross. Hgpflow: extending hypergraph particle flow to collider event reconstruc- tion. The European Physical Journal C, 85(8):1-18, 2025.
R. E. Kalman. A new approach linear filtering and prediction problems. Journal of Ba- sic Engineering, 82(1))35 45, 03 1960. ISSN 0021-9223. doi: 10.1115/1.366255.. URL https://bh.org/00.111//....6
12

---

Published as a conference paper at ICLR 2026
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
Henry Klest. Overview and design of the sPHENIX TPC. J. Phys. Conf. Ser, 1498:012025, 2020. doi: 10.1088/1742-6596 he////19////// 19011/0
Maxim Kolodiazhnyi, Anna Vorontsova, Anton Komushin, and Danila Rukhovich. Oneformer3d: One transformer for unified point cloud segmentation. In Proceedings of the IEEE/CVF Confer- on Computer Vision and Pattern Recognition, pp. 20943 -0993,, 2024.
Qing Li, Zhihang Hu, Yixuan Wang, Lei Li, Yimin Fan, Irwin King, Le Song, and Yu Li. Progress and opportunities of foundation models in bioinformatics. Briefings in Bioinformatics, 2024. doi: 10. 10.109//b//bbbd Survey of foundation models in bioinformatics.
Shuhang Li, Yi Huang, David Park, Xihaier Luo, Haiwang Yu, Yeonju Go, Christopher Pinkenburg, Yuewei Lin, Shinjae Yoo, Joseph Osborn, Christof Roland, Jin Huang, and Yihui Ren. Tpepp- 10m: Simulated proton-proton collisions in a time projection chamber for ai foundation models, 2025. URL httpss//arri..og///b225500.5.1.0
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 1rXi0:1711.0500,, 2017.
Qin Ma, Yi Jiang, Hao Cheng, and Dong Xu. Hamesssing the deep learning power of foundation models in single-cell omics. Nature Reviews Molecular Cell Biology, 25(8):593-594, 2024.
Siqi Miao, Zhiyuan Lu, Mia Liu, Javier Duarte, and Pan Li. Locality-sensitive bashing--bseed effi- cient point transformer with applications in high-energy physics. International Conference on Machine Learning, pp. 15546935669. PMLR, 2024.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99-106, 2021.
Farouk Mokhtar, Joosep Pata, Dolores Garcia, Eric Wulff, Mengke Zhang, Michael Kagan, and Javier Duarte. Fine-tuning machine-learned particle-flow reconstruction for new detector geome- tries in future colliders. Physical Review D) 111(9):092015, 2025.
Clara Moskowitz. Tiny bubbles of primordial soup re-create early universe. Sci- entific American, 328(3), March 2023. URL https//www..sieeiiaaeeriaa..mm
oony-uubbees-o--rrmoooo www.rddiaa-ssuu---errrreeee wwwreetee-aary--ueeeeee
Accessed: 2025-07-19.
Tung Johanees Brandstetter, Ashish Kapoor, Jayesh K. Gupta, and Aditya Grover. Climax: foundation model for weather and climate. In Proceedings of the 40th International Conference on Machine Learning. ICML'23, JMLR.org. 2023.
Johannes
Roman Novak, Yasaman Bahri, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Sensitivity and generalization in neural networks: an empirical study. In International Conference on Learning Representations, 2018.
Joseph D. Osborn, Anthony D. Frawley, Jin Huang, Sookhyun Lee, Hugo Pereira Costa, Michael Peters, Christopher Pinkenburg, Christof Roland, and Haiwang Yu. Implementation of ACTS into sPHENIX Track Reconstruction. Comput. Softw. Big Sci., 5(1):23, 2021. doi: 10.1007/
$41781-021-00068-ww
Edward O. Pyzer-Knapp, Matteo Manica, Peter Staar, Lucas Morin, Patrick Ruch, Teodoro Laino, John R. Smith, and Alessandro Curioni. Foundation models for materials discovery - cur- rent state and future directions. npj Computational Materials, 11:15, 2025a. doi: 10.1038/
0415524-025-1533--
Edward O Pyzer-Knane Matteo Manica, Peter Staar, Lucas Morin, Patrick Ruch, Teodoro Laino, John R Smith, and Alessandro Curioni. Foundation models for materials discovery- current state and future directions. Npj Computational Materials, (1()),,1, 2025b.
13

---

Published as a conference paper at ICLR 2026
Alec Radford, Karthik Narasimhan, Tim Salimans, and llya Sutskever. Improving language under- standing by generative pre-training Technical report, OpenAI, 2018.
T Konstantin Rusch, Michael M Bronstein, and Siddhartha Mishra. A survey on oversmoothing in graph neural networks. arXiv preprini arXiv:2303. 2tt::2303.109933 10993, 2023.
Aäron van den Oord, Nal Kalchbrenner, and Koray Kavukeuoglu. Pixel recurrent neural networks. In Proceedings of the 33rd International Conference on Machine Learning (ICML), 2016.
Leena Chennuru Vankadara, Jin Xu, Moritz Haas, and Volkan Cevher. On feature learning in struc- tured state space models. In The Thirty-eighth Annual Conference on Neural Information Pro- cessing Systems, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Wenhai Wang, Jifeng Dai, Zhe Chen, Zhenhang Huang, Zhiqi Li, Xizhou Zhu, Xiaowei Hu, and et al. Internimage: Exploring large-scale vision foundation models with deformable convolutions. In CVPR 2023, pp. 3591-3600,, 2023. doi: 10.1 1..10///22P22688..... 2023. 00356.
Greg Yang, Edward Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ry- der, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs Tuming large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022.
Yukun Zhou, Mark A Chia, Siegfried K Wagner, Murat S Ayhan, Dominic J Williamson, Robbert R Struyven, Timing Liu, Moucheng Xu, Mateo G Lozano, Peter Woodward-Courf et founda- tion model for generalizable disease detection from retinal images. Nature, 622(798 (001199911116--11 1):116-1113,, 2023.
14

---

Published as a conference paper at ICLR 2026
A DATASET
The dataset used in this work is based simulated proton-proton (p+p) collisions a center-of- mass energy of v8 = 200 GeV, corresponding to conditions of the sPHENIX experiment at the Relativistic Heavy Ion Collider (RHIC). Charged-particle trajectories are recorded with the Time Projection Chamber (TPC). p+p collisions serve as precision workhorse for testing quantum chro- modynamies (QCD) and nucleon structure and provide the baseline for quantifying how particle production in heavy-ion collisions, viewed as a superposition of p+p interactions, is modified by the sourk-gluon plasma (QGP) Busza al. (2018).
Relativistic Heavy Ion Collider
Gele beand
RHIC
1.2 kmn
beam
SPHENIX
STAR
LINAC
NSR
EBIS
BOOSTER
AGS
Time Projection Chamber
Figure 6: Relativistic Heavy Ion Collider Brookhaven National Lab and sPHENIX Experiment.
A.1 SIMULATION AND PROVENANCE
Minimum-bias p+p collisions are generated with PYTIIAA-.337 Sjöstrand et al. (2015) "Detroit" tune Aguilar et al. (2022) then propagated through full GEANT4 Agostinelli et al. (2003) simu- lation of the bu-built sPHENIX detector, including its detailed CAD geometry and measured 1.4T field map. The 'FTFP BERT HP' physics list used for high- treatment neutron and hadron interactions. The simulation chain models continuous energy loss, multiple scattering, ondary particle production, and decay processes with the true material budget, as well as supports space-charge distortion and its data-driven correction and carries signals through the full front-end electronics (shaping, digitization, zero suppression, and channel-by-channe gain/noise).
The simulated TPC response, so-called G4HITS, emulates raw ionization signals from charged par- ticles traversing the TPC volume, which reconstructed into spacepoints reflecting the true spatial resolution and distortions. Each spacepoint is then matched to the Monte Carlo truth particle that produced it, and the particle's properties identity, momentum, and track association recorded as ground-truth labels for our downstream tasks (track finding, particle identification (PID), and noise identification).
A.2 CONTENTS AND STRUCTURE
Each event contains:
Reconstructed spacepoints from the TPC, including position and ionization energy.
Monte Carlo truth particles with their PDG identity, momentum at production, and vertex location at production.
Associations between spacepoints and truth particles.
15

---

Published as a conference paper at ICLR 2026
A.3 DATASET STATISTICS
The event-level complexity in the dataset varies widely. As shown in Figure the number of constructed TPC spacepoints per event ranges from a few hundred to tens of thousands, reflecting low-multiplicity relatively busy collision topologies. Correspondingly, the number of truth tracks per event spans from under 10 up to nearly 100.
Figure summarizes the class composition for the noise-tagging and PID downstream tasks.
(a) Number of spacepoints distribution
(b) Number of tracks distribution
800
mean= 856
400
mean= 15.6
Comnt
600
300
400
200
200
100
0
2,000
4,000
6,000
8,000
20
40
60
80
Number
Number
Figure Distributions of number of spacepoints (a) and tracks (b) per event.
Noise-tagging. Noise spacepoints are defined operationally based on the truth-level kinematics of their progenitor particles. Specifically, any spacepoint associated with a Monte Carlo truth track whose momentum below 60 MeV / labeled noise. Particles produced in the primary p+p collision with such low momentum kinematically unable reach the active TPC volume due to the magnetic field. Therefore, spacepoints matched to these low-momentum tracks arise predomi- nantly from secondary interactions with detector material (e.g., delta electrons, conversion products, or other material- naterril-innuued processes). These secondary oorggin spacepoints are not part the primary signal topology of interest and treated as noise for the purposes of the corresponding downstream classification task.
Noise tagging
(b) Particle
69%
250
23.8%
38.2%
7%
pion
loson
proton
noise
signal
electron
others
Figure 8: Class ratio of noise tagging (a) and particle identification (b).
The PID task uses five coarse -rrrine- target classes, grouping charge-conjugate species to- gether reduce sparsity while preserving physics relevance:
PID.
Pion: ग., π
Kaon: K+,K** K
Proton: proton and anti-proton
Electron: electron and positron
Other: all remaining particle species.
16

---

Published as a conference paper at ICLR 2026
The class ratios shown in Figure reflect the inherent imbalance in these labels, driven by the underlying particle production spectra and the noise definition.
17

---

Published as a conference paper at ICLR 2026
B METHODOLOGY
B.1 PRELIMINARIES
This section outlines a compact mathematical way to express the hierarchical relationship between events, tracks, and spacepoints in particle detector, such as Time Projection Chamber (TPC). A collision event E represented as a set tracks (7,), where each track T3 an ordered sequence of spacepoints (sx) and each spacepoint is vector (Eácp, _, y, 2.....) containing its physical properties. Concretely, we express a single event, E, as follows:
E = {T
This states that an event (E) is set containing m individual tracks (T₃). The number of tracks, m, varies for each event. Each track, in turn, is defined by its constituent spacepoints
T = (
This expresses that a single track (T²) is an ordered sequence of n, spacepoints ($y.k). The sequence ordered because particles follow specific path through the detector, and the number of space- points per track, "z. is also variable. Finally, each individual spacepoint is a vector of its properties, which can be represented abstractly as:
$j,k R
A spacepoint (s) is a vector in a D-dimensional feature space. A Spacepoint (s3.2) is now explicitly defined as a vector containing its primary physical properties:
$j = (E, x, y, 2)j.k
where & is the energy deposited by the particle at that point in the detector and (ar, y, 2) is the spatial coordinates of the spacepoint.
B.2 COORDINATE TRANSFORMATION
We transform spacepoint coordinates from Cartesian (x. y, 2) to cylindrical-polar system (r, d, n) that better reflects the geometry and symmetries of collider experiments. The radial distance is defined
/ 22 + y², measuring how far point lies from the beamline in the transverse plane, and is essential for evaluating transverse momentum and energy. The azimuthal angle á is given by o atan2(y, z). describing the orientation of the spacepoint in the x-y plane and ex- ploiting the detector's s cylindrical symmetry around the beam axis. The pseudorapidity vj is defined as 17 In [tan (0/2)), where 0 atan2(r; is the polar angle. This coordinate is used instead of 0 because particle production tends to be uniform in 7), and, for highly relativistic particles, approximates the Lorentz-invariant rapidity. Finally, ensure consistent feature scaling, we apply a min max normalization to the spatial coordinates, transforming the pseudorapidity (17 [-2, 2), azimuthal angle (ő [- T),, and radial distance (r E [30, 78), centimeters) into the interval [0, 1). The transformed 8; (E, r, o. 17); are used for all analyses described in this paper.
B.3 SERIALIZATION
Our objective is to perform self-supervised pretraining on the raw three- diree-dimensional (3D) point cloud of particle spacepoints from a collision event, S = {s sy}. To leverage the power of sequential models like MAMBA, which have excelled in learning rich representations, first solve the fundamental problem of transforming the unordered 3D set into an ordered one- dimensional sequence. This serialization process is not merely technical step. The choice of ordering scheme is critical for preserving the data's underlying physical structure.
must
An ideal serialization must satisfy two competing demands: it must respect the global physics of the event (i.e., particles flying outwards) while simultaneously preserving the local continuity of individual particle tracks.
We first analyze and dismiss naive approaches. For example, a space-filling curve excels at preserv- ing 3D locality but completely disregards the concept of a track. Its path erratically jumps between physically distinct trajectories, creating a chaotic signal. Conversely, a simple global raster scan on the spacepoints' cylindrical coordinates, s (rr, 6; n1), respects the outward propagation along the radius but fails on local continuity. The initial hits of track (at low r) become "context-starved their preceding elements in the sequence belong to entirely different tracks.
18

---

Published as a conference paper at ICLR 2026
Proposed Solution: Hierarchical Raster Scan To resolve this dichotomy, introduce Hier- archical Raster Scan. This method balances global structure with local context by operating on two levels:
1. Partitioning: The entire detector volume is partitioned into grid of smaller 3D "boxes."
2. Ordering: A raster scan using the physically motivated order (r, 6.) 17) is applied twice. First, it orders the spacepoints within each box (intra-box ordering). Second, it orders the boxes themselves based on their geometric centers (inter-box ordering).
This strategy ensures the sequence progresses globally outwards but maintains local contiguity within each partitioned region. However, even with this optimal serialization, profound challenge remains. the learning objective simply is to predict the next hit in this sequence, the model would be forced learn the arbitrary artifacts of the serialization itself, particularly the artificial jumps box boundaries.
Therefore, designing a robust serialization scheme a necessary but insufficient step. The learning objective must be intelligently designed to be independent of these serialization artifacts, challenge we address the subsequent section.
The division of the detector volume into grid is not uniform. Instead, is a physics-inforned partitioning designed to align with both the detector's physical geometry and the observed distribution of particle hits. This ensures the partitioning itself provides a meaningful structural prior for the learning task.
Physics-informed Partitioning
For the azimuthal angle (ø) and pseudorapidity (n) dimensions, the binning data-driven. The boundaries are specifically chosen to create bins with roughly uniform density of hits. This strat- egy balances the information content across partitions, preventing high-occupancy regions from disproportionately influencing the model. A detailed number of bins and illustration of this binning strategy provided in Figure
Distribution
200000
373000
150000
cals with Number of bins: (Eta, Phi, Radius) = (8, 8, 6)
125000
50000
25000
Distribution Radius
200000
3.75
175000
1.50
300000
15000
8.75
50000
8:50
23000
8.00
Figure 9: Physics-informed Partitioning. Top-left graph shows the binning of the data space into 384 bins (8 8 x 6). The other plots show the distribution of spacepoint values in normalized Phi, Eta, and Radius dimensions, respectively, computed using 50,000 events.
For the radial dimension (r), the partitioning mirrors the sPHENIX detector's physical construction. The detector's ' 48 layers are arranged three major groups. Thus, we create six radial bins, allo- cating two bins to each major detector group. By embedding the detector's 's known layered structure into the partitioning scheme, we further ground the serialization process in the experiment' 's physical reality.
19

---

Published as a conference paper at ICLR 2026
B.4 MAMBA: SELECTIVE STATE SPACE MODELS
Mamba represents a significant advancement in sequence modeling, challenging the dominance of the Transformer architecture, particularly for long sequences. It a selective state space model (SSM) that combines the strengths of recurrent neural networks (RNNs) and convolutional neural networks (CNNs) to offer linear-time complexity and constant- time inference.
Mamba's foundation is the SSM, a continuous- time system described by the following linear ordi- nary differential equation:
dh(()) = Ah(t) + Bx(t), dt
y(t) - Ch(() + Dz((t)
Here, h(t) is the latent state, xxt) is the input, and y(t) is the output. A, B. C, and D are matrices that are typically learned from data.
For use in deep learning, this continuous system is discretized. A crucial step in Mamba is making the key matrices, particularly the transition matrix A and the input projection matrix B, selective and input-dependent. This is achieved by having dedicated neural networks that predict these matrices based on the current input token.
The discretized formulation of the state transition is:
he = Ah + Bx,
(1)
Where À and B are the discretized, input-dependent matrices. This selectivity allows Mamba to modulate its recurrent state, effectively controlling how much of the past to retain and how to incor- porate the current input. The model can be unrolled for efficient parallel training, similar a CNN, or used in recurrent manner for constant-time inference.
Architectural Principles The core innovation of Mamba is its selective mechanism, which allows the model to dynamically adapt its parameters based on the input. This enables it to focus on relevant information and filter out noise, crucial capability for processing long and complex sequences. Unlike traditional SSMs, which are time-invariant, Mamba's parameters are functions of the input, making it a time-varying system. Key components the Mamba architecture include:
Selective State Space Layer:
This is the fundamental building block of Mamba. It replaces the attention mechanism and feed-forward network of a Transformer block.
Hardware-Award Algorithm: Mamba employs a parallel scan algorithm that is optimized for modern hardware (GPUs), enabling efficient training and inference. This algorithm avoids the materialization of the full state sequence, a significant memory bottleneck in traditional SSMs.
Mamba-2 Mamba-2 is direct successor Mamba, designed further improve upon its effi- ciency and performance. It introduces a new theoretical framework called Structured State Space Duality (SSD), which provides a deeper understanding of the relationship between SSMs and other architectures like Transformers. The primary motivation behind Mamba-2 was to address some of the hardware utilization inefficiencies of the original Mamba. While Mamba offered linear- time complexity, its performance on modern GPUs could still be optimized. Key improvements Mamba-2 include:
SSD: This framework establishes formal equivalence between class of structured SSMs and a form of global convolution. This duality allows for the design of more efficient algorithms by leveraging insights from both perspectives.
Architectural Simplifications: Mamba-2 simplifies the Mamba block by replacing the complex selective scan with a more structured and hardware- wewear--rreedlyyyyyeool formulation derived from the SSD framework. This often involves a multi-headed Mamba block, analogous to the multi-head attention in Transformers.
Improved Hardware Utilization: The redesigned architecture of Mamba-2 2 is more amenable to parallelization on modern hardware, leading to significant speedups in both training and inference compared the original Mamba.
20

---

Published as a conference paper at ICLR 2026
B.5 FM4NPP: ARCHITECTURE
Positional Embedding The model first transforms the raw input data into high-dimensional space suitable for sequence processing. An input batch of serialized collision events is represented tensor of shape (B, S, 4), where B is the batch size, S is the sequence length, and each space- point is a nour-dienssioal (4D) vector comprising its deposited energy and 3D spatial coordinates [Baep 7); 6, r). This tensor is processed by embedding module that projects the 4D spacepoint features into the model's 's latent space, Dmodel. It also computes a positional encoding from the 3D spatial coordinates using function 7()) inspired by Neural Radiance Fields (NeRF), defined
(insin(22) (p) = (p. sin (2p), cos (2p)..... sin (2'p). cos (2'p))
where p is the coordinate vector and the frequencies 2 are sampled from geometric progression. This encoding, also mapped to Dmedal is combined with the feature representation via element-wise addition. The output of this stage is a single tensor of shape (B, S, modeel), where Dmotel is the model width.
Network Architecture and k-Next Nearest Neighbor Prediction Head The core architecture consists of a stack of Mamba blocks that sequentially process the embedded hits. The input to the first block is the (B, S. Dmoka) tensor from the embedding stage. Each block operates as follows:
:re-Normalization: The input tensor is first passed through Root Mean Square Normal- ization (RMSNor)) layer. This layer normalizes the feature vector of each spacepoint independently.
Sequence Modeling: The normalized (B, S, Dimdss) tensor is then processed by the Mamba-2 layer.
Residual Connection: A residual or "skip" connection is applied around the normalization and Mamba-2 layers. The original input to the block is added element-wise to the output of the Mamba-2 layer.
After passing through the final Mamba block, the sequence is processed by one last RMSNorm layer. The resulting (B, S, Dmedal) tensor is then fed into the prediction head. This head is a single linear layer that projects the -indea-dimensional representation of each hit yielding a final output tensor of shape (B, S, 3k). Here, k = 30 is the number of neighbors to be predicted. This output format designed specifically for the Causal k-Nearest Neighbor (kNN) objective.
3l-dimensional vector,
B.6 MAXIMAL UPDATE PARAMETERIZATION
Challenge in Scaling Models Imagine building with LEGOs. If you build small car, it is stable. Yet, if you try to build a life-sized car using the exact same small-brick techniques, it will be flimsy and fall apart. Modern AI models face a similar problem. When we try to make them bigger and more powerful by adding more "width" or digital neurons, their internal mathematics can become unstable during training. The signals inside can either "explode" into uselessly large numbers or "vanish" to zero, making it impossible for the model to learn. μ- r-PParameterzztion (μPP is groundbreaking set of rules that solves this problem. It is much like a master blueprint for building Al models, explaining exactly how to adjust the initial settings and learning rate based on the model's size. This ensures that as the model scales up, internal signals stay perfectly balanced, allowing it to train stably and effectively. A major benefit is the best training settings found on a small, cheap model can be directly transferred to a massive, expensive one, saving enormous amounts of time and computational cost.
Concretely, standard infinite-width network analyses, such as those based on the Neural Tangent Kernel (NTK), predict that wide networks operate in a "lazy regime," where they fail to learn meaningful features from data. μP was introduced to overcome this limitation by defining a spe- cific scaling of model initializations and learning rates that guarantee nontrivial feature evolution in the infin limit. A significant practical advantage of μP is that it enables zero-shot hy- perparameter transfer, allowing optimal settings found on small-scale models to be directly applied to their large-scale counterparts. This mitigates the often prohibitive computational costs associated with tuning large models.
21

---

Published as a conference paper at ICLR 2026
Applications in Modern Architectures The principles of juP have been successfully extended be- yond simple multilayer perceptrons (MLPs) to a range of complex architectures. In Transformers, uP facilitates hyperparameter transfer, although achieving stable scaling with respect to both model width and depth. The framework also has been adapted for sci- entific machine learning models, such as Fourier Neural Operators (FNOs), where specific μ-FNO parameterization ensures stable training as the model size and number of Fourier modes are scaled. Recently, μP has been applied stabilize training of large Diffusion Models, again enabling hy- perparameter transfer for these computationally intensive generative systems. This body of research highlights both the generality of the uP framework and the necessity of deriving architecture speccfffc scaling laws.
limit requires careful
μP for MAMBA
To address this, a corrected scaling for SSMs, termed μP-SSM (Maximal Update Parameterization for SSMs), was derived by analyzing signal propagation directly within the Mamba architecture. This analysis yielded specific scaling rules for initialization variances (σ), which con- trol the scale of the model's initial random weights, and learning rates (n) that determine the step size during training. The key formulas dictate how these parameters for Mamba's 's weight matrices (Wn n. We) should be scaled relative to the model's 's latent state dimension (N₂) and input dimension (N.). Using asymptotic Big-Thete (+) notation, the rules are:
Initialization Variances: dB E (
No N,
and oc 00 (NNN..
Learning Rates: 1/8 E O and IJc E 0( NEVNN.
We have integrated this μP-SSM methodology into our own Mamba- based model. The effectiveness of this approach is evidenced by the stable scaling of layer-wise activation norms across different model sizes as empirically verified our experiments. Unlike standard parameterizations that lead to exploding signals or heuristic μP which leads to vanishing signals, model's activations and their updates remain correctly scaled, confirming that the model is operating in a stable feature- learning regime.
B.7 ADDITIONAL DETAILS FOR PRETRAINING
The model trained using the Adam optimizer, which incorporates weight decay for regulariza- tion against overfitting. To manage the learning rate dynamics, employ cosine decay schedule, which is preceded by a brief linear warmup period at the beginning of training to ensure initial stabil- ity. To further prevent training instabilities arising from large gradients, we apply gradient clipping. The learning objective is to minimize a Mean Squared Error (MSE) loss function. This loss quan- tifies the Euclidean distance between the model's predicted coordinates for the kNN and the truth coordinates. These truth neighbors are pre-computed for each particle spacepoint during the data loading phase ensure efficient throughput during training.
Loss Rescaling by Event Difficulty We identified nuisance structure in the training data related event spacepoint density: events with a larger number of spacepoints are inherently easier to predict as the average distance between neighboring spacepoints is smaller. This variance in diffi- culty can lead to training instability, manifesting as loss spikes. To mitigate this, we introduce a loss rescaling strategy based event binning. Events are first grouped into discrete bins based on their average KNN distance, which serves as a proxy for prediction difficulty. Let gi) be the function that maps event i to its corresponding difficulty bin. The loss objective is then modified as follows: 1) the MSE for each event re- weighted by a factor Wg(t), corresponding to the average difficulty of its bin, and 2) the total batch loss is calculated by averaging these re-weighted individual losses. This formulated as: C
B W 0
B
S S Sij
1
C.
Σ gg
Σ j=l
Viyl1
B
B
2=1
Here, B the number of events in the batch; L; is the standard MSE for event with Sn spacepoints; Si and yij are the predicted and truth coordinates, respectively: and Wg(s) is the pre-computed weight for the difficulty bin to which the event belongs. This ensures a single batch-averaged loss is computed only after accounting for the inherent difficulty of each event in the batch.
22

---

Published as a conference paper at ICLR 2026
C ADDITIONAL RESULTS
C.I. DOWNSTREAM MODEL
C.I.I TRACKING (INSTANCE SEGMENTATION)
Our lightweight downstream model for track finding formulated per-point instance seg- mentation task is inspired by image panoptic segmentation models, such as MASKFORMER and MASK2FORMER, adapted to point cloud data.
Let X = {x denote the input set of N points, where each x; E Rd is d-dimensional point- level feature (either raw input, pretrained representation, or from a randomly initialized encoder). These are first projected into latent embedding space via linear layer:
e,
W p
E R
We denote the of projected spacepoint embeddings E {e:
To represent candidate tracks, use K learnable queries (track queries) Q {4k (0) j where (0) each I E Rdd. These prototypes are refined over transformer decoder layers. Each decoder
layer consists of:
Cross-attention: updates qu by attending to point embeddings E.
Self-attention: refines interaction among the K prototypes.
Feed-forward network (FFN): standard Transformer update.
After L decoder layers, we obtain the refined track queries Q(4) = =1 1 Jkm Each refined query vector is then processed by two MLPs:
mg mLP (9()) (L)),
UK MLP (
where m. E Id is the track embedding for the k-th prototype and track instance prediction EN [0, 1] is the probability of corresponding to a real track (versus a "no-object" class).
Each track embedding m2 used compute point-to-prototype assignment logits:
Zik e, mi; Pix o (zsk);
where o-) denotes the sigmoid function. The predicted assignment probability P.x represents the likelihood that point i belongs to prototype k.
To encourage each track query to focus on the subset of points it is likely responsible for, we apply an additive attention mask during cross-altention. The attention mask is defined as:
Au - log(pi +
with small constant added for numerical stability. This mask added to the attention logits before the softmax operation in the cross-attentios layer. This dynamic masking suppresses contributions from low-probability points and improves localization by making each prototype attend selectively to its likely constituent points.
M be the set of M ground-truth -runh tracks (instance labels). We com- pute bipartite matching between the M ground-truth truth tracks and the K refined track queries using the Hungarian algorithm. The matching minimizes a cost function combining:
Training Loss. Let T
Dice loss Laice on the per-point predicted versus ground-truth track
Focal loss Lrocal on point-wise assignment probabilities
Classification loss Lets the track/no-object prediction.
For each matched pair (T), qx), the total loss is:
c (j.k) maich
Adice
C ((.h) dice
+ Atecal focal + Aets (((h)
For unmatched prototypes, only compute (k) with the ground truth label being "no-object."
23

---

Published as a conference paper at ICLR 2026
well as The final training loss includes auxiliary losses from each decoder layer l 1, .,.... ,,, L, from the initial prototype vectors:
Lotal Σ
2=0
During inference, we assign each spacepoint to the track whose combined mask and classification score is maximal. Concretely, we compute k = arg max (pix Uts and label point as belonging to track kj-
This formulation enables end-to-end training of the instance segmentation model, while allowing the pretrained or learned point embeddings to guide track-level grouping.
C.1.2 PARTICLE IDENTIFICATION AND NOISE IDENTIFICATION
For both PID and noise classification, we use a simple lightweight adapter:
Embedding: A linear layer projects each point feature E R to d, dimensnsio em- bedding.
A single self-attention layer aggregates global information across all point em- beddings.
Context:
Prediction: An MLP with softmax over C' output classes.
C.2
COMPARATIVE METHODS FOR DOWNSTREAM TASKS
C.2.1 ADAPT EXA. TRKX PIPELINE FOR SPHENIX TRACKING-FINDING
In this section, we discuss several adaptions done to the Exa. TrkX pipeline for it to work well on the sPHENIX data. We need to apply adaptions to the first four stages data pre-processing. hit embedding, edge filtering, and GNN edge classification out of six stages of the Exa TrkX pipeline.
Pre-processing. The Exa. TrkX's study was based on the TrackML dataset Amrouche et al. 2020). The dataset provides two sources for the construction of the neural network input the 3D location of the spacepoints and the directional information and summary statistics from the charge deposited in each spacepoint (eight-dimensional)) The second source of information is called cell features in the paper. The hit feature the concatenation of the location and cell features. Because sPHENIX data do not provide cell features, we only used the location of hits in the HEP-coordinate construct the input. More precisely, let (f) o, f) be the location of hit (normalized pscudorapid- ity, angle, and normalized radius), the features of this hit is eie-dimensional vector
(E, fj cos( fi)
where E is the energy. We use (cos(o), sin Φ) instead of á to overcome the discontinuity of Φ at 2m. We normalize the pscudorapidity by 1.96 to get the normalized pscudorapidity n (-1, 1). To normalize a radius, we first match it to the closest one of the 48 radius bins and use the bin index to replace the radius. Then, we divide the index by 48 to normalize the radius to number between (0, We do this because the distance between the sPHENIX TPC layers are not uniform with outer layers, spacing farther apart than the inner ones. This may problem for distance- based edge set construction for GNN model same-track hits toward the end of the track may be less likely to be connected by the model.
Embedding and Filtering. The Exa TrkX pipeline embeds the spacepoints and filters the edges as two separate steps. To adapt them for sPHENIX, we modified the procedure in the following aspects: 1) how determine whether pair of hits is connected, 2) how candidate hit pairs generated, 3) how to traine the models, and 4) how to construct the models.
In the embedding stage, Exa. TzkX trained an MLP network to embed each hit into a latent repre- sentation, so pairs of neighboring hits from the same track closer in the latent space than pairs that are not (e.g., from different tracks or not neighbors on the same track). The embedding network trained by first passing the two hits through the same embedding network then minimizing the hinge loss of the distance between the two embeddings.
Because sPHENIX data do not provide information to determine if two same-track hits are direct neighbors (although this information could be inferred for high-energy tracks), we decide not to
24

---

Published as a conference paper at ICLR 2026
distinguish whether two same-track hits are neighboring. This approach also was recommended by the Exa. Trkx research team as a valid alternative.
In the filtering stage, Exa. Trkx takes a pair of hits, passes them through the embedding network, concatenates the two embeddings, and passes the concatenation through MLP filtering network to predict if the two hits are connected. The prediction is optimized by a binary cross entropy loss.
For both the embedding and filtering models, we need to provide candidate hit pairs. For the embed- ding stage, Exa. TrkX uses two types of pairs, random and kNN pairs, as a form of hard negative mining- As random pair has an extremely low chance to be connected, Exa. also trains on pairs formed by hit and its closest neighbors in the latent representation space.
We follow the pipeline as closely as possible. However, because of the differences between sPHENIX and TrackML input features and the fact we treat all pairs from the same track as being connected (in contrast to Exa TrkX's is approach where only immediate neighbors are connected), we have to choose different cutoffs in both embedding and filtering. Specifically, we set a threshold 2. for distance in the embedding space with pairs less than the threshold apart classified as having an edge between them. The threshold was chosen as ensures we have an over .8 recall (efficiency in the Exa. TrkX terminology) in identifying pairs from the same track. Of note, here we did not select a threshold that ensures close to a 100% recall because we can afford the model to fail to recognize faraway points from the same track as being connected.
For the filtering step, we choose a threshold of .675 for probability of a true edge with pairs exceed- ing the threshold considered as being connected. The threshold is selected because it ensures the false positive rate in edge identification will dip below 1%.
GNN edge classification. For the final GNN step, we also employ the Interaction Network Battaglia et al. 2016) architecture with the same hyperparameters used by the Exa. TrkX pipeline. For edge classification, we choose a threshold of .9 as probability of a true edge. With this choice, we achieved 91%79% tracking efficiency (recall) 194. for tracks with pr > 1 GeV) and track purity (pre- cision) of 66.46% 46. With a threshold of .8, the tracking efficiency drops slightly to 90.01% (92.60% for pT
GeV) with notable improvement in purity 76.72%
C.2.2 ADAPT EGGNET FOR SPHENIX TRACK-FINDIIG
The EggNet. study also is based on the TrackML dataset Amrouche et al. (2020) and shares the same data pre-processing approach with Exa. TrkX. To partially compensate for the lack of cell features from sP HENIX data, we use the following approach to augment the input. Let (ño, 6o. fo) be the location of a hit (normalized pseudorapidity, angle, and normalized radius). The features of this hit is a 12 dimensional vector
(oo cos( sin()v), fo; i. cos((1)) (in())), A1; f22 cos( sin((2), f2),
where (m, Φ1, f1) and (f2 Φ2 f22) are the locations of the two closest neighbors of the hit in the (1) (os( (on( f) space. The motivation for augmenting the hit with the two closest neighbors is that for the majority of the hits in a high-energy track, the two closest neighbors are most likely from the same track. In which case, the augmented hit features can provide information about the track's direction.
For the GNN model, EggNet adopted similar approach to GravNet Qasim et al. (2019). The outstanding feature of a GravNet-lype model that the edge set not predetermined but con- structed dynamically. More precisely, EggNet will run N normal GNN iterations, but, before each GNN iteration, the edge set will be constructed via kNN based on the current node embeddings. To adapt EggNet to sPHENIX data, set GNN iterations to be 4 and used four pressage-passing rounds for each GNN iteration. The nearest 10 hits in the embedding space are used to form the neighborhood of hit. Different from the original GravNet (but similar to the interaction GNN used by Exx..TkX) EggNet also has an edge network for calculating edge messages. More- over, EggNet also uses a dedicated node decoding network to produce the node embeddings for the KNN. All sub-networks of EggNet (node encoding/decoding networks and edge network) are MLPs, each with two hidden layers and 64 hidden features. The embedding dimension of the node (i.c.. the number output features from the node decoding network) is 24.
The network is trained with a hinge loss of margin 1, aimed at reducing the Euclidean distance in the embedding space between pair of hits from the same track and enlarging the distance between
25

---

Published as a conference paper at ICLR 2026
pair from different tracks. The model was trained for 300 epochs, and the final clustering was done using DBSCAN with c = 1 and minimum number samples 2.
C.2.3 ADAPT HEPT FOR SPHENIX
HEPT Miao et al. 2024) is locality-sensitive hashing-basel efficient point Transformer designed for large-scale point cloud processing in high energy physics. Unlike Exa. Trkx and EggNet, which rely on graph neural networks, HEPT uses self-attention mechanisms with LSH-based ap- proximation to achieve ncar-linear computational complexity.
To adapt HEPT for SPHENTX TPC tracking, we made the following modifications to the model and training procedure:
Pre-processing, We use the same normalized high energy physics coordinates as discussed C.2.1 (f) (o), sin()), f). where r) is the normalized pseudorapidity and f is the normalized and binned radius. For the input features, we concatenate the energy of each hit with its Cartesian coordinates (r, y, and the normalized high energy physics coordinates, resulting in an eight-dimensional fea- ture vector per spacepoint. Unlike some baseline methods, we do not filter tracks by transverse momentum pT and considered particles across all momentum ranges.
Contrastive Learning. HEPT is trained using a contrastive learning objective that brings embed- dings of same-track hits closer together while pushing embeddings from different tracks apart. For negative sampling, we form negative examples from at most 64 neighboring hits from different tracks (reduced from the original 256 to accommodate the lower spacepoint density in sPHENIX TPC data).
Training Configuration. We use the most recent model architecture from the HEPT example folder. For optimization, we set the initial learning rate to 0.0001 (instead of the original 0.01) and switched the AdamW optimizer (from Adam) for better regularization. These adjustments were necessary to achieve stable training convergence on the sP X dataset.
Track Formation. Because НEРТ produces per-point embeddings rather than end-to- end tracking predictions, we apply HDBSCAN McInnes et al. (207)) clustering on the learned embeddings to form track candidates. We use the following HDBSCAN hyper- parameters: metric= cluster _eteerrsllllle ect on www.methhddd/ae "eom". These parameters are tuned to maximize the average per-event Adjusted Rand Index (ARI) on the test set, balancing cluster granularity and noise robust- ness.
dean", min cluster s size12 min samples=15, 5, and
C.2.4 ADAPT GNNS FOR SPHENIX PARTICLE IDENTIFICATION AND NOISE TAGGING
We selected four GNN models, GATConv, GCNConv, GraphConv, and SAGEConv, as benchmarking algorithms for the PID and noise tagging downstream tasks. We used the torch.geomet ri implementations for the models. We used the same data pre-processing pro- tocol as discussed in C.2.1 To generate the edge set, for hit location (17 cos( (in(),))())))))))))))))))) sin()), (d), f), we connect it 50 nearest neighbor hits with distance < 1. We allowed the edges to be directed. The node features to the GNNs are the energy E of the hit together with its 4D location. For the node encoding network, uniformly six GNN layers for each GNN model. For the hit classification network, we use an MLP of two hidden layers with 128 and 64 hidden features. The GNNs are trained with cross entropy loss. Each GNN is trained for 200 epochs.
an MLP, each with two hidden layers and 256 hidden features. We
In general, GNNs' performance on the two downstream tasks are suboptimal. We hypothesize that the failure of GNNs is result of their difficulty in capturing and communicating more global pat- terns of the tracks as solving both particle identification and noise tagging requires a model to un- derstand the general shape of tracks that span significant space in TPC.
C.2.5 ADAPT ONEF 3R FOR SP PEHNIX PARTICLE IDENTIFICATION AND NOISE TAGGING
OneFormer3D is state-of-ft www.of--tt--tttte object detection algorithm for 3D point cloud data that can solve semantic and instance segmentation task in one run. The model architecture of OneF ormer 3D is U-Net backbone followed by transformer decoder.
26

---

Published as a conference paper at ICLR 2026
To rum OneFormer on point cloud data, we first need to get the so-called "super points" (a grouping of raw points) either by a clustering algorithm or voxelization. To adapt OneF ormee 3D the same pre- prepeeocssing approach as discussed in C.2.1 and voxelized the resulting point cloud to a grid of shape (64, 64, 48) in i), é i, respectively.
to sPHENIX data,
The super points first pass through the sparse convolution -poeedord--oedddd U-Net backbone to be featur- ized. Then, the super point features serve as the keys and values in the Transformer encoder. The learnable queries output from the transformer decoder are then used to produce instance/semantic segmentation predictions on the super points. In the final step, the prediction on the super points will be broadcast to their constituent raw points. As both partiele identification and noise tagging can be considered semantic segmentation tasks, we separate the part of the code (primarily in prediction and loss function) for semantic segmentation from OneF o ormer 3D while keeping the neural archi- tecture identical. We use the same network parameters the example of CneF Formereer33 3D on the S3DIS dataset.
Table 5: Noise tagging per-class recall and precision.
Accuracy
Macro
Non-noise
Noise
Recall
Precision
Recall
Precision
Recall
Precision
GATConV
0.9099
0.6730
0.8060
0.9788
0.9242
0.3672
0.6878
GCNConv
0.9095
0.6728
0.8037
0.9784
0.9241
0.3672
0.6832
GraphConv
0.9190
0.7213
0.8252
0.9764
0.9351
0.4661
0.7152
SAGEConv
0.9174
0.7227
0.8165
0.9740
0.9355
0.4714
0.6975
OneFormer3D
0.9646
0.9404
0.8948
0.9716
0.9884
0.9092
0.8012
AdapterOnly
0.9111
0.6215
0.8359
0.9901
0.9169
0.2528
0.7548
FM4NPP (m6)
0.9708
0.9122
0.9114
0.9809
0.9812
0.8435
0.8416
Table 6: Particle Identification per-class recall and precision.
Accuracy
Others
Proton
Roc.
Pre.
Rec.
Pre.
Rec.
Pre.
Rec.
Ptc.
Rec.
Pre.
Rec.
Pre.
GATConv
0.6922
0.3973
0.6368
0.0947
0.5709
0.9106
0.7014
0.0057
0.6146
0.4567
0.6117
0.5190
0.6854
GCNConv
0.6892
0.3911
0.6319
0.0782
0.5762
0.9144
0.6966
0.0073
0.5871
0.4501
0.6140
0.5058
0.6858
GraphCoav
0.7079
0.1176
0.6125
0.1301
0.5739
0.9133
0.71416
0.0080
0.5791
0.1766
0.6272
0.5597
0.7178
SAGBConv
0.7262
0.4563
0.6502
0.1085
0.5790
0.9126
0.7382
0.0338
0.5239
0.6242
0.7071
0.6024
0.7028
OneFormer3D
0.7701
0.1897
0.5767
0.3029
0.5758
0.9207
0.7658
0.0000
0.0000
0.1859
0.6991
0.7389
0.8427
AdapterOnly
0.6631
0.7714
0.9511
0.0002
0,2872
0.4120
PM4NPF (m6)
0.904
0.6623
0.8328
0.4449
0.7647
0.9651
0.8484
0.2712
0.7829
0.8068
0.8763
0.8336
0.8919
Table 7: Diagnostic metrics for tracking performance.
model
ARI
overall spacepoint efficiency
overall spacepoint purity
no. parameters
EggNet
0.7256
93.01% %
92.34%
0.16M
Exa.TrkX
0.8765
94.47%
98.83%
3.86M
AdapterOnly
0.7243
89.34%
92.09%
2.39M
FM4NP? (m6)
0.9448
97.56%
98.34%
188M + 2.39M
C.3 ADDITIONAL RESULT ON SHARED VERSUS INDEPENDENT ADAPTER ARCHITECTURES FOR DOWNSTREAM TASKS
To address whether the downstream tasks can benefit from end-to-end joint training with shared adapter layers, conduct ablation study comparing multitask learning against isolated task- specific adapters. We focus on the two most architecturally similar tasks: PID and noise tagging, both of which operate on point-level representations.
27

---

Published as a conference paper at ICLR 2026
Noise Tagging Ground Truth
AdapterOnly Prediction
PНАNPP Prediction
sigual (1028)
(0.97, 0.68)
(0.98, (1, 0.99)
noise (1476)
(0.68, 0.97)
(a) Event 6937
signal (3044)
(1,0.55) (0:03, 0.67)
(0.97, 0.99)
(530)
(0:93, 0.85)
(b) Event $520
signal (3094) mbiise (28)
(0.769, 0.992)
(0.978, 0.997)
(0.357, 0.014)
(0.643, 0. 0.209)
(c) Event 192
Figure 10: Performance of AdapterOnly and FMANPP on Noise Tagging, The numbers in the parentheses in the target sub-figures the number of signal and noise spacepoints. The numbers in the parentheses the prediction sub-figures are the recall and precision of the class.
Multitask Architecture. The multitask model shares a common input projection layer followed by two self-attention (SA) and feed-forward network (FFN) layers across both tasks with separate task-specific classification heads for PID and NID. The training objective combines both task losses with manual weighting: Lioul UewDCrD + CND We the pretrained m5 backbone (1536- dim, frozen) with weights WNID = 2.5 and wPID = 0.5 to prioritize the simpler binary noise tagging task.
Results. Table compares the best validation losses achieved by multitask learning against iso- lated training, where each task its own dedicated adapter layers. Despite sharing representations through common SA+FFN layers, the multitask model exhibits negative transfer: both tasks perform worse than when trained independently. The NID task degrades by 7.6 7.61% %, while the PID task suffers more severe 30.69% increase in validation loss.
Table 8: Comparison of multitask versus isolated adapter training for PID and noise tagging tasks. Both configurations use the frozen m5 backbone with two shared SA+FEN layers. Lower validation loss better.
Task
Isolated Training Multi-Task (Shared)
A (%)
Noise Tagging (NID)
0.0513
0.0552
+7.61%
Particle ID (PID)
0.2377
0.3107
+30.69%
Interpretation. These results suggest that despite conceptual overlap between PID and noise tag- ging. both classify individual spacepoints, the learned representations required for optimal perfor- mance differ substantially between tasks. The negative transfer likely arises from conflicting gradient
28

---

Published as a conference paper at ICLR 2026
Particle Identification Ground Truth
AdapterOnly Prediction
PMANPP Prediction
pion (1212)
(0.97,0.57) 0.57)
(0.93, 0.94)
lnon (94)
[0, ()
10.81,0.93)
proton (0)
(0,0)
(0, 0)
electron (316)
(0.15, 0.54
(0.91. 0.94
others (514)
[0.0]
(0.91,
(a) Evens 5
pion (1339)
(0.98, 0.69)
(0.99, 0.94)
knon (109)
(0.0)
(0.79.1)
proton (134)
[0.38, 0.8)
(0.84, 0.93)
electron (110)
(0. (0.01.0 16
(0.92. 0.99
others (17)
(0.0)
(..59.0.67)
(b) Event. 1966
pion (632)
(0.76, 0.63)
(0.88, 0.74)
knon (d1)
(0,0)
(0.59.4)
proton (59)
(0,0)
(0.58.0.3 0.97)
electron (1691).
(0.88, 0.87)
(0.89, 0.97)
others (61)
(0.0)
(0, 0)
(c) Event 6937
Figure 11: Performance of AdapterOnly and FM4NPP on particle identification. The numbers the parentheses in the target sub-figures are the number of spacepoints each particle ID class. The numbers in the parentheses in the prediction sub-figures are the recall and precision of the class.
signals: noise tagging requires distinguishing signal from detector noise based on energy deposition patterns, while PID must differentiate between particle species using ionization profiles. The more severe degradation in PID performance indicates the shared adapter prioritizes the simpler, more heavily weighted NID task at the expense of the more complex PID classification.
For the tracking task, which employs a fundamentally different DETR (DEtection TRansformer)- style set prediction architecture rather than point classification, joint training with PID/NID is even less suitable. Consequently, we retain independent task spec-sp adapters for each downstream task, allowing each to specialize its learned representations without interference while still leveraging the shared pretrained backbone.
C.4 ADDITIONAL RESULT ON ADAPTER HEAD CAPACITY VERSUS FOUNDATION MODEL QUALITY TRADE-OFF
To understand whether adapter capacity or foundation model quality is the limiting factor for down- stream performance, conduct ablation study on the PID task by sweeping adapter depth while keeping the pretrained m5 backbone frozen.
We varied the number of self-attention (SA) layers in the adapter head from 0 to 4 while maintaining the frozen m6 backbone. The configurations tested include:
Experimental Setup.
OL (0.56M params): No SA layers, only linear projection + MLP
1L (1.09M params): SA layer
2L (1.62M params): SA layers
4L (2.67M params): SA layers
29

---

Published as a conference paper at ICLR 2026
Track Finding Ground Truth
AdapterOnly Prediction
FNANPP Prediction
(a) Event
(b) Track 14052 and
(c) Track
(d) Track 3832 and
Track -3619 and
Figure 12: Performance AdapterOnly and FM4NPP on track finding. In panel (a). we show the ground-trute truth tracks, the AdapterOnis track candidates, and the FM4NPP track candidates (note that two different tracks might have the same color as the length of the color cycle used may be smaller than the number of tracks). In panel (b)-(e), we show four pairs of close-by ground-truth tracks that the AdapterOnl, erönly model fails to separate while the FMANPP model does.
For each configuration, train two variants: one with the pretrained backbone (frozen) and one without pretraining (replacing the backbone with a learnable linear layer) to isolate the effect of pretraining quality.
Results. Figure D shows validation loss as a function of adapter capacity. Performance improves substantially when adding the first SA layer (OI. 1I: 45.7% loss reduction) with diminishing
30

---

Published as a conference paper at ICLR 2026
returns thereafter. The optimal configuration uses two SA layers, achieving validation loss of 0.238. Adding more layers (4L) slightly degrades performance to 0.244, suggesting overfitting or that the adapter capacity exceeds what the frozen backbone effectively support.
PID Classification: Validation Loss vs Adapter Size
Pretraining
1.0
Pretraining
Pretraining improvement (% loss)
50
Best Validation Loss
Impr (%)
34.4%
0.7
+5
0.5
-1
02
Adapter Parameters (Milllons)
Figure 13: Adapter capacity versus performance on PID task. Validation loss decreases as adapter depth increases from zero to two self-attention layers then plateaus or slightly degrades at four layers. The gap between pretrained and non-pretrained models shrinks with larger adapters, indicating that small adapter heads are the performance bottleneck rather than foundation model quality.
Interpretation. The results reveal two key insights: 1) performance plateaus after two SA layers, indicating the frozen FM representation is already sufficiently rich small adapter heads are the bottleneck rather than backbone quality. 2) While the relative pretraining benefit diminishes with larger adapters (54.2% for OL decreasing to 19.8% for 4L), the pretrained backbone still provides substantial absolute improvements across all configurations. Even with the largest 4I. adapter, the model with pretrained features achieves 0.244 validation loss compared to 0.304 without pretraining significant gap that underscores the effectiveness of the FM representation.
C.5 ADDITIONAL LEARNED EMBEDDINGS RESULTS
FM4NPP (frozen)
FM features
Linear map
Post-lincar features
Adapter (trainable)
PCA Reduction
PCA Reduction
Figure 14: Visualization of learned embeddings from the foundation model (left) and the post- lost-linear map features (right), projected via PCA reduction. Each marker corresponds to a spacepoint, colored by its associated track identity.
31

---

Published as a conference paper at ICLR 2026
We analyze the neural embeddings from the frozen FM and their transformation after simple linear projection, which precedes the lightweight adapter used for downstream tasks. To probe task speci- ficity, we apply dimensionality reduction techniques (.g.,.CAA to both the raw FM embeddings and linearly projected features, focusing on a representative downstream task: track finding. As shown in Figure [1] the raw FM embeddings exhibit no clear separation among particle tracks, indicating the representations are task-agnostic. However, after applying a single linear projection, distinct and well-separated clusters emerge, corresponding to different particle tracks.
In Figure 15| present results obtained by applying various dimensionality reduction techniques (including PCA, t-SNE, and UMAP) to both FM features and downstream adapter features. For illustrative clarity, we randomly select two data samples. The results demonstrate consistent improvement, clearly showcasing the FM features' adaptability. Even after a single linear projection, the FM embeddings exhibit substantial clustering and separability, indicating rapid adaptation to the downstream track-findingg task. Adapter features consistently provide superior discrimination, yielding distinctly well-separated clusters corresponding different track categories. Because linear transformation alone cannot create separabilityy where none exists, this demonstrates the FM encodes rich, general- purppsse information that only requires minimal alignment to become task- specific. It also explains why lightweight adapters, when built atop FM embeddings, outperform non-FM baselines by leveraging semantically meaningful input features.
PCA Reduction
t-NNE Hedrcctio
UMAP Reduction
Fleer
Figure 15: Dimensionality reduction results using PCA, t-SNE, and UMAP on randomly selected test data samples.
To further validate the robustness and generalizability of the FM features, we systematically in- vestigate the impact of varying dimensionality reduction parameters using t-SNE. Specifically, we conduct experiments by setting the reduced dimensionality to three, four, and five and visualized the results by plotting the first two t-SNE components (See Figure I6]. Across all tested dimensional configurations, the FM features consistently demonstrate pronounced clustering patterns and clear separability, highlighting their intrinsic adaptability and effectiveness in supporting diverse down- stream classification tasks.
32

---

Published as a conference paper at ICLR 2026
Redused dimension
Reduced dimennsion
Reduced dimeasion
motem stett
Faonr 2
Figure 16: t-SNE visualizations for randomly selected test instances across various reduced dimen- sions.
In Figure 17 we extende analysis to multiple downstream tasks, again using randomly selected test data instances and employing t-SNE for visualization. The FM features' separability was no- tably effective for the track-finding task, slightly diminished for particle identification, and consid- erably reduced for noise tagging. The limited performance observed in noise tagging is attributed to the inherent imbalance of the binary classification data, making separability challenging due to the dominant prevalence of a single label. Overall, our analyses confirm a hierarchy of effectiveness FM embeddings across downstream tasks. Track-findigg demonstrates the strongest separability, followed by particle identification, and lastly noise tagging. These findings align well with the FM's pretraining objective, neighbor identification, and are consistent with task relevance from a physics perspective.
Partide
Notse
Thoging
Figure 17: 1-SNE visualizations for randomly selected test instances across various downstream tasks.
33

---

Published as a conference paper at ICLR 2026
REFERENCES
Agostinelli al. GEANT4- A Simulation Toolkit, Nucl. Instrum. Meth. A, 506:250-303, 2003. doi: 10.1016/50168-9000000 (00000688880000000)06888884
Manny Rosales Aguilar, Zilong Chang, Raghav Kunnawalkam Elayavalli, Renee Fatemi, Yang He, Yuanjing Ji, Dmitry Kalinkin, Matthew Kelsey, Isaac Mooney, and Veronica Verkest. pythia8 underlying event tune for RHIC energies. Phys. Rev. D, 105(1):01001, 2022. PhysRevD.100.01600
doi: 10.1 103/
Sabrina Amrouche, Laurent Basara, Paolo Calafiura, Victor Estrade, Steven Farrell, Diogo R. Fer- reira, Liam Finnie, Nicole Finnie, Cécile Germain, Vladimir Vava Gligorov, Tobias Golling. Sergey Gorbunov, Heather Gray, Isabelle Guyon, Mikhail Hushchyn, Vincenzo Innocente, Moritz Kiehn, Edward Moyse, Jean-Franee SaanFFraçcois Puget, Yuval Reina, David Rousseau, Andreas Salzburger, Andrey Ustyuzhanin, Jean-Roch Vlimant, Johan Sokrates Wind, Trian Xylouris, and Yetkin Yil- maz. The tracking machine learning challenge: Accuracy phase. In Sergio Escalera and Ralf Herbrich (eds.), The NeurlPS 18 Competition, pp. 231-266- Cham, 2020. Springer International Publishing.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 29, 2016.
Wit Busza, Krishna Rajagopal, and Wilke van der Schee. Heavy Ion Collisions: The Big Pic- ture, and the Big Questions. Ann. Rev. Nucl. Part. Sci., 68:339-376, 2018. doi: 10.1146/ canurev-nucl-lee 1011--00117--205522
Leland McInnes, John Healy, and Steve Astels. hdbscan: Hierarchical density based clustering. Journal of Open Source Software, 2(11))205, 2017.
Siqi Miao, Zhiyuan Lu, Mia Liu, Javier Duarte, and Pan Li. Locality hashing- based effi- cient point transformer with applications in high-energy physics. In International Conference Machine Learning, pp. 35546-35569. PMIR, 2024.
Shah Rukh Qasim, Jan Kieseler, Yutaro liyama, and Maurizio Pierini. Learning representations of irregular particle-detector geometry with distance-weighted graph networks. The European Physical Journal C, 79(7):1-11, 2019.
Torbjörn Sjöstrand, Stefan Ask, Jesper R. Christiansen, Richard Corke. Nishita Desai, Philip Ilten, Stephen Mrenna, Stefan Prestel, Christine O. Rasmussen, and Peter Z. Skands. An introduction to PYTHIA 8.2. Comput. Phys. Commun., 191:159- 177, 2015. doi: 10.10166..pc22015.1.0444
34


Reviews:
[{"review_id": "CST81ZyUgi", "reviewer": "Reviewer_QgWZ", "summary": "The manuscript presents a foundation model for nuclear and particle physics, trained on proton-proton (p+p) collision data at a center-of-mass energy √s = 200 GeV. The proposed architecture utilizes the Mamba model as its backbone and is pretrained using a novel k-Next-Nearest-Neighbor prediction objective. The authors evaluate the model's performance on three downstream tasks: Track Finding, Particle Identification, and Noise Tagging.", "questions": "- The model is positioned as a foundation for particle physics, yet it is trained only on RHIC data. Given that the LHC represents the forefront of particle physics, why was the model not trained or validated on LHC data to demonstrate its utility as a true foundation model for the entire field? The authors must justify this fundamental limitation in scope and discuss the model's generalizability to higher-energy, more complex LHC environments.\n-  Can the authors quantify the computational overhead of the k-NNN objective compared to a baseline, and justify how the performance benefits warrant this increased expense?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The authors present detailed technical considerations for building the foundation model, though there are shortcomings on specific design and experiment choices.", "weaknesses": "- The manuscript does not adequately situate its contribution within the emerging landscape of foundation models for high-energy physics. It fails to discuss or compare against other recent efforts (e.g., models based on transformers, graph networks, or other architectures applied to similar data).\n- The datasets used appear limited in scope, focusing on a specific collision system and energy. How does this restricted training data affect the model's validity and utility as a general-purpose foundation model for the broader field of nuclear and particle physics?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qkyLRiQplR", "reviewer": "Reviewer_9giK", "summary": "The paper introduces FM4NPP, a scientific foundation model for particle physics that attempts to address existing scaling challenges outside of natural language due to, e.g., the sparse nature of detector data, difficulties in serializing 3D data. The model is trained on simulated collisions under conditions of the sPHENIX experiment and evaluated on three core downstream tasks, mostly outperforming non-FM baselines on key metrics. FM4NPP is also shown to observe scaling laws at high parameter counts, and may present a useful paradigm for scientific FMs in other domains.", "questions": "- Were any evaluations run using the adapter models under a different (non-FM4NPP) spacepoint embedding model? This might be a distinct way to isolate the performance of the proposed FM architecture.\n- As presented with the three downstream tasks, can this approach be trained end-to-end? The explicit adapter-based approach seems to be the most flexible route and keeps the backbone in more of an embedding role, but given the overlap in the three tasks it seems plausible that these could fully captured end-to-end (potentially internalizing/sharing knowledge across tasks previously isolated within the adapter models).\n- Were any other kinds of downstream tasks tested (e.g., forecasting)?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is well-organized, includes a comprehensive literature review, and provides high quality figures/diagrams.\n- The paper's stated contributions are clear and address a difficult, high-impact problem in the NPP domain. The approach potentially lays the groundwork for tackling broader challenges relevant to scientific foundation models beyond particle physics settings.\n- The reported evaluation is extensive, covering a variety of important dimensions that help position the model's utility and behavior. For instance, the model is compared to several baseline methods on the downstream tasks (highlighting its comparative advantages), key ablations are reported (justifying architectural decisions), and different model sizes are evaluated, highlighting scaling laws.\n- I appreciate the discussion style of Section 5.3, which sets up several key questions relevant to practitioners seeking to leverage FM4NPP on new tasks.", "weaknesses": "- I find the core positioning of the paper to be somewhat difficult to pin down. There is a clear focus on particle physics as the primary domain of interest, but the proposed training scheme for scaling and adapting FMs to downstream tasks is very general. To fully demonstrate the generality of the scheme, more (scientific) domains would need to be evaluated to verify the approach transfers well to other settings (helping substantiate the stated hypothesis that FMs encode task-agnostic representations useful for downstream tasks).\n\n  Note that this isn't to take away from the model's performance on this task, but simply that the methodology leans more general-purpose while the evaluation focuses entirely on a single domain/dataset. \n- The downstream tasks used for evaluation appear limited in their variety. In particular, they seem to test for similar kinds of model understanding, which is to say classifying spatiotemporally connected groups of points in the particle space. For instance, given that the model performs well on track finding, it's not surprising that it would perform similarly well on particle identification or noise tagging (as these tasks, naively, simply re-group the tracks). Including downstream tasks like forecasting would be useful for demonstrating broader model understanding, more explicitly highlighting its ability to understand/extrapolate particle dynamics.\n- It is ultimately quite difficult to asses how well the trained model fills the role of a foundation model. This point is made in tandem with the first two, i.e., that it's hard to gauge general model understanding provided the single dataset and task suite. What's concretely demonstrated is that FM4NPP is an effective embedding model for particles under conditions of the sPHENIX experiment, and large parameter counts can help with downstream classification tasks. But the jump to labeling FM4NPP as a foundation model is hard to justify without stronger evidence of its generality, such as more diverse tests of model understanding and/or evaluation under different conditions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QFwhIazxzv", "reviewer": "Reviewer_2w6N", "summary": "This work introduces a foundation model tailored towards nuclear and particle physics. The authors develop a large scale open dataset with 10 million events and benchmark on three downstream tasks. The authors develop a self-supervised training methodology for their foundation models, and present benchmarks on downstream tasks in addition to scaling behaviors for their foundation models. They demonstrate improved performance of their model over baselines across all tasks.", "questions": "* Why choose mamba over other linear transformers (deltanet, GLA..) and state space models? \n* How long are the sequence lengths for the input? State space models struggle with extremely long context. \n* Does the hierarchical raster scan have a strong impact on the amount of compute?\n* What is the full model architecture of the foundation model?\n* Can the code and models be made public?\n* Do you actually claim that self-supervised pre-training + fine-tuning outperforms a dedicated model trained from scratch on the downstream task in Fig. 5 (b)? Or are the architectures different? If you extend the axis, would Adaptor only eventually converge to the FM performance, as is seen in other FM papers in particle physics?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The authors present a well motivated problem that they address with a multi-step comprehensive framework \n* The self-supervised learning objective is physics motivated, and intuitively works with the detector setup\n* Very thorough benchmarking and ablation studies show improved performance for all pieces of the foundation model\n* Performance on downstream task show greatly increased performance over baselines\n* Scaling curves show increased performance for increased flops/parameter count, indicating scalability of the model\n* Shows greatly increased performance against state of the art models\n* Shows that the self-supervised learning technique improves performance on downstream tasks", "weaknesses": "* Does not justify choice of mamba over linear transformer models, or other state space model alternatives\n* Unclear how the hierarchical raster scan would impact the efficiency\n* No benchmarks against efficient transformer baselines such as HEP-T (locality sensitive hashing transformer for high energy physics applications, https://arxiv.org/abs/2402.12535), which performs very well on trackML datasets\n* For PID, also missing comparisons to dedicated physics ML algorithms, like MLPF (https://arxiv.org/abs/2503.00131) or HGPF (https://arxiv.org/abs/2410.23236)\n* Does not define the model architecture/construction (number of layers, parameters per layer etc)\n* There are no code or pretrained weights available\n* Unclear of pretraining + fine-tuning outperforms dedicated models trained from scratch; Most other literature on FMs in particle physics do not claim this and instead focus on data efficiency (i.e. better performance when training on smaller datasets)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FOxlutr8sf", "reviewer": "Reviewer_Soo5", "summary": "The paper proposes FM4NPP, a self-supervised foundation model for nuclear & particle physics (NPP) built on a Mamba (SSM) backbone. It introduces (i) a >10M-event simulated RHIC/sPHENIX TPC dataset and three downstream tasks (track finding, particle ID, noise tagging), (ii) a Hierarchical Raster Scan to serialize sparse detector spacepoints, and (iii) a k-Next-Nearest-Neighbor (k-NNN) autoregressive pretraining objective that predicts outward (in radius) neighbors to avoid order-leakage. With frozen FM features and lightweight adapters, the largest model (188M params) reportedly outperforms GNN/point-cloud baselines across tasks and shows clean scaling trends w.r.t. model/data/compute.", "questions": "1. Have you evaluated FM4NPP embeddings on real sPHENIX runs (even small samples)? How does sim-trained FM perform zero-shot/adapter-shot on real data, and what domain-gap mitigation (calibration, augmentation, feature normalization) is most effective?\n2. How do different grid resolutions $r \\times \\eta \\times \\phi$, binning strategies, or non-radial intra-box sortings affect pretraining loss and downstream ARI? Could a learnable/local ordering, or permutation-invariant SSM input layer, match or exceed Hierarchical Raster?\n3. Why constrain neighbors strictly to $r_j > r_i$ rather than a cone or helix-aware neighborhood? Did you try multi-horizon targets mixing local and mid-range neighbors, or contrastive variants that avoid explicit regression?\n4. For tracking, why a DETR-style query decoder over bipartite clustering in embedding space? For PID/noise, did you test slightly larger heads or lightweight decoders (e.g., 2–4 self-attention layers) to probe head-capacity vs. FM-quality trade-offs?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper targets an important gap: general, reusable representations for sparse detector data rather than bespoke task models. The two-stage “pretrain FM -> freeze ->  task adapters” design is well-motivated.\n2. The Hierarchical Raster Scan balances global outward geometry with local track continuity; k-NNN aligns the self-supervised target with physics causality (increasing radius) rather than sequence order. The ablation table isolates benefits over next-token prediction and Hilbert curves.\n3. Clear power-law trends for model/data/compute; the paper also documents hardware/iterations, optimizer, and $\\mu$-param transfer, useful for reproducibility. Also, FM4NPP + small adapters beats GNN and 3D point-cloud baselines on ARI and efficiency.", "weaknesses": "1. The approach based on the proposed ordering. Although ablations compare Hierarchical Raster vs. Hilbert, it’s unclear how sensitive performance is to bin sizes, box aspect ratios, or detector layout changes (e.g., silicon layers, calorimeters), and whether learned positional encodings could reduce this hand-crafted dependency.\n2. k-NNN chooses neighbors with a larger radius (future in $r$). That choice encodes a physics prior (outward propagation) but may under-represent curling/looping tracks, secondaries, or back-scattering. Ablations vary $k$ but not the neighborhood definition (e.g., cone in $\\eta,\\phi, r$).\n3. All tasks are TPC-centric, point-level segmentations. For a foundation model, it would be compelling to include fusion tasks (TPC+EMCal), rare-event tagging, vertexing, or particle-flow style reconstruction, even as small-scale probes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)