OmniLearn: A Method to Simultaneously Facilitate All Jet Physics Tasks
Vinicius Mikuni Benjamin Nachman





vmikuni@lbl.gov
National Energy Research Scientific Computing Center, Berkeley Lab, Berkeley, CA 94720, USA


bpnachman@lbl.gov
Physics Division, Lawrence Berkeley National Laboratory, Berkeley, CA 94720, USA
Berkeley Institute for Data Science, University of California, Berkeley, CA 94720, USA


    Machine learning has become an essential tool in jet physics.  Due to their complex, high-dimensional nature, jets can be explored holistically by neural networks in ways that are not possible manually.  However, innovations in all areas of jet physics are proceeding in parallel.  We show that specially constructed machine learning models trained for a specific jet classification task can improve the accuracy, precision, or speed of all other jet physics tasks.  This is demonstrated by training on a particular multiclass classification task and then using the learned representation for different classification tasks, for datasets with a different (full) detector simulation, for jets from a different collision system (pp versus ep), for generative models, for likelihood ratio estimation, and for anomaly detection.  Our OmniLearn approach is thus a foundation model and is made publicly available for use in any area where state-of-the-art precision is required for analyses involving jets and their substructure.


OmniLearn: A Method to Simultaneously Facilitate All Jet Physics Tasks
    Benjamin Nachman
    25 April 2024
======================================================================










§ INTRODUCTION



The study of high-energy hadronic final states – jet physics – is seeing a paradigm shift from modern machine learning (ML).  Jets are composed of many particles, each with properties of their own.  This means that jets are represented in high-dimensional spaces and are thus difficult to analyze manually.  There has been incredible progress over the last two decades in classical jet physics to develop observables and other techniques using direct physics reasoning <cit.>, but the deep learning revolution of the last few years has shown that automation and indirect physics reasoning (e.g. through simulations and general physics considerations) can significantly improve performance in many tasks.  For example, classifying a jet as originating from a top quark or a generic quark/gluon jet is over 20 times more effective with the latest deep learning solutions compared with classical methods <cit.>.  Deep learning has also enabled new studies that were unimaginable before, like unbinned differential cross sections in tens or hundreds of dimensions simultaneously <cit.>. 
However, one feature these innovations have in common is that they are all advancing in parallel.  Our question is simple: can we make progress on all jet physics tasks at the same time?

One answer to this question has been inspired by recent progress in large language modeling (LLMs).  Tools like ChatGPT (<https://chat.openai.com>) and others <cit.> are called foundation models because they are able to approach many downstream tasks either with little or no finetuning to the specific problem.  They have many millions (or billions/trillions) parameters and are trained using many millions (or much more) of examples.  

Recently, LLMs have also been adapted for particle physics-specific queries <cit.>.  Such tools complement our approach: they target the process of particle physics research while we focus on processing particle physics data.

Foundation models are usually trained using a form of self-supervision (e.g. mask data and learn to fill in the blanks) in order to learn to represent the structure of data.  This sort of representation learning has been studied recently in particle physics using a number of interesting approaches <cit.>.  As desired, the learned representations are also demonstrated to improve downstream tasks. However, the way these models are trained is not aligned with any actual analysis goal (we do not actually want to fill in blanks).  Our intuition is that foundation models are useful because they increase the effective size of the training dataset for a downstream task.  The closer the foundation model training is to the downstream task, the larger the increase in the effective size for a fixed sample size used to train the foundation model.  A key difference between foundation models in (particle) physics and foundation models in society at large is the existence of ab initio simulations.  These simulations provide large datasets that can be used to target specific tasks.  Since we are in a privileged situation in which machine learning models can be constructed for dedicated problems, can such a ML model dedicated to a specific jet physics task act as a foundation model?  

The process of (pre)training an ML model for one task and then applying it elsewhere, usually with some fine-tuning to the downstream task, is called transfer learning.  This strategy has been shown to be effective in a number of applications across particle physics <cit.>.  In all of these cases, there was only one downstream task.  Furthermore, either the pretraining was very far from the target task (e.g. pretraining with generic non-physics images) or very close to the target task (e.g. classify jet of type A and transfer to type B).  These papers were not trying to build foundation models, but we hypothesize that a significantly scaled up version of the transfer learning task can form the basis for a foundation model in particle physics. As such, our goal is to see if we can build a foundation model for jet physics by using a supervised instead of self-supervised learning.  Our neural network will have many millions of parameters and will be trained on 100 million jets.  We will achieve success if the implicit representation learned by our model improves (in training speed, accuracy, or both) essentially all downstream tasks in jet physics.  

An approach with a similar philosophy to ours is OmniJet-α <cit.>.  The authors of OmniJet-α use a generative model trained on a specific jet physics task in order to build a foundation model.  A key difference is that OmniJet-α is based on a language model and so jet constituents are discretized (into `tokens') and the generation is autoregressive (like generating a sentence from left to right).  On the application side, OmniJet-α was trained on one generative task and applied on one classification task using the same dataset, whereas we try to explore most areas of jet physics across a variety of datasets with our model. 

Our paper is organized as follows.  Section <ref> describes how jets can be represented as point clouds and the neural network architecture at the center of our foundation model. Just how OmniFold <cit.> can unfold all dimensions simultaneously, our approach – called OmniLearn – can learn useful representations for all jet physics tasks simultaneously.  This is demonstrated in subsequent sections by showing that our model can enhance and/or accelerate tasks other than the one it was trained on (multi-class jets from fast simulation) to binary classification on a different dataset (Sec. <ref>), to binary classification using a full detector simulation (Sec. <ref>), to jets originating from a different collision system (Sec. <ref>), to generative models (Sec. <ref>), to likelihood ratio estimation (Sec. <ref>), and to anomaly detection (Sec. <ref>).  The paper ends with conclusions and outlook in Sec. <ref>.



§ POINT CLOUDS FOR JET PHYSICS





The interpretation of jets as point clouds motivated the development of new neural network architectures to naturally address the data structure represented by an unordered set of particles with varying number of constituents. Successful neural network models using sets <cit.>, graph neural networks (GNNs) <cit.>, and more recently transformers <cit.>, were all able to successfully improve upon previous deep learning approaches in collider physics. The first transformer model in high energy physics <cit.> used graph-attention networks <cit.>, combining the advantages of both GNNs and transformers. The rise in popularity of transformer models is also attributed to their strong scaling properties <cit.>, making it the backbone of almost all modern foundational models and the choice of neural network architecture used to build OmniLearn.

For collider physics applications, we expect that in the presence of large datasets and suitable tasks, a foundational model for jets will be able to leverage the information to learn a general representation of jets, thus providing a stronger inductive bias that will be transferable to different datasets and tasks. We also notice that the overwhelming majority of machine learning applications for jet physics can be summarized as a classification or generation tasks. These include the applications in jet tagging and jet generation as well as complex tasks such as event reweighting, unfolding, and anomaly detection. Motivated by this observation, we conjecture that a flexible model, capable of learning both to generate and classify jets, will also learn a useful and general data representation that can be employed to quickly adapt to new downstream datasets and applications. 

While multiple generative models have been proposed for point cloud generation in particle physics based on generative adversarial networks <cit.>, variational autoencoders <cit.>, normalizing flows <cit.>, and auto-regressive transformers <cit.>,  we choose to focus on diffusion generative models <cit.>. Diffusion generative models use time-dependent perturbations applied to the data to learn an approximation of the score function, or gradients of the logarithm of the probability density of the data. This choice also aligns with the classification task. Since the perturbation process is designed such that no perturbation is applied at the initial time t=0, conditioning the model over the time parameter ensures the network is able to accommodate both perturbed and unperturbed data simultaneously. To build OmniLearn, we design the architecture with a shared representation whose outputs are then fed to task-specific neural networks. This approach enables flexibility and efficient design, since downstream task applications only need to load the shared representation and relevant task-specific network, reducing the overall model. The main building blocks of the network are summarized in Fig. <ref>. In the following subsections we will provide a detailed description of the model and the core design choices.



 §.§ Point-Edge Transformer



The shared representation of the network takes as inputs the particles clustered inside the jets and is conditioned on the diffusion time parameter. The time information, following previous diffusion models for collider physics <cit.>, is encoded to a higher dimensional space using a time embedding layer. The time embedding consists Fourier features <cit.> followed by two multi-layer perceptrons (MLPs) with GELU activation function <cit.>. Unless otherwise stated, all MLP layers used in this work are followed by a GELU non-linear activation. Contrary to previous diffusion models, we modify the time embedding by multiplying the output of the Fourier features by the time parameter, such that the output of the time embedding is zero when the input time is also zero. This choice ensures the time embedding is effectively turned off when the model is evaluated in classifier mode. The next step is to combine the time information with the input particle information. Datasets store different levels of information for each particle. The most basic information, described by the kinematic information of each particle, is always stored. However, additional information such as particle identification (PID) and vertex information for charged particles is only available in specific datasets, with the latter only provided in the JetClass dataset <cit.> among benchmark, public jet datasets. To avoid training multiple models to accommodate each dataset, and thus defeating the purpose of a generalized model, we instead adopt a feature drop approach. During training, we consider as inputs both the kinematic information for each particle and their respective PID[The vertex information while present in the JetClass dataset is not used during training. While that could also be included in the training methodology, our focus is the application to multiple datasets without this information, hence for simplicity we skip these features. ]. With a probability p = 0.2 we drop the PID information by replacing it with zeros. This approach is similar to dropout layers <cit.> that encourage the network to learn a useful representation both in the presence and absence of these features. After the feature drop, the inputs are encoded to a higher dimensional space using a feature embedding consisting of two MLP layers. The outputs of the feature embedding are then combined with the time information though a shift and scaling operation. Before the transformer block we introduce a positional token to encode the geometrical information of the neighborhood surrounding each particle inside the jet. Even though transformers are capable of learning general correlations between particles, the addition of local information can generally improve performance <cit.>, creating a better latent representation that is aware of the distances between particles. We create the local encoding using dynamic graph convolution (DGCNNs) layers where the neighborhood is defined using a k-nearest neighbor algorithm with number of neighbors fixed to 10. The distances are calculated in the pseudorapidity-azimuthal angle space. For each of the k-neighbors, edge features are defined based on the particle features concatenated with the subtraction between the particle features and each of the respective neighbors. An MLP is used over all edges before an average pooling operation over the neighbor dimension. The result of the operation represents the new particle features. A second DGCNN layer is then created  with distances calculated based on the Euclidean distances between particles after the updated features. This second operation allows the dynamic construction of edges between particles in a learned geometrical space that pulls relevant particles closer together and far apart otherwise.  The particle embedding equipped with local information is then passed through multiple transformer blocks. The transformer block closely follows the original proposal <cit.>, combining multi-head attention modules with additional skip connections. Our main modification is the addition of LayerScale <cit.> layers. The LayerScale operation introduces a learnable multiplicative factor to each skip connection in the transformer block. Using a small initial value (10^-5 in our implementation), the operation improves the stability and convergence of the model by allowing each transformer block to learn the relevance of each attention block in the transformer architecture. The outputs of the transformer blocks are then added to the original particle embedding after the combination with the local tokens, improving the information flow over the entire model. The outputs are then used as inputs of the task-specific blocks. Since our proposed backbone architecture combines edge creation with transformer modules we refer to it as Point-Edge Transformer (PET), noting that OmniLearn refers to the join training strategy, while PET refers to the specific neural network architecture used to implement OmniLearn.




 §.§ Classifier Head

The classification task across different jet types is accomplished through a dedicated classification head that takes as inputs the particles after the shared network, now referred to as the PET body. Additionally, we also include the overall information from the jet kinematics, including the jet mass, transverse momentum , pseudorapidity η, and particle multiplicity. Even though the additional information is partially redundant compared to the initial particle features, the addition of the jet kinematic information helps the model converge faster when evaluated over datasets covering different fiducial regions than the ones used during training. The jet information is embedded in a higher dimensional space using a jet embedding layer that includes two MLP layers of same size as the current particle embedding dimensionality. This information is then combined with the outputs of the PET body through a scaling and shift operations. A trainable class token <cit.> is then used to summarize the information of the particle embeddings before the classification output. The class token is essentially interpreted as an additional particle, concatenated to the true particle inputs. Inside the transformer block, the outputs of the PET body are not updated but only the class token is allowed to change at the end of each transformer block. The output predictions are then determined by passing the updated class tokens over one last MLP with output size determined by the number of classes in the dataset.



 §.§ Generator Head

Similarly to the classification head, the generator head takes as inputs both particle embeddings and jet kinematic information. Additionally, we include the time information and the set of class labels to condition the generator over the jet types to be simulated. The time and jet information are embedded in a higher dimensional space using the same encoding blocks used in the PET body and classifier head, respectively. This information is then combined through an addition operation followed by an MLP. The classification labels are also mapped to a higher dimensional space using a single MLP. The outputs of the MLP are then passed through a layer drop operation that, similarly to the feature drop, has a probability p = 0.1 to replace the entire output of the MLP with zeros. This choice is motivated by two observations: when the model is used during downstream tasks, the classes used to condition the PET architecture are hardly going to be the same as the ones used during training. Randomly ignoring the class labels encourages the entire architecture to learn both a general and specialized representation, leading to quicker convergence when adapted to other datasets. Second, this technique is similar to classifier-free guidance, observed to improve the generation quality of diffusion models <cit.>.
The results of the layer drop operation are then added to the  outputs of the combined jet and time embeddings. The result of this combination is then used as a diffusion token, where similarly to the classification token, is tasked to summarize the particle embedding information inside the transformer block. However, while the classification token is interpreted as an additional particle, we interpret the diffusion token as a conditional shift of the particle embeddings produced by the PET body. Initially, all particles are simultaneously shifted by the diffusion token created from the combined class labels, time, and jet information. The diffusion tokens are then updated with every transformer layer. The diffusion prediction is then the sum of the original PET body outputs with the learned diffusion tokens. 



 §.§ Loss Function


The loss function of OmniLearn consists of multiple terms designed to combine both the classification and generation tasks. Each of the terms is shown in Equation <ref>.


    ℒ   = ℒ_class + ℒ_gen + ℒ_class smear
       = CE(y,y_pred) +  𝐯 - 𝐯_pred^2 + α^2CE(y,ŷ_pred) .


For input data x, the cross-entropy (CE) loss is calculated using the output of the PET classifier y_pred and true class labels y. The PET generator takes as inputs perturbed data x̂ = α(t)x + σ(t)ϵ with time-dependent perturbation parameters α(t) and σ(t) and predicts a velocity parameter v_pred that is compared with the true velocity value v = α(t)ϵ - σ(t)x. Additionally, the perturbed inputs can also be interpreted as a form of data augmentation that can further improve the classifier performance. We evaluate the PET classifier over the perturbed inputs to get the predictions ŷ_pred used in the calculation of the cross entropy loss. A weight of α(t)^2 is applied to ensure that at t=0, where α(0) = 1 and σ(0) = 0, we recover the classifier loss over clean inputs and at t=1, where α(1) = 0 and σ(1) = 1, the completely corrupted data does not negatively impact the classification performance.



 §.§ Training Details



We train OmniLearn using the JetClass dataset <cit.>. A total of 10 different jet classes are provided, simulated using MADGRAPH5_aMC@NLO <cit.> for the matrix element calculation and Pythia <cit.> to perform the parton showering and hadronization. Detector effects are simulated using Delphes <cit.> with the CMS detector configuration. Jets are clustered using the anti-k_T algorithm with radius parameter of R=0.8. <cit.>. Jets with transverse momentum between 500-1000 GeV and pseudorapidity |η| < 2.0 are saved. The training dataset consists of 100M jets, divided equally between each of the 10 classes. The kinematic information and PID for each particle is used as input features for the training and listed in Appendix <ref>.

Up to 150 particles are saved per jet to be used during training. The training is carried out on the Perlmutter Supercomputer <cit.> using 128 GPUs simultaneously with Horovod <cit.> package for data distributed training. A local batch size of size 256 is used with model training up to 200 epochs. OmniLearn is implemented in TensorFlow <cit.> with Keras <cit.> backend. The cosine learning rate schedule <cit.> is used with an initial learning rate of 3×10^-5, increasing to 3√(128)×10^-5 after three epochs and decreasing to 10^-6 until the end of the training. The Lion optimizer <cit.> is used with parameters β_1 = 0.95 and β_2 = 0.99. The fine-tuning of OmniLearn across different datasets and tasks is performed by setting the learning rate of the PET body to be a factor 10 smaller than the learning rate used in the rest of the PET architecture. The PET body model has 1.3M trainable weights, while the classifier and generator heads have 268k and 416k trainable parameters, respectively. 

In every application, all pre-trained weights are loaded unless there is a change in input dimensions. This is only the case for the output layer of the classification model and the input class labels in the generator model, where the number of classes used in each downstream task changes. In these cases, a new layer with correct input and output sizes replaces the trained weights and is initialized using random weights. We evaluate OmniLearn across 10 different datasets with results described in the following sections.




§ GENERALIZATION ACROSS JET TYPES











We first evaluate OmniLearn on two widely-used benchmark datasets for jet tagging: the top quark tagging <cit.> and quark/gluon <cit.> classification. In the top quark tagging dataset, events are simulated using Pythia 8 and Delphes with the ATLAS configuration.  The background consists of dijets produced via QCD and the signal consists of top quark pair production with all-hadronic decays.  The default energy flow algorithm in Delphes is used to create jet constituents, which are clustered using the anti-k_T algorithm with R=0.8.  All jets in the range 550 GeV < p_T < 650 GeV and |η| < 2 are saved.  Note that while top quark and QCD categories are present in the JetClass dataset, the Delphes detector configuration and p_T ranges are different. The quark/gluon dataset consists of stable particles clustered into jets, excluding neutrinos, using the anti-k_T algorithm with radius R = 0.4. The quark-initiated sample (signal) is generated using a Z(νν) + (u,d,s) while the gluon-initiated data (background) are generated using Z(νν) +g processes. Both samples are generated using Pythia8 without detector effects. Jets are required to have transverse momentum ∈ [500,550] GeV and rapidity |y|<1.7 for the reconstruction. For each dataset we evaluate the results of adapting OmniLearn on each dataset and compare with the training carried out with the PET classifier architecture from scratch. The results are compared with other models using the same datasets in Tables <ref> and <ref>. 

In both cases we observe the performance obtained by OmniLearn to be significantly better than other models trained from scratch, while matching and sometimes surpassing the performance observed by the fine-tuned version of the state-of-the-art model PartT <cit.>. In Figure. <ref>, we show the loss curve obtained during training in the validation set. After a single epoch, OmniLearn already reaches the performance from the PET classifier with convergence observed after only three epochs in both datasets, reducing the overall training time by a factor 3.




§ GENERALIZATION ACROSS DETECTORS













A more realistic scenario is to consider a complete simulation of the detector response. In this scenario, the generalization power from OmniLearn trained on fast simulations could  greatly reduce the need for large simulation samples with full detector simulation for the training of particle taggers at the LHC or be quickly deployed using open data released by the LHC Experiments. We investigate this scenario using the public CMS Open Data <cit.> and ATLAS Top tagging dataset <cit.>. We use the simulations from the CMS Open Data release for 2011A run period <cit.> processed by the MIT Open Data software <cit.> to select simulated QCD jets in proton-proton collisions  at a center-of-mass energy of √(s) = 7 TeV produced with Pythia6 <cit.> and Geant4 <cit.> for detector effects. Particle-Flow objects <cit.> are used to define the objects that are clustered into jets using the anti-k_T algorithm with radius parameter 0.5 with additional requirement to have > 375 GeV to achieve high trigger efficiency. The jet flavor is defined by the hard parton associated to the jet. We define quark-initiated jets as jets associated to uds partons and save these jets for the classification task against gluon-initiated jets. For simplicity, we fix the number of quark and gluon jets to be the same and around 20M events in total. We train OmniLearn using 70% of the simulated jets and report the results compared to the PET classifier trained from scratch using a separate test set consisting of 10% of the events in Table <ref>. We observe the performance achieved by OmniLearn to be better in all metrics compared to a classifier trained from scratch. This observation is encouraging since OmniLearn shows that classification performance can be enhanced even when available datasets are of similar size as the ones used to train OmniLearn. Moreover, OmniLearn is able to converge 2 times faster and to a better minimum than training from scratch as shown by the progression of the validation loss in Figure <ref>.

Next, we use the data released by the ATLAS Collaboration consisting of top quarks and QCD jets simulated using the ATLAS detector simulation based on Geant4. Events are generated with Pythia8 using the NNPDF2.3LO <cit.> set of parton distribution functions and the A14 <cit.> set of tuned parameters. Pileup effects are simulated by overlaying inelastic interactions on top of the underlying hard scattering process based on the 2017 data taking period. Hadronic boosted top quarks are obtained in simulated events containing the decay of a heavy Z' boson with mass of 2 TeV. The cross section of this process is reweighted to produce an approximately flat jet p_T distribution to efficiently populate the full kinematic region. Background QCD jets are obtained in simulated events containing pairs of light quarks or gluons. Unified Flow Objects <cit.> are used to determine the jet constituents. Jets are clustered using anti-k_T algorithm with R=1.0 with additional pileup mitigation algorithms <cit.> applied. The Soft-Drop algorithm <cit.>  is also applied to remove soft and wide-angle radiation.
We investigate two scenarios. One where we fine-tune OmniLearn using all 40M training events available in the dataset and one where we evaluate the fine-tuning performance using only 4M events. Results are shown in Table. <ref> compared to other architectures trained using the same dataset and with the PET classifier trained from scratch. 

Once again, OmniLearn is able to achieve state-of-the-art performance, surpassing the results reported by all previous models trained using the same dataset. The difference between the OmniLearn and PET classifiers trained over the entire dataset is less significant since the 40M training events are only a factor 2 smaller than the entire JetClass dataset. Conversely, OmniLearn fine-tuned on 4M events already matches the performance of all previous models trained from scratch using the entire dataset, requiring less examples to achieve the same performance. In Figure. <ref> we show the validation loss curves obtained during the training. In both cases, OmniLearn already starts from a lower loss and converges quicker to a lower value than the classifier trained from scratch.



§ GENERALIZATION ACROSS COLLISION SYSTEMS







We also evaluate the generalization capability of OmniLearn for jets measured in different collision systems, covering completely different fiducial regions of the phase space compared to the LHC. We use simulations of neutral-current deep inelastic scattering (DIS) generated using the Rapgap 3.1 <cit.> generator for electron-proton collisions with electron and proton beam energies of 27.6 GeV and 920 GeV, respectively. The simulations are provided by the H1 Collaboration using the Heracles routines <cit.> for QED radiation, CTEQ6L PDF set <cit.>, and the Lund hadronization model <cit.>. The detector simulation is performed using the Geant3 <cit.> package. An energy-flow algorithm <cit.> is then used to reconstruct the particles clustered into jets using the k_T algorithm with R=1.0. A second simulation using the Djangoh 1.4 <cit.> generator is used during the classification. The task is to separate jets between the two different simulations. Notice that both simulations target the description of DIS events, hence their differences are more subtle than the previous classification tasks investigated thus far. This choice of classification problem is motivated by previous studies on the unfolding of jet substructure observables <cit.>. One of the leading uncertainties, the closure test performed between two different simulations, is carried out using the same simulation routines and classification task used in this study, thus any improvements driven by OmniLearn could also lead to better unfolding algorithms to be developed for current and future unfolding analyses. A total of 2.5M events are generated for each simulation and used during the training. The performance obtained is listed in Table <ref>.


We observe the same level of performance between OmniLearn and the PET classifier. Compared to the previous classification tasks, both accuracy and AUC values are much lower, evidencing the challenge of distinguishing jets from the two simulations apart. While the final performance is the same, we observe a much quicker convergence from OmniLearn as shown in the validation loss curve in Figure <ref>, resulting in a factor 3.5 faster training compared to starting from scratch. As described in <cit.>, the full determination of the uncertainties require the training of thousands of classifiers, making any improvements in training speed an important asset for the applicability of the algorithm. These results also highlight the capability of OmniLearn to quickly adapt to new datasets, providing a strong asset for future experimental facilities such as the EIC. 





§ CONDITIONAL GENERATION





Next, we evaluate the generation quality obtained by OmniLearn using the JetNet <cit.> datasets consisting of jets initiated by light-quarks, gluons, top quarks, W and Z bosons. The jets are generated with transverse momenta  around 1 TeV and are clustered using the anti-k_t algorithm with a radius parameter of 0.4. Each jet has a maximum number of particles stored fixed to 30 <cit.> or 150 <cit.>. For each jet, the four-momentum information (_jet,η_jet,ϕ_jet, m_jet) is provided, as well as the particle multiplicity. We adopt the two-model strategy presented in <cit.>, training a model that only learns the kinematic information of the jets and using that information as conditional information for the diffusion model trained using particles as inputs. Note that in this case, OmniLearn is used only to learn the particle information while the overall jet kinematic information model is always trained from scratch. The performance of the generation is evaluated using multiple physics-based metrics proposed in <cit.> and listed in Tables <ref> and <ref>.
We also provide the results of OmniLearn and the PET generator in the ideal case, where the jet information, used to condition the particle generation model, is taken directly from the validation set of the JetNet dataset, effectively separating the impact of the jet generation, that does not benefit from OmniLearn, from the particle generation process.






In all metrics investigated in this study, OmniLearn shows similar or improved performance compared to previous models for both datasets consisting of 30 and 150 particles. We also notice that differences between the idealized version of OmniLearn and PET generator are more significant than at full generation level where the jet generation quality also affects the overall model performance. Additionally, statistical uncertainties determined from the bootstrapping method are often above the 20% level for some metrics, requiring access to larger evaluation sample sizes to provide a more precise comparison between models. In Figure <ref> we show the validation loss curve for the PET generator and OmniLearn training in the JetNet dataset. Similarly to previous results, the OmniLearn training starts from a lower value of the loss function and is able to converge quicker, requiring roughly 20% and 30% fewer training epochs than a model trained from scratch. 




§ REWEIGHING AND UNFOLDING





Unfolding, or correcting observables for detector effects, is a fundamental task in collider physics to enable efficient comparisons between measurements and theory predictions. Traditional unfolding methods use histograms to estimate the unfolded response through a regularized inversion of the detector response matrix <cit.>. Machine learning greatly increase the flexibility and potential of unfolding by allowing the simultaneous correction of multiple distributions without the use of histograms <cit.>. The OmniFold algorithm <cit.> introduced an iterative approach for unfolding based on reweighting functions. In the first iteration of the method, a reweighting function that corrects simulations towards measured data, is determined using a classifier. We examine the potential of OmniLearn applied to high-dimensional reweighting and complete unfolding using  all available features and the same dataset introduced in Ref. <cit.>, available on Zenodo <cit.>. The dataset consists of proton-proton collisions producing a Z boson, generated at a center-of-mass energy of √(s)=14 TeV. A sample used as the `data' representative is simulated using particle collisions with the default tune of Herwig 7.1.5 <cit.>. A second dataset, representative of the `simulation' we want to correct, is simulated using Pythia8 with Tune 26 <cit.>. Detector distortions are simulated with Delphes and the CMS tune that uses a particle flow reconstruction. Jets are clustered using all particle flow objects at detector level and all stable non-neutrino truth particles at particle level.  They are defined by the anti-k_T algorithm with radius parameter R=0.4 as implemented in FastJet 3.3.2. The Z bosons are required to have >200 GeV in order to mitigate acceptance effects. While all clustered particles are used during the training, we report the performance obtained by different algorithms using the same set of observables reported in the original OmniFold publication. 

First, we investigate the capability of OmniLearn to determine the full phase space reweighting function by fine-tuning the classifier. We also compare with the results obtained from a classifier trained from scratch (PET classifier). The results of the reweighted distributions at reconstruction level are shown in Figure <ref>. From a visual inspection we see that both OmniLearn and the PET classifier are able to correctly reweight the distributions of all high-level observables we investigated, however OmniLearn shows a better agreement with the `data'. We quantify the improvement brought by OmniLearn by calculating the triangular discriminator <cit.> for each observable with results reported in Table <ref>. In all cases, the values obtained by OmniLearn are significantly better than the baseline training.



In Figure <ref>, we show the validation loss curve obtained during the reweighting training for both OmniLearn and PET classifier. Similarly to previous classifier results, OmniLearn starts from a lower validation and and converges faster to a lower minimum compared to a classifier trained from scratch.






Next, we investigate the impact of using OmniLearn during the entire unfolding process. We follow the OmniFold methodology training the model for five iterations, similarly to the original OmniFold publication. In the first iteration, we either start from a model trained from scratch, in the case of the PET classifier baseline, or use OmniLearn as the starting point. Each following iteration starts from the trained model of the previous iteration. The results of the unfolded distributions are shown in Figure <ref> with numerical comparison of the triangular discriminator shown in Table <ref>. Results are compared with the values reported in the original OmniFold publication, where particle-level unfolding was accomplished using a DeepSets architecture <cit.>.




Once again OmniLearn shows an improvement in performance compared to other methods and same classifier architecture trained from scratch. These results are encouraging since data availability is strictly limited at collider experiments, presenting a limiting factor to the applicability of OmniFold. However, using OmniLearn as the starting point, we can mitigate this issue to achieve better performance. 




§ WEAK SUPERVISION AND RESONANT ANOMALY DETECTION



The search for new particle interactions is a challenging task at collider experiments. Expected to be rare, signs for new physics might be subtle, affecting only very specific observables. While traditional approaches rely on theory predictions to narrow down possible new physics scenarios, the current lack of new physics discoveries at the LHC motivates new strategies to search for new phenomena. Anomaly detection brings a change in this paradigm <cit.>. Unexpected data structures can be automatically identified by algorithms as possible hints for new physics. One well-studied approach designed for resonant anomaly detection is based on Classification Without Labels (CWoLa) framework <cit.>, where weakly-supervised learning enables training directly on (unlabeled) data. In the CWoLa approach, samples with mixed fractions of a possible signal and background are used to train a classifier whose goal is to identify the origin of the sample. In the best case, we can imagine a sample consisting of only the  background process, possibly as part of a background simulation or derived from a control region, that is then used to train the classifier against data possibly containing new particle interactions in addition to the background process. We examine the benefits of using OmniLearn during the classification process and evaluate the performance using the R&D dataset from the LHC Olympics data challenge <cit.>.  The background consists of dijet final states from QCD production while the signal is a resonant boson production A → B (→ q q') C(→ q q') with masses m_A, m_B, m_C = 3.5, 0.5, 0.1 TeV, respectively.  Signal and background vents are generated with Pythia8 interfaced with Delphes3.4.1 for detector simulation. Jets are defined using the anti-k_T algorithm as implemented in FastJet with R=1. We focus on the two leading jets in transverse momentum space and require the leading jet to have > 1.2 TeV. After selection, we save all particles associated to the two most energetic jets, resulting in a maximum particle multiplicity of 279 particles per jet. 

During weakly-supervised training, we follow previous studies <cit.> and define the signal region of interest for events with dijet mass 3300 GeV < m_jj < 3700 GeV. In the region of interest, we call `data' the combination of 100k background events with varying amounts of signal events. The background-only distribution consists of 350k independently simulated background events. Differently to previous applications of OmniLearn described in this paper, here, we need to include the information of both jets in the classifier. We modify the PET classifier to accommodate the changes in the dataset while also preserving the permutation equivariance of the complete network. This is achieved by first passing the particles present in each jet through the PET body independently, such that particles belonging to different jets do not interact with each other. The outputs of the PET body are then shifted and scaled by the outputs of the jet embedding block that takes as input the kinematic information of each jet. This strategy allows us to maintain the permutation equivariance of the model while giving jet specific information to each particle. The shifted and scaled particles are then passed to the classifier head, reshaped as if all particles belonged to the same jet. Since the shift and scaling operations are jet-dependent, the reshaping operation allows all particles to be conditionally mapped to the same space without loss of information. The classifier head is unchanged, with a class token used to summarize the information of all particles before the classification output. We use the output of the classifier as the anomaly score to determine the sensitivity to this specific new physics scenario. We quantify the performance based on the maximum value of the significance improvement characteristic curve (SIC) defined as the signal efficiency divided by the square root of the background efficiency versus the signal efficiency. The SIC represents a multiplicative factor by which the initial significance of a signal present in the data would increase when a particular threshold of the classifier output is chosen.  Maximum SIC values above unity indicate value added. We show the results in Figure <ref> and compare the results obtained by OmniLearn and PET classifier with the results reported in <cit.>. Since the reference background process is statistically identical to the background presented in the `data' construction, we call this scenario idealized[The authors of Ref. <cit.> show that it is possible to achieve even better performance if the functional form is known - it would be interesting to see such a strategy combined with OmniLearn in the future.].  




OmniLearn shows non-negligible signal sensitivity for signals injections above 600, corresponding to an initial significance S/√(B)∼ 2, representing a large increase in sensitivity compared to  previous results where signal sensitivity was only achieved for signal injections above 1500 (S/√(B)∼ 5). Compared to previous results, we also observe the performance of PET classifier to be similar at lower signal injections to the results reported in <cit.> and performing worse at higher signal injections. The reason for this difference is due to the limited amount of data in the signal region (around 100k). Even with the larger generated background of 350k events, the dataset size of this application is at least 2 times smaller than all previous datasets investigated so far. In the low data regimes, data efficient models are often observed to perform better than large transformer models, a limitation that is mitigated by OmniLearn. 

In general, a pure background process in the region of interest is often not readily available, requiring alternative strategies to produce a background-only description. In the resonant case, a natural strategy is to use sidebands around the region of interest to determine the properties of the background process.  As in a number of previous weakly supervised studies <cit.>, a generative model conditioned on the resonant variable and trained only in the sidebands can be used to interpolate the background description in the region of interest. We use this strategy to train OmniLearn in the sidebands using the dijet mass value as a conditional variable replacing the class labels as inputs to the  PET generator head. After training of the generative model, predicted background events in the region of interest are created by generating 350k background events. The same classifier used in the weakly supervision case is used to separate events from the generated background and `data', created using true background events and different amounts of signal injected. Even though the training of the generative model is carried out using all clustered particles, we require particles in the generated and data samples to have > 1 GeV, to reduce the impact of low energy particles driving the classifier performance. Results are shown in Figure <ref> with idealized results also included for completion.



Similarly to the idealized scenario, OmniLearn increases the reach in signal sensitivity to signal injections above 700 (S/√(B)∼ 2.2) and greatly improving upon previous results. While the Cathode method, using 6 high level observables <cit.>, is still more sensitive to the specific new physics scenario at lower signal injection values, we expect the use of all particles enhanced by the OmniLearn model to attain sensitivity to a more general class of possible new physics scenarios, that would otherwise not be identified by the specific set of high level observables chosen. This application highlights the potential of OmniLearn for anomaly detection, where the data used to train the generative model and classifiers is limited by the experimental setup and thus cannot be increased arbitrarily.

Finally, we also investigate the changes in validation loss as the training of the generative model progresses. The results are shown in Figure <ref>.



Not only do we observe OmniLearn starting from a considerably lower loss, but OmniLearn also reaches an overall lower validation loss and matches the minimum validation loss from the PET generator in half the number of epochs. Reducing the training time for the generator is important for realistic applications at the LHC, where a scan using multiple signal windows is necessary to cover the entire dijet mass spectrum, requiring the training of several generative models for each signal region of interest.



§ CONCLUSION AND OUTLOOK



In this paper, we have introduced a foundation model for jet physics called OmniLearn.  This machine learning approach is a neural network capable of advancing a wide variety of research areas within jet physics.  We have shown that even though this model was trained on one specific dataset, it accelerates and/or improves the accuracy of other classification tasks and generation quality evaluated over nine additional datasets across initial states, final states, and simulation levels as well as of tasks in jet generation, reweighting/unfolding, and anomaly detection.  While our neural network has millions of trainable parameters, it is much smaller than many other foundation models (i.e. it is not `large' in the sense of an LLM).  We view this as a strength for usability and we hypothesize that this is possible because our model had a training task close to the application tasks. 

In all tasks, OmniLearn is nearly always as good or (much) better than previous dedicated models.  This is potentially transformative in a number of applications and we highlight three to illustrate the immense potential of this new approach:


    
  * For tasks with expensive simulations, we have shown that OmniLearn can effectively amplify the training statistics to improve the performance.  For example, on a full detector simulation dataset for top tagging from ATLAS <cit.>, we show that we can achieve (or exceed) state-of-the-art performance with only 10% of the training dataset.  This could lead to significant computational savings for the experiments when developing new taggers.
    
  * For full phase space unfolding, OmniLearn is not only more precise, but also much faster - converging in about half the time as a dedicated approach.  This is critical for practical applications of unfolding, where ensembling and statistical uncertainties require the training of a computationally expensive number of models <cit.>.    
    
  * For anomaly detection, we have shown for the first time that full phase space methods are capable of non-trivial discoveries.  Previous work in resonant anomaly detection <cit.> had shown that full phase-space methods could significantly amplify injected signals, but only if their starting significance was well above 2.  With OmniLearn, this is pushed down to 2.


Our methodology and trained model are publicly available.  While we have focused on jet physics, it would be exciting to see OmniLearn applied in other, related areas with similar challenges and thus also similar potential rewards.



§ CODE AVAILABILITY


The code for this paper can be found at <https://github.com/ViniciusMikuni/OmniLearn>.



§ ACKNOWLEDGMENTS

We thank F. Dreyer, Ming Fong, R. Grabarczyk, K. Greif, P.F. Monni, and D. Whiteson for interesting discussions related to transfer learning.  We also thank J. Birk,  A. Hallin , and G. Kasieczka for comments on the manuscript

Additionally, we thank our colleagues from the H1 Collaboration for allowing us to use the simulated MC event samples. We also thank DESY-IT and the MPI für Physik for providing computing infrastructure and supporting the data preservation project of the HERA experiments.

VM, and BN are supported by the U.S. Department of Energy (DOE), Office of Science under contract DE-AC02-05CH11231.  This research used resources of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02-05CH11231 using NERSC awards HEP-ERCAP0021099 and HEP-ERCAP0028249.  






§ INPUT VARIABLES


The input features used to train OmniLearn are described in Table <ref>.




apsrev4-1







