fixing inference

Browse files

Files changed (4) hide show

README.md +64 -65
root_gnn_dgl/README.md +9 -3
root_gnn_dgl/scripts/inference.py +48 -32
root_gnn_dgl/setup/Dockerfile +4 -7

README.md CHANGED Viewed

@@ -31,11 +31,11 @@ We present foundation models that adopt an architecture similar to that used for
 To provide a diverse set of physics processes for the pretraining, we use Madgraph@NLO 2.7.3 [Alwall et al., 2014](#ref-alwall-2014hca) to generate proton-proton collision events at next-to-leading order (NLO) in Quantum Chromodynamics (QCD). We generate 12 distinct Standard Model (SM) physics processes, including six major Higgs boson production mechanisms: gluon fusion production \\(ggF\\), vector boson fusion \\(VBF\\), associated production of the Higgs boson with a W boson \\(WH\\) or a Z boson \\(ZH\\), associated production of the Higgs boson with a top-quark pair \\(t\bar{t}H\\), and associated production of the Higgs boson with a single top quark and a forward quark \\(tHq\\). Additionally, we simulate six top quark production processes: single top production, top-quark pair production \\(t\bar{t}\\), top quark pair production in association with a pair of photons \\(t\bar{t}\gamma\gamma\\), associated production of a top-quark pair with a W boson \\(t\bar{t}W\\), simultaneous production of three top quarks \\(t\bar{t}t\\), and simultaneous production of four top quarks \\(t\bar{t}t\bar{t}\\). In these samples, the Higgs boson and top quarks decay inclusively. These 12 Higgs and top quark production processes constitute the pretraining dataset.
-To test the pretrained model, we further generated four processes including three beyond Standard Model (SM) processes: a SM \\(t\\bar{t}H\\) production where the Higgs boson decays exclusively to a pair of photons, a \\(t\\bar{t}H\\) production with the Higgs boson decaying to a pair of photons, where the top-Yukawa coupling is CP-odd, implemented using the Higgs Characterization model [Artoisenet et al., 2013](#ref-artoisinet-2013puc), the production of a pair of superpartners of the top quark (s-top) using the Minimal Supersymmetric Standard Model (MSSM) [Rosiek, 1990](#ref-rosiek-1990), [Allanach et al., 2009](#ref-allanach-2009), and flavor changing neutral current (FCNC) processes [Degrande et al., 2015](#ref-degrande-2015), [Durieux et al., 2015](#ref-durieux-2015). For the s-top process, we simulate the production of heavier s-top pairs \\(t_2\\bar{t_2}\\), where each heavier s-top (mass 582 GeV) decays into a lighter s-top \\(t_1\\) or \\(\bar{t_1}\\), mass 400 GeV) and a Higgs boson. The FCNC process involves \\(t\\bar{t}\\) production where one top quark decays to a Higgs boson and a light quark. We generate 10 million events for each process, except for \\(tHq\\) and \\(t\\bar{t}t\\bar{t}\\), where 5 million events were produced.
-In all simulation samples, the center of mass energy of the proton-proton collision is set to 13 TeV. The Higgs boson, top quarks, and vector bosons are set to decay inclusively (except the \\(t\\bar{t}H \\rightarrow \\gamma\\gamma\\) samples), with MadSpin [Artoisenet et al., 2012](#ref-artoisinet-2012st) handling the decays of top quarks and W bosons. The generated events are processed through Pythia 8.235 [Sjostrand et al., 2015](#ref-sjostrand-2015) for parton showering and heavy particle decays, followed by Delphes 3.4.2 [de Favereau et al., 2014](#ref-defavereau-2014) configured to emulate the ATLAS detector [ATLAS Collaboration, 2008](#ref-atlas-2008) for fast detector simulation.
-The detector-level object selection criteria are defined to align with typical experimental conditions. Photons are required to have transverse momentum \\(p_T \\geq 20~\\mathrm{GeV}\\) and pseudorapidity \\(|\\eta| \\leq 2.37\\), excluding the electromagnetic calorimeter crack region \\(1.37 < |\\eta| < 1.52\\). Electrons must have \\(p_T \\geq 10~\\mathrm{GeV}\\) and \\(|\\eta| \\leq 2.47\\) (excluding the same crack region), while muons are selected with \\(p_T \\geq 10~\\mathrm{GeV}\\) and \\(|\\eta| \\leq 2.7\\). Jets are reconstructed using the anti-\\(k_t\\) algorithm [Cacciari et al., 2008](#ref-cacciari-2008gp) with radius parameter \\(\\Delta R=0.4\\), where \\(\\Delta R\\) is defined as \\(\\sqrt{\\Delta\\eta ^2 + \\Delta\\phi^2}\\), with \\(\\Delta\\eta\\) being the difference in pseudorapidity and \\(\\Delta\\phi\\) the difference in azimuthal angle. Jets must satisfy \\(p_T \\geq 25~\\mathrm{GeV}\\) and \\(|\\eta| \\leq 2.5\\). To avoid double-counting, jets are removed if they are within \\(\\Delta R < 0.4\\) of a photon or lepton. The identification of jets originating from b-quark decays (b-tagging) is performed by matching jets within \\(\\Delta R = 0.4\\) of a b-quark, with efficiency corrections applied to match the performance of the ATLAS experiment's b-tagging algorithm [ATLAS Collaboration, 2019](#ref-atlas-2019bwq).
 ## Methods
@@ -45,11 +45,11 @@ We present a methodology for developing and evaluating a foundation model for pa
 Given the prevalence of classification problems in particle physics data analysis, we evaluate the model's efficacy through a systematic assessment across five binary classification tasks:
-- $t\bar{t}H(\rightarrow \gamma\gamma)$ with CP-even versus CP-odd t-H interaction
-- $t\bar{t}$ with FCNC top quark decays versus $tHq$ processes
-- $t\bar{t}W$ versus $ttt$ processes
-- Stop pair production with Higgs bosons in the decay chain versus $t\bar{t}H$ processes
-- $WH$ versus $ZH$ production modes
 Our evaluation metrics encompass classification performance, computational efficiency, and model interpretability. The investigation extends to analyzing the model's scaling behavior with respect to training dataset size, benchmarked against models trained without pretraining. Although we explored transfer learning through parameter freezing of pretrained layers, this approach did not yield performance improvements, leading us to focus our detailed analysis on fine-tuning strategies.
@@ -59,7 +59,7 @@ This methodological framework demonstrates the potential of foundation models to
 ### GNN Architecture
-We implement a Graph Neural Network (GNN) architecture that naturally accommodates the point-cloud structure of particle physics data, employing the DGL framework with a PyTorch backend [Wang et al., 2019][ref-dgl-2019], [Paszke et al., 2019][ref-pytorch-2019]. A fully connected graph is constructed for each event, with nodes corresponding to reconstructed jets, electrons, muons, photons, and $\vec{E}_T^{\text{miss}}$. The features of each node include the four-momentum $(p_T, \eta, \phi, E)$ of the object with a massless assumption ($E = p_T \cosh \eta$), the b-tagging label (for jets), the charge (for leptons), and an integer labeling the type of object represented by the node. We use a placeholder value of 0 for features which are not defined for every node type such as the b-jet tag, lepton charge, or the pseudorapidity of $\vec{E}_T^{\text{miss}}$. We assign the angular distances ($\Delta \eta, \Delta \phi, \Delta R$) as edge features and the number of nodes $N$ in the graph as a global feature. We denote the node features $\{\vec x_i\}$, edge features $\{\vec y_{ij}\}$, and global features $\{\vec z\}$.
 The GNN model is based on the graph network architecture described in [Battaglia et al., 2018][ref-graphnets-2018] using simple multilayer perceptron (MLP) feature functions and summation aggregation. The model is comprised of three primary components: an encoder, the graph network, and a decoder. In the encoder, three MLPs embed the nodes, edges, and global features into a latent space of dimension 64. The graph network block, which is designed to facilitate message passing between different domains of the graph, performs an edge update $f_e$, followed by a node update $f_n$, and finally a global update $f_g$, all defined below. The inputs to each update MLP are concatenated.
@@ -77,7 +77,7 @@ $$
 This graph block is iterated four times with the same update MLPs. Finally, the global features are passed through a decoder MLP and a final layer linear to produce the desired model outputs. Each MLP consists of 4 linear layers, each with an output width of 64, with the `ReLU` activation function. The output of the MLP is then passed through a `LayerNorm` layer [Ba et al., 2016][ref-layernorm-2016]. The total number of trainable parameters in this model is about 400,000.
-As a performance benchmark, a baseline GNN model is trained from scratch for each classification task. The initial learning rate is set to $10^{-4}$ with an exponential decay following $LR(x) = LR_{\text{initial}}\cdot(0.99)^x$, where $x$ represents the epoch number.
 ---
@@ -95,23 +95,22 @@ This approach combines both classification and regression tasks to characterize
 We develop a comprehensive set of 41 labels that capture both particle multiplicities and kinematic properties. This approach increases prediction granularity and enhances model interpretability. By training the model to predict event kinematics rather than event identification, we create a task-independent framework that can potentially generalize better to novel scenarios not seen during pretraining.
-The particle multiplicity labels count the number of Higgs bosons ($n_{\text{higgs}}$), top quarks ($n_{\text{tops}}$), vector bosons ($n_V$), $W$ bosons ($n_W$), and $Z$ bosons ($n_Z$). The kinematic labels characterize the transverse momentum ($p_T$), pseudorapidity ($\eta$), and azimuthal angle ($\phi$) of Higgs bosons and top quarks through binned classifications.
-For Higgs bosons, $p_T$ is categorized into three ranges: (0, 30) GeV, (30, 200) GeV, and (200, $\infty$) GeV, with the upper range particularly sensitive to potential BSM effects. Similarly, both leading and subleading top quarks have $p_T$ classifications spanning (0, 30) GeV, (30, 300) GeV, and (300, $\infty$) GeV. When no particle exists within a specific $p_T$ range, the corresponding label is set to $[0, 0, 0]$. For all particles, $\eta$ measurements are divided into 4 bins with boundaries at $[-1.5, 0, 1.5]$, while $\phi$ measurements use 4 bins with boundaries at $[-\frac{\pi}{2}, 0, \frac{\pi}{2}]$. As with $p_T$, both $\eta$ and $\phi$ labels default to $[0, 0, 0, 0]$ in the absence of a particle. This comprehensive labeling schema enables fine-grained learning of kinematic distributions and particle multiplicities, essential for characterizing complex collision events.
 The loss function combines individual losses from all 41 labels through weighted averaging. Binary cross-entropy is applied to classification labels, while mean-squared error is used for regression labels. The model generates predictions for all labels simultaneously, with individual losses calculated according to their respective types. The final loss is computed as an equally-weighted average across all labels, with weights set to 1 to ensure uniform contribution to the optimization process. The output layer of the multilabel model has 2,688 trainable parameters.
 #### Pretraining
-During pre-training, the initial learning rate is $10^{-4}$, and the learning rate decays by 1% each epoch following the power law function $LR(x) = 10^{-4}\cdot(0.99)^x$, where $x$ is the number of epochs. Both pre-trained models reach a plateau in loss by epoch 50, at which point the training is stopped.
 ---
 ### Fine-tuning Methodology
 For downstream tasks, we adjust the model architecture for fine-tuning by replacing the original output layer (final linear layer) with a newly initialized linear layer while retaining the pre-trained weights for all other layers. This modification allows the model to specialize in the specific downstream task while leveraging the general features learned during pretraining.
-The fine-tuning process begins with distinct learning rate setups for different parts of the model. The newly initialized linear layer is trained with an initial learning rate of $10^{-4}$, matching the rate used for models trained from scratch. Meanwhile, the pre-trained layers are fine-tuned more cautiously with a lower initial learning rate of $10^{-5}$. This approach ensures that the pre-trained layers adapt gradually without losing their general features, while the new layer learns effectively from scratch. Both learning rates decay over time following the same power law function, $LR(x) = LR_{initial} \cdot (0.99)^x$, to promote stable convergence as training progresses.
 We also evaluated a transfer learning setup in which either the decoder MLP or the final linear layer was replaced with a newly initialized component. During this process, all other model parameters remained frozen, leveraging the pre-trained features without further updating them. However, we did not observe performance improvements using the transfer learning setup. Consequently, we focus on reporting results obtained with the fine-tuning approach.
@@ -123,28 +122,28 @@ We assess model performance using two figures of merit: the classification accur
 To obtain reliable performance estimates and uncertainties, we employ an ensemble training approach where 5 independent models are trained for each configuration with random weight initialization and random subsets of the training dataset. This enables us to evaluate both the models' sensitivity to initial parameters and to quantify uncertainties in their performance.
-To investigate how model performance scales with training data, we conducted training runs using sample sizes ranging from $10^3$ to $10^7$ events per class ($10^3$, $10^4$, $10^5$, $10^6$, and $10^7$) for each model setup: the from-scratch baseline and models fine-tuned from multi-class or multi-label pretrained models. For the $10^7$ case, only the initialization was randomized due to dataset size limitations. All models were evaluated on the same testing dataset, consisting of 2 million events per class, which remained separate from the training process.
-| **Name of Task**     | **Pretraining Task** | $10^3$         | $10^4$         | $10^5$         | $10^6$         | $10^7$         |
-|----------------------|----------------------|----------------|----------------|----------------|----------------|----------------|
-| **ttH CP Even vs Odd** | Baseline Accuracy     | 56.5 $\pm$ 1.1 | 62.2 $\pm$ 0.1 | 64.3 $\pm$ 0.0 | 65.7 $\pm$ 0.0 | 66.2 $\pm$ 0.0 |
-|                        | Multiclass (%)        | +4.8 $\pm$ 1.1 | +3.4 $\pm$ 0.1 | +1.3 $\pm$ 0.0 | +0.2 $\pm$ 0.0 | -0.0 $\pm$ 0.0 |
-|                        | Multilabel (%)        | +2.1 $\pm$ 1.2 | +1.9 $\pm$ 0.1 | +0.8 $\pm$ 0.1 | +0.0 $\pm$ 0.0 | -0.1 $\pm$ 0.0 |
-| **FCNC vs tHq**        | Baseline Accuracy     | 63.6 $\pm$ 0.7 | 67.8 $\pm$ 0.4 | 68.4 $\pm$ 0.3 | 69.3 $\pm$ 0.3 | 67.9 $\pm$ 0.0 |
-|                        | Multiclass (%)        | +5.8 $\pm$ 0.8 | +1.2 $\pm$ 0.4 | +1.4 $\pm$ 0.3 | +0.5 $\pm$ 0.3 | -0.0 $\pm$ 0.0 |
-|                        | Multilabel (%)        | -5.3 $\pm$ 0.8 | -1.3 $\pm$ 0.4 | +0.9 $\pm$ 0.4 | +0.3 $\pm$ 0.3 | +0.4 $\pm$ 0.1 |
-| **ttW vs ttt**         | Baseline Accuracy     | 75.8 $\pm$ 0.1 | 77.6 $\pm$ 0.1 | 78.9 $\pm$ 0.0 | 79.8 $\pm$ 0.0 | 80.3 $\pm$ 0.0 |
-|                        | Multiclass (%)        | +3.7 $\pm$ 0.1 | +2.7 $\pm$ 0.1 | +1.3 $\pm$ 0.0 | +0.4 $\pm$ 0.0 | +0.0 $\pm$ 0.0 |
-|                        | Multilabel (%)        | +2.2 $\pm$ 0.1 | +1.1 $\pm$ 0.1 | +0.5 $\pm$ 0.0 | +0.0 $\pm$ 0.0 | -0.1 $\pm$ 0.0 |
-| **stop vs ttH**        | Baseline Accuracy     | 83.0 $\pm$ 0.2 | 86.3 $\pm$ 0.1 | 87.6 $\pm$ 0.0 | 88.5 $\pm$ 0.0 | 88.8 $\pm$ 0.0 |
-|                        | Multiclass (%)        | +0.4 $\pm$ 0.2 | +1.9 $\pm$ 0.1 | +1.0 $\pm$ 0.0 | +0.3 $\pm$ 0.0 | +0.0 $\pm$ 0.0 |
-|                        | Multilabel (%)        | +2.8 $\pm$ 0.2 | +1.0 $\pm$ 0.1 | +0.5 $\pm$ 0.0 | +0.0 $\pm$ 0.0 | -0.0 $\pm$ 0.0 |
-| **WH vs ZH**           | Baseline Accuracy     | 51.4 $\pm$ 0.1 | 53.9 $\pm$ 0.1 | 55.8 $\pm$ 0.0 | 57.5 $\pm$ 0.0 | 58.0 $\pm$ 0.0 |
-|                        | Multiclass (%)        | +5.2 $\pm$ 0.1 | +5.3 $\pm$ 0.1 | +3.1 $\pm$ 0.0 | +0.6 $\pm$ 0.0 | +0.1 $\pm$ 0.0 |
-|                        | Multilabel (%)        | -1.1 $\pm$ 0.1 | -0.9 $\pm$ 0.2 | +0.5 $\pm$ 0.1 | +0.1 $\pm$ 0.0 | -0.1 $\pm$ 0.0 |
 > **Table 1**: Accuracy of the traditional model versus the accuracy increase due to fine-tuning from various pretraining tasks.
-> The accuracies are averaged over 5 independently trained models with randomly initialized weights and trained on a random subset of the data. One exception is the $10^7$ training where all models use the same dataset due to limitations on our dataset size. The random subsets are allowed to overlap, but this overlap should be very minimal because all models take an independent random subset of $10^7$ events. The testing accuracy is calculated from the same testing set of 2 million events per class across all models for a specific training task. The errors are the propagated errors (root sum of squares) of the standard deviation of accuracies for each model.
 ## Results
@@ -152,9 +151,9 @@ To investigate how model performance scales with training data, we conducted tra
 Since the observations of AUC and accuracy show similar trends, we focus the presentation of the results using accuracy here for conciseness in Table 1.
-In general, the fine-tuned pretrained model achieves at least the same level of classification performance as the baseline model. Notably, there are significant improvements, particularly when the sample size is small, ranging from \$10^3$ to \$10^4$ events. In some cases, the accuracy improvements exceed five percentage points, demonstrating that pretrained models provide a strong initial representation that compensates for limited data. The numerical values of the improvements in accuracy may not fully capture the impact on the sensitivity of the measurements for which the neural network classifier is used, and the final sensitivity improvement is likely to be greater.
-As the training sample size grows to \$10^5$, \$10^6$, and eventually \$10^7$ events, the added benefit of pretraining diminishes. With abundant data, models trained from scratch approach or even match the accuracy of fine-tuned pretrained models. This suggests that large datasets enable effective learning from scratch, rendering the advantage of pretraining negligible in such scenarios.
 Although both pretraining approaches offer benefits, multiclass pretraining tends to provide more consistent improvements across tasks, especially in the low-data regime. In contrast, multilabel pretraining can sometimes lead to neutral or even slightly negative effects for certain tasks and data sizes. This highlights the importance of the pretraining task design, as the similarity between pretraining and fine-tuning tasks in the multiclass approach appears to yield better-aligned representations.
@@ -172,45 +171,45 @@ The similarity is evaluated using a 64-dimensional latent representation after t
 To provide an intuitive understanding of CKA values, we construct a table of the CKA scores for various transformations performed on a set of dummy data.
-- **A:** randomly initialized matrix with shape (1000, 64), following a normal distribution ($\sigma = 1, \mu=0$)
-- **B:** matrix with shape (1000, 64) constructed via various transformations performed on $A$
-- **Noise:** randomly initialized noise matrix with shape (1000, 64), following a normal distribution ($\sigma = 1, \mu=0$)
 | Dataset | CKA Score |
 |---------|-----------|
-| $A, B = A$ | 1.00 |
-| $A, B =$ permutation on columns of $A$ | 1.00 |
-| $A, B = A + \mathrm{Noise}(0.1)$ | 0.99 |
-| $A, B = A + \mathrm{Noise}(0.5)$ | 0.80 |
-| $A, B = A + \mathrm{Noise}(0.75)$ | 0.77 |
-| $A, B = A\cdot \mathrm{Noise}(1)$ (Linear Transformation) | 0.76 |
-| $A, B = A + \mathrm{Noise}(1)$ | 0.69 |
-| $A, B = A + \mathrm{Noise}(2)$ | 0.51 |
-| $A, B = A + \mathrm{Noise}(5)$ | 0.39 |
-**Table 2:** CKA scores for a dummy dataset $A$ and $B$, where $B$ is created via various transformations performed on $A$.
 As seen in Table 2 and in the definition of the CKA, the CKA score is permutation-invariant. We will use the CKA score to evaluate the similarity between various models and gain insight into the learned representation of detector events in each model (i.e., the information that each model learns).
 We train ensembles of models for each training task to observe how the CKA score changes due to the random initialization of our models. The CKA score between two models is then defined to be:
-$$
-CKA(A, B) = \frac{1}{n^2}\sum_i^n \sum_j^n CKA(A_i, B_j)
-$$
-where $A_i$ is the representation learned by the $i^{th}$ model in an ensemble with $n$ total models. The error in CKA is the standard deviation of $CKA(A_i, B_j)$.
 Here we present results for the CKA similarity between the final model in each setup with the final model in the baseline, shown in Table 3.
 | Training Task         | Baseline         | Multiclass      | Multilabel      |
-|----------------------|------------------|-----------------|-----------------|
-| ttH CP Even vs Odd   | 0.94 ± 0.05      | 0.82 ± 0.01     | 0.77 ± 0.06     |
-| FCNC vs tHq          | 0.96 ± 0.03      | 0.76 ± 0.01     | 0.81 ± 0.01     |
-| ttW vs ttt           | 0.91 ± 0.08      | 0.75 ± 0.10     | 0.72 ± 0.05     |
-| stop vs ttH          | 0.87 ± 0.11      | 0.79 ± 0.12     | 0.71 ± 0.08     |
-| WH vs ZH             | 0.90 ± 0.07      | 0.53 ± 0.03     | 0.44 ± 0.06     |
-**Table 3:** CKA Similarity of the latent representation before the decoder with the baseline model, averaged over 3 models per training setup, and all models trained with the full dataset (\$10^7$). The baseline column is not guaranteed to be 1.0 because of the random initialization of the model. Each baseline model converges to a slightly different representation as seen in the CKA values in that column.
 The baseline models with different initializations exhibit high similarity values, ranging from approximately 0.87 to 0.96, which indicates that independently trained baseline models tend to converge on similar internal representations despite random initialization. Across the considered tasks, models trained as multi-class or multi-label classifiers exhibit noticeably lower CKA similarity scores when compared to the baseline model. For example, in the WH vs ZH task, the baseline model and another baseline trained model have a high similarity of 0.90, whereas the multi-class and multi-label models show significantly reduced similarities (0.53 and 0.44, respectively). This pattern suggests that the representational spaces developed by multi-class or multi-label models differ substantially from those learned by the baseline model that was trained directly on the downstream classification task.
@@ -221,7 +220,7 @@ To estimate the computational resources required for each approach, we measured
 ![The ratio of the fine-tuning time required to achieve 99% of the baseline model's final classification accuracy to the total time spent training the baseline model.](training_time.png)
 *Fig. 1: The ratio of the fine-tuning time required to achieve 99% of the baseline model's final classification accuracy to the total time spent training the baseline model.*
-Figure 1 shows the fine-tuning time for the model pretrained with multiclass classification, relative to the time required for the baseline model, as a function of training sample size. In general, the fine-tuning time is significantly shorter than the training time required by the baseline model approach. For smaller training sets, on the order of $10^5$ events, tasks such as FCNC vs. tHq and ttW vs. ttt benefit substantially from the pretrained model’s “head start,” achieving their final performance in only about 1% of the baseline time. For large training datasets, the fine-tuning time relative to the baseline training time becomes larger; however, given that the large training sample typically requires longer training time, fine-tuning still yields much faster training convergence. The ttH CP-even vs. ttH CP-odd task, with a training sample size of $10^7$ events, is an exception where the fine-tuning time exceeds the training time required for the baseline model. This is likely because the processes involved in this task include photon objects in the final states, which are absent from the events used during pretraining.
 To accurately evaluate the total time consumption, it is necessary to include the pretraining time required for the foundation model approach. The pretraining times are as follows:
@@ -230,9 +229,9 @@ To accurately evaluate the total time consumption, it is necessary to include th
 The GPU hours recorded for the multi-label model represent the total time required when training the model in parallel on 16 GPUs. This includes a model synchronization step, which results in higher GPU hours compared to the multi-class pretraining model.
-The foundation model approach becomes increasingly efficient when a large number of tasks are fine-tuned using the same pretrained model, compared to training each task independently from scratch. To illustrate this, we evaluate the computational time required for a scenario where the training sample contains $10^7$ events. For the five tasks tested in this study, the baseline training time (training from scratch) ranges from 1.68 GPU hours (WH vs. ZH) to 5.30 GPU hours (ttW vs. ttt), with an average baseline training time of 2.94 GPU hours. In contrast, the average fine-tuning time for the foundation model approach, relative to the baseline, is 38% of the baseline training time for $10^7$ events. Based on these averages, we estimate that the foundation model approach becomes more computationally efficient than the baseline approach when fine-tuning is performed for more than 41 tasks.
-As a practical example, the ATLAS measurement of Higgs boson couplings using the $H \rightarrow \gamma\gamma$ decay channel [ATLAS Collaboration, 2023][ref-atlas-2023-higg] involved training 42 classifiers for event categorization. This coincides with our estimate, suggesting that the foundation model approach can reduce computational costs even for a single high-energy physics measurement.
 ## Conclusions

 To provide a diverse set of physics processes for the pretraining, we use Madgraph@NLO 2.7.3 [Alwall et al., 2014](#ref-alwall-2014hca) to generate proton-proton collision events at next-to-leading order (NLO) in Quantum Chromodynamics (QCD). We generate 12 distinct Standard Model (SM) physics processes, including six major Higgs boson production mechanisms: gluon fusion production \\(ggF\\), vector boson fusion \\(VBF\\), associated production of the Higgs boson with a W boson \\(WH\\) or a Z boson \\(ZH\\), associated production of the Higgs boson with a top-quark pair \\(t\bar{t}H\\), and associated production of the Higgs boson with a single top quark and a forward quark \\(tHq\\). Additionally, we simulate six top quark production processes: single top production, top-quark pair production \\(t\bar{t}\\), top quark pair production in association with a pair of photons \\(t\bar{t}\gamma\gamma\\), associated production of a top-quark pair with a W boson \\(t\bar{t}W\\), simultaneous production of three top quarks \\(t\bar{t}t\\), and simultaneous production of four top quarks \\(t\bar{t}t\bar{t}\\). In these samples, the Higgs boson and top quarks decay inclusively. These 12 Higgs and top quark production processes constitute the pretraining dataset.
+To test the pretrained model, we further generated four processes including three beyond Standard Model (SM) processes: a SM \\(t\bar{t}H\\) production where the Higgs boson decays exclusively to a pair of photons, a \\(t\bar{t}H\\) production with the Higgs boson decaying to a pair of photons, where the top-Yukawa coupling is CP-odd, implemented using the Higgs Characterization model [Artoisenet et al., 2013](#ref-artoisinet-2013puc), the production of a pair of superpartners of the top quark (s-top) using the Minimal Supersymmetric Standard Model (MSSM) [Rosiek, 1990](#ref-rosiek-1990), [Allanach et al., 2009](#ref-allanach-2009), and flavor changing neutral current (FCNC) processes [Degrande et al., 2015](#ref-degrande-2015), [Durieux et al., 2015](#ref-durieux-2015). For the s-top process, we simulate the production of heavier s-top pairs \\(t_2\bar{t_2}\\), where each heavier s-top (mass 582 GeV) decays into a lighter s-top \\(t_1\\) or \\(\bar{t_1}\\), mass 400 GeV) and a Higgs boson. The FCNC process involves \\(t\bar{t}\\) production where one top quark decays to a Higgs boson and a light quark. We generate 10 million events for each process, except for \\(tHq\\) and \\(t\bar{t}t\bar{t}\\), where 5 million events were produced.
+In all simulation samples, the center of mass energy of the proton-proton collision is set to 13 TeV. The Higgs boson, top quarks, and vector bosons are set to decay inclusively (except the \\(t\bar{t}H \rightarrow \gamma\gamma\\) samples), with MadSpin [Artoisenet et al., 2012](#ref-artoisinet-2012st) handling the decays of top quarks and W bosons. The generated events are processed through Pythia 8.235 [Sjostrand et al., 2015](#ref-sjostrand-2015) for parton showering and heavy particle decays, followed by Delphes 3.4.2 [de Favereau et al., 2014](#ref-defavereau-2014) configured to emulate the ATLAS detector [ATLAS Collaboration, 2008](#ref-atlas-2008) for fast detector simulation.
+The detector-level object selection criteria are defined to align with typical experimental conditions. Photons are required to have transverse momentum \\(p_T \geq 20~\mathrm{GeV}\\) and pseudorapidity \\(|\eta| \leq 2.37\\), excluding the electromagnetic calorimeter crack region \\(1.37 < |\eta| < 1.52\\). Electrons must have \\(p_T \geq 10~\mathrm{GeV}\\) and \\(|\eta| \leq 2.47\\) (excluding the same crack region), while muons are selected with \\(p_T \geq 10~\mathrm{GeV}\\) and \\(|\eta| \leq 2.7\\). Jets are reconstructed using the anti-\\(k_t\\) algorithm [Cacciari et al., 2008](#ref-cacciari-2008gp) with radius parameter \\(\Delta R=0.4\\), where \\(\Delta R\\) is defined as \\(\sqrt{\Delta\eta ^2 + \Delta\phi^2}\\), with \\(\Delta\eta\\) being the difference in pseudorapidity and \\(\Delta\phi\\) the difference in azimuthal angle. Jets must satisfy \\(p_T \geq 25~\mathrm{GeV}\\) and \\(|\eta| \leq 2.5\\). To avoid double-counting, jets are removed if they are within \\(\Delta R < 0.4\\) of a photon or lepton. The identification of jets originating from b-quark decays (b-tagging) is performed by matching jets within \\(\Delta R = 0.4\\) of a b-quark, with efficiency corrections applied to match the performance of the ATLAS experiment's b-tagging algorithm [ATLAS Collaboration, 2019](#ref-atlas-2019bwq).
 ## Methods
 Given the prevalence of classification problems in particle physics data analysis, we evaluate the model's efficacy through a systematic assessment across five binary classification tasks:
+- \\(t\bar{t}H(\rightarrow \gamma\gamma)\\) with CP-even versus CP-odd t-H interaction
+- \\(t\bar{t}\\) with FCNC top quark decays versus $tHq$ processes
+- \\(t\bar{t}W\\) versus $ttt$ processes
+- Stop pair production with Higgs bosons in the decay chain versus \\(t\bar{t}H\\) processes
+- \\(WH\\) versus \\(ZH\\) production modes
 Our evaluation metrics encompass classification performance, computational efficiency, and model interpretability. The investigation extends to analyzing the model's scaling behavior with respect to training dataset size, benchmarked against models trained without pretraining. Although we explored transfer learning through parameter freezing of pretrained layers, this approach did not yield performance improvements, leading us to focus our detailed analysis on fine-tuning strategies.
 ### GNN Architecture
+We implement a Graph Neural Network (GNN) architecture that naturally accommodates the point-cloud structure of particle physics data, employing the DGL framework with a PyTorch backend [Wang et al., 2019][ref-dgl-2019], [Paszke et al., 2019][ref-pytorch-2019]. A fully connected graph is constructed for each event, with nodes corresponding to reconstructed jets, electrons, muons, photons, and \\(\vec{E}_T^{\text{miss}}\\). The features of each node include the four-momentum \\((p_T, \eta, \phi, E)\\) of the object with a massless assumption (\\(E = p_T \cosh \eta\\)), the b-tagging label (for jets), the charge (for leptons), and an integer labeling the type of object represented by the node. We use a placeholder value of 0 for features which are not defined for every node type such as the b-jet tag, lepton charge, or the pseudorapidity of \\(\vec{E}_T^{\text{miss}}\\). We assign the angular distances (\\(\Delta \eta, \Delta \phi, \Delta R\\)) as edge features and the number of nodes $N$ in the graph as a global feature. We denote the node features \\(\{\vec x_i\}\\), edge features \\(\{\vec y_{ij}\}\\), and global features \\(\{\vec z\}\\).
 The GNN model is based on the graph network architecture described in [Battaglia et al., 2018][ref-graphnets-2018] using simple multilayer perceptron (MLP) feature functions and summation aggregation. The model is comprised of three primary components: an encoder, the graph network, and a decoder. In the encoder, three MLPs embed the nodes, edges, and global features into a latent space of dimension 64. The graph network block, which is designed to facilitate message passing between different domains of the graph, performs an edge update $f_e$, followed by a node update $f_n$, and finally a global update $f_g$, all defined below. The inputs to each update MLP are concatenated.
 This graph block is iterated four times with the same update MLPs. Finally, the global features are passed through a decoder MLP and a final layer linear to produce the desired model outputs. Each MLP consists of 4 linear layers, each with an output width of 64, with the `ReLU` activation function. The output of the MLP is then passed through a `LayerNorm` layer [Ba et al., 2016][ref-layernorm-2016]. The total number of trainable parameters in this model is about 400,000.
+As a performance benchmark, a baseline GNN model is trained from scratch for each classification task. The initial learning rate is set to \\(10^{-4}\\) with an exponential decay following \\(LR(x) = LR_{\text{initial}}\cdot(0.99)^x\\), where \\(x\\) represents the epoch number.
 ---
 We develop a comprehensive set of 41 labels that capture both particle multiplicities and kinematic properties. This approach increases prediction granularity and enhances model interpretability. By training the model to predict event kinematics rather than event identification, we create a task-independent framework that can potentially generalize better to novel scenarios not seen during pretraining.
+The particle multiplicity labels count the number of Higgs bosons (\\(n_{\text{higgs}}\\)), top quarks (\\(n_{\text{tops}}\\)), vector bosons (\\(n_V\\)), \\(W\\) bosons (\\(n_W\\)), and \\(Z\\) bosons (\\(n_Z\\)). The kinematic labels characterize the transverse momentum (\\(p_T\\)), pseudorapidity (\\(\eta\\)), and azimuthal angle (\\(\phi\\)) of Higgs bosons and top quarks through binned classifications.
+For Higgs bosons, $p_T$ is categorized into three ranges: (0, 30) GeV, (30, 200) GeV, and (200, \\(\infty\\)) GeV, with the upper range particularly sensitive to potential BSM effects. Similarly, both leading and subleading top quarks have $p_T$ classifications spanning (0, 30) GeV, (30, 300) GeV, and (300, \\(\infty\\)) GeV. When no particle exists within a specific \\(p_T\\) range, the corresponding label is set to \\([0, 0, 0]\\). For all particles, \\(\eta\\) measurements are divided into 4 bins with boundaries at \\([-1.5, 0, 1.5]\\), while \\(\phi\\) measurements use 4 bins with boundaries at \\([-\frac{\pi}{2}, 0, \frac{\pi}{2}]\\). As with \\(p_T\\), both \\(\eta\\) and \\(\phi\\) labels default to \\([0, 0, 0, 0]\\) in the absence of a particle. This comprehensive labeling schema enables fine-grained learning of kinematic distributions and particle multiplicities, essential for characterizing complex collision events.
 The loss function combines individual losses from all 41 labels through weighted averaging. Binary cross-entropy is applied to classification labels, while mean-squared error is used for regression labels. The model generates predictions for all labels simultaneously, with individual losses calculated according to their respective types. The final loss is computed as an equally-weighted average across all labels, with weights set to 1 to ensure uniform contribution to the optimization process. The output layer of the multilabel model has 2,688 trainable parameters.
 #### Pretraining
+During pre-training, the initial learning rate is \\(10^{-4}\\), and the learning rate decays by 1% each epoch following the power law function \\(LR(x) = 10^{-4}\cdot(0.99)^x\\), where \\(x\\) is the number of epochs. Both pre-trained models reach a plateau in loss by epoch 50, at which point the training is stopped.
 ---
 ### Fine-tuning Methodology
 For downstream tasks, we adjust the model architecture for fine-tuning by replacing the original output layer (final linear layer) with a newly initialized linear layer while retaining the pre-trained weights for all other layers. This modification allows the model to specialize in the specific downstream task while leveraging the general features learned during pretraining.
+The fine-tuning process begins with distinct learning rate setups for different parts of the model. The newly initialized linear layer is trained with an initial learning rate of \\(10^{-4}\\), matching the rate used for models trained from scratch. Meanwhile, the pre-trained layers are fine-tuned more cautiously with a lower initial learning rate of \\(10^{-5}\\). This approach ensures that the pre-trained layers adapt gradually without losing their general features, while the new layer learns effectively from scratch. Both learning rates decay over time following the same power law function, \\(LR(x) = LR_{initial} \cdot (0.99)^x\\), to promote stable convergence as training progresses.
 We also evaluated a transfer learning setup in which either the decoder MLP or the final linear layer was replaced with a newly initialized component. During this process, all other model parameters remained frozen, leveraging the pre-trained features without further updating them. However, we did not observe performance improvements using the transfer learning setup. Consequently, we focus on reporting results obtained with the fine-tuning approach.
 To obtain reliable performance estimates and uncertainties, we employ an ensemble training approach where 5 independent models are trained for each configuration with random weight initialization and random subsets of the training dataset. This enables us to evaluate both the models' sensitivity to initial parameters and to quantify uncertainties in their performance.
+To investigate how model performance scales with training data, we conducted training runs using sample sizes ranging from \\(10^3\\) to \\(10^7\\) events per class (\\(10^3\\), \\(10^4\\), \\(10^5\\), \\(10^6\\), and \\(10^7\\)) for each model setup: the from-scratch baseline and models fine-tuned from multi-class or multi-label pretrained models. For the \\(10^7\\) case, only the initialization was randomized due to dataset size limitations. All models were evaluated on the same testing dataset, consisting of 2 million events per class, which remained separate from the training process.
+| **Name of Task**     | **Pretraining Task** | \\(10^3\\)         | \\(10^4\\)         | \\(10^5\\)         | \\(10^6\\)         | \\(10^7\\)         |
+|----------------------|----------------------|--------------------|--------------------|--------------------|--------------------|--------------------|
+| **ttH CP Even vs Odd** | Baseline Accuracy     | 56.5 ± 1.1 | 62.2 ± 0.1 | 64.3 ± 0.0 | 65.7 ± 0.0 | 66.2 ± 0.0 |
+|                        | Multiclass (%)        | +4.8 ± 1.1 | +3.4 ± 0.1 | +1.3 ± 0.0 | +0.2 ± 0.0 | −0.0 ± 0.0 |
+|                        | Multilabel (%)        | +2.1 ± 1.2 | +1.9 ± 0.1 | +0.8 ± 0.1 | +0.0 ± 0.0 | −0.1 ± 0.0 |
+| **FCNC vs tHq**        | Baseline Accuracy     | 63.6 ± 0.7 | 67.8 ± 0.4 | 68.4 ± 0.3 | 69.3 ± 0.3 | 67.9 ± 0.0 |
+|                        | Multiclass (%)        | +5.8 ± 0.8 | +1.2 ± 0.4 | +1.4 ± 0.3 | +0.5 ± 0.3 | −0.0 ± 0.0 |
+|                        | Multilabel (%)        | −5.3 ± 0.8 | −1.3 ± 0.4 | +0.9 ± 0.4 | +0.3 ± 0.3 | +0.4 ± 0.1 |
+| **ttW vs ttt**         | Baseline Accuracy     | 75.8 ± 0.1 | 77.6 ± 0.1 | 78.9 ± 0.0 | 79.8 ± 0.0 | 80.3 ± 0.0 |
+|                        | Multiclass (%)        | +3.7 ± 0.1 | +2.7 ± 0.1 | +1.3 ± 0.0 | +0.4 ± 0.0 | +0.0 ± 0.0 |
+|                        | Multilabel (%)        | +2.2 ± 0.1 | +1.1 ± 0.1 | +0.5 ± 0.0 | +0.0 ± 0.0 | −0.1 ± 0.0 |
+| **stop vs ttH**        | Baseline Accuracy     | 83.0 ± 0.2 | 86.3 ± 0.1 | 87.6 ± 0.0 | 88.5 ± 0.0 | 88.8 ± 0.0 |
+|                        | Multiclass (%)        | +0.4 ± 0.2 | +1.9 ± 0.1 | +1.0 ± 0.0 | +0.3 ± 0.0 | +0.0 ± 0.0 |
+|                        | Multilabel (%)        | +2.8 ± 0.2 | +1.0 ± 0.1 | +0.5 ± 0.0 | +0.0 ± 0.0 | −0.0 ± 0.0 |
+| **WH vs ZH**           | Baseline Accuracy     | 51.4 ± 0.1 | 53.9 ± 0.1 | 55.8 ± 0.0 | 57.5 ± 0.0 | 58.0 ± 0.0 |
+|                        | Multiclass (%)        | +5.2 ± 0.1 | +5.3 ± 0.1 | +3.1 ± 0.0 | +0.6 ± 0.0 | +0.1 ± 0.0 |
+|                        | Multilabel (%)        | −1.1 ± 0.1 | −0.9 ± 0.2 | +0.5 ± 0.1 | +0.1 ± 0.0 | −0.1 ± 0.0 |
 > **Table 1**: Accuracy of the traditional model versus the accuracy increase due to fine-tuning from various pretraining tasks.
+> The accuracies are averaged over 5 independently trained models with randomly initialized weights and trained on a random subset of the data. One exception is the \\(10^7\\) training where all models use the same dataset due to limitations on our dataset size. The random subsets are allowed to overlap, but this overlap should be very minimal because all models take an independent random subset of \\(10^7\\) events. The testing accuracy is calculated from the same testing set of 2 million events per class across all models for a specific training task. The errors are the propagated errors (root sum of squares) of the standard deviation of accuracies for each model.
 ## Results
 Since the observations of AUC and accuracy show similar trends, we focus the presentation of the results using accuracy here for conciseness in Table 1.
+In general, the fine-tuned pretrained model achieves at least the same level of classification performance as the baseline model. Notably, there are significant improvements, particularly when the sample size is small, ranging from \\(10^3\\) to \\(10^4\\) events. In some cases, the accuracy improvements exceed five percentage points, demonstrating that pretrained models provide a strong initial representation that compensates for limited data. The numerical values of the improvements in accuracy may not fully capture the impact on the sensitivity of the measurements for which the neural network classifier is used, and the final sensitivity improvement is likely to be greater.
+As the training sample size grows to \\(10^5\\), \\(10^6\\), and eventually \\(10^7\\) events, the added benefit of pretraining diminishes. With abundant data, models trained from scratch approach or even match the accuracy of fine-tuned pretrained models. This suggests that large datasets enable effective learning from scratch, rendering the advantage of pretraining negligible in such scenarios.
 Although both pretraining approaches offer benefits, multiclass pretraining tends to provide more consistent improvements across tasks, especially in the low-data regime. In contrast, multilabel pretraining can sometimes lead to neutral or even slightly negative effects for certain tasks and data sizes. This highlights the importance of the pretraining task design, as the similarity between pretraining and fine-tuning tasks in the multiclass approach appears to yield better-aligned representations.
 To provide an intuitive understanding of CKA values, we construct a table of the CKA scores for various transformations performed on a set of dummy data.
+- **A:** randomly initialized matrix with shape (1000, 64), following a normal distribution (\\(\sigma = 1, \mu = 0\\))
+- **B:** matrix with shape (1000, 64) constructed via various transformations performed on \\(A\\)
+- **Noise:** randomly initialized noise matrix with shape (1000, 64), following a normal distribution (\\(\sigma = 1, \mu = 0\\))
 | Dataset | CKA Score |
 |---------|-----------|
+| \\(A, B = A\\) | 1.00 |
+| \\(A, B =\\) permutation on columns of \\(A\\) | 1.00 |
+| \\(A, B = A + \mathrm{Noise}(0.1)\\) | 0.99 |
+| \\(A, B = A + \mathrm{Noise}(0.5)\\) | 0.80 |
+| \\(A, B = A + \mathrm{Noise}(0.75)\\) | 0.77 |
+| \\(A, B = A \cdot \mathrm{Noise}(1)\\) (Linear Transformation) | 0.76 |
+| \\(A, B = A + \mathrm{Noise}(1)\\) | 0.69 |
+| \\(A, B = A + \mathrm{Noise}(2)\\) | 0.51 |
+| \\(A, B = A + \mathrm{Noise}(5)\\) | 0.39 |
+**Table 2:** CKA scores for a dummy dataset \\(A\\) and \\(B\\), where \\(B\\) is created via various transformations performed on \\(A\\).
 As seen in Table 2 and in the definition of the CKA, the CKA score is permutation-invariant. We will use the CKA score to evaluate the similarity between various models and gain insight into the learned representation of detector events in each model (i.e., the information that each model learns).
 We train ensembles of models for each training task to observe how the CKA score changes due to the random initialization of our models. The CKA score between two models is then defined to be:
+\\[
+CKA(A, B) = \frac{1}{n^2} \sum_i^n \sum_j^n CKA(A_i, B_j)
+\\]
+where \\(A_i\\) is the representation learned by the \\(i^{\text{th}}\\) model in an ensemble with \\(n\\) total models. The error in CKA is the standard deviation of \\(CKA(A_i, B_j)\\).
 Here we present results for the CKA similarity between the final model in each setup with the final model in the baseline, shown in Table 3.
 | Training Task         | Baseline         | Multiclass      | Multilabel      |
+|-----------------------|------------------|-----------------|-----------------|
+| ttH CP Even vs Odd    | 0.94 ± 0.05      | 0.82 ± 0.01     | 0.77 ± 0.06     |
+| FCNC vs tHq           | 0.96 ± 0.03      | 0.76 ± 0.01     | 0.81 ± 0.01     |
+| ttW vs ttt            | 0.91 ± 0.08      | 0.75 ± 0.10     | 0.72 ± 0.05     |
+| stop vs ttH           | 0.87 ± 0.11      | 0.79 ± 0.12     | 0.71 ± 0.08     |
+| WH vs ZH              | 0.90 ± 0.07      | 0.53 ± 0.03     | 0.44 ± 0.06     |
+**Table 3:** CKA Similarity of the latent representation before the decoder with the baseline model, averaged over 3 models per training setup, and all models trained with the full dataset (\\(10^7\\)). The baseline column is not guaranteed to be 1.0 because of the random initialization of the model. Each baseline model converges to a slightly different representation as seen in the CKA values in that column.
 The baseline models with different initializations exhibit high similarity values, ranging from approximately 0.87 to 0.96, which indicates that independently trained baseline models tend to converge on similar internal representations despite random initialization. Across the considered tasks, models trained as multi-class or multi-label classifiers exhibit noticeably lower CKA similarity scores when compared to the baseline model. For example, in the WH vs ZH task, the baseline model and another baseline trained model have a high similarity of 0.90, whereas the multi-class and multi-label models show significantly reduced similarities (0.53 and 0.44, respectively). This pattern suggests that the representational spaces developed by multi-class or multi-label models differ substantially from those learned by the baseline model that was trained directly on the downstream classification task.
 ![The ratio of the fine-tuning time required to achieve 99% of the baseline model's final classification accuracy to the total time spent training the baseline model.](training_time.png)
 *Fig. 1: The ratio of the fine-tuning time required to achieve 99% of the baseline model's final classification accuracy to the total time spent training the baseline model.*
+Figure 1 shows the fine-tuning time for the model pretrained with multiclass classification, relative to the time required for the baseline model, as a function of training sample size. In general, the fine-tuning time is significantly shorter than the training time required by the baseline model approach. For smaller training sets, on the order of \\(10^5\\) events, tasks such as FCNC vs. tHq and ttW vs. ttt benefit substantially from the pretrained model’s “head start,” achieving their final performance in only about 1% of the baseline time. For large training datasets, the fine-tuning time relative to the baseline training time becomes larger; however, given that the large training sample typically requires longer training time, fine-tuning still yields much faster training convergence. The ttH CP-even vs. ttH CP-odd task, with a training sample size of \\(10^7\\) events, is an exception where the fine-tuning time exceeds the training time required for the baseline model. This is likely because the processes involved in this task include photon objects in the final states, which are absent from the events used during pretraining.
 To accurately evaluate the total time consumption, it is necessary to include the pretraining time required for the foundation model approach. The pretraining times are as follows:
 The GPU hours recorded for the multi-label model represent the total time required when training the model in parallel on 16 GPUs. This includes a model synchronization step, which results in higher GPU hours compared to the multi-class pretraining model.
+The foundation model approach becomes increasingly efficient when a large number of tasks are fine-tuned using the same pretrained model, compared to training each task independently from scratch. To illustrate this, we evaluate the computational time required for a scenario where the training sample contains \\(10^7\\) events. For the five tasks tested in this study, the baseline training time (training from scratch) ranges from 1.68 GPU hours (WH vs. ZH) to 5.30 GPU hours (ttW vs. ttt), with an average baseline training time of 2.94 GPU hours. In contrast, the average fine-tuning time for the foundation model approach, relative to the baseline, is 38% of the baseline training time for \\(10^7\\) events. Based on these averages, we estimate that the foundation model approach becomes more computationally efficient than the baseline approach when fine-tuning is performed for more than 41 tasks.
+As a practical example, the ATLAS measurement of Higgs boson couplings using the \\(H \rightarrow \gamma\gamma\\) decay channel [ATLAS Collaboration, 2023][ref-atlas-2023-higg] involved training 42 classifiers for event categorization. This coincides with our estimate, suggesting that the foundation model approach can reduce computational costs even for a single high-energy physics measurement.
 ## Conclusions

root_gnn_dgl/README.md CHANGED Viewed

@@ -1,15 +1,19 @@
 # root_gnn_dgl
-Pretrained DGL-based ROOT graph neural network.
 ## Overview
 - Stable release with pretrained model weights.
-Pretrained model location: ``
 ## Conda setup
 ```bash
 cd setup
 conda env create -f environment.yml
@@ -23,6 +27,8 @@ python setup/test_setup.py
 - NERSC Perlmutter environment with `podman-hpc` available.
 - Access to `joshuaho/pytorch:1.0` on Docker Hub [https://hub.docker.com/r/joshuaho/pytorch](https://hub.docker.com/r/joshuaho/pytorch)
 ### Pull the Prebuilt Image
 ```bash

 # root_gnn_dgl
+Pretrained DGL-based ROOT graph neural network.
+Pretrained model location: `/global/cfs/projectdirs/atlas/joshua/Pretrained_GNN/multiclass_pretrained_model_12/`
+To use the pretrained model, take a look at a finetuning config in `configs`.
+Replace `pretraining_path:` with `/global/cfs/projectdirs/atlas/joshua/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt`.
 ## Overview
 - Stable release with pretrained model weights.
 ## Conda setup
+The conda environment is required for the inference step: applying the GNN onto root files and saving GNN scores as an additional branch. This is because the infereces script uses pyROOT.
 ```bash
 cd setup
 conda env create -f environment.yml
 - NERSC Perlmutter environment with `podman-hpc` available.
 - Access to `joshuaho/pytorch:1.0` on Docker Hub [https://hub.docker.com/r/joshuaho/pytorch](https://hub.docker.com/r/joshuaho/pytorch)
+The inference step requires the conda environment, since the container does not contain ROOT.
 ### Pull the Prebuilt Image
 ```bash

root_gnn_dgl/scripts/inference.py CHANGED Viewed

@@ -1,6 +1,5 @@
 import sys
-import os
-file_path = os.getcwd()
 sys.path.append(file_path)
 import os
 import argparse
@@ -135,6 +134,7 @@ def main():
     import time
     start  = time.time()
     import torch
     from array import array
     import numpy as np
@@ -186,6 +186,7 @@ def main():
     lend = time.time()
     print('Loader finished in {:.2f} seconds'.format(lend - lstart))
     sample_graph, _, _, global_sample = loader[0]
     print('dset length =', len(dset))
     print('loader length =', len(loader))
@@ -197,6 +198,7 @@ def main():
         for config_file, branch in zip(args.config, args.branch_name):
             config = load_config(config_file)
             model = utils.buildFromConfig(config['Model'], {'sample_graph' : sample_graph, 'sample_global': global_sample}).to(device)
             if args.ckpt < 0:
                 ep, checkpoint = utils.get_best_epoch(config, var=args.var, mode='max', device=device)
             else:
@@ -246,41 +248,55 @@ def main():
             all_labels[branch] = labels
             all_tracking[branch] = tracking_info
     if args.write:
-        import uproot
-        import awkward as ak
-        # Open the original ROOT file and get the tree
-        infile = uproot.open(args.target)
-        tree = infile[dset_config['args']['tree_name']]
-        # Read the original tree as an awkward array
-        original_data = tree.arrays(library="ak")
-        # Prepare new branches as dicts of arrays
-        new_branches = {}
-        n_entries = len(original_data)
-        for branch, scores in all_scores.items():
-            # Ensure the scores array is the right length
-            scores = np.asarray(scores)
-            if scores.shape[0] != n_entries:
-                raise ValueError(f"Branch '{branch}' has {scores.shape[0]} entries, but tree has {n_entries}")
-            new_branches[branch] = scores
-        # Merge all arrays (original + new branches)
-        # Convert awkward to dict of numpy arrays for uproot
-        out_dict = {k: np.asarray(v) for k, v in ak.to_numpy(original_data).items()}
-        out_dict.update(new_branches)
-        # Write to new ROOT file
         os.makedirs(os.path.split(args.destination)[0], exist_ok=True)
-        with uproot.recreate(args.destination) as outfile:
-            outfile.mktree(dset_config['args']['tree_name'], {k: v.dtype for k, v in out_dict.items()})
-            outfile[dset_config['args']['tree_name']].extend(out_dict)
-        print(f"Wrote new ROOT file {args.destination} with new branches {list(new_branches.keys())}")
     else:
         os.makedirs(os.path.split(args.destination)[0], exist_ok=True)
         np.savez(args.destination, scores=all_scores, labels=all_labels, tracking_info=all_tracking)

 import sys
+file_path = "/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl"
 sys.path.append(file_path)
 import os
 import argparse
     import time
     start  = time.time()
+    import ROOT
     import torch
     from array import array
     import numpy as np
     lend = time.time()
     print('Loader finished in {:.2f} seconds'.format(lend - lstart))
     sample_graph, _, _, global_sample = loader[0]
+    global_sample = []
     print('dset length =', len(dset))
     print('loader length =', len(loader))
         for config_file, branch in zip(args.config, args.branch_name):
             config = load_config(config_file)
             model = utils.buildFromConfig(config['Model'], {'sample_graph' : sample_graph, 'sample_global': global_sample}).to(device)
             if args.ckpt < 0:
                 ep, checkpoint = utils.get_best_epoch(config, var=args.var, mode='max', device=device)
             else:
             all_labels[branch] = labels
             all_tracking[branch] = tracking_info
     if args.write:
+        from ROOT import std
+        # Open the original ROOT file
+        infile = ROOT.TFile.Open(args.target)
+        tree = infile.Get(dset_config['args']['tree_name'])
+        # Create the destination directory if it doesn't exist
         os.makedirs(os.path.split(args.destination)[0], exist_ok=True)
+        # Create a new ROOT file to write the modified tree
+        outfile = ROOT.TFile.Open(args.destination, 'RECREATE')
+        # Clone the original tree structure
+        outtree = tree.CloneTree(0)
+        # Create branches for all scores
+        branch_vectors = {}
+        for branch, scores in all_scores.items():
+            if isinstance(scores[0], (list, tuple, np.ndarray)) and len(scores[0]) > 1:
+                # Create a new branch for vectors
+                branch_vectors[branch] = std.vector('float')()
+                outtree.Branch(branch, branch_vectors[branch])
+            else:
+                # Create a new branch for single floats
+                branch_vectors[branch] = array('f', [0])
+                outtree.Branch(branch, branch_vectors[branch], f'{branch}/F')
+        # Fill the tree
+        for i in range(tree.GetEntries()):
+            tree.GetEntry(i)
+            for branch, scores in all_scores.items():
+                branch_data = branch_vectors[branch]
+                if isinstance(branch_data, array):  # Check if it's a single float array
+                    branch_data[0] = float(scores[i])
+                else:  # Assume it's a std::vector<float>
+                    branch_data.clear()
+                    for value in scores[i]:
+                        branch_data.push_back(float(value))
+            outtree.Fill()
+        # Write the modified tree to the new file
+        print(f'Writing to file {args.destination}')
+        print(f'Input entries: {tree.GetEntries()}, Output entries: {outtree.GetEntries()}')
+        print(f'Wrote scores to {args.branch_name}')
+        outtree.Write()
+        outfile.Close()
+        infile.Close()
     else:
         os.makedirs(os.path.split(args.destination)[0], exist_ok=True)
         np.savez(args.destination, scores=all_scores, labels=all_labels, tracking_info=all_tracking)

root_gnn_dgl/setup/Dockerfile CHANGED Viewed

@@ -7,19 +7,16 @@ LABEL maintainer.email="ho22joshua@berkeley.edu"
 ENV LANG=C.UTF-8
-# Install system dependencies: vim, OpenMPI, and build tools
 RUN apt-get update -qq \
  && apt-get install -y --no-install-recommends \
-    wget lsb-release gnupg software-properties-common \
     vim \
     g++-11 gcc-11 libstdc++-11-dev \
     openmpi-bin openmpi-common libopenmpi-dev \
  && rm -rf /var/lib/apt/lists/*
-# Install Python packages: mpi4py and jupyter
 RUN pip install --no-cache-dir mpi4py jupyter uproot
-# (Optional) Expose Jupyter port
-EXPOSE 8888

 ENV LANG=C.UTF-8
+# System deps (with CA certs for HTTPS downloads)
 RUN apt-get update -qq \
  && apt-get install -y --no-install-recommends \
+    wget curl ca-certificates lsb-release gnupg software-properties-common \
     vim \
     g++-11 gcc-11 libstdc++-11-dev \
     openmpi-bin openmpi-common libopenmpi-dev \
  && rm -rf /var/lib/apt/lists/*
+# Python packages
 RUN pip install --no-cache-dir mpi4py jupyter uproot
+EXPOSE 8888