rewriting codebase

.github/workflows/ci.yml

@@ -0,0 +1,75 @@
+name: CI
+
+on:
+  push:
+  pull_request:
+
+concurrency:
+  group: ci-${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  lint:
+    name: lint
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev
+      - run: uv run ruff check .
+      - run: uv run ruff format --check .
+
+  test:
+    name: test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev
+      - run: uv run pytest
+
+  package:
+    name: package
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev
+      - run: uv build
+      - run: uv run python -m twine check dist/*
+      - name: Wheel install smoke test
+        shell: bash
+        run: |
+          set -euo pipefail
+          smoke_dir="$(mktemp -d)"
+          trap 'rm -rf "$smoke_dir"' EXIT
+          uv venv "$smoke_dir/venv"
+          uv pip install --python "$smoke_dir/venv/bin/python" dist/*.whl
+          (
+            cd "$smoke_dir"
+            "$smoke_dir/venv/bin/python" -c 'import gnn4colliders; print(gnn4colliders.__version__)'
+            "$smoke_dir/venv/bin/gnn4colliders" --help
+            "$smoke_dir/venv/bin/gnn4colliders" train --help
+            "$smoke_dir/venv/bin/gnn4colliders" export --help
+          )
+
+  onnx:
+    name: onnx
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v6
+        with:
+          version: "0.8.x"
+          enable-cache: true
+      - run: uv sync --dev --extra root-gnn --extra onnx
+      - run: uv run pytest tests/unit/export -v

.gitignore

@@ -4,3 +4,29 @@ scores/
 slurm/
 .onnx
 .png
+
+# Python tooling
+*.py[cod]
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+*.egg-info/
+.venv/
+venv/
+htmlcov/
+.coverage
+profiles/
+*.trace.json
+dist/
+build/
+benchmark-results/
+validation_output/
+
+# Local data and generated outputs
+data/raw/*
+!data/raw/.gitkeep
+data/processed/
+outputs/
+checkpoints/
+
+# Keep the legacy implementation immutable by convention; do not add generated files there.

.python-version

@@ -0,0 +1 @@
+3.12

AGENTS.md

@@ -0,0 +1,477 @@
+# GNN4Colliders Agent Instructions
+
+## Project purpose
+
+GNN4Colliders is a collider machine-learning package intended to support multiple model architectures over shared collider-physics data infrastructure.
+
+The first model family being rewritten is:
+
+```text
+root_gnn
+```
+
+Future model families may include architectures such as:
+
+```text
+root_transformer
+```
+
+Do not design shared infrastructure around assumptions that only apply to GNNs.
+
+---
+
+## Source of truth
+
+Before making non-trivial changes, inspect:
+
+```text
+docs/architecture.md
+docs/migration.md
+```
+
+and the relevant existing source and tests.
+
+The `legacy/` tree is a behavioral reference for the rewrite.
+
+Unless explicitly instructed otherwise:
+
+* do not modify legacy code
+* do not reorganize legacy code
+* do not mechanically copy legacy architecture into the new package
+
+When legacy behavior and documentation disagree, identify the discrepancy rather than silently choosing one.
+
+---
+
+## Package architecture
+
+Production Python code belongs under:
+
+```text
+src/gnn4colliders/
+```
+
+The intended boundaries are:
+
+```text
+data/
+    Data access, datasets, caching, batching, folds, and generic
+    representation-independent data infrastructure.
+
+features/
+    Collider-domain feature transformations, selections, scaling,
+    and derived physics quantities.
+
+graphs/
+    Graph representation, topology, edge construction, and other
+    graph-specific transformations.
+
+models/
+    Architecture-specific neural network implementations.
+
+models/root_gnn/
+    ROOT-GNN-specific model components.
+
+training/
+    Training lifecycle, losses, metrics, checkpointing,
+    reproducibility, and distributed-training utilities where
+    architecture-independent.
+
+inference/
+    Prediction, evaluation, output writing, and model export.
+
+cli/
+    Thin command-line entry points only.
+```
+
+Do not place ROOT file-reading logic inside a specific model family unless it is genuinely architecture-specific.
+
+Do not place reusable physics-feature logic inside `root_gnn`.
+
+Do not place core implementation in notebooks or shell scripts.
+
+---
+
+## Shared infrastructure vs model-specific code
+
+Use this decision rule:
+
+**Shared collider or experiment behavior → shared package module.**
+
+Examples:
+
+```text
+ROOT reading
+event selections
+physics features
+dataset splits
+metrics
+checkpoint orchestration
+```
+
+**Representation-specific behavior → representation module.**
+
+Examples:
+
+```text
+graph topology
+edge construction
+sequence/token construction
+```
+
+**Architecture-specific behavior → models/<model_family>/.**
+
+Examples:
+
+```text
+message-passing network
+transformer encoder
+architecture-specific layers
+```
+
+The package should allow future architectures to reuse the same data and physics infrastructure where practical.
+
+---
+
+## Configuration
+
+Experiment configuration should describe intent rather than expose implementation details.
+
+Prefer:
+
+```yaml
+model:
+  type: root_gnn
+```
+
+over configuration that directly names Python module paths and class names.
+
+Configuration should eventually support composition across concerns such as:
+
+```text
+data
+model
+task
+trainer
+environment
+```
+
+Do not hardcode site-specific filesystem paths, CUDA settings, Slurm settings, or machine configuration into model or task definitions.
+
+Environment-specific configuration belongs in an environment configuration layer.
+
+---
+
+## CLI design
+
+The intended user-facing interface is a single project CLI with subcommands conceptually similar to:
+
+```bash
+gnn4colliders prepare
+gnn4colliders train
+gnn4colliders evaluate
+gnn4colliders predict
+gnn4colliders export
+```
+
+CLI modules should be thin.
+
+They may:
+
+* parse/compose configuration
+* construct application objects
+* invoke library functions
+* handle user-facing errors
+
+They should not contain core training, data-processing, graph-building, or model logic.
+
+---
+
+## Rewrite strategy
+
+This repository is being rewritten incrementally.
+
+Do not attempt to rewrite the entire legacy repository in one task.
+
+For substantial migrations:
+
+1. inspect the relevant legacy implementation
+2. identify externally observable behavior
+3. inspect existing characterization/parity tests
+4. state or infer the intended new interface
+5. implement the smallest coherent unit
+6. add or update tests
+7. run relevant validation
+8. report intentional differences and unresolved ambiguity
+
+Prefer vertical, testable migration steps over large speculative refactors.
+
+---
+
+## Behavioral parity
+
+The legacy implementation defines important behavior that may need to remain compatible during migration.
+
+Important compatibility areas include, where applicable:
+
+* input feature definitions
+* feature ordering
+* tensor shapes
+* tensor dtypes
+* graph topology
+* edge feature definitions
+* label semantics
+* event weights
+* fold semantics
+* loss calculations
+* metric definitions
+* model outputs
+* checkpoint compatibility
+* inference outputs
+
+Do not alter behavior merely because the legacy implementation appears unusual.
+
+If behavior seems incorrect or ambiguous:
+
+1. document it
+2. characterize it with a test when possible
+3. separate compatibility from proposed improvement
+
+Improvements can be made deliberately after the behavior is understood.
+
+---
+
+## Testing
+
+Tests belong under:
+
+```text
+tests/unit/
+tests/integration/
+tests/parity/
+```
+
+### Unit tests
+
+Use unit tests for isolated transformations and components.
+
+They should be:
+
+* fast
+* deterministic
+* focused
+* independent of large external datasets
+
+### Integration tests
+
+Use integration tests for small end-to-end workflows such as:
+
+```text
+input fixture
+→ features
+→ representation
+→ model
+→ output
+```
+
+### Parity tests
+
+Use parity tests to compare the rewrite with the legacy implementation.
+
+Prefer small deterministic fixtures and reference outputs.
+
+Do not weaken or delete parity tests simply to make new code pass.
+
+If a parity difference is intentional, document the reason.
+
+---
+
+## Test data
+
+Large production datasets do not belong in the repository.
+
+Small deterministic fixtures may live under:
+
+```text
+tests/fixtures/
+```
+
+Fixtures should be only large enough to exercise relevant behavior.
+
+Where possible, create reference outputs for deterministic legacy behavior before replacing that behavior.
+
+---
+
+## Reproducibility
+
+Reproducibility is a project requirement.
+
+Randomness should be explicit and controllable.
+
+Where relevant, account for:
+
+* Python random state
+* NumPy random state
+* PyTorch random state
+* CUDA random state
+* DataLoader workers
+* shuffling
+* sampling
+* data augmentation
+* distributed execution
+
+Do not introduce hidden global random-state mutation.
+
+Seeds should be passed or configured explicitly.
+
+Do not claim bitwise deterministic training unless the execution environment actually guarantees it.
+
+---
+
+## Code quality
+
+Domain meaning must be represented by named fields or schemas, not only by
+positional column indices. New APIs must not require callers to know that a
+particular tensor column means fold, weight, or another domain field.
+
+Prefer:
+
+* small modules with clear responsibilities
+* typed function signatures
+* explicit inputs and outputs
+* dataclasses or typed configuration where appropriate
+* composition over hidden global state
+* readable names over abbreviations
+* dependency injection over implicit filesystem/environment assumptions
+
+Avoid:
+
+* repository-relative `sys.path` modifications
+* mutable global configuration
+* hidden singleton state
+* wildcard imports
+* giant utility modules
+* giant training scripts
+* model-specific behavior in generic data code
+* duplicated preprocessing logic
+* unnecessary abstraction introduced before a second use case exists
+
+Keep public interfaces intentionally small.
+
+---
+
+## Dependencies
+
+Do not introduce a dependency merely to simplify a small amount of code.
+
+Before adding a major framework or runtime dependency, explain why it is necessary.
+
+In particular, do not introduce architectural frameworks such as:
+
+```text
+PyTorch Lightning
+Kedro
+```
+
+unless the task explicitly calls for evaluating or adopting them.
+
+Scientific dependencies should be introduced intentionally and with environment compatibility in mind.
+
+PyTorch is shared infrastructure for the expected model families and belongs
+in the base package dependencies. Architecture-specific dependencies belong in
+named extras, for example `root-gnn` for DGL. The canonical development
+environment for the active ROOT-GNN rewrite is:
+
+```bash
+uv sync --extra root-gnn
+```
+
+---
+
+## Documentation
+
+Update documentation when a change alters:
+
+* architecture
+* configuration
+* public interfaces
+* expected workflow
+* compatibility guarantees
+
+Do not document speculative functionality as if it already exists.
+
+Use:
+
+```text
+docs/architecture.md
+```
+
+for the target system structure and important architectural decisions.
+
+Use:
+
+```text
+docs/migration.md
+```
+
+for rewrite progress and migration sequencing.
+
+---
+
+## Scope discipline
+
+Do not make unrelated changes.
+
+When implementing a task:
+
+* modify only the modules necessary for that task
+* do not opportunistically refactor unrelated code
+* do not rename public concepts without a clear reason
+* do not remove compatibility behavior without explicit instruction
+* do not modify generated artifacts unless required
+
+If a larger architectural issue is discovered, report it rather than expanding the task automatically.
+
+---
+
+## Validation before finishing
+
+After code changes, run the relevant available checks.
+
+As the project matures, the expected baseline should include:
+
+```bash
+pytest
+ruff check .
+```
+
+ROOT-GNN parity validation requires DGL. A package-wide environment where DGL
+is intentionally absent may skip DGL-dependent tests, but that is incomplete
+ROOT-GNN validation. The required parity gate is:
+
+```bash
+GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 pytest
+```
+
+The gate must fail when DGL cannot be imported.
+
+Run more focused tests first when appropriate.
+
+If the full suite is expensive, run the relevant subset and clearly report what was and was not run.
+
+Do not claim validation succeeded unless the commands actually succeeded.
+
+---
+
+## Completion report
+
+For non-trivial changes, finish with a concise report containing:
+
+* files changed
+* behavior implemented
+* tests/checks run
+* parity status where relevant
+* intentional deviations from legacy behavior
+* unresolved questions or ambiguities
+
+Do not hide failed checks or incomplete behavior.

CHANGELOG.md

@@ -0,0 +1,13 @@
+# Changelog
+
+## Unreleased
+
+### Added
+
+- Release validation and CI coverage for the rewritten package.
+- Installed-package Hydra configuration discovery and public import smoke tests.
+
+### Compatibility
+
+- The `0.1.0` package version remains a static, single-source version exposed
+  as `gnn4colliders.__version__`.

MANIFEST.in

@@ -0,0 +1,3 @@
+include README.md
+include CHANGELOG.md
+recursive-include docs *.md

README.md

@@ -1,358 +1,252 @@
----
-license: mit
-tags: 
-- arXiv:2412.10665
----
-
-This is a demo is of the approach described in the paper, ["Pretrained Event Classification Model for High Energy Physics Analysis"](https://arxiv.org/abs/2412.10665)
-```
-@misc{ho2024pretrained,
-  title={Pretrained Event Classification Model for High Energy Physics Analysis},
-  author={Joshua Ho, Benjamin Ryan Roberts, Shuo Han, Haichen Wang},
-  year={2024},
-  eprint={2412.10665},
-  archivePrefix={arXiv}
-}
-```
-
-## Abstract
-
-We introduce a foundation model for event classification in high-energy physics, built on a **Graph Neural Network** architecture and trained on **120 million simulated proton-proton collision events** spanning 12 distinct physics processes. The model is *pretrained* to learn a general and robust representation of collision data using challenging multiclass and multilabel classification tasks.
-
-Its performance is evaluated across five event classification tasks, which include both physics processes used during pretraining and new processes not encountered during pretraining. Fine-tuning the pretrained model significantly improves classification performance, particularly in scenarios with limited training data, demonstrating gains in both accuracy and computational efficiency.
-
-To investigate the underlying mechanisms behind these performance improvements, we employ a representational similarity evaluation framework based on *Centered Kernel Alignment*. This analysis reveals notable differences in the learned representations of fine-tuned pretrained models compared to baseline models trained from scratch.
-
-## Introduction
-
-Machine learning has become a ubiquitous tool in particle physics, employed in a variety of tasks including triggering, simulation, reconstruction, and offline analysis. While its utility spans classification, regression, and generative tasks, the current paradigm of developing machine learning models from scratch for each specific application presents several challenges. This approach not only demands specialized expertise and substantial computing resources but can also result in suboptimal performance due to limited training data. The from-scratch development of models necessitates individual validation studies to ensure that neural networks utilize well-modeled information from training samples, whether derived from Monte Carlo simulations or control samples from experimental data.
-
-Foundation models offer a promising direction to address these limitations. These models, pre-trained on large, diverse datasets across various tasks, provide robust and general representations of underlying data structures. Notable examples in other fields include GPT-4 [OpenAI et al., 2024](#ref-openai-2024-gpt4) and BERT [Devlin et al., 2018](#ref-devlin-2018-bert) in natural language processing, Stable Diffusion [Rombach et al., 2021](#ref-rombach-2021-latentdiffusion) in image processing, and AlphaFold [Jumper et al., 2021](#ref-jumper-2021-alphafold) in structural biology. The foundation model approach offers several advantages for particle physics applications: reduced computing resources for fine-tuning [Yosinski et al., 2014](#ref-yosinski-2014-transfer) compared to training from scratch, superior performance on specific tasks (particularly with limited training data), and potentially simplified validation procedures as downstream tasks inherit verified representations from the pre-trained model.
-
-Current literature on pretrained models for particle physics can be categorized based on the data representation they handle. Models operating on particle- or event-level numerical data use features like particle four momenta or jets, leveraging self-supervised or generative methods to learn versatile representations. Detector-focused models operate on high-dimensional responses such as calorimeter deposits or pixel hits, employing geometry-aware techniques for accurate simulation and analysis. Finally, models using textual or code representations apply large language model architectures to integrate domain knowledge, enabling tasks like question answering and code generation.
-
-Recent studies have begun exploring foundation models tailored to particle physics data, which has a variety of distinct structures and properties across many experiments and data processing stages, including:
-
-- particle-level & event-level numeric data [Wildridge et al., 2024](#ref-wildridge-2024-bumblebee), [Katel et al., 2024](#ref-katel-2024-jet), [Golling et al., 2024](#ref-golling-2024-maskedset), [Mikuni & Nachman, 2024](#ref-mikuni-2024-omnilearn), [Harris et al., 2024](#ref-harris-2024-resimulation), [Birk et al., 2024](#ref-birk-2024-omnijet), [Vigl et al., 2024](#ref-vigl-2024-finetune),
-- detector-level & geometry-aware data [Araz et al., 2024](#ref-araz-2024-pointcloud), [Liu et al., 2023](#ref-liu-2023-gaam), [Hashemi et al., 2024](#ref-hashemi-2024-gen), [Huang et al., 2024](#ref-huang-2024-lmtracking),
-- textual or code data [Zhang et al., 2024](#ref-zhang-2024-xiwu).
-
-This paper presents a foundation model designed specifically for collider event-level data. In modern collider experiments, final-stage analysis processes information from reconstructed objects that either directly correspond to particles in collision final states (such as leptons and photons) or serve as proxies (such as jets and missing transverse energy). While traditional approaches often relied on "high-level" variables calculated from object features, recent trends favor direct input of event objects and their features into neural networks for analysis tasks. A notable example is [ATLAS Collaboration, 2023](#ref-atlas-2023-4top), which established the observation of simultaneous production of four top quarks with the ATLAS experiment by employing a graph neural network (GNN) architecture to process event-level object information.
-
-We present foundation models that adopt an architecture similar to that used for [ATLAS Collaboration, 2023](#ref-atlas-2023-4top). Our models are pre-trained using either multiclass classification or multi-label learning tasks across 12 distinct physics processes. We evaluate these models through fine-tuning and testing on five classification tasks, including both familiar and novel processes not seen during pre-training. Our analysis benchmarks the models' performance improvements, their scaling behavior with training sample size, and computational efficiency, representing the first prototype of a foundation model operating on collider final-state object data.
-
-## Data Samples
-
-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
-
-### Overview
-
-We present a methodology for developing and evaluating a foundation model for particle collision event analysis. The approach centers on pretraining a Graph Neural Network (GNN) architecture using a comprehensive dataset that spans multiple physics tasks, enabling the model to learn robust and transferable features. For task-specific applications, we employ a fine-tuning strategy that combines output layer adaptation with carefully calibrated learning rates for updating the pretrained parameters.
-
-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.
-
-This methodological framework demonstrates the potential of foundation models to enhance the efficiency of particle physics analyses while improving task-specific performance, offering a promising direction for future high-energy physics research.
-
----
-
-### 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.
-
-$$
-\vec {y'}_{ij} = f_e\left(\{\vec x_k\},\vec y_{ij},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\vec x_j,\vec y_{ij},\vec z\right)
-$$
-
-$$
-\vec{x'}_{i} = f_n\left(\vec x_i,\{\vec{y'}_{jk}\},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\sum_j\vec{y'}_{ij},\vec z\right)
-$$
-
-$$
-\vec{z'} = f_g\left(\{\vec{x'}_i\},\{\vec{y'}_{ij}\},\vec z\right) = \mathrm{MLP}\left(\sum_i\vec{x'}_i,\sum_{i,j}\vec{y'}_{ij},\vec z\right)
-$$
-
-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.
-
----
-
-### Pretraining Strategy
-
-We explore two complementary pretraining approaches to develop robust representations of collision events: (1) multi-class classification, which trains the model to distinguish between different physics processes, and (2) multi-label classification, which predicts the existence and kinematics of heavy particles with prompt decays. The pretraining dataset consists of approximately 120 million events, evenly distributed across 12 distinct physics processes, including all major Higgs boson production mechanisms and top quark processes as described in [Data Samples](#sec-data). This large-scale pretraining effort was conducted on the Perlmutter supercomputer at NERSC.
-
-#### Multi-class Classification
-
-For Monte Carlo simulated events, the underlying physics process that generated each event is known precisely, providing natural labels for supervised learning. However, the challenge lies in the complexity of collision events: different physics processes can produce similar kinematics and event topologies, particularly in certain regions of phase space. No single observable can unambiguously identify the underlying process. By training the model to distinguish between 12 different processes simultaneously, we challenge it to learn subtle differences in kinematics and topology that collectively characterize each process. The model is trained using categorical cross entropy as the loss function. The output layer of the multiclass classification model has 832 trainable parameters.
-
-#### Multi-label Classification
-
-This approach combines both classification and regression tasks to characterize collision events. For discrete properties like particle presence in specific kinematic regions, we employ classification labels with binary cross-entropy loss. For continuous quantities like particle multiplicities, we use regression labels with mean-squared error loss. This hybrid approach enables the model to learn both categorical and continuous aspects of the physics processes simultaneously.
-
-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.
-
----
-
-### Performance Evaluation
-
-We assess model performance using two figures of merit: the classification accuracy and the Area Under the Curve (AUC) of the Receiver Operating Characteristic (ROC) curve. The accuracy is defined as the fraction of correctly classified events when applying a threshold of 0.5 to the neural network output score. Both metrics demonstrate consistent trends in our analysis.
-
-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
-
-### Classification Performance
-
-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.
-
-Finally, the spread of accuracy across the five tasks for the baseline model is quite large, offering a robust test of fine-tuning across tasks of varying difficulty. The consistent observation of these trends across tasks confirms the reliability and robustness of the findings.
-
----
-
-### Model Interpretability
-
-We aim to understand whether pretrained and baseline models learn the same underlying representations. If the two models exhibit high similarity, a plausible interpretation is that pretraining provides the pretrained model with an advantageous initialization, allowing it to converge to a similar state as the baseline model more efficiently. Conversely, significant differences between the models would indicate that pretraining facilitates the development of a more general and robust latent space, which serves as a foundation for fine-tuning to effectively adapt to the downstream task. To investigate this, we analyzed the representational similarity between a pretrained model fine-tuned for the downstream task and a baseline model trained directly on the downstream task without pretraining.
-
-We use Centered Kernel Alignment (CKA) [Kornblith et al., 2019][ref-kornblith-2019-cka] to analyze model similarity and interpretability. CKA is a robust metric that quantifies the similarity between the internal representations of neural networks by comparing their feature matrices in a manner that is invariant to scaling, rotation, and alignment. This invariance makes CKA particularly effective for studying relationships between network layers, even across networks of different sizes or those trained from varying initializations.
-
-The similarity is evaluated using a 64-dimensional latent representation after the decoder stage of the GNN model. This choice allows us to compare the internal states of the models at a fine-grained level and understand how training strategies impact the representations directly used for the output task.
-
-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.
-
-### Computational Efficiency
-
-To estimate the computational resources required for each approach, we measured the wall time needed for a model to reach its final performance. For baseline models, this is defined as the wall time from the start of training until the loss of the model plateaus. For the foundation model approach, the estimate includes both the pretraining time and the fine-tuning time, each measured from the start of training until the loss plateaus. This approach ensures a consistent and comprehensive evaluation of the computational demands.
-
-![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:
-
-- **Multi-class pretraining:** 45.5 GPU hours
-- **Multi-label pretraining:** 60.0 GPU hours
-
-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
-
-We presented an in-depth study of a particle physics foundation model designed to operate on the four-momentum and identification properties of event final-state objects. This model is built on a Graph Neural Network (GNN) architecture and trained on a dataset comprising 120 million simulated proton-proton collision events across 12 distinct physics processes. The pretraining phase explored both multiclass and multilabel classification tasks, providing a robust foundation for downstream applications. Notably, the pretrained models demonstrated significant improvements in event classification performance when fine-tuned, particularly for tasks with limited training samples.
-
-The foundation model approach also offers substantial computational advantages. By leveraging fine-tuning, this methodology reduces the computational resources required for large-scale applications across multiple tasks. Our estimates indicate that significant resource savings can be achieved even for single particle physics measurements, making this approach both scalable and efficient.
-
-To better understand the learned representations of the pretrained model and guide future optimization efforts, we employed a representational similarity evaluation framework using Centered Kernel Alignment (CKA). This metric allowed us to investigate the source of the performance gains observed in the foundation model. Our analysis revealed notable differences in the learned representations between the fine-tuned pretrained model and a baseline model trained from scratch. In deep learning, it is well-established that multiple equally valid solutions can exist. Future studies are necessary to determine whether the low similarity in latent representations reflects complementary information uniquely captured by the foundation and baseline models, or if it can simply be attributed to connected local minima in the loss landscape.
-
-## Acknowledgments
-
-This work is supported by the U.S. National Science Foundation under the Award No. 2046280, and by U.S. Department of Energy, Office of Science under contract DE-AC02-05CH11231.
-
-## References
-
-- <span id="ref-openai-2024-gpt4"></span> **OpenAI et al.** GPT-4 Technical Report. arXiv:2303.08774 (2024). [https://arxiv.org/abs/2303.08774](https://arxiv.org/abs/2303.08774)
-
-- <span id="ref-yosinski-2014-transfer"></span> **Jason Yosinski, Jeff Clune, Yoshua Bengio, Hod Lipson.** How transferable are features in deep neural networks? CoRR abs/1411.1792 (2014). [http://arxiv.org/abs/1411.1792](http://arxiv.org/abs/1411.1792)
-
-- <span id="ref-rombach-2021-latentdiffusion"></span> **Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer.** High-Resolution Image Synthesis with Latent Diffusion Models. CoRR abs/2112.10752 (2021). [https://arxiv.org/abs/2112.10752](https://arxiv.org/abs/2112.10752)
-
-- <span id="ref-podell-2023-sdxl"></span> **Dustin Podell, Zion English, Kyle Lacey et al.** SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis. arXiv:2307.01952 (2023). [https://arxiv.org/abs/2307.01952](https://arxiv.org/abs/2307.01952)
-
-- <span id="ref-jumper-2021-alphafold"></span> **John Jumper, Richard Evans, Alexander Pritzel et al.** Highly accurate protein structure prediction with AlphaFold. Nature 596, 583-589 (2021). [https://doi.org/10.1038/s41586-021-03819-2](https://doi.org/10.1038/s41586-021-03819-2)
-
-- <span id="ref-devlin-2018-bert"></span> **Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova.** BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. CoRR abs/1810.04805 (2018). [http://arxiv.org/abs/1810.04805](http://arxiv.org/abs/1810.04805)
-
-- <span id="ref-atlas-2023-higg"></span> **ATLAS Collaboration.** Measurement of the properties of Higgs boson production at \\(\sqrt{s} = 13\,\text{TeV}\\) in the \\(H \to \gamma\gamma\\) channel using \\(139\,\text{fb}^{-1}\\) of \\(pp\\) collision data with the ATLAS experiment. JHEP 07 (2023) 088. [arXiv:2207.00348](https://arxiv.org/abs/2207.00348), [https://doi.org/10.1007/JHEP07(2023)088](https://doi.org/10.1007/JHEP07(2023)088)
-
-- <span id="ref-atlas-2023-4top"></span> **ATLAS Collaboration.** Observation of four-top-quark production in the multilepton final state with the ATLAS detector. Eur. Phys. J. C 83 (2023) 496. [arXiv:2303.15061](https://arxiv.org/abs/2303.15061), [https://doi.org/10.1140/epjc/s10052-023-11573-0](https://doi.org/10.1140/epjc/s10052-023-11573-0)
-
-- <span id="ref-kornblith-2019-cka"></span> **Simon Kornblith, Mohammad Norouzi, Honglak Lee, Geoffrey Hinton.** Similarity of Neural Network Representations Revisited. CoRR abs/1905.00414 (2019). [http://arxiv.org/abs/1905.00414](http://arxiv.org/abs/1905.00414)
-
----
-
-<!-- Historical/General Physics foundational texts -->
-
-- <span id="ref-birell-1982-qfields"></span> **N. D. Birell, P. C. W. Davies.** Quantum Fields in Curved Space. Cambridge Univ. Press (1982).
-
-- <span id="ref-feynman-1954"></span> **R. P. Feynman.** Phys. Rev. 94, 262 (1954).
-
-- <span id="ref-einstein-1935-epr"></span> **A. Einstein, Yu. Podolsky, N. Rosen.** Phys. Rev. 47, 777 (1935).
-
-- <span id="ref-berman-1983-stability"></span> **G. P. Berman, Jr., F. M. Izrailev, Jr.** Stability of nonlinear modes. Physica D 88, 445 (1983).
-
-- <span id="ref-davies-1988-trapped"></span> **E. B. Davies, L. Parns.** Trapped modes in acoustic waveguides. Q. J. Mech. Appl. Math. 51, 477–492 (1988).
-
-- <span id="ref-witten-2001"></span> **Edward Witten.** hep-th/0106109 (2001). [https://arxiv.org/abs/hep-th/0106109](https://arxiv.org/abs/hep-th/0106109)
-
----
-
-<!-- Particle physics/data science foundational models -->
-
-- <span id="ref-beutler-1994-hem"></span> **E. Beutler.** Williams Hematology, 5th Edition, Chapter 7, pp. 654–662. McGraw-Hill, New York (1994).
-
-- <span id="ref-knuth-1973-fa"></span> **Donald E. Knuth.** The Art of Computer Programming vol. 1: Fundamental Algorithms, 2nd Ed., Addison-Wesley (1973).
-
-- <span id="ref-smith-2005-philos"></span> **J. S. Smith, G. W. Johnson.** Philos. Trans. R. Soc. London, Ser. B 777, 1395 (2005).
-
-- <span id="ref-smith-2010-jap-unpub"></span> **W. J. Smith, T. J. Johnson, B. G. Miller.** Surface chemistry and preferential crystal orientation on a silicon surface. J. Appl. Phys. (unpublished, 2010).
-
-- <span id="ref-smith-2010-jap-sub"></span> **V. K. Smith, K. Johnson, M. O. Klein.** Surface chemistry and preferential crystal orientation on a silicon surface. J. Appl. Phys. (submitted, 2010).
-
-- <span id="ref-underwood-1988-lowerbounds"></span> **Ulrich Underwood, Ned Net, Paul Pot.** Lower Bounds for Wishful Research Results. Talk at Fanstord University (1988).
-
-- <span id="ref-johnson-2007-comm"></span> **M. P. Johnson, K. L. Miller, K. Smith.** Personal communication (Jan-May 2007).
-
----
-
-<!-- Prototypical collider software and tools -->
-
-- <span id="ref-pytorch-2019"></span> **Adam Paszke et al.** PyTorch: An Imperative Style, High-Performance Deep Learning Library. arXiv:1912.01703 (2019). [http://arxiv.org/abs/1912.01703](http://arxiv.org/abs/1912.01703)
-
-- <span id="ref-dgl-2019"></span> **Minjie Wang et al.** Deep Graph Library: Towards Efficient and Scalable Deep Learning on Graphs. arXiv:1909.01315 (2019). [http://arxiv.org/abs/1909.01315](http://arxiv.org/abs/1909.01315)
-
-- <span id="ref-graphnets-2018"></span> **Peter W. Battaglia et al.** Relational inductive biases, deep learning, and graph networks. arXiv:1806.01261 (2018). [http://arxiv.org/abs/1806.01261](http://arxiv.org/abs/1806.01261)
-
-- <span id="ref-layernorm-2016"></span> **Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton.** Layer Normalization. arXiv:1607.06450 (2016). [https://arxiv.org/abs/1607.06450](https://arxiv.org/abs/1607.06450)
-
----
-
-<!-- Recent & foundation models in HEP ML -->
-
-- <span id="ref-wildridge-2024-bumblebee"></span> **Andrew J. Wildridge et al.** Bumblebee: Foundation Model for Particle Physics Discovery. arXiv:2412.07867 (2024). [https://arxiv.org/abs/2412.07867](https://arxiv.org/abs/2412.07867)
-
-- <span id="ref-katel-2024-jet"></span> **Subash Katel et al.** Learning Symmetry-Independent Jet Representations via Jet-Based Joint Embedding Predictive Architecture. arXiv:2412.05333 (2024). [https://arxiv.org/abs/2412.05333](https://arxiv.org/abs/2412.05333)
-
-- <span id="ref-araz-2024-pointcloud"></span> **Jack Y. Araz et al.** Point cloud-based diffusion models for the Electron-Ion Collider. arXiv:2410.22421 (2024). [https://arxiv.org/abs/2410.22421](https://arxiv.org/abs/2410.22421)
-
-- <span id="ref-leigh-2024-maskedparticle"></span> **Matthew Leigh et al.** Is Tokenization Needed for Masked Particle Modelling? arXiv:2409.12589 (2024). [https://arxiv.org/abs/2409.12589](https://arxiv.org/abs/2409.12589)
-
-- <span id="ref-mikuni-2024-omnilearn"></span> **Vinicius Mikuni, Benjamin Nachman.** OmniLearn: A Method to Simultaneously Facilitate All Jet Physics Tasks. arXiv:2404.16091 (2024). [https://arxiv.org/abs/2404.16091](https://arxiv.org/abs/2404.16091)
-
-- <span id="ref-zhang-2024-xiwu"></span> **Zhengde Zhang et al.** Xiwu: A Basis Flexible and Learnable LLM for High Energy Physics. arXiv:2404.08001 (2024). [https://arxiv.org/abs/2404.08001](https://arxiv.org/abs/2404.08001)
-
-- <span id="ref-harris-2024-resimulation"></span> **Philip Harris et al.** Re-Simulation-based Self-Supervised Learning for Pre-Training Foundation Models. arXiv:2403.07066 (2024). [https://arxiv.org/abs/2403.07066](https://arxiv.org/abs/2403.07066)
-
-- <span id="ref-birk-2024-omnijet"></span> **Joschka Birk, Anna Hallin, Gregor Kasieczka.** OmniJet-$\alpha$: the first cross-task foundation model for particle physics. Machine Learning: Science and Technology. 5(3), 035031 (Aug 2024). [https://doi.org/10.1088/2632-2153/ad66ad](https://doi.org/10.1088/2632-2153/ad66ad)
-
-- <span id="ref-huang-2024-lmtracking"></span> **Andris Huang et al.** A Language Model for Particle Tracking. arXiv:2402.10239 (2024). [https://arxiv.org/abs/2402.10239](https://arxiv.org/abs/2402.10239)
-
-- <span id="ref-golling-2024-maskedset"></span> **Tobias Golling et al.** Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models. arXiv:2401.13537 (2024). [https://arxiv.org/abs/2401.13537](https://arxiv.org/abs/2401.13537)
-
-- <span id="ref-liu-2023-gaam"></span> **Junze Liu et al.** Generalizing to new geometries with Geometry-Aware Autoregressive Models (GAAMs) for fast calorimeter simulation. Journal of Instrumentation 18(11), P11003 (Nov 2023). [https://doi.org/10.1088/1748-0221/18/11/p11003](https://doi.org/10.1088/1748-0221/18/11/p11003)
-
-- <span id="ref-hashemi-2024-gen"></span> **Baran Hashemi et al.** Ultra-high-granularity detector simulation with intra-event aware generative adversarial network and self-supervised relational reasoning. Nature Communications 15(1) (June 2024). [https://doi.org/10.1038/s41467-024-49104-4](https://doi.org/10.1038/s41467-024-49104-4)
-
-- <span id="ref-vigl-2024-finetune"></span> **Matthias Vigl et al.** Finetuning Foundation Models for Joint Analysis Optimization. arXiv:2401.13536 (2024). [https://arxiv.org/abs/2401.13536](https://arxiv.org/abs/2401.13536)
-
-- <span id="ref-li-2024-refine"></span> **Chen Li, Hao Cai, Xianyang Jiang.** Refine neutrino events reconstruction with BEiT-3. Journal of Instrumentation 19(6), T06003 (Jun 2024). [https://doi.org/10.1088/1748-0221/19/06/t06003](https://doi.org/10.1088/1748-0221/19/06/t06003)
+# GNN4Colliders
+
+GNN4Colliders is a collider-machine-learning toolkit. The repository name
+reflects its first production model family, ROOT-GNN; the Python package is
+`gnn4colliders`, and the configuration identifier is `root_gnn`. Shared ROOT
+ingestion, collider features, metadata, tasks, training, inference, and
+distributed utilities are designed so that a future sequence model can reuse
+them without requiring every event to be a graph.
+
+```text
+ROOT files -> EventSample -> shared collider features
+                              ├── GraphSample -> ROOT-GNN
+                              └── future SequenceSample -> ROOT-Transformer
+```
+
+The new implementation lives under [`src/gnn4colliders`](src/gnn4colliders/).
+[`legacy/`](legacy/) is a frozen behavioral reference for parity work and
+historical checkpoint investigation, not a supported runtime backend.
+
+## Installation
+
+The supported development environment is Python 3.12 (`>=3.12,<3.13`), with
+PyTorch 2.2.2 and the optional ROOT-GNN stack DGL 2.4.0. The canonical setup
+is:
+
+```bash
+uv sync --dev --extra root-gnn
+```
+
+The core package can be installed without DGL when only shared data or task
+code is needed. ROOT-GNN models, graph construction, and ROOT-GNN parity tests
+require the `root-gnn` extra. The DGL extra uses the validated CUDA 12.1 wheel
+source configured in `pyproject.toml`; a compatible NVIDIA driver is still
+required. Do not add site-specific CUDA, Slurm, or filesystem paths to model
+or task configuration.
+
+## Quick start
+
+Prepare a graph cache from a ROOT tree. The feature specifications below are
+illustrative placeholders; replace them with the branches in the input tree.
+The full preparation interface is documented in
+[`docs/configuration.md`](docs/configuration.md).
+
+```bash
+uv run gnn4colliders prepare \
+  data.files=[data/events.root] \
+  data.tree_name=Events \
+  data.cache.path=cache/events.pt \
+  'data.feature_branches=[["jet_pt"],["jet_eta"],["jet_phi"],CALC_E,[1.0],[0.0],NODE_TYPE]' \
+  data.object_types=[vector] \
+  data.scales=[1,1,1,1,1,1,1]
+```
+
+Train, evaluate, and predict from that cache:
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=cache/events.pt \
+  trainer.max_epochs=1 \
+  environment.output_root=outputs/pretraining_multiclass
+
+uv run gnn4colliders evaluate \
+  data.cache.path=cache/events.pt \
+  inference.checkpoint=outputs/pretraining_multiclass/checkpoints/epoch_0000.pt
+
+uv run gnn4colliders predict \
+  data.cache.path=cache/events.pt \
+  inference.checkpoint=outputs/pretraining_multiclass/checkpoints/epoch_0000.pt \
+  inference.output=outputs/pretraining_multiclass/predictions.npz
+```
+
+For a dependency-complete, temporary-data version of this flow, run
+`uv run python scripts/dev/smoke_end_to_end.py`.
+
+## Core concepts
+
+`EventSample` is the architecture-neutral event boundary. It contains the
+selected `objects`, `label`, `global_features`, and named `EventMetadata`.
+Metadata includes `fold`, `weight`, and stable `sample_id`; callers should not
+interpret public `tracking[:, N]` columns. Legacy tracking mappings exist only
+at compatibility boundaries.
+
+The ROOT-GNN adapter converts shared features to a directed, fully connected
+graph with no self-loops: an event with `N` nodes has `N * (N - 1)` edges.
+Node columns are, in order, `pt`, `eta`, `phi`, `energy`, `btag`, `charge`,
+and `node_type`. Edge columns are `deta`, wrapped `dphi`, and `dR`.
+Object collections are concatenated in configured object-type order. The
+compatibility energy is `pt * cosh(eta)` before per-column scaling.
+
+`GraphSampleCache` stores processed graph samples and schema metadata. It is a
+Level-2 graph cache, not the universal event cache. Feature, graph, and cache
+schema versions are checked when loading; incompatible versions fail before
+training.
+
+## ROOT-GNN training and transfer
+
+`EdgeNetwork` encodes node, edge, and global features, performs iterative
+edge/node/global message passing, decodes a graph representation, and applies
+the classifier. Its output is raw logits; sigmoid or softmax is task-owned.
+
+Multiclass pretraining uses the semantic `model=root_gnn/edge_network` and
+`task=pretraining_multiclass` groups:
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=cache/events.pt \
+  model=root_gnn/edge_network task=pretraining_multiclass \
+  trainer.max_epochs=20 data.batch_size=64 \
+  environment.output_root=outputs/pretraining_multiclass
+```
+
+Fine-tuning is a separate workflow. It loads a pretrained backbone, replaces
+the classifier, and creates a new task/head optimizer:
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=cache/target.pt \
+  model=root_gnn/fine_tuned_edge_network \
+  task=binary_classification \
+  checkpoint.pretrained=/path/to/pretrained.pt \
+  model.freeze_backbone=true \
+  trainer.max_epochs=10
+```
+
+Set `model.freeze_backbone=false` to train the reused backbone as well.
+Transfer learning is not resume training:
+
+| Workflow | Meaning | Restored state |
+| --- | --- | --- |
+| Resume | Continue the same task/run | model, optimizer, scheduler, trainer, early stopping, and RNG state when present |
+| Transfer | Start a new task from a pretrained backbone | model weights only; new classifier and optimizer |
+
+Resume example:
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=cache/events.pt \
+  checkpoint.resume=outputs/pretraining_multiclass/checkpoints/epoch_0000.pt \
+  trainer.max_epochs=20
+```
+
+Validation is evaluated each epoch and drives scheduling/early stopping;
+`test` remains held out. Evaluation computes task metrics over the complete
+selected split, including weighted ROC AUC where defined:
+
+```bash
+uv run gnn4colliders evaluate \
+  data.cache.path=cache/events.pt \
+  inference.split=test \
+  inference.checkpoint=/path/to/checkpoint.pt
+```
+
+Prediction writes a named compressed NPZ. Labeled data includes `labels`;
+`fold` and `weight` are included when available. Every result includes
+`sample_id`, `logits`, `scores`, and `predictions`:
+
+```bash
+uv run gnn4colliders predict \
+  data.cache.path=cache/events.pt \
+  inference.checkpoint=/path/to/checkpoint.pt \
+  inference.output=outputs/predictions.npz
+```
+
+Optional Python-level ROOT writing is provided by
+`gnn4colliders.inference.write_root_scores`. It clones the selected tree,
+adds `score` (or `score_class_N`), and writes `selection_pass`; IDs ending in
+`:<entry>` preserve alignment and unselected entries receive NaN scores. The
+CLI currently exposes NPZ output only.
+
+The supported legacy checkpoint, metadata, and output boundary is documented
+in [`docs/compatibility.md`](docs/compatibility.md). New code should use named
+metadata fields; positional tracking is accepted only by the explicit
+compatibility adapter.
+
+### ONNX export
+
+Install the optional export dependencies and export a prepared graph-cache
+checkpoint with numerical ONNX validation:
+
+```bash
+uv sync --extra root-gnn --extra onnx
+uv run gnn4colliders export \
+  export.checkpoint=/path/to/checkpoint.pt \
+  export.output=model.onnx \
+  data.cache.path=/path/to/graph-cache.pt
+```
+
+The model accepts processed graph tensors and returns raw logits. See
+[`docs/export.md`](docs/export.md) for the tensor contract and limitations.
+
+## Configuration and environments
+
+Hydra groups are `data`, `model`, `task`, `trainer`, `checkpoint`,
+`inference`, `environment`, and `distributed`. Use configuration for a new
+experiment and Python for new behavior. Examples:
+
+```bash
+uv run gnn4colliders train trainer.max_epochs=50 data.batch_size=64
+uv run gnn4colliders train environment=perlmutter environment.device=cuda
+uv run gnn4colliders train distributed=ddp environment=perlmutter
+```
+
+Each run writes a resolved configuration to
+`<environment.output_root>/resolved_config.yaml`. See
+[`docs/configuration.md`](docs/configuration.md) for the group reference and
+[`docs/perlmutter.md`](docs/perlmutter.md) for launch examples.
+
+## Distributed execution and reproducibility
+
+Launch DDP with `torchrun` or the provided Slurm wrappers. `data.batch_size`
+and `data.num_workers` are per process, so the ordinary effective batch size
+is `batch_size * world_size`. Training shards may be padded for equal steps;
+validation and prediction are unpadded. Rank 0 writes shared checkpoints,
+configs, and predictions, and metrics/results are gathered across ranks.
+
+The configured seed controls initialization and deterministic local loader
+ordering; distributed process seeds are rank-offset and samplers use
+`set_epoch`. CPU runs are reproducible for fixed inputs and environment. GPU
+kernels, DGL, and distributed scheduling can remain nondeterministic, so the
+project does not promise bitwise GPU reproducibility.
+
+## Development and validation
+
+```bash
+uv run pytest
+uv run pytest tests/unit
+GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest tests/parity -v
+uv run ruff check .
+uv run ruff format --check .
+uv run python benchmarks/benchmark_preprocessing.py
+uv run python benchmarks/benchmark_training.py --device cpu
+```
+
+Unit tests cover isolated components, integration tests cover small workflows,
+and parity tests compare deterministic behavior with the frozen legacy
+reference. Performance guidance and measured caveats are in
+[`docs/performance.md`](docs/performance.md) and
+[`benchmarks/README.md`](benchmarks/README.md).
+See [`docs/testing.md`](docs/testing.md) for test layers, optional dependency
+markers, and package smoke validation.
+
+## Architecture and migration status
+
+See [`docs/architecture.md`](docs/architecture.md) for responsibility
+boundaries and the future sequence-model extension point. See
+[`docs/migration.md`](docs/migration.md) for the migration matrix,
+intentional redesigns, compatibility limits, and deferred work.
+
+ROOT-GNN v1 covers ROOT preparation, validated feature/graph/model/task
+behavior, training, fine-tuning, checkpoint resume, evaluation, prediction,
+single-process/DDP execution, and validated ONNX export. Streaming distributed
+output, legacy cleanup, and ROOT-Transformer remain follow-up work.

benchmarks/README.md

@@ -0,0 +1,48 @@
+# Performance benchmarks
+
+These scripts use fixed synthetic inputs, explicit seeds, warmup iterations,
+and JSON-lines output. GPU timings synchronize after every measured operation;
+warmup and one-time setup costs are excluded from steady-state numbers.
+
+```bash
+uv run python benchmarks/benchmark_preprocessing.py
+uv run python benchmarks/benchmark_dataloader.py
+uv run python benchmarks/benchmark_training.py --device cpu
+uv run python benchmarks/benchmark_inference.py --device cpu
+uv run python benchmarks/benchmark_training.py --device cuda --profile
+```
+
+DGL-dependent scripts report a structured `skipped` result when the optional
+`root-gnn` extra is absent. Profiler traces go under `profiles/` and are not
+committed.
+
+## Baseline measurements
+
+The portable baseline in this checkout uses Python 3.12.13, PyTorch 2.2.2+cu121,
+DGL 2.4.0+cu121, CPU, `nodes=32`, `iterations=10`, and `warmup=3`. Exact timings are machine
+dependent; the JSON output from a local run is authoritative.
+
+| Measurement | Mean | Median |
+| --- | ---: | ---: |
+| feature construction | 0.47 ms/event | 0.46 ms/event |
+| edge features | 0.11 ms/graph | 0.11 ms/graph |
+| graph-sample loader | 1.65 ms/iteration | 1.65 ms/iteration |
+| ROOT-GNN training step | 21.19 ms/step | 5.32 ms/step |
+| ROOT-GNN inference | 1.48 ms/graph | 1.49 ms/graph |
+
+On the available NVIDIA A100-PCIE-40GB with CUDA 12.1 runtime, the same small
+synthetic ROOT-GNN benchmark measured 6.41 ms/step (17.3 MiB peak allocated)
+and 2.00 ms/graph. These are microbenchmarks, not production-workload claims.
+
+DGL is available in this environment. The training mean is skewed by one CPU
+warmup-adjacent outlier; median is the more useful steady-state indicator. DDP
+measurements require a multi-process run and are not inferred from CPU numbers.
+
+The topology cache is bounded to 32 node-count/device/policy entries and only
+stores reusable index tensors, never event-specific edge features.
+
+On the same CPU, a direct topology microbenchmark for 64 nodes measured
+201.6 microseconds per cold construction versus 10.5 microseconds for a warm
+cache lookup (about 19x for this isolated operation). This is a targeted
+index-construction result, not an end-to-end training speedup; graph feature
+construction and DGL message passing remain separate costs.

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Standalone, reproducible performance measurements for GNN4Colliders."""

benchmarks/_common.py

@@ -0,0 +1,70 @@
+"""Small helpers shared by benchmark entry points."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import platform
+import statistics
+import time
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+
+def common_parser(description: str) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=description)
+    parser.add_argument("--iterations", type=int, default=20)
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--seed", type=int, default=1234)
+    return parser
+
+
+def synchronize(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def measure(
+    operation: Callable[[], Any], *, iterations: int, warmup: int, device: torch.device
+) -> dict[str, float]:
+    for _ in range(warmup):
+        operation()
+    synchronize(device)
+    durations = []
+    for _ in range(iterations):
+        start = time.perf_counter()
+        operation()
+        synchronize(device)
+        durations.append(time.perf_counter() - start)
+    return {
+        "mean_ms": statistics.mean(durations) * 1000,
+        "median_ms": statistics.median(durations) * 1000,
+        "min_ms": min(durations) * 1000,
+        "max_ms": max(durations) * 1000,
+        "stdev_ms": statistics.stdev(durations) * 1000 if len(durations) > 1 else 0.0,
+    }
+
+
+def metadata(device: torch.device) -> dict[str, Any]:
+    result: dict[str, Any] = {
+        "python": platform.python_version(),
+        "torch": torch.__version__,
+        "device": str(device),
+        "cuda_available": torch.cuda.is_available(),
+    }
+    if device.type == "cuda" and torch.cuda.is_available():
+        result["gpu"] = torch.cuda.get_device_name(device)
+    try:
+        import dgl
+
+        result["dgl"] = dgl.__version__
+    except ImportError:
+        result["dgl"] = None
+    return result
+
+
+def report(name: str, values: dict[str, Any]) -> None:
+    print(json.dumps({"benchmark": name, **values}, sort_keys=True))

benchmarks/benchmark_dataloader.py

@@ -0,0 +1,63 @@
+"""Measure deterministic loader/batch construction when ROOT-GNN is installed."""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--samples", type=int, default=64)
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.data import (
+            EventMetadata,
+            GraphDataLoader,
+            GraphDataset,
+            GraphSample,
+        )
+    except ImportError:
+        report(
+            "dataloader",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    samples = []
+    for index in range(args.samples):
+        graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2)
+        graph.ndata["features"] = torch.ones(2, 7)
+        graph.edata["features"] = torch.ones(2, 3)
+        samples.append(
+            GraphSample(
+                graph, torch.tensor(index % 2), None, EventMetadata(0, 1.0, str(index))
+            )
+        )
+    loader = GraphDataLoader(GraphDataset(samples), args.batch_size)
+    result = measure(
+        lambda: list(loader),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    report(
+        "dataloader",
+        {
+            **metadata(device),
+            **result,
+            "batch_size": args.batch_size,
+            "batches_per_second": 1000 / result["mean_ms"],
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_inference.py

@@ -0,0 +1,52 @@
+"""Measure ROOT-GNN inference with fixed synthetic graph input."""
+
+from __future__ import annotations
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.models.root_gnn import EdgeNetwork
+    except ImportError:
+        report(
+            "inference",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2).to(
+        device
+    )
+    graph.ndata["features"] = torch.ones(2, 7, device=device)
+    graph.edata["features"] = torch.ones(2, 3, device=device)
+    model = (
+        EdgeNetwork(graph, None, hid_size=16, out_size=2, n_layers=1, n_proc_steps=1)
+        .to(device)
+        .eval()
+    )
+
+    def infer() -> None:
+        with torch.inference_mode():
+            model(graph)
+
+    result = measure(
+        infer, iterations=args.iterations, warmup=args.warmup, device=device
+    )
+    report(
+        "inference",
+        {**metadata(device), **result, "graphs_per_second": 1000 / result["mean_ms"]},
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_preprocessing.py

@@ -0,0 +1,61 @@
+"""Measure shared feature and graph preprocessing on fixed synthetic events."""
+
+from __future__ import annotations
+
+import torch
+
+from gnn4colliders.features import build_node_features
+from gnn4colliders.graphs import build_edge_features, fully_connected_edges
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--nodes", type=int, default=32)
+    args = parser.parse_args()
+    torch.manual_seed(args.seed)
+    event = {
+        "pt": torch.arange(args.nodes, dtype=torch.float32) + 1,
+        "eta": torch.linspace(-2, 2, args.nodes),
+        "phi": torch.linspace(-3.0, 3.0, args.nodes),
+    }
+    branches = [["pt"], ["eta"], ["phi"], "CALC_E", [1.0], [0.0], "NODE_TYPE"]
+    object_types = ["vector"]
+    scales = [1.0] * 7
+    device = torch.device(args.device)
+
+    feature_result = measure(
+        lambda: build_node_features(event, branches, object_types, scales),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    nodes = build_node_features(event, branches, object_types, scales)[0]
+    src, dst = fully_connected_edges(args.nodes)
+    graph_result = measure(
+        lambda: build_edge_features(nodes, src, dst, eta_index=1, phi_index=2),
+        iterations=args.iterations,
+        warmup=args.warmup,
+        device=device,
+    )
+    base = {**metadata(device), "nodes": args.nodes, "iterations": args.iterations}
+    report(
+        "feature_construction",
+        {
+            **base,
+            **feature_result,
+            "events_per_second": 1000 / feature_result["mean_ms"],
+        },
+    )
+    report(
+        "edge_features",
+        {**base, **graph_result, "graphs_per_second": 1000 / graph_result["mean_ms"]},
+    )
+
+
+if __name__ == "__main__":
+    main()

benchmarks/benchmark_training.py

@@ -0,0 +1,73 @@
+"""Measure a short ROOT-GNN training section and optionally emit a profiler trace."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+try:
+    from ._common import common_parser, measure, metadata, report
+except ImportError:
+    from _common import common_parser, measure, metadata, report
+
+
+def main() -> None:
+    parser = common_parser(__doc__)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--profile", action="store_true")
+    parser.add_argument("--profile-dir", default="profiles")
+    args = parser.parse_args()
+    device = torch.device(args.device)
+    try:
+        import dgl
+
+        from gnn4colliders.models.root_gnn import EdgeNetwork
+    except ImportError:
+        report(
+            "training",
+            {**metadata(device), "status": "skipped: install root-gnn extra"},
+        )
+        return
+    torch.manual_seed(args.seed)
+    graph = dgl.graph((torch.tensor([0, 1]), torch.tensor([1, 0])), num_nodes=2).to(
+        device
+    )
+    graph.ndata["features"] = torch.randn(2, 7, device=device)
+    graph.edata["features"] = torch.randn(2, 3, device=device)
+    model = EdgeNetwork(
+        graph, None, hid_size=16, out_size=2, n_layers=1, n_proc_steps=1
+    ).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+
+    def step() -> None:
+        optimizer.zero_grad(set_to_none=True)
+        loss = model(graph).square().mean()
+        loss.backward()
+        optimizer.step()
+
+    if args.profile:
+        profile_dir = Path(args.profile_dir)
+        profile_dir.mkdir(parents=True, exist_ok=True)
+        with torch.profiler.profile(record_shapes=True, profile_memory=True) as prof:
+            for _ in range(args.warmup + args.iterations):
+                step()
+        prof.export_chrome_trace(str(profile_dir / "training.trace.json"))
+    result = measure(
+        step, iterations=args.iterations, warmup=args.warmup, device=device
+    )
+    report(
+        "training_step",
+        {
+            **metadata(device),
+            **result,
+            "steps_per_second": 1000 / result["mean_ms"],
+            "peak_memory_mb": torch.cuda.max_memory_allocated(device) / 2**20
+            if device.type == "cuda"
+            else 0.0,
+        },
+    )
+
+
+if __name__ == "__main__":
+    main()

configs/.gitkeep

@@ -0,0 +1 @@
+

configs/data/.gitkeep

@@ -0,0 +1 @@
+

configs/environment/.gitkeep

@@ -0,0 +1 @@
+

configs/model/.gitkeep

@@ -0,0 +1 @@
+

configs/model/root_gnn/.gitkeep

@@ -0,0 +1 @@
+

configs/task/.gitkeep

@@ -0,0 +1 @@
+

configs/trainer/.gitkeep

@@ -0,0 +1 @@
+

data/fixtures/.gitkeep

@@ -0,0 +1 @@
+

data/raw/.gitkeep

@@ -0,0 +1 @@
+

docs/architecture.md

@@ -0,0 +1,375 @@
+# GNN4Colliders architecture
+
+## Current v1 architecture
+
+The supported rewrite is layered around an architecture-neutral event boundary:
+
+```text
+ROOT/Awkward
+    ↓
+EventSample + EventMetadata
+    ↓
+shared collider feature construction
+    ├── GraphSample -> versioned graph cache -> GraphBatch
+    └── future SequenceSample -> ROOT-Transformer (not implemented)
+                                      ↓
+                               ROOT-GNN EdgeNetwork
+                                      ↓
+                               raw logits -> Task
+                                  ├── loss
+                                  ├── predictions
+                                  └── full-split metrics
+                                      ↓
+                          Trainer / Predictor / outputs
+```
+
+| Layer | Responsibility |
+| --- | --- |
+| `data` | ROOT/Awkward ingestion, event samples, metadata, graph caches, folds, and batching |
+| `features` | Shared collider-object features and derived physics quantities |
+| `graphs` | Topology, edge features, and the DGL representation adapter |
+| `models/root_gnn` | ROOT-GNN encoders, message passing, classifier, and transfer boundary |
+| `tasks` | Loss, score/prediction, labels, weights, and metrics |
+| `training` | Optimizer lifecycle, validation, early stopping, checkpointing, and reproducibility |
+| `inference` | Ordered prediction/evaluation and NPZ/ROOT output adapters |
+| `distributed` | Rank-local devices, sharding, DDP, and cross-rank collection |
+| `config` / `cli` | Semantic Hydra composition and thin user-facing commands |
+
+`EventSample` is shared infrastructure, not a ROOT-GNN object. `GraphSample`
+is the current representation-specific adapter. This separation is the
+extension point for a future sequence/token representation.
+
+The new public metadata contract is named `EventMetadata(fold, weight,
+sample_id, extra)`. The legacy positional tracking tensor is accepted only by
+compatibility-facing ingestion code. A `GraphSampleCache` is deliberately a
+Level-2 graph cache; replacing it with a universal cache would couple future
+model families to DGL.
+
+For deployment, a prepared `GraphBatch` can pass through the isolated
+`RootGNNExportAdapter` into an ONNX model. This is an inference boundary only;
+the native model continues to consume DGL graphs and ONNX does not read ROOT
+or construct collider features.
+
+## Current public workflow
+
+`prepare` reads ROOT through `RootEventDataset`, builds shared features and
+DGL graphs, and saves a schema-checked cache. `train` creates a model/task and
+`Trainer`; validation is the model-selection split and test is held out.
+`evaluate` computes metrics after collecting the complete split. `predict`
+returns detached CPU tensors in loader order and writes named NPZ fields.
+`write_root_scores` is an optional Python adapter with explicit entry alignment;
+the CLI currently exposes NPZ output.
+
+Compatibility responsibilities are isolated in `gnn4colliders.compat`.
+Supported historical checkpoint prefixes, classifier names, and the two-column
+tracking conversion are listed in [`compatibility.md`](compatibility.md).
+The modern pipeline does not propagate positional tracking or historical NPZ
+fields.
+
+Checkpoints are independent of the model implementation: they carry model and
+task metadata, lifecycle state, schema versions, and optional RNG state.
+Prefix normalization supports DDP `module.` and compiled `_orig_mod.` weights,
+plus the active ROOT-GNN historical classifier-name compatibility path.
+
+---
+
+## Historical behavioral reference
+
+This document covers the target system in `legacy/root_gnn_dgl/`. The sibling
+`legacy/physicsnemo/` tree is a prior rewrite attempt and is not a behavioral
+target.
+
+## 1. High-level system description
+
+The active rewrite exposes `gnn4colliders.models.root_gnn.EdgeNetwork`. Its
+encoders and message-passing blocks form a reusable backbone whose decoded
+graph representation is passed to an explicit classifier. `FineTunedEdgeNetwork`
+reuses that backbone and replaces only the task-specific classifier, with
+explicit frozen or trainable-backbone control.
+
+`root_gnn_dgl` is a ROOT-to-DGL graph classification system. YAML selects
+dataset, model, loss, and finish-function classes by import path. The dataset
+reads ROOT trees, converts collider objects to fully connected DGL graphs, and
+saves graph chunks. Training loads those chunks, applies fold selection and
+optional pre-batching/padding, trains a graph network, writes one PyTorch
+checkpoint per epoch, and reports weighted loss, accuracy, and ROC AUC.
+
+The primary model is `models.GCN.Edge_Network` ([`GCN.py:182-251`](../legacy/root_gnn_dgl/models/GCN.py)).
+It encodes node, edge, and global features, repeats edge -> node -> global
+message passing `n_proc_steps` times, decodes the global state, and applies
+`classify`. Fine-tuning uses `models.GCN.Transferred_Learning_Finetuning`
+([`GCN.py:884-997`](../legacy/root_gnn_dgl/models/GCN.py)), which loads a
+pretrained `Edge_Network`, removes its final classifier, and applies a new one.
+The active configs use output size 12 for multiclass pretraining and output
+size 1 for binary tasks ([`configs/stats_100K/pretraining_multiclass.yaml:1-45`](../legacy/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml),
+[`configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml:1-45`](../legacy/root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml)).
+
+### Entry points and flows
+
+- `scripts/training_script.py:main` and its CLI parser load YAML, create
+  loaders, construct the model, and call `train`; `--evaluate` calls
+  `evaluate` ([`training_script.py:638-843`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+- `scripts/prep_data.py:main` creates configured graph caches
+  ([`prep_data.py:68-110`](../legacy/root_gnn_dgl/scripts/prep_data.py)).
+- `scripts/inference.py:main` reconstructs an unlazy dataset, loads one or
+  more checkpoints, and writes `.npz` or ROOT scores
+  ([`inference.py:163-387`](../legacy/root_gnn_dgl/scripts/inference.py)).
+- `scripts/export_onnx.py:main` exports an ONNX-friendly model
+  ([`export_onnx.py:979-1035`](../legacy/root_gnn_dgl/scripts/export_onnx.py)).
+- `selections.py:main`, `check_dataset_files.py:main`, and
+  `plot_config_distributions.py:main` are diagnostic entry points. `run_demo.sh`
+  sequences pretraining, binary training, fine-tuning, and inference
+  ([`run_demo.sh:3-59`](../legacy/root_gnn_dgl/run_demo.sh)).
+
+The training flow is:
+
+```text
+YAML -> load_config/buildFromConfig -> RootDataset/LazyDataset
+     -> ROOT/Awkward -> DGL graph + labels/tracking/globals -> .bin cache
+     -> fold_selection -> prebatch/padding -> GraphDataLoader
+     -> Edge_Network or transfer model -> weighted loss/metrics
+     -> model_epoch_N.pt, logs, evaluation/inference output
+```
+
+`training_script.train` is the lifecycle implementation
+([`training_script.py:143-614`](../legacy/root_gnn_dgl/scripts/training_script.py));
+distributed paths use NCCL/DDP ([`training_script.py:616-839`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Inference uses `CustomPreBatchedDataset`, applies a configured finish function,
+and collects `scores`, `labels`, and `tracking_info`
+([`inference.py:20-76`](../legacy/root_gnn_dgl/scripts/inference.py),
+[`inference.py:223-325`](../legacy/root_gnn_dgl/scripts/inference.py)).
+
+## 2. Dependency and data-flow map
+
+```text
+ROOT files (raw_dir/file_names, tree_name)
+  -> gnn4colliders.data.RootEventDataset (Uproot/Awkward)
+  -> node feature construction (dataset.py:15-50)
+  -> full_connected_graph (dataset.py:52-59)
+  -> EdgeDataset.make_graph: [deta, dphi, dR] (dataset.py:471-482)
+  -> DGL .bin cache / LazyDataset / PreBatchedDataset
+  -> GraphDataLoader -> models.GCN -> loss/metrics -> outputs
+```
+
+`RootDataset` provides `process`, `save`, `load`, `__getitem__`, and `__len__`
+([`dataset.py:160-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
+`LazyDataset` loads one chunk through a ring buffer
+([`dataset.py:525-578`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
+`PreBatchedDataset.process` selects, shuffles, batches, pads, and caches
+([`batched_dataset.py:34-146`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+
+`load_config` uses PyYAML `FullLoader` and shallow `include` merging, while
+`buildFromConfig` dynamically imports `module`, resolves `class`, merges extra
+keys into `args`, converts list-valued weights to tensors, and injects runtime
+arguments ([`utils.py:10-43`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
+This reflection shape is a de facto interface for configured components.
+
+### Data and preprocessing
+
+The active node schema is seven columns: `pt`, `eta`, `phi`, `energy`, `btag`,
+`charge`, and `node_type`. `CALC_E` is `pt*cosh(eta)`, constants are broadcast
+per object type, `NODE_TYPE` is an integer type code, and feature scales are
+applied columnwise ([`dataset.py:15-50`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
+`full_connected_graph` makes directed all-pairs edges; `EdgeDataset` requests
+no self-loops and stores `[deta, dphi, dR]`
+([`dataset.py:52-59`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`dataset.py:471-482`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
+
+Selections are strings evaluated with builtins disabled or
+`(variable, cut, operator)` triples (`check_selection`, `selection_mask`;
+[`dataset.py:75-145`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)). Fold
+selection uses `tracking[:,0] % n_folds`; tracking column 0 is fold and column 1
+is weight ([`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py),
+[`dataset.py:176-182`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
+`hash_partition` and a seeded Torch generator control pre-batch order
+([`batched_dataset.py:27-99`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+Padding modes are `NONE`, `STEPS`, `FIXED`, and `NODE`; `FIXED` is hardcoded to
+16,000 nodes and 104,000 edges ([`batched_dataset.py:100-125`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+
+### Model, losses, and metrics
+
+`Make_MLP` builds linear/ReLU/dropout blocks followed by LayerNorm
+([`GCN.py:18-35`](../legacy/root_gnn_dgl/models/GCN.py)). Each `Edge_Network`
+step encodes inputs, copies source/destination states to edges, updates edges,
+sums edge messages into nodes, updates nodes, then mean-pools nodes/edges to
+update globals ([`GCN.py:195-249`](../legacy/root_gnn_dgl/models/GCN.py)). It
+returns logits `[graphs, out_size]` without sigmoid/softmax.
+
+The default objective is elementwise `BCEWithLogitsLoss`, multiplied by
+`tracking[:,1]`, averaged separately per unique label, then averaged across
+labels ([`training_script.py:143-185`](../legacy/root_gnn_dgl/scripts/training_script.py),
+[`training_script.py:320-359`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+`--abs` makes weights positive. Binary metrics use sigmoid threshold 0.5 and
+weighted ROC AUC; multiclass metrics use argmax and one-vs-rest ROC AUC
+([`training_script.py:438-510`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Additional configurable losses/finishers live in `models/loss.py`
+([`loss.py:6-310`](../legacy/root_gnn_dgl/models/loss.py)).
+
+### Checkpoints and outputs
+
+Training writes `Training_Directory/model_epoch_<epoch>.pt` containing `epoch`,
+`model_state_dict`, `optimizer_state_dict`, and serialized `early_stop`
+([`training_script.py:565-604`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Keys strip `module.` and compiled models save the underlying `_orig_mod` state;
+`get_last_epoch`, `get_specific_epoch`, and `get_best_epoch` load the files
+([`utils.py:145-248`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
+`evaluate` writes `evaluation_<epoch>.npz`; inference writes `.npz` fields
+`scores`, `labels`, `tracking_info`, or adds score branches and `selection_pass`
+to a cloned ROOT tree ([`training_script.py:57-140`](../legacy/root_gnn_dgl/scripts/training_script.py),
+[`inference.py:328-385`](../legacy/root_gnn_dgl/scripts/inference.py)).
+
+### Training lifecycle boundary
+
+The active rewrite keeps lifecycle orchestration architecture-independent:
+
+```text
+GraphDataLoader -> GraphBatch -> Model -> Task -> Trainer
+                                      loss/metrics
+```
+
+`gnn4colliders.training.Trainer` owns device placement, train/evaluation mode,
+gradient and optimizer steps, epoch aggregation, optional scheduler stepping,
+early stopping, and in-memory history. Tasks own loss and metric semantics;
+the trainer does not inspect positional tracking columns or collider-specific
+features. Evaluation concatenates detached outputs across the complete split
+before calling task metrics, so ROC AUC is not computed per mini-batch.
+
+Checkpoint persistence and the Python inference/output layer are implemented
+as separate adapters. The semantic CLI and distributed application boundary
+are implemented in the current stack. `gnn4colliders.inference.Predictor` accumulates detached CPU
+logits, task-defined scores/predictions, labels, and named event metadata in
+loader order; `write_npz` is the primary named-field format and ROOT score
+writing is an optional alignment-aware adapter.
+
+The new lifecycle uses conventional split semantics. `train` updates model
+parameters, `validation` is evaluated after every epoch and drives scheduler,
+early-stopping, and later model selection, and `test` is held out. The trainer
+does not accept a test loader in `fit`; callers evaluate the held-out test set
+separately after training. This deliberately corrects the legacy convention
+where a loader named `test` was used for model selection and `val` represented
+held-out testing.
+
+## Distributed execution
+
+`gnn4colliders.distributed` contains the small DDP boundary used by the
+application layer. `DistributedContext` reads the standard `torchrun`
+environment (`RANK`, `LOCAL_RANK`, and `WORLD_SIZE`), selects the rank-local
+device, and owns process-group cleanup. Graph samples are sharded before
+batching; training may pad rank shards for equal step counts, while validation
+and prediction use unpadded shards so events are not counted twice.
+
+The configured graph `batch_size` is per process. DDP wraps an otherwise
+ordinary model after device placement, and checkpoint state is normalized to
+the underlying model keys. Loss gradients are synchronized by DDP; epoch
+metrics and evaluation outputs are gathered across ranks. Rank 0 writes
+resolved configuration, checkpoints, and NPZ predictions. Moderate-size
+prediction gathering is in-memory; streaming/sharded output is a future
+extension.
+
+
+## Randomness, external services, and coupling
+
+The CLI exposes `--seed`, but `main` passes it to model construction rather
+than globally seeding Python, NumPy, or Torch
+([`training_script.py:638-753`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Pre-batching has explicit seeds, but `AugmentedDataset` mutates the process-wide
+NumPy seed ([`dataset.py:716-827`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
+clustering uses unseeded `torch.randperm`/`randint` (`loss.py:259-295`), and
+model reset/fine-tuning hardcodes `torch.manual_seed(2)`
+([`GCN.py:58-65`](../legacy/root_gnn_dgl/models/GCN.py),
+[`GCN.py:900-915`](../legacy/root_gnn_dgl/models/GCN.py)). CUDA kernels, DDP,
+and DataLoader behavior are not made deterministic.
+
+Implicit coupling includes repository-relative `sys.path` insertion
+([`training_script.py:14-20`](../legacy/root_gnn_dgl/scripts/training_script.py)),
+dynamic imports, mutable default lists/dicts, global `FEATURE_DTYPE`
+([`dataset.py:13`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
+in-place `tracking_info` mutation ([`dataset.py:176-181`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
+and in-place DGL graph mutation during forward.
+
+The legacy environment assumes Python 3.8, PyTorch 2.0.1, CUDA 11.8, DGL
+1.1.1, ROOT, Awkward, Uproot, PyYAML, and scikit-learn
+([`setup/environment.yml:1-10`](../legacy/root_gnn_dgl/setup/environment.yml),
+[`setup/environment.yml:240-295`](../legacy/root_gnn_dgl/setup/environment.yml)).
+The active Linux development environment is intentionally separate: Python
+3.12, PyTorch 2.2.2/CUDA 12.1, and DGL 2.4.0 from the official DGL wheel
+repository. CUDA runtime wheels do not replace the compatible host NVIDIA
+driver and do not encode Perlmutter module settings.
+Standard configs assume `/global/cfs/` and `/pscratch/` paths, CUDA/NCCL,
+Slurm, and optionally Podman-HPC. `setup/download_data.sh` downloads the
+external Hugging Face dataset `HWresearch/Delphes`
+([`download_data.sh:13-67`](../legacy/root_gnn_dgl/setup/download_data.sh)).
+
+## Apparent unused or secondary code
+
+Not selected by the standard stats/Delphes configs, or only reachable from
+optional workflows, are `GCN_global`, `GCN_global_2way`, most transfer variants,
+attention models, `MultiModel`, and `Clustering`
+([`GCN.py:122-1933`](../legacy/root_gnn_dgl/models/GCN.py)); `UprootDataset`,
+`tHbbEdgeDataset`, `AugmentedDataset`, and photon-ID paths
+([`uproot_dataset.py:10-31`](../legacy/root_gnn_dgl/root_gnn_base/uproot_dataset.py),
+[`dataset.py:484-827`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`photon_ID_dataset.py:1-33`](../legacy/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py));
+optional loss and similarity utilities; and the no-op
+`root_gnn_base.utils.graph_augmentation` ([`utils.py:393-395`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
+The main path evaluates `test_loaders`; validation loaders are only assembled
+when a config has a validation fold ([`training_script.py:682-747`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+
+## 3. De facto interfaces to preserve
+
+### Metadata-aware dataset boundary
+
+The rewrite uses named `EventMetadata` (`fold`, `weight`, and stable
+`sample_id`) instead of exposing the legacy positional tracking tensor.
+`GraphSample`, `GraphBatch`, `SplitDefinition`, and `GraphDataLoader` form the
+ROOT-GNN orchestration boundary. Graph caches carry feature, graph, and cache
+schema versions and reject incompatible artifacts before loading.
+
+The current cache implementation stores processed `GraphSample` values (the
+Level-2 cache). The separation from `RootEventDataset` is intentional: a
+future Level-1 cache can store normalized `EventSample`/feature data for
+sequence or transformer representations without requiring DGL graph caches.
+
+The rewrite's shared data boundary is `gnn4colliders.data`: it reads selected
+ROOT/Awkward branches and returns architecture-neutral event samples. Feature
+construction and graph building remain separate downstream boundaries, so the
+same samples can be reused by non-graph model families.
+
+### Configuration and CLI boundary
+
+Hydra composes semantic YAML groups under `configs/` and passes the resolved
+configuration to explicit application factories in `gnn4colliders.config`.
+Those factories allow-list supported models, tasks, and trainer components;
+YAML is never treated as an arbitrary Python import specification. The thin
+`gnn4colliders` CLI selects `prepare`, `train`, `evaluate`, or `predict` and
+delegates to the stable data, training, checkpoint, and inference APIs. A new
+experiment should generally be a YAML change; new behavior belongs in Python.
+
+1. YAML `module`, `class`, `args`, plus runtime `sample_graph` and
+   `sample_global` injection.
+2. Dataset items `(DGLGraph, label, tracking, global_features)`.
+3. `ndata['features']`, `edata['features']`, seven node columns, and three edge
+   columns in `[deta, dphi, dR]` order.
+4. Tracking column 0 fold and column 1 weight semantics.
+5. `model(graph, global_feats)`, logits shape `[batch, out_size]`, and
+   `representation` where used.
+6. Weighted per-label loss, metric thresholds, checkpoint keys/prefix cleanup,
+   epoch filenames, and `.npz`/ROOT output fields.
+
+## 4. Ambiguous behavior
+
+- Historical edge order/self-loop expectations; empty and padding graph inputs.
+- Whether negative weights are meaningful or should always be absolute.
+- Whether “validation” is intended to differ from the active test-loader path.
+- Shape semantics of multi-label finishers and experimental transfer classes.
+- Whether chunk IDs must match historical `np.array_split` boundaries.
+- Required behavior for missing branches and dynamic selection expressions.
+
+## 5. Recommended rewrite boundaries
+
+Separate typed configuration; ROOT/Awkward I/O; selections/folds/features/
+edges; DGL cache/lazy loading/batching; active models and checkpoint adapters;
+objectives/metrics; training lifecycle; and inference/ONNX applications.
+Establish parity for the active `LazyDataset -> PreBatchedDataset ->
+Edge_Network` binary/multiclass path first. Add experimental classes only when
+a config or consumer proves they are required.

docs/compatibility.md

@@ -0,0 +1,22 @@
+# Compatibility boundary
+
+The package's canonical APIs use named `EventMetadata` fields and the version
+1 checkpoint schema. Compatibility is explicit and one-way: old artifacts are
+adapted into the new representation and are never rewritten implicitly.
+
+| Artifact or behavior | Supported | Boundary | Notes |
+| --- | :---: | --- | --- |
+| New checkpoint (`schema_version: 1`) | yes | `CheckpointManager` | Full model/lifecycle resume when state is present |
+| Legacy `model_epoch_N.pt` checkpoint | yes | `gnn4colliders.compat.load_legacy_checkpoint` | Model state and active early-stop fields are adapted |
+| Legacy DDP/compiled prefixes | yes | `normalize_legacy_state_dict_keys` | Supports `module.` and `_orig_mod.` |
+| Legacy ROOT-GNN classifier name | yes | `map_legacy_edge_network_state_dict` | Maps `classify` to `classifier` |
+| Legacy optimizer state | partial | legacy checkpoint adapter | Loaded when present; scheduler state is not available in the historical format |
+| Legacy scheduler state | no | — | Historical checkpoints do not carry a supported scheduler state |
+| Positional tracking rows | yes, at ingestion boundary | `EventMetadata.from_legacy_tracking` | Exactly `tracking[0] = fold`, `tracking[1] = weight`; shorter rows fail |
+| Generic/unknown tracking layouts | no | — | The package does not guess historical column meanings |
+| Modern NPZ output | yes | `inference.write_npz` | Named fields: `sample_id`, `logits`, `scores`, `predictions`, and available metadata |
+| Historical `tracking_info` NPZ output | no | — | No active consumer remains; positional output is intentionally unsupported |
+| Legacy YAML `module`/`class`/`args` | compatibility only | configuration boundary | Accepted only where the semantic factory can safely interpret it; new configs use semantic model names |
+
+The frozen `legacy/` tree remains available to parity tests and historical
+investigation. Production modules do not import executable code from it.

docs/configuration.md

@@ -0,0 +1,94 @@
+# Configuration guide
+
+The CLI composes the installed `gnn4colliders.configs/config.yaml` with semantic
+Hydra groups. A config
+describes experiment intent; it does not contain arbitrary Python module or
+class import paths.
+
+## Groups
+
+| Group | Purpose |
+| --- | --- |
+| `data` | ROOT source or graph-cache path, batch settings, and fold splits |
+| `model` | Model family and supported ROOT-GNN constructor settings |
+| `task` | Binary or multiclass loss, score, and metric semantics |
+| `trainer` | Device, seed, epochs, optimizer, scheduler, and early stopping |
+| `checkpoint` | Output directory, resume checkpoint, or pretrained checkpoint |
+| `inference` | Checkpoint, split, output path, and output format |
+| `environment` | Device and experiment output root |
+| `distributed` | Single-process/DDP selection and process-group backend |
+
+The active model groups are `root_gnn/edge_network` and
+`root_gnn/fine_tuned_edge_network`. The active task groups are
+`pretraining_multiclass`, `binary_classification`, and `tth_cp_finetune`.
+
+## Preparation
+
+`prepare` requires `data.files`, `data.cache.path`, `data.feature_branches`,
+`data.object_types`, and `data.scales`. `feature_branches` follows the shared
+seven-column feature contract: one branch/constant specification per output
+column and one entry per configured object type. `CALC_E` and `NODE_TYPE` are
+reserved derived specifications. `object_types` entries are `vector` or
+`single`. Preparation reads the configured tree in file order and writes a
+versioned `GraphSampleCache`.
+
+Example overrides are easiest to maintain in a YAML file for real datasets:
+
+```yaml
+# project-local example: data/my_events.yaml
+files: [data/events.root]
+tree_name: Events
+cache:
+  path: cache/events.pt
+feature_branches:
+  - [jet_pt]
+  - [jet_eta]
+  - [jet_phi]
+  - CALC_E
+  - [1.0]
+  - [0.0]
+  - NODE_TYPE
+object_types: [vector]
+scales: [1, 1, 1, 1, 1, 1, 1]
+fold_var: eventNumber
+weight_var: weight
+```
+
+Then compose it with `data=my_events`. The cache stores processed graph
+samples, labels, globals, named metadata, and feature/graph/cache schema
+versions. Changing the feature or graph schema requires a new compatible cache;
+loading a mismatched schema raises an error.
+
+## Common overrides
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=cache/events.pt \
+  data.batch_size=64 \
+  trainer.max_epochs=50 \
+  trainer.seed=123 \
+  environment.output_root=outputs/my_run
+```
+
+`data.batch_size` is per process. `data.splits.train_folds`,
+`validation_folds`, and `test_folds` define conventional train/validation/test
+selection and must be disjoint. Model/task mismatches are rejected during
+config validation; binary tasks require `model.out_size=1`, while multiclass
+tasks require `model.out_size=task.num_classes`.
+
+## Checkpoints and resolved configuration
+
+Training writes `epoch_####.pt` and the fully resolved configuration at
+`<environment.output_root>/resolved_config.yaml`. A checkpoint includes schema
+version, model weights/config, task config, trainer/optimizer/scheduler state,
+early stopping state, metadata, and optional RNG state. Set
+`checkpoint.resume=/path/to/epoch_####.pt` to continue a run. Set
+`checkpoint.pretrained=/path/to/epoch_####.pt` with the fine-tuned model group
+to load weights into a new task head; these options are mutually exclusive.
+
+## Environment profiles
+
+`environment=local` selects CPU by default. `environment=perlmutter` selects
+the CUDA device and a conventional output-root pattern. Profiles should hold
+device/output policy only; site-specific module loads and filesystem paths
+belong in a launcher or shell environment.

docs/end_to_end_validation.md

@@ -0,0 +1,30 @@
+# End-to-end legacy/rewrite validation
+
+Task 21 compares staged event identity, labels, folds, weights, globals, node
+features, topology, edge features, batching, fixed-weight forward, loss and
+metrics, one optimizer step, short training, checkpoint reload, and inference.
+The first failing stage is retained in JSON and Markdown reports.
+
+Legacy and rewrite extraction may use separate environments. Each writes the
+version-1 artifact described in `validation/README.md`; the comparator never
+imports legacy modules or DGL graph objects. Record ROOT file/tree, selection,
+event count, file size, SHA-256, branches, seeds, versions, device and source
+identities in the artifact manifest. Do not use an absolute personal path as a
+scientific event identity.
+
+The captured runtime versions and seed policy are recorded in
+`validation/manifests/environments.json`; the HF composite fixture provenance
+is recorded in `validation/manifests/multiclass_fixture.json`.
+
+Strict parity applies to preprocessing, topology, fixed forward, loss, and
+one-step CPU updates. Multi-epoch, GPU and DDP comparisons are scientific:
+compare curves, metrics, distributions and event-level correlations. Named
+metadata, modern NPZ names and the version-1 checkpoint schema are expected
+interface differences and must be reported explicitly.
+
+Run the existing gates separately:
+
+```bash
+GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest
+uv run ruff check .
+```

docs/export.md

@@ -0,0 +1,36 @@
+# ROOT-GNN ONNX export
+
+ONNX export is an inference-only adapter for prepared ROOT-GNN graph batches.
+It does not read ROOT files or move feature construction into the deployment
+model. Install the optional dependencies with:
+
+```bash
+uv sync --extra root-gnn --extra onnx
+```
+
+Export a checkpoint with the project CLI:
+
+```bash
+uv run gnn4colliders export \
+  export.checkpoint=/path/to/checkpoint.pt \
+  export.output=model.onnx \
+  data.cache.path=/path/to/graph-cache.pt
+```
+
+The exported model accepts six tensor inputs: `node_features`,
+`edge_features`, `edge_src`, `edge_dst`, `node_batch`, and `global_features`.
+`node_batch` identifies the graph for each node; edge membership is derived
+from `node_batch[edge_src]`. Graph, node, edge, and batch dimensions are
+dynamic. The model returns raw `logits`; task postprocessing and metadata stay
+outside ONNX.
+
+The exporter uses ONNX opset 17, validates the structure with `onnx.checker`,
+and writes compact provenance/schema metadata to `model.onnx.json`. It
+supports `EdgeNetwork` multiclass models and `FineTunedEdgeNetwork` binary
+models. Empty graphs are invalid; single-node graphs (zero edges) are
+represented and pooled with a zero edge contribution.
+
+Direct DGL export was not retained: the active model uses DGL graph mutation
+and reductions (`apply_edges`, `update_all`, and graph pooling) that are not a
+portable ONNX contract. The adapter expresses those operations with standard
+tensor indexing, `index_add`, and per-graph means.

docs/migration.md

@@ -0,0 +1,288 @@
+# Incremental migration plan: `root_gnn_dgl`
+
+The target is `legacy/root_gnn_dgl/`. `legacy/physicsnemo/` is a prior rewrite
+attempt and may inspire abstractions, but it is not a parity target. Neither
+legacy tree should be modified during migration.
+
+Each phase should add focused unit tests, a deterministic fixture in
+`data/fixtures/`, and parity tests under `tests/parity/` before moving upward.
+Record intentional differences and checkpoint consequences here.
+
+## Phase 0 — freeze observations and fixtures
+
+Capture a small representative ROOT-equivalent fixture containing the seven
+active node features, three edge features, labels, fold values, weights, and
+globals. Record outputs of `node_features_from_tree`, `full_connected_graph`,
+`EdgeDataset.make_graph`, and `fold_selection`
+([`dataset.py:15-59`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`dataset.py:471-482`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Preserve
+one `.bin`, one `model_epoch_N.pt`, one evaluation `.npz`, and one inference
+`.npz` fixture if available.
+
+Task 3 characterization records the active-path observations. Node rows are
+concatenated by object type in the configured order (jets, electrons, muons,
+photons, MET), and the seven columns are `[pt, eta, phi, energy, btag, charge,
+node_type]`. `CALC_E` is `pt*cosh(eta)` before the configured column scale is
+applied. The graph is directed and uses all ordered pairs except self-loops
+for graphs with more than one node; edge order is source-major. A one-node
+graph is a special case: the no-self-loop branch retains its sole self-loop.
+Edge columns are `[deta, dphi, dR]`, with `dphi` wrapped into `[-pi, pi]`.
+Dataset items expose `(graph, label, tracking, global_features)`; tracking
+column 0 is the fold identifier and column 1 is the event weight. These are
+compatibility observations, not proposed fixes.
+
+## Phase 1 — configuration boundary
+
+Implement a typed configuration layer that reads `Training`, `Model`,
+optional `Loss`, and `Datasets`. Initially retain a compatibility adapter for
+`module`/`class`/`args` and runtime injection of `sample_graph` and
+`sample_global`, matching `buildFromConfig`
+([`utils.py:10-43`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Keep
+dynamic imports isolated at this boundary rather than spreading reflection
+through new code.
+
+## Phase 2 — pure preprocessing parity
+
+Port and test, in isolation:
+
+- branch-to-node conversion, `CALC_E`, `NODE_TYPE`, constants, scaling, empty
+  objects, and dtypes (`node_features_from_tree`,
+  [`dataset.py:15-50`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
+- string/tuple selections and cutflow (`check_selection`, `selection_mask`,
+  `compute_cutflow`, [`dataset.py:75-158`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
+- fold masks and cache suffixes (`fold_selection`, `fold_selection_name`,
+  [`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py));
+- deterministic chunk partitioning (`hash_partition`,
+  [`batched_dataset.py:27-31`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+
+This is the highest-value parity layer: model parity is invalid if graph inputs
+differ.
+
+Task 4 implements the shared branch-to-node feature builder under
+`gnn4colliders.features`. It preserves the active seven-column schema,
+object-type ordering, explicit scales, derived `CALC_E`, node-type codes,
+float32 output, and supported empty vector collections. Selection, fold, and
+chunk helpers remain deferred to later data-infrastructure work.
+
+Task 6 implements the shared ROOT/Awkward ingestion boundary under
+`gnn4colliders.data`. `RootEventDataset` returns immutable, architecture-neutral
+`EventSample` values with selected branch data, labels, tracking, and globals;
+events are ordered by input file order with a global zero-based index. Fold
+filtering, caching, batching, and model-specific conversion remain deferred.
+
+## Phase 3 — graph construction and cache format
+
+Implement graph construction with tests for node/edge counts, directed edge
+ordering, self-loop policy, `[deta, dphi, dR]` order, metadata, and empty graphs.
+Preserve the dataset item contract `(graph, label, tracking, global_features)`
+from `RootDataset.__getitem__`
+([`dataset.py:465-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
+
+Then implement DGL `.bin` serialization, lazy chunk loading, pre-batching, and
+padding. Compare against `RootDataset.save/load`, `LazyDataset`, and
+`PreBatchedDataset` ([`dataset.py:396-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
+[`batched_dataset.py:129-174`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
+Treat `NONE`, `STEPS`, `FIXED`, and `NODE` as explicit features; do not hide
+the hardcoded fixed padding sizes.
+
+Task 7 establishes the metadata-aware orchestration boundary around this
+phase: `EventMetadata`, `GraphSample`, `GraphBatch`, fold-based split
+selection, deterministic batching, and a version-checked graph-sample cache.
+The cache is deliberately Level 2; normalized event caching remains a future
+extension so non-graph model families can reuse ROOT preprocessing.
+
+## Phase 4 — active model parity
+
+Task 8 adds the active `EdgeNetwork` and `FineTunedEdgeNetwork` under
+`gnn4colliders.models.root_gnn`. The update order and MLP ordering follow the
+legacy active path. The rewrite uses an explicit backbone/classifier boundary,
+local DGL graph scope, and does not mutate global RNG state in constructors.
+Model parity now covers fixed-weight pretraining and transfer paths, including
+historical checkpoint prefixes. The legacy transfer implementation has an
+active bug when nonempty globals are supplied (`Pretrained_Output` ignores its
+argument); parity therefore characterizes its supported no-global path, while
+the rewritten model supports both global and fallback modes.
+
+Port `models.GCN.Edge_Network` first. Preserve constructor parameters,
+`forward(graph, global_feats)`, feature keys, processor order, MLP LayerNorm
+placement, and logits shape. Compare intermediate and final tensors on fixed
+graphs using the legacy architecture
+([`GCN.py:18-35`](../legacy/root_gnn_dgl/models/GCN.py),
+[`GCN.py:182-251`](../legacy/root_gnn_dgl/models/GCN.py)).
+
+Next port `Transferred_Learning_Finetuning`, including pretrained
+`model_state_dict` loading, removal of the final classifier, and new classifier
+initialization ([`GCN.py:884-997`](../legacy/root_gnn_dgl/models/GCN.py)). Test
+both frozen and unfrozen modes. Defer other model classes until an active
+config or consumer proves they are needed.
+
+## Phase 5 — objectives and metrics
+
+Implement the default objective exactly: elementwise configured loss,
+tracking-column weights, per-unique-label normalization, and averaging across
+labels ([`training_script.py:320-359`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Add parity cases for positive, zero, and negative weights and binary versus
+multiclass shapes.
+
+Port metric behavior from
+[`training_script.py:438-510`](../legacy/root_gnn_dgl/scripts/training_script.py):
+sigmoid threshold 0.5, argmax, weight masking, weighted ROC AUC, one-vs-rest
+multiclass AUC, and NaN behavior when AUC is undefined. Add `models/loss.py`
+classes only with dedicated tests; do not substitute their reductions.
+
+## Phase 6 — checkpoint and lifecycle
+
+Task 10 implemented the in-memory single-process training lifecycle before the
+checkpoint portion of this phase: `Trainer`, explicit optimizer/scheduler
+builders, `EarlyStopping`, reproducibility seeding, `GraphBatch.to`, and
+epoch/history result types. Checkpoint persistence/resume and the Python
+inference/evaluation and named NPZ/ROOT output layers are now implemented.
+Distributed execution and CLI wiring were completed in the later phases.
+
+Task 10 also establishes corrected split semantics: validation is evaluated
+every epoch and is the only split used for model selection or early stopping;
+the test split remains held out and is evaluated separately after fitting. The
+legacy loader naming inversion (`test` used for selection and `val` held out)
+is not carried into the rewrite.
+
+Create a checkpoint adapter preserving `model_epoch_<epoch>.pt` and keys
+`epoch`, `model_state_dict`, `optimizer_state_dict`, and `early_stop`
+([`training_script.py:565-604`](../legacy/root_gnn_dgl/scripts/training_script.py)).
+Support legacy DDP/compiled prefixes (`module.` and `_orig_mod.`) as exercised
+by checkpoint lookup and inference
+([`utils.py:145-248`](../legacy/root_gnn_dgl/root_gnn_base/utils.py),
+[`inference.py:274-290`](../legacy/root_gnn_dgl/scripts/inference.py)). Port
+`EarlyStop` state and log parsing separately
+([`utils.py:325-390`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Verify
+resume, restart, early termination, and `.npz` fields before distributed work.
+
+## Phase 7 — CLI, inference, and export
+
+Task 12 implemented ordered prediction/evaluation, task-owned score
+semantics, checkpoint weight-only loading, named metadata retention, NPZ
+output, and explicit ROOT entry alignment. The semantic CLI and the validated
+ROOT-GNN ONNX export adapter are implemented.
+
+Task 13 adds Hydra composition and a single-process CLI around those existing
+APIs. The current application data boundary is a versioned
+`GraphSampleCache`; ROOT preparation converts events through the shared
+feature and graph builders before writing that cache.
+
+Build thin new applications around tested library interfaces in this order:
+
+1. preprocessing/cache generation (`scripts/prep_data.py`);
+2. training/evaluation (`scripts/training_script.py`);
+3. inference to `.npz` and ROOT (`scripts/inference.py`);
+4. ONNX export after PyTorch parity (`gnn4colliders export`).
+
+Use subprocess integration tests with tiny fixtures. Preserve CLI options only
+where they serve an active workflow; document removed diagnostic/cluster-only
+options.
+
+## Phase 8 — reproducibility and deployment
+
+Task 14 adds the initial deployment boundary: CPU/GPU DDP through standard
+`torchrun` variables, rank-local graph-sample sharding, global metric/output
+gathering, rank-0 checkpoint/config writing, and Perlmutter-oriented Slurm
+examples. Evaluation deliberately avoids sampler padding duplicates. The
+remaining follow-up is a streaming or sharded output path for very large
+distributed inference jobs.
+
+The seed policy remains explicit: the configured seed is offset by rank for
+process-local randomness, while distributed sample assignment is derived from
+the configured seed, world size, and epoch. GPU kernel nondeterminism and
+exact per-rank RNG checkpoint replay remain environment-dependent. Slurm/NCCL,
+Podman-HPC, ROOT, and Hugging Face integrations stay in launcher/adapters
+rather than package code.
+
+## Checkpoint compatibility checklist
+
+- [x] Load a checked-in or generated multiclass pretrained checkpoint.
+- [x] Load a legacy fine-tuning checkpoint after prefix normalization.
+- [x] Resume optimizer and early-stop state.
+- [x] Produce equivalent logits on a deterministic graph fixture.
+- [x] Produce equivalent `.npz` score, label, and metadata fields.
+- [x] Preserve ROOT scalar/vector score branch conventions in the Python adapter.
+
+Known risks are documented in [`architecture.md`](architecture.md): edge order,
+self-loops, weight semantics, validation/test naming, padding, dynamic
+selection evaluation, reproducibility, and the experimental model/loss surface.
+
+## Migration closure status
+
+### Task 18 compatibility closure
+
+The compatibility boundary is now explicit in `gnn4colliders.compat`.
+Production ingestion stores named `EventMetadata`; legacy two-column tracking
+is converted only at the compatibility boundary. Checkpoint prefix cleanup and
+the historical ROOT-GNN `classify` to `classifier` mapping have one canonical
+implementation. The new checkpoint schema and named NPZ output remain
+canonical. See [`compatibility.md`](compatibility.md) for the supported and
+intentionally unsupported historical artifacts.
+
+The following matrix describes the supported new stack, rather than every
+class that exists in `legacy/`:
+
+| Legacy area | New-stack status | Notes |
+| --- | --- | --- |
+| ROOT/Awkward ingestion | migrated | `RootEventDataset` returns `EventSample` in file/event order |
+| node features | migrated + parity-tested | seven-column schema, `CALC_E`, ordering, scales, float32 |
+| edge construction | migrated + parity-tested | directed source-major topology and `[deta,dphi,dR]` |
+| graph cache | migrated | versioned `GraphSampleCache`; graph-level cache only |
+| folds and weights | migrated | named `EventMetadata.fold` and `.weight` |
+| batching | migrated | deterministic local loader and DDP sharding |
+| legacy padding modes | deferred | no active new-stack consumer |
+| `Edge_Network` | migrated + parity-tested | `EdgeNetwork`, raw logits |
+| transfer/fine-tuning | migrated + parity-tested | frozen or trainable backbone |
+| loss and metrics | migrated + parity-tested | task-owned weighted reductions and full-split AUC |
+| training lifecycle | migrated | `Trainer`, validation semantics, scheduler, early stopping |
+| checkpoints/resume | migrated | schema v1; historical weight/prefix adapter |
+| inference/NPZ | migrated | named output fields and ordered accumulation |
+| ROOT score output | compatibility adapter | Python API supported; CLI currently NPZ-only |
+| DDP | migrated | torchrun boundary, rank-0 artifacts, gathered metrics |
+| Slurm/Perlmutter | launcher examples | site policy remains outside package code |
+| ONNX export | migrated for ROOT-GNN | tensor-only adapter, ONNX Runtime validation, and `export` CLI; raw graph tensors are the input contract |
+
+### Intentional redesigns
+
+These are deliberate new-stack contracts, not accidental parity failures:
+
+* `tracking[:, 0]` and `tracking[:, 1]` become named `metadata.fold` and
+  `metadata.weight`; public consumers do not depend on positional columns.
+* Dynamic legacy YAML `module`/`class` construction becomes allow-listed
+  semantic Hydra configuration.
+* The monolithic training script becomes `Task` + `Trainer` + checkpoint and
+  inference adapters.
+* Graph state is scoped to the forward pass rather than relying on persistent
+  mutation of shared graph state.
+* Model constructors do not mutate global RNG state; seeding is explicit in
+  the training/application boundary.
+* Validation is the selection/early-stopping split and test is held out. This
+  corrects the legacy loader-name inversion.
+
+Compatibility preserves externally observable scientific behavior where it is
+validated; it does not promise to preserve every legacy implementation bug.
+The characterized legacy transfer path had a nonempty-global handling bug;
+the rewrite supports named globals. Negative weights, rare empty graphs,
+historical checkpoint variants, and legacy padding edge cases remain areas to
+audit when a supported consumer requires them.
+
+## ROOT-GNN v1 completion checklist
+
+- [x] active ROOT data path and graph cache
+- [x] validated feature, graph, model, task, and metric behavior
+- [x] train from scratch and fine-tune a pretrained backbone
+- [x] resume new-stack checkpoints and load supported historical weights
+- [x] evaluate and predict named outputs
+- [x] single-process and DDP application boundaries
+- [x] Perlmutter/Slurm launcher examples and profiling guidance
+- [x] ROOT-GNN ONNX export and CPU Runtime parity
+- [ ] streaming/sharded large-scale prediction output
+- [ ] removal of frozen legacy reference
+- [ ] ROOT-Transformer representation/model
+
+ROOT-GNN v1 is complete when the checked-in new stack can prepare active data,
+reproduce validated legacy behavior, train, transfer, resume, evaluate,
+predict, and run single-process or DDP workflows. The remaining unchecked
+items are intentionally deferred rather than undocumented promises.

docs/performance.md

@@ -0,0 +1,20 @@
+# Performance notes
+
+The scripts in [`../benchmarks/`](../benchmarks/) separate setup, warmup, and
+steady-state timings and emit JSON lines with execution metadata. Float32 eager
+execution remains the correctness reference path.
+
+`fully_connected_edges` has a bounded 32-entry cache keyed by node count,
+self-loop policy, and device. It preserves source-major ordering and only
+reuses topology indices; event-dependent edge features are always recomputed.
+The cache is graph-specific and does not alter scientific behavior.
+
+Training already uses `zero_grad(set_to_none=True)` and inference already uses
+`torch.inference_mode()` with detached CPU accumulation.
+
+Mixed precision, `torch.compile`, custom kernels, aggressive worker defaults,
+and cache-format replacement were not retained without target-machine
+measurements. The main known bottleneck is the quadratic graph workload
+`N * (N - 1)` and associated DGL message passing; size-aware batching and
+streaming prediction remain follow-up work because they affect ordering or
+output semantics.

docs/perlmutter.md

@@ -0,0 +1,58 @@
+# Perlmutter execution
+
+The package does not encode Perlmutter paths, modules, or allocation policy.
+Create the uv environment in the project location appropriate for your
+account, select a site-compatible GPU/driver environment, and use the same
+Hydra configuration as on a workstation.
+
+## Single GPU
+
+```bash
+uv run gnn4colliders train \
+  data.cache.path=/path/to/graphs.pt \
+  environment=perlmutter \
+  trainer.device=cuda \
+  trainer.max_epochs=10
+```
+
+The equivalent Slurm wrapper is:
+
+```bash
+sbatch scripts/slurm/train_single_gpu.sh \
+  data.cache.path=/path/to/graphs.pt trainer.max_epochs=10
+```
+
+## Single-node DDP
+
+```bash
+GPUS_PER_NODE=4 sbatch scripts/slurm/train_multi_gpu.sh \
+  data.cache.path=/path/to/graphs.pt trainer.max_epochs=10
+```
+
+The wrapper uses `torchrun`; `batch_size` and `num_workers` are per GPU.
+Effective batch size is `data.batch_size * number_of_processes`, and only rank
+0 writes the shared checkpoint/config/prediction artifacts.
+
+## Multi-node DDP
+
+```bash
+GPUS_PER_NODE=4 sbatch scripts/slurm/train_multi_node.sh \
+  data.cache.path=/path/to/graphs.pt trainer.max_epochs=10
+```
+
+Use the provided script as a template and adapt only allocation/account
+settings required by the site. Evaluation and prediction can use
+`scripts/slurm/evaluate.sh`; prediction currently gathers moderate-size
+results in memory.
+
+## Checks and common failures
+
+```bash
+uv run python -c "import torch, dgl; print(torch.__version__, dgl.__version__, torch.cuda.is_available())"
+uv run gnn4colliders --help
+```
+
+An unavailable DGL wheel, incompatible driver, missing cache, or mismatched
+cache schema should be fixed in the environment/input rather than hidden with
+package-level path changes. GPU kernels and distributed execution are not
+promised bitwise deterministic.

docs/releasing.md

@@ -0,0 +1,23 @@
+# Releasing
+
+The project uses a static version defined once in
+`src/gnn4colliders/__init__.py`; setuptools reads that value for package
+metadata. The current package is not published to PyPI.
+
+Before a release:
+
+1. Update `__version__`, `CHANGELOG.md`, and this release checklist.
+2. Run `uv sync --dev` (and `--extra root-gnn` when ROOT-GNN validation is
+   available), then run lint and tests.
+3. Run `bash scripts/dev/check_release.sh` to build and inspect the sdist and
+   wheel and smoke-test a wheel installation outside the checkout.
+4. Review the generated artifacts and `git diff`, then create a version tag
+   according to the repository's release policy.
+
+The release-validation script never publishes artifacts. Publishing, if
+adopted later, must use repository-managed credentials or trusted publishing.
+
+The direct DGL wheel source is retained in `pyproject.toml` because the
+validated ROOT-GNN environment requires the CUDA 12.1 DGL wheel. Core imports
+do not import DGL or ONNX; install the corresponding extras for those
+workflows.

docs/testing.md

@@ -0,0 +1,29 @@
+# Testing
+
+The suite is layered by dependency and purpose:
+
+```bash
+# Fast shared-package tests
+uv run pytest tests/unit
+
+# Full CPU suite, including integration and parity tests
+uv run pytest
+
+# Required ROOT-GNN parity gate (must fail if DGL is unavailable)
+GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest
+
+# Optional layers
+uv run pytest -m distributed -v
+uv run pytest -m onnx -v
+uv run pytest -m gpu -v
+GNN4COLLIDERS_ROOT_FIXTURE=/path/to/reduced.root uv run pytest -m real_data -v
+```
+
+Unit tests use deterministic, small tensors and generated ROOT files. DGL,
+ONNX, CUDA, distributed execution, and reduced real-data fixtures remain
+optional layers. Tests should assert public contracts and scientific
+invariants rather than private call sequences. New regression tests should
+use `tmp_path`, explicit seeds, and justified numerical tolerances.
+
+Release validation additionally builds a wheel and runs the CLI help command
+in a clean environment; it is not part of the ordinary pytest suite.

examples/README.md

@@ -0,0 +1,18 @@
+# Examples
+
+The supported experiment interface is the root CLI with Hydra configuration;
+examples do not define a second application framework. The repository's
+reference groups cover multiclass pretraining, binary fine-tuning, resume,
+evaluation, and NPZ prediction. Start with the commands in the top-level
+[`README.md`](../README.md), then copy a config group into a local config file
+when a dataset needs more than command-line overrides.
+
+For a complete CPU smoke workflow using generated temporary ROOT data:
+
+```bash
+uv run python scripts/dev/smoke_end_to_end.py
+```
+
+Values such as input ROOT paths, cache locations, and checkpoint paths are
+deliberately local placeholders. Production datasets and generated outputs do
+not belong in this directory.

legacy/README.md

@@ -0,0 +1,358 @@
+---
+license: mit
+tags: 
+- arXiv:2412.10665
+---
+
+This is a demo is of the approach described in the paper, ["Pretrained Event Classification Model for High Energy Physics Analysis"](https://arxiv.org/abs/2412.10665)
+```
+@misc{ho2024pretrained,
+  title={Pretrained Event Classification Model for High Energy Physics Analysis},
+  author={Joshua Ho, Benjamin Ryan Roberts, Shuo Han, Haichen Wang},
+  year={2024},
+  eprint={2412.10665},
+  archivePrefix={arXiv}
+}
+```
+
+## Abstract
+
+We introduce a foundation model for event classification in high-energy physics, built on a **Graph Neural Network** architecture and trained on **120 million simulated proton-proton collision events** spanning 12 distinct physics processes. The model is *pretrained* to learn a general and robust representation of collision data using challenging multiclass and multilabel classification tasks.
+
+Its performance is evaluated across five event classification tasks, which include both physics processes used during pretraining and new processes not encountered during pretraining. Fine-tuning the pretrained model significantly improves classification performance, particularly in scenarios with limited training data, demonstrating gains in both accuracy and computational efficiency.
+
+To investigate the underlying mechanisms behind these performance improvements, we employ a representational similarity evaluation framework based on *Centered Kernel Alignment*. This analysis reveals notable differences in the learned representations of fine-tuned pretrained models compared to baseline models trained from scratch.
+
+## Introduction
+
+Machine learning has become a ubiquitous tool in particle physics, employed in a variety of tasks including triggering, simulation, reconstruction, and offline analysis. While its utility spans classification, regression, and generative tasks, the current paradigm of developing machine learning models from scratch for each specific application presents several challenges. This approach not only demands specialized expertise and substantial computing resources but can also result in suboptimal performance due to limited training data. The from-scratch development of models necessitates individual validation studies to ensure that neural networks utilize well-modeled information from training samples, whether derived from Monte Carlo simulations or control samples from experimental data.
+
+Foundation models offer a promising direction to address these limitations. These models, pre-trained on large, diverse datasets across various tasks, provide robust and general representations of underlying data structures. Notable examples in other fields include GPT-4 [OpenAI et al., 2024](#ref-openai-2024-gpt4) and BERT [Devlin et al., 2018](#ref-devlin-2018-bert) in natural language processing, Stable Diffusion [Rombach et al., 2021](#ref-rombach-2021-latentdiffusion) in image processing, and AlphaFold [Jumper et al., 2021](#ref-jumper-2021-alphafold) in structural biology. The foundation model approach offers several advantages for particle physics applications: reduced computing resources for fine-tuning [Yosinski et al., 2014](#ref-yosinski-2014-transfer) compared to training from scratch, superior performance on specific tasks (particularly with limited training data), and potentially simplified validation procedures as downstream tasks inherit verified representations from the pre-trained model.
+
+Current literature on pretrained models for particle physics can be categorized based on the data representation they handle. Models operating on particle- or event-level numerical data use features like particle four momenta or jets, leveraging self-supervised or generative methods to learn versatile representations. Detector-focused models operate on high-dimensional responses such as calorimeter deposits or pixel hits, employing geometry-aware techniques for accurate simulation and analysis. Finally, models using textual or code representations apply large language model architectures to integrate domain knowledge, enabling tasks like question answering and code generation.
+
+Recent studies have begun exploring foundation models tailored to particle physics data, which has a variety of distinct structures and properties across many experiments and data processing stages, including:
+
+- particle-level & event-level numeric data [Wildridge et al., 2024](#ref-wildridge-2024-bumblebee), [Katel et al., 2024](#ref-katel-2024-jet), [Golling et al., 2024](#ref-golling-2024-maskedset), [Mikuni & Nachman, 2024](#ref-mikuni-2024-omnilearn), [Harris et al., 2024](#ref-harris-2024-resimulation), [Birk et al., 2024](#ref-birk-2024-omnijet), [Vigl et al., 2024](#ref-vigl-2024-finetune),
+- detector-level & geometry-aware data [Araz et al., 2024](#ref-araz-2024-pointcloud), [Liu et al., 2023](#ref-liu-2023-gaam), [Hashemi et al., 2024](#ref-hashemi-2024-gen), [Huang et al., 2024](#ref-huang-2024-lmtracking),
+- textual or code data [Zhang et al., 2024](#ref-zhang-2024-xiwu).
+
+This paper presents a foundation model designed specifically for collider event-level data. In modern collider experiments, final-stage analysis processes information from reconstructed objects that either directly correspond to particles in collision final states (such as leptons and photons) or serve as proxies (such as jets and missing transverse energy). While traditional approaches often relied on "high-level" variables calculated from object features, recent trends favor direct input of event objects and their features into neural networks for analysis tasks. A notable example is [ATLAS Collaboration, 2023](#ref-atlas-2023-4top), which established the observation of simultaneous production of four top quarks with the ATLAS experiment by employing a graph neural network (GNN) architecture to process event-level object information.
+
+We present foundation models that adopt an architecture similar to that used for [ATLAS Collaboration, 2023](#ref-atlas-2023-4top). Our models are pre-trained using either multiclass classification or multi-label learning tasks across 12 distinct physics processes. We evaluate these models through fine-tuning and testing on five classification tasks, including both familiar and novel processes not seen during pre-training. Our analysis benchmarks the models' performance improvements, their scaling behavior with training sample size, and computational efficiency, representing the first prototype of a foundation model operating on collider final-state object data.
+
+## Data Samples
+
+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
+
+### Overview
+
+We present a methodology for developing and evaluating a foundation model for particle collision event analysis. The approach centers on pretraining a Graph Neural Network (GNN) architecture using a comprehensive dataset that spans multiple physics tasks, enabling the model to learn robust and transferable features. For task-specific applications, we employ a fine-tuning strategy that combines output layer adaptation with carefully calibrated learning rates for updating the pretrained parameters.
+
+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.
+
+This methodological framework demonstrates the potential of foundation models to enhance the efficiency of particle physics analyses while improving task-specific performance, offering a promising direction for future high-energy physics research.
+
+---
+
+### 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.
+
+$$
+\vec {y'}_{ij} = f_e\left(\{\vec x_k\},\vec y_{ij},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\vec x_j,\vec y_{ij},\vec z\right)
+$$
+
+$$
+\vec{x'}_{i} = f_n\left(\vec x_i,\{\vec{y'}_{jk}\},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\sum_j\vec{y'}_{ij},\vec z\right)
+$$
+
+$$
+\vec{z'} = f_g\left(\{\vec{x'}_i\},\{\vec{y'}_{ij}\},\vec z\right) = \mathrm{MLP}\left(\sum_i\vec{x'}_i,\sum_{i,j}\vec{y'}_{ij},\vec z\right)
+$$
+
+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.
+
+---
+
+### Pretraining Strategy
+
+We explore two complementary pretraining approaches to develop robust representations of collision events: (1) multi-class classification, which trains the model to distinguish between different physics processes, and (2) multi-label classification, which predicts the existence and kinematics of heavy particles with prompt decays. The pretraining dataset consists of approximately 120 million events, evenly distributed across 12 distinct physics processes, including all major Higgs boson production mechanisms and top quark processes as described in [Data Samples](#sec-data). This large-scale pretraining effort was conducted on the Perlmutter supercomputer at NERSC.
+
+#### Multi-class Classification
+
+For Monte Carlo simulated events, the underlying physics process that generated each event is known precisely, providing natural labels for supervised learning. However, the challenge lies in the complexity of collision events: different physics processes can produce similar kinematics and event topologies, particularly in certain regions of phase space. No single observable can unambiguously identify the underlying process. By training the model to distinguish between 12 different processes simultaneously, we challenge it to learn subtle differences in kinematics and topology that collectively characterize each process. The model is trained using categorical cross entropy as the loss function. The output layer of the multiclass classification model has 832 trainable parameters.
+
+#### Multi-label Classification
+
+This approach combines both classification and regression tasks to characterize collision events. For discrete properties like particle presence in specific kinematic regions, we employ classification labels with binary cross-entropy loss. For continuous quantities like particle multiplicities, we use regression labels with mean-squared error loss. This hybrid approach enables the model to learn both categorical and continuous aspects of the physics processes simultaneously.
+
+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.
+
+---
+
+### Performance Evaluation
+
+We assess model performance using two figures of merit: the classification accuracy and the Area Under the Curve (AUC) of the Receiver Operating Characteristic (ROC) curve. The accuracy is defined as the fraction of correctly classified events when applying a threshold of 0.5 to the neural network output score. Both metrics demonstrate consistent trends in our analysis.
+
+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
+
+### Classification Performance
+
+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.
+
+Finally, the spread of accuracy across the five tasks for the baseline model is quite large, offering a robust test of fine-tuning across tasks of varying difficulty. The consistent observation of these trends across tasks confirms the reliability and robustness of the findings.
+
+---
+
+### Model Interpretability
+
+We aim to understand whether pretrained and baseline models learn the same underlying representations. If the two models exhibit high similarity, a plausible interpretation is that pretraining provides the pretrained model with an advantageous initialization, allowing it to converge to a similar state as the baseline model more efficiently. Conversely, significant differences between the models would indicate that pretraining facilitates the development of a more general and robust latent space, which serves as a foundation for fine-tuning to effectively adapt to the downstream task. To investigate this, we analyzed the representational similarity between a pretrained model fine-tuned for the downstream task and a baseline model trained directly on the downstream task without pretraining.
+
+We use Centered Kernel Alignment (CKA) [Kornblith et al., 2019][ref-kornblith-2019-cka] to analyze model similarity and interpretability. CKA is a robust metric that quantifies the similarity between the internal representations of neural networks by comparing their feature matrices in a manner that is invariant to scaling, rotation, and alignment. This invariance makes CKA particularly effective for studying relationships between network layers, even across networks of different sizes or those trained from varying initializations.
+
+The similarity is evaluated using a 64-dimensional latent representation after the decoder stage of the GNN model. This choice allows us to compare the internal states of the models at a fine-grained level and understand how training strategies impact the representations directly used for the output task.
+
+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.
+
+### Computational Efficiency
+
+To estimate the computational resources required for each approach, we measured the wall time needed for a model to reach its final performance. For baseline models, this is defined as the wall time from the start of training until the loss of the model plateaus. For the foundation model approach, the estimate includes both the pretraining time and the fine-tuning time, each measured from the start of training until the loss plateaus. This approach ensures a consistent and comprehensive evaluation of the computational demands.
+
+![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:
+
+- **Multi-class pretraining:** 45.5 GPU hours
+- **Multi-label pretraining:** 60.0 GPU hours
+
+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
+
+We presented an in-depth study of a particle physics foundation model designed to operate on the four-momentum and identification properties of event final-state objects. This model is built on a Graph Neural Network (GNN) architecture and trained on a dataset comprising 120 million simulated proton-proton collision events across 12 distinct physics processes. The pretraining phase explored both multiclass and multilabel classification tasks, providing a robust foundation for downstream applications. Notably, the pretrained models demonstrated significant improvements in event classification performance when fine-tuned, particularly for tasks with limited training samples.
+
+The foundation model approach also offers substantial computational advantages. By leveraging fine-tuning, this methodology reduces the computational resources required for large-scale applications across multiple tasks. Our estimates indicate that significant resource savings can be achieved even for single particle physics measurements, making this approach both scalable and efficient.
+
+To better understand the learned representations of the pretrained model and guide future optimization efforts, we employed a representational similarity evaluation framework using Centered Kernel Alignment (CKA). This metric allowed us to investigate the source of the performance gains observed in the foundation model. Our analysis revealed notable differences in the learned representations between the fine-tuned pretrained model and a baseline model trained from scratch. In deep learning, it is well-established that multiple equally valid solutions can exist. Future studies are necessary to determine whether the low similarity in latent representations reflects complementary information uniquely captured by the foundation and baseline models, or if it can simply be attributed to connected local minima in the loss landscape.
+
+## Acknowledgments
+
+This work is supported by the U.S. National Science Foundation under the Award No. 2046280, and by U.S. Department of Energy, Office of Science under contract DE-AC02-05CH11231.
+
+## References
+
+- <span id="ref-openai-2024-gpt4"></span> **OpenAI et al.** GPT-4 Technical Report. arXiv:2303.08774 (2024). [https://arxiv.org/abs/2303.08774](https://arxiv.org/abs/2303.08774)
+
+- <span id="ref-yosinski-2014-transfer"></span> **Jason Yosinski, Jeff Clune, Yoshua Bengio, Hod Lipson.** How transferable are features in deep neural networks? CoRR abs/1411.1792 (2014). [http://arxiv.org/abs/1411.1792](http://arxiv.org/abs/1411.1792)
+
+- <span id="ref-rombach-2021-latentdiffusion"></span> **Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer.** High-Resolution Image Synthesis with Latent Diffusion Models. CoRR abs/2112.10752 (2021). [https://arxiv.org/abs/2112.10752](https://arxiv.org/abs/2112.10752)
+
+- <span id="ref-podell-2023-sdxl"></span> **Dustin Podell, Zion English, Kyle Lacey et al.** SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis. arXiv:2307.01952 (2023). [https://arxiv.org/abs/2307.01952](https://arxiv.org/abs/2307.01952)
+
+- <span id="ref-jumper-2021-alphafold"></span> **John Jumper, Richard Evans, Alexander Pritzel et al.** Highly accurate protein structure prediction with AlphaFold. Nature 596, 583-589 (2021). [https://doi.org/10.1038/s41586-021-03819-2](https://doi.org/10.1038/s41586-021-03819-2)
+
+- <span id="ref-devlin-2018-bert"></span> **Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova.** BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. CoRR abs/1810.04805 (2018). [http://arxiv.org/abs/1810.04805](http://arxiv.org/abs/1810.04805)
+
+- <span id="ref-atlas-2023-higg"></span> **ATLAS Collaboration.** Measurement of the properties of Higgs boson production at \\(\sqrt{s} = 13\,\text{TeV}\\) in the \\(H \to \gamma\gamma\\) channel using \\(139\,\text{fb}^{-1}\\) of \\(pp\\) collision data with the ATLAS experiment. JHEP 07 (2023) 088. [arXiv:2207.00348](https://arxiv.org/abs/2207.00348), [https://doi.org/10.1007/JHEP07(2023)088](https://doi.org/10.1007/JHEP07(2023)088)
+
+- <span id="ref-atlas-2023-4top"></span> **ATLAS Collaboration.** Observation of four-top-quark production in the multilepton final state with the ATLAS detector. Eur. Phys. J. C 83 (2023) 496. [arXiv:2303.15061](https://arxiv.org/abs/2303.15061), [https://doi.org/10.1140/epjc/s10052-023-11573-0](https://doi.org/10.1140/epjc/s10052-023-11573-0)
+
+- <span id="ref-kornblith-2019-cka"></span> **Simon Kornblith, Mohammad Norouzi, Honglak Lee, Geoffrey Hinton.** Similarity of Neural Network Representations Revisited. CoRR abs/1905.00414 (2019). [http://arxiv.org/abs/1905.00414](http://arxiv.org/abs/1905.00414)
+
+---
+
+<!-- Historical/General Physics foundational texts -->
+
+- <span id="ref-birell-1982-qfields"></span> **N. D. Birell, P. C. W. Davies.** Quantum Fields in Curved Space. Cambridge Univ. Press (1982).
+
+- <span id="ref-feynman-1954"></span> **R. P. Feynman.** Phys. Rev. 94, 262 (1954).
+
+- <span id="ref-einstein-1935-epr"></span> **A. Einstein, Yu. Podolsky, N. Rosen.** Phys. Rev. 47, 777 (1935).
+
+- <span id="ref-berman-1983-stability"></span> **G. P. Berman, Jr., F. M. Izrailev, Jr.** Stability of nonlinear modes. Physica D 88, 445 (1983).
+
+- <span id="ref-davies-1988-trapped"></span> **E. B. Davies, L. Parns.** Trapped modes in acoustic waveguides. Q. J. Mech. Appl. Math. 51, 477–492 (1988).
+
+- <span id="ref-witten-2001"></span> **Edward Witten.** hep-th/0106109 (2001). [https://arxiv.org/abs/hep-th/0106109](https://arxiv.org/abs/hep-th/0106109)
+
+---
+
+<!-- Particle physics/data science foundational models -->
+
+- <span id="ref-beutler-1994-hem"></span> **E. Beutler.** Williams Hematology, 5th Edition, Chapter 7, pp. 654–662. McGraw-Hill, New York (1994).
+
+- <span id="ref-knuth-1973-fa"></span> **Donald E. Knuth.** The Art of Computer Programming vol. 1: Fundamental Algorithms, 2nd Ed., Addison-Wesley (1973).
+
+- <span id="ref-smith-2005-philos"></span> **J. S. Smith, G. W. Johnson.** Philos. Trans. R. Soc. London, Ser. B 777, 1395 (2005).
+
+- <span id="ref-smith-2010-jap-unpub"></span> **W. J. Smith, T. J. Johnson, B. G. Miller.** Surface chemistry and preferential crystal orientation on a silicon surface. J. Appl. Phys. (unpublished, 2010).
+
+- <span id="ref-smith-2010-jap-sub"></span> **V. K. Smith, K. Johnson, M. O. Klein.** Surface chemistry and preferential crystal orientation on a silicon surface. J. Appl. Phys. (submitted, 2010).
+
+- <span id="ref-underwood-1988-lowerbounds"></span> **Ulrich Underwood, Ned Net, Paul Pot.** Lower Bounds for Wishful Research Results. Talk at Fanstord University (1988).
+
+- <span id="ref-johnson-2007-comm"></span> **M. P. Johnson, K. L. Miller, K. Smith.** Personal communication (Jan-May 2007).
+
+---
+
+<!-- Prototypical collider software and tools -->
+
+- <span id="ref-pytorch-2019"></span> **Adam Paszke et al.** PyTorch: An Imperative Style, High-Performance Deep Learning Library. arXiv:1912.01703 (2019). [http://arxiv.org/abs/1912.01703](http://arxiv.org/abs/1912.01703)
+
+- <span id="ref-dgl-2019"></span> **Minjie Wang et al.** Deep Graph Library: Towards Efficient and Scalable Deep Learning on Graphs. arXiv:1909.01315 (2019). [http://arxiv.org/abs/1909.01315](http://arxiv.org/abs/1909.01315)
+
+- <span id="ref-graphnets-2018"></span> **Peter W. Battaglia et al.** Relational inductive biases, deep learning, and graph networks. arXiv:1806.01261 (2018). [http://arxiv.org/abs/1806.01261](http://arxiv.org/abs/1806.01261)
+
+- <span id="ref-layernorm-2016"></span> **Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton.** Layer Normalization. arXiv:1607.06450 (2016). [https://arxiv.org/abs/1607.06450](https://arxiv.org/abs/1607.06450)
+
+---
+
+<!-- Recent & foundation models in HEP ML -->
+
+- <span id="ref-wildridge-2024-bumblebee"></span> **Andrew J. Wildridge et al.** Bumblebee: Foundation Model for Particle Physics Discovery. arXiv:2412.07867 (2024). [https://arxiv.org/abs/2412.07867](https://arxiv.org/abs/2412.07867)
+
+- <span id="ref-katel-2024-jet"></span> **Subash Katel et al.** Learning Symmetry-Independent Jet Representations via Jet-Based Joint Embedding Predictive Architecture. arXiv:2412.05333 (2024). [https://arxiv.org/abs/2412.05333](https://arxiv.org/abs/2412.05333)
+
+- <span id="ref-araz-2024-pointcloud"></span> **Jack Y. Araz et al.** Point cloud-based diffusion models for the Electron-Ion Collider. arXiv:2410.22421 (2024). [https://arxiv.org/abs/2410.22421](https://arxiv.org/abs/2410.22421)
+
+- <span id="ref-leigh-2024-maskedparticle"></span> **Matthew Leigh et al.** Is Tokenization Needed for Masked Particle Modelling? arXiv:2409.12589 (2024). [https://arxiv.org/abs/2409.12589](https://arxiv.org/abs/2409.12589)
+
+- <span id="ref-mikuni-2024-omnilearn"></span> **Vinicius Mikuni, Benjamin Nachman.** OmniLearn: A Method to Simultaneously Facilitate All Jet Physics Tasks. arXiv:2404.16091 (2024). [https://arxiv.org/abs/2404.16091](https://arxiv.org/abs/2404.16091)
+
+- <span id="ref-zhang-2024-xiwu"></span> **Zhengde Zhang et al.** Xiwu: A Basis Flexible and Learnable LLM for High Energy Physics. arXiv:2404.08001 (2024). [https://arxiv.org/abs/2404.08001](https://arxiv.org/abs/2404.08001)
+
+- <span id="ref-harris-2024-resimulation"></span> **Philip Harris et al.** Re-Simulation-based Self-Supervised Learning for Pre-Training Foundation Models. arXiv:2403.07066 (2024). [https://arxiv.org/abs/2403.07066](https://arxiv.org/abs/2403.07066)
+
+- <span id="ref-birk-2024-omnijet"></span> **Joschka Birk, Anna Hallin, Gregor Kasieczka.** OmniJet-$\alpha$: the first cross-task foundation model for particle physics. Machine Learning: Science and Technology. 5(3), 035031 (Aug 2024). [https://doi.org/10.1088/2632-2153/ad66ad](https://doi.org/10.1088/2632-2153/ad66ad)
+
+- <span id="ref-huang-2024-lmtracking"></span> **Andris Huang et al.** A Language Model for Particle Tracking. arXiv:2402.10239 (2024). [https://arxiv.org/abs/2402.10239](https://arxiv.org/abs/2402.10239)
+
+- <span id="ref-golling-2024-maskedset"></span> **Tobias Golling et al.** Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models. arXiv:2401.13537 (2024). [https://arxiv.org/abs/2401.13537](https://arxiv.org/abs/2401.13537)
+
+- <span id="ref-liu-2023-gaam"></span> **Junze Liu et al.** Generalizing to new geometries with Geometry-Aware Autoregressive Models (GAAMs) for fast calorimeter simulation. Journal of Instrumentation 18(11), P11003 (Nov 2023). [https://doi.org/10.1088/1748-0221/18/11/p11003](https://doi.org/10.1088/1748-0221/18/11/p11003)
+
+- <span id="ref-hashemi-2024-gen"></span> **Baran Hashemi et al.** Ultra-high-granularity detector simulation with intra-event aware generative adversarial network and self-supervised relational reasoning. Nature Communications 15(1) (June 2024). [https://doi.org/10.1038/s41467-024-49104-4](https://doi.org/10.1038/s41467-024-49104-4)
+
+- <span id="ref-vigl-2024-finetune"></span> **Matthias Vigl et al.** Finetuning Foundation Models for Joint Analysis Optimization. arXiv:2401.13536 (2024). [https://arxiv.org/abs/2401.13536](https://arxiv.org/abs/2401.13536)
+
+- <span id="ref-li-2024-refine"></span> **Chen Li, Hao Cai, Xianyang Jiang.** Refine neutrino events reconstruction with BEiT-3. Journal of Instrumentation 19(6), T06003 (Jun 2024). [https://doi.org/10.1088/1748-0221/19/06/t06003](https://doi.org/10.1088/1748-0221/19/06/t06003)