Self-supervised learning (SSL) has been a key driver in recent machine learning (ML) advancements; through the use of data augmentation, SSL algorithms rely on pseudo labels to create representations of the data that can be highly useful and efficiently adapted for downstream tasks. Accordingly, the resulting SSL trained model can act as a powerful foundation model [1], defined here as a model that is "pretrained" on a generic task, in this case one defined with SSL, and capable of being fine-tuned for a variety of purposes. Foundation models have quickly arisen as a strategy for describing complex tasks, especially in computer vision (e.g., [2]), natural language processing (e.g., [3]), and speech recognition (e.g., [4]). More recently, foundation models for machine learning in science have become an important and active area of research, for instance in biology [5], chemistry [6][7][8][9], differential equation solving [10,11], cosmology [12,13], and, finally, high-energy physics [14], not all of which, however are necessarily SSL-based. SSL benefits from unlabeled training and from the creation of a relation between data augmentations. Leveraging vast unlabeled datasets can enable SSL to outperform supervised learning. In this work, we focus on contrastive learning, a variant of SSL. We propose a new strategy for data augmentation and study how it performs within the SimCLR framework, which was first introduced in [15]. In SimCLR, the learning objective is to map a data point and its augmentation(s) to similar representations, while pushing different data points toward differing representations. In contrastive learning, the quality of the learning task is highly dependent on the set of augmentations. Domain completeness, i.e., when the augmentation set covers all possible variations in the data, can lead to better representation learning because the model can gain a full view of plausible data during learning.
To that end, we propose a simulation-based strategy to obtain an augmented data set that is as domain complete as possible and only limited by the knowledge encapsulated in a simulator. More specifically, in settings with stochastic simulators, we propose a strategy whereby we intervene in the middle of a simulation process, fix the upstream generated latent state of an event, and resimulate downstream components to sample augmentations from the set of all plausible observations for a given latent state. From a programmatic standpoint, the first step fixes all initial conditions and stochastic latent variables in the simulation up to the intervention point, while the second step resamples the stochastic latent variables in the simulation after the intervention point and thus generates a new stochastic output of the simulation conditioned on the fixed latent state at the intervention point. Hence, this allows us to interpret augmentations in self-supervised learning in a novel way: the point at which we intervene in the simulation chain is the point at which further information downstream is considered an augmentation, above is information. The goal is then to learn a robust representation that contains as much information as possible while integrating out the variability due to the augmentations. We dub our new strategy RS3L ([0rIzə], "Resimulation-based self-supervised representation learning"). We will show that this method allows for powerful contrastive pretraining that, through altering the simulator settings for resimulation, also aids in the mitigation of uncertainties arising from domain shift between potentially imperfect simulations and real data.
To test our method, we will focus our experiments on high energy physics (HEP), where high-fidelity simulators are readily available (see Ref. [16] for an overview). In our case, the fixed latent state is represented by the elementary particle generated in a hard-scattering process, which can be described with perturbative quantum field theory. We then resimulate the steps in which secondary particles are created through radiation off of the particles produced in the hard scattering, the subsequent hadronization of secondary particles into stable, bound particle states, and then the stable particle interaction with the detector material. In the following, this step is referred to as parton showering.
In particular, we will focus on learning representations of jets, which are highly energetic, collimated streams of secondary particles. Jets are core objects in HEP; they are produced from the showering and hadronization of highenergy quarks and gluons or through the decay of heavy elementary particles like the Higgs boson, the top quark, or heavy weak gauge bosons. Jets are featured in many theories for new physics beyond the Standard Model. Hence, the ability to "tag" a jet, i.e., to correctly and robustly classify the particle which initiated the jet, is crucial for the success of the physics program of experiments like ATLAS [17] and CMS [18] at the Large Hadron Collider (LHC) [19] and future high energy collider experiments. For instance, one of the canonical classification tasks in HEP and a prerequisite for the precise understanding of Higgs boson couplings to bottom quarks [20][21][22][23][24] is the discrimination between jets originating from Higgs bosons and jets originating from QCD. Our choice of resimulating the parton shower is motivated by a desire to characterize the progenitor of a jet, while capturing the large variability that results for each parton shower. That is, the resimulated shower defines an augmentation to the existing shower that we claim captures all the physics that we desire downstream.
By sampling from our high-fidelity simulator to obtain augmented versions of jets, our work presents an evolution of the methods developed in JetCLR [25,26], which focused on empirical methods of augmentation generation (such as rotations), whereas our work develops a domain-complete augmentation set based concretely on the physics in the simulation. As a result, we outline a strategy through augmentations within SSL to pretrain a latent space that captures the most salient, physically motivated, features of a simulation. Within high energy physics, other SSL methods have also been explored for foundation model pretraining. For instance, the masked modeling strategies developed in BERT [27] and BEiT [28] have been adapted for set-type data in order to mask and predict particles in a jet [29], and to mask and predict particle-type information [30]. Supervised pretraining and fine-tuning strategies have also been explored for high energy physics contexts [31]. A similar strategy has been also pursued within astrophysics [32].
Accordingly, the contributions of this paper are as follows:
(i) The development of a methodology for resimulation through intervening in the simulation chain. This has the benefit of generating plausible and physically motivated downstream outcomes that serve as augmentations in a contrastive pretraining, enabling learning of the salient components of physics simulations. (ii) The creation of a large open dataset for common development in the community https://doi.org/ 10.5281/zenodo.10633815.
(iii) A systematic study of the gains of resimulationdriven contrastive learning against fully supervised learning strategies. While the focus of this work will be on jets produced at the LHC, related strategies can be extended to other domains where simulation is present.
Our resimulation strategy aims to develop a physically motivated pretraining of a model that we dub the RS3L backbone. This model represents data in a multidimensional space rich in information that is common between an object and its augmentations, leading to a localization of the object within the space. Through this strategy, we can further reduce the dimensionality of the overall space, leading to a low-dimensional representation that aims to capture the key properties of the high-level objects. This strategy is illustrated in Fig. 1, which is described in the following.
Events are produced by a high-fidelity stochastic simulation of the hard scattering of elementary particles (in our case, outgoing particles can be Higgs bosons, quarks, or gluons). Different configurations for the parton showering and hadronization are employed during resimulation, each giving an augmented realization, i.e., a different jet, of the same initial parton. In the following, the PYTHIA8 [33] FIG. 1. Illustration of the RS3L setup, including downstream resimulation, sampling, graph computation, and the construction of positive and negative pairs. These are then used in a contrastive loss function aiming to align positive pairs and push negative pairs apart.
parton shower and hadronization simulation with standardized settings derived from tuning to CMS experimental data [34] is used to generate the nominal scenario.
Domain completeness can be achieved by considering a set of augmentations which represents the full knowledge embedded in our high-fidelity simulator. Accordingly, we perform augmentations that span an array of plausible experimental outcomes, namely (1) In-domain augmentation by keeping the simulator settings fixed but resampling with a different numerical seed. (2) Out-of-domain augmentation by either (i) varying the primary simulator settings within well-motivated bounds, in our case altering the probability for finalstate radiation (FSR) branchings in the parton showering step, or (ii) using a different simulator, in our case a different parton shower model (HER-WIG7 [35]). The above set captures both (1) implicit uncertainty in the definition of a quark or gluon parton shower, which results in a cascade of particles that can vary in number of particles and distribution of energy across the particles, and (2) the variations within parton shower simulations that are known to cover the observed variations in the data. The augmentations themselves can be interpreted physically as a construction of the parton "wave function" before its collapse into a measured state. The resulting space thus aims to "project away," i.e., disregard, the imperfect modeling of the parton shower process and the statistical variability in our simulator, including quantum effects and ill-defined hidden parameters.
The RS3L space is trained with 5M events using a 50%/ 50% admixture of QCD jets and Higgs jets. Given a nominal jet, we randomly sample one of its augmentations to build a positive pair. All other possible jet pairs in a minibatch, which consists of 100 nominal and 100 augmented jets, are considered negative pairs. We then process each jet with the backbone, which we define using graphbased architectures with built-in graph convolutions and message-passing. Graph neural networks (GNNs) have seen large success over other ML algorithms like dense or convolutional neural networks in the analysis of LHC data due to the point-cloud nature of the particles that comprise a jet or an event (for an extensive review, see Ref. [36]). Finally, each jet is embedded into an Ndimensional latent space. We use N ¼ 8 dimensions for the latent space as a balance between expressiveness and computational complexity, although higher-dimensional spaces are possible. 1 The 8D jet representations are then used in the contrastive loss function based on SimCLR. The loss function aims at aligning the 8D vectors of jets building a positive pair, and pushes apart jets from negative pairs. A temperature parameter τ regulates the relative importance of these two simultaneous training objectives and is set to 0.1.
Hard-scattering events for pp → Z þ jet and pp → HZ, with H → b b and Z → νν, from which we sample the jets for training the networks, are generated at a center-of-mass energy of ffiffi ffi s p ¼ 13 TeV using MADGRAPH5_AMC@NLOv3.4.0 [37] with the NNPDF30_nlo_nf_5_pdfas parton distribution function [38] at leading-order accuracy in QCD. The j and, respectively, the Higgs boson are generated within a p T range of 400-600 GeV and within jηj < 0.1. The hard-scattering events are then passed through a parton shower simulation and are reconstructed using the DELPHESv3.4.3pre1 detector simulation [39] with a CMS detectorlike geometry. For gluons and the Higgs, the first splitting or, respectively, decay is generated at matrixelement level with a minimum ΔR ¼ ffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffiffi ffi ðΔηÞ 2 þ ðΔϕÞ 2 p > 0.001. All further splittings are generated by the subsequent parton shower.
In the nominal configuration, PYTHIA8v8.244 with the CP5 tune [34] is used for parton showering, hadronization, and modeling of the underlying event. No effects from pile-up are simulated. Four augmentations of each nominal jet are created via resimulation: (i) the same settings as in the nominal configuration are used, but the partons are passed through PYTHIA8 using a different random initial seed; (ii) the value for the parton shower renormalization scale used in the determination of final-state radiation is multiplied by a factor ffiffi ffi 2 p ; (iii) the same scale is multiplied by a factor 1= ffiffi ffi 2 p ; (iv) HERWIG7 with the default tune is used for parton showering and hadronization. In all parton shower configurations, initial-state radiation has been disallowed to focus solely radiation effects of the final state.
Jets are finally clustered from the DELPHES E-Flow candidates using the anti-k t algorithm [40] with a radius parameter R ¼ 0.8. The 100 highest-p T E-Flow candidates per jet are used for training the networks. For jets to be considered in the training and evaluation, one QCD parton (gluon or quark) or a Higgs boson along with its two decay b quarks have to be matched to the jets, i.e., fulfill ΔR < 0.8, with Δη and Δϕ being the difference in pseudorapidity η and, respectively, in azimuthal angle ϕ between the parton and the jet axis. Jets with a transverse momentum of p T > 450 GeV, a mass of at least 10 GeV, and jηj < 0.1 are kept. These requirements are also imposed for all augmentations of a given nominal jet.
We use a graph neural network that employs a stack of DYNAMICEDGECONV [41] network layers for enriching the particle features with information from neighboring particles. The k nearest neighbors are determined dynamically, i.e., in the latent feature space obtained after the initial embedding or, respectively, after each graph convolution. A BERT-like transformer architecture [42] adapted from [43] has been used to cross-check the results and yields comparable performance. We point out that an optimization of the architecture used for the demonstration of the contrastive approach goes beyond a proof-of-concept and, thus, beyond the scope of this article. The architecture used in the RS3L pretraining is:
The input feature set X is given by 15 features per particle, listed in Table I. These input features are not normalized over a minibatch, but their magnitudes are comparable among the features and Oð1Þ, thereby regularizing the network and leading to a similar effect as a pretransformation that enforces a mean of 0 and standard deviation of 1.
For embedding the 15 particle features into a higherdimensional latent space, two fully connected layers with dimensions (128=128) are used in MLP embed . After a series of three convolutions with k ¼ 24, MLP enc consists of four fully connected layers with dimensions (64=32=32=8). Outputs of all neurons in both MLPs are activated with an exponential linear unit [44]. The final building block, GLOBALSUMPOOL, aggregates the eight features of each particle in the jet from the last layer of MLP enc to obtain the final eight features z characterizing the jet in the RS3L space.
The loss function employed is based on the SimCLR framework introduced in [15] and has first been used for jet physics in [25]. It is given by
with sðz i ; z j Þ being the cosine similarity between negative jet pairs i and j or, in the case of positive pairs, the nominal jet i and the augmented jet i 0
The network will try to push apart dissimilar jets (negative pairs) on the N -1-dimensional hypersphere with unit radius given by the N input features, while trying to align similar jets (positive pairs) and characterize them by similar features. The competition between these two effects is regulated by a temperature parameter τ. We studied this parameter based on the performance obtained when using the RS3L space to distinguish Higgs from QCD jets. Specifically, we trained RS3L spaces with τ ranging from 0.05 to 0.5 and fine-tuned these networks on the Higgs vs. QCD classification task with a fixed backbone. A fixed backbone offers the best insight
TABLE I. Description of the particle input features. Particle feature Description log p T Logarithm of the transverse momentum log ðp T =p Jet T Þ Logarithm of p T normalized with respect to the jet log E Logarithm of the energy log ðE=E Jet Þ Logarithm of E normalized with respect to the jet Δη Difference in pseudorapidity with respect to the jet Δϕ Difference in ϕ with respect to the jet ΔR Distance in η -ϕ plane from the jet axis tanh d 0 Hyperbolic tangent of the transverse impact parameter tanh d z Hyperbolic tangent of the longitudinal impact parameter q Particle charge isPhoton 1 if particle is a photon, else 0 isMuon 1 if particle is a muon, else 0 isElectron 1 if particle is an electron, else 0 isCH 1 if particle is a charged hadron, else 0 isNH 1 if particle is a neutral hadron, else 0 FIG. 2. Metrics pertaining to the convergence of the RS3L training as a function of epoch. Top panel: the average of cosine similarity between the positive pairs (anchor jet and augmented jet). Bottom panel: the cosine similarity between the average Higgs vector and average QCD vector. The variation, indicated by the error bands, is computed over three RS3L trainings.
into the performance of the RS3L pretraining step because it projects jets into a one-dimensional classification space using only the information derived during the self-supervised step. We observed that the Higgs vs. QCD binary classification performance was optimal with τ ¼ 0.1. Figure 2 illustrates the convergence of the RS3L space as a function of training time. In the top panel, the cosine similarity is shown for nominal and augmented jets (averaged over each class). This is analogous to the positive pair component of the loss function, the pairs that the network will attempt to make parallel in the space. The network makes augmentations nearly parallel to their augmentations, with average angles between 0.74 and 0.89. We observe that the network further aligns Higgs vectors over time, while slightly misaligning QCD vectors over time. In the bottom panel, the cosine similarity between the average vectors of the two summary classes (Higgs and QCD) is shown. At the beginning of the training, Higgs and QCD occupy roughly the same region of the space, however as the network converges the populations are pushed to lie on opposing ends of the space, despite not knowing about their different class nature. In all plots, the bands are indicative of averaging over 3 repeated RS3L trainings with random initialization.
The MLP placed on top of the RS3L architecture is given by one additional fully connected layer with 8 input dimensions and a single sigmoid activation function at the output. The backbone network architecture used in the finetuning is the exact same as the fully supervised. The only difference between the approaches is that in the fine-tuning method the base graph trained with RS3L is loaded as a hotstart for the classification and the MLP is added on top of the backbone network. We consider two configurations of finetuning the RS3L space: a "fixed" graph (where only the MLP parameters are floating) and a "floating" graph (where all the MLP and graph parameters are floating). The admixture of QCD/H/W, as well as augmentations composition, are always consistent between the RS3L fine-tuning and full supervision. The RS3L space and classification steps are trained on independent datasets.
In each case, the network is trained repeatedly (at least three times) to gain a measure of whether a result is statistically significant. The numbers in ROCs and tables indicate the averages of these repeated trainings. The difference between subsequent training runs is observed to be quite small.
Of the eight dimensions, the most discriminating dimension between Higgs bosons jets and quarks and gluons is shown on the left of Fig. 3. As a qualitative statement on the robustness of the space we derive through RS3L, we also show in the same figure an equivalent plot for the variable N 2 , which is a theory-motivated function of the energy flow of particles in the jet [45]. The latter represents a ratio between the compatibilities of the jet to have one or, respectively, two prongs, which are regions of collimated radiation. As a powerful jet tagging variable, N 2 is commonly used in searches for processes involving hadronically decaying Higgs bosons [46][47][48] or (additional) weak gauge bosons FIG. 3. Left: one of eight features derived in the RS3L pretraining. Right: jet substructure variable N 2 . The main (upper) panel shows the distributions for the nominal parton shower scenario. The ratio panels show the difference between the respective varied distributions and the nominal distribution. For FSR, the up and down variations form a band around the nominal distribution. [49,50] to reduce the overwhelming background where jets arise from quantum chromodynamics (QCD) interactions. both plots, the distributions of Higgs jets are shown in magenta and the distributions of QCD jets (jets initiated from quarks and gluons) are shown in other colors (see also Sec. II). In particular, QCD jets are divided by their flavor content and whether the initial parton was a single quark or a gluon. The main (upper) panel shows the RS3L or, respectively, N 2 feature for the nominal parton shower scenario, while the bottom panels show the ratios of the varied distributions to the nominal distribution. In these plots, ratios closer to 1 indicate smaller differences in the learned representation for jets with different simulator configurations, and therefore smaller systematic uncertainties arising from the parton showering model chosen. Large trends are observed in these ratio when comparing different simulator configurations in the case of N 2 , while the trends are mitigated in the RS3L space thanks to the employed selfsupervision strategy of aligning nominal jets with their augmented versions.
Both the RS3L feature and N 2 provide discrimination power between Higgs and QCD jets. This class separation in RS3L comes about despite the network not receiving any information about the true class of the jet. Rather, it emerges dynamically through self-supervision resulting from fulfilling the contrastive training objective. The corner plots in Fig. 4 show the correlation between the eight features derived through RS3L, split up into Higgs and QCD classes. The Pearson correlation coefficient [51] is provided for each feature pair. The largest (anti)correlation for feature pairs describing Higgs boson jets is -0.16, while for QCD jets it is 0.49. Overall, the eight features are largely uncorrelated, indicating a near orthonormal feature set to characterize the jets was found, in particular the ones coming from Higgs bosons.
Finally, we visually probe the entire jet representation contained in the 8D RS3L space by performing a reduction to two dimensions using the t-SNE algorithm [52]. The plot on the top of Fig. 5 shows this 2D space for gluons, quarks, and Higgs jets. Again, strong separation between Higgs and QCD jets as well as strong clustering for a given class is visible. Notably, the model tries to push apart jets from two different Higgs bosons. However, as all Higgs bosons share common characteristics, this objective is hard to meet, leading to the observed clustering. On the bottom left, we show the space divided into jets showered with PYTHIA8 and, respectively, HERWIG7. At the population level, it is clear that different showering configurations occupy similar regions in the RS3L space and, hence, our approach provides decent robustness against domain shift. The bottom right of Fig. 5 shows the t-SNE 2D space of a RS3L training that did not include HERWIG7 augmentations. Compared to the RS3L space obtained with all augmentations, this space has visibly smaller overlap (congruence) FIG. 5. 2D visualization of the 8D RS3L space, derived via t-SNE dimensionality reduction. Top: a good class separation is seen between Higgs jets and QCD (quark and gluon) jets. Bottom left: jets shown by parton shower model for PYTHIA8 and HERWIG7 for a RS3L space trained with HERWIG7 augmentations. Bottom right: The same for a RS3L space trained without HERWIG7 augmentations. The congruence of the different parton shower models is visibly worse in the right-hand scenario. between PYTHIA8-showered and HERWIG7-showered jets, speaking to larger uncertainties incurred by domain shift.
We now study the use of RS3L as a backbone for indistribution and out-of-distribution classification tasks, and we quantify performance and uncertainties. For all finetunings, we consider a sample comprising the same types of augmentations used for training the respective RS3L space. We then compare the RS3L pretraining + fine-tuning performance to the performance of a fully supervised network with the same architecture. For a fair comparison, we also train the fully supervised network on a sample comprising the same types of augmentations as were used to train the RS3L space. The fully supervised network, therefore, learns to marginalize over the augmentations, not explicitly exploiting the information that the nominal and augmented jet come from the same initial particle. We point out that the standard approach for fully supervised algorithms at experiments such as ATLAS and CMS is to train on samples comprising only jets from a nominal parton shower configuration.
The most straightforward application of RS3L is a binary classification fine-tuning to separate the same jet classes FIG. 6. Tagging performance on the Higgs vs. QCD classification task. All numbers are obtained on a sample comprising only jets from the default parton shower scenario. Fine-tunings on top of a RS3L space that was trained with 5M jets are shown in light and dark red, for graphs that have floating and fixed weights, respectively. The number in parentheses indicates the fine-tuning dataset size. Fully supervised networks with various training sizes are shown in shades of blue. Three repeated trainings with different random initialization are used to compute an average and standard deviation for each ROC. These are shown as bands. Note that fixed-weight trainings do not have any appreciable uncertainty.
TABLE II. QCD rejection rates for various training configurations and Higgs efficiencies. The bolded numbers indicate the network with the best performance at a given Higgs efficiency. The uncertainties are calculated using the standard deviation of repeated trainings with different random initializations. Higgs efficiency 0.3 0.5 0.7 N 2 29 16 8 Fine-tuned (3M, fixed) 140 AE 1 64AE 0 39AE 0 Fine-tuned (3M, floating) 1325 AE 31 380 AE 4 135 AE 1 Fully supervised (1M) 834 AE 11 257 AE 6 96AE 2 Fully supervised (3M) 1070 AE 29 317 AE 3 115 AE 1 Fully supervised (8M) 1312 AE 15 381 AE 7 134 AE 1 FIG. 7. ROC curves showing the tagging performance of RS3L fine-tunings on the in-domain Higgs vs. QCD classification task. Top: graph fine-tuned with floating weights. Bottom: graph finetuned with fixed weights.
that are used to train the RS3L space. Thus, we first consider the aforementioned canonical classification task of discriminating between jets originating from Higgs bosons and jets originating from QCD. Two variants are considered for fine-tuning the RS3L space: (i) The weights of the backbone model are kept fixed, and a task-specific head consisting of a one-layer multilayer perceptron (MLP) is added, whose parameters are trained on the task at hand. (ii) The weights of the backbone model are used as initial parameters but are allowed to be adjusted ("floated") together with the ones of the task head.
A comparison of RS3L spaces fine-tuned for a Higgs vs. QCD classification task is presented in Fig. 6, which shows the receiver operating characteristic (ROC) curve calculated from nominal jets only. The QCD rejection rate (1/ false positive rate) is shown in Table II for various Higgs tagging efficiencies, again computed from only nominal jets. Compared to a fully supervised network trained on 8M jets, a fine-tuning using 3M jets on top of a 5M RS3L space yields a similar performance. If we instead match the yields of the training dataset of the fully supervised approach and the finetuning (i.e., both training on 3M labeled samples), a gain in performance is observed for the RS3L + fine-tuning strategy. This indicates that a common, large pretraining, for which RS3L can be one possible strategy, allows for efficient finetuning on comparably much smaller dataset sizes than a fully supervised setup due to the "hot start." Furthermore, it is important to note that all training configurations that use floating networks are more performant than networks trained with fixed weights on top of the RS3L backbone. This is to be expected because in the fixed weight case the model is severely restricted, as the only trainable parameters belong to the MLP head. Lastly, all networks outperform N 2 by a wide margin.
The performance of the RS3L space fine-tuned on the indomain Higgs vs. QCD classification task is shown in ROC curves Fig. 7 with separate ROC curves for each individual QCD subprocess. Note that the training sample consists of an admixture with similar weightings of all QCD subprocesses. The top (bottom) shows a graph with floating (fixed) weights. Superior performance is observed for a graph with floating weights in particular for light flavor QCD jets. This is an interesting observation, as it seems that the handles on flavor tagging such as the impact parameters of charged jet constituents get exploited during the finetuning step relative to the pretraining stage. For the finetuning with floating weights, all light-flavor subprocesses end up with roughly the same ROC curves, while for finetuning with fixed weights, the discrimination against single-quark light flavor jets (labeled "q") is much better than against light-flavor gluon splittings (labeled "g(qq)"). In addition, the discrimination power against g → b b to first order does not stem from flavor information, as the Higgs boson also decays into a pair of b quarks. Hence, a direct comparison with the ROC curve for N 2 from Fig. 6 allows an understanding of the substructure usage, with even fixed-weight RS3L improving over N 2 as much as 50% for a Higgs acceptance of 0.2.
TABLE III. Wasserstein distances between distributions of tagger output for nominal and augmented jets. Each column indicates the distances between the nominal and the respective augmentation. The number of training events is included for each training setup. Finetunings are performed on a 5M RS3L pretraining with either fixed or floating backbone weights. The inclusion of specific augmentations (e.g. "seed" or "FSR") in the training setup indicates that the pretraining and fine-tuning or supervision is performed on that set of augmentations. Uncertainties are derived from three independent trainings. The bolded numbers indicate the network with the smallest Wasserstein distances. Training setup Seed FSR up FSR down Herwig N 2 1.51 × 10 -4 2.49 × 10 -3 2.98 × 10 -3 1.22 × 10 -2 Fine-tuned (3M, fixed, all) (8.63 AE 0.23) ×10 -4 (4.85 AE 0.11Þ × 10 -4 (1.02 AE 0.02Þ × 10 -3 ð1.08 AE 0.01Þ × 10 -3 Fine-tuned (3M, floating, all) (8.23 AE 0.24) ×10 -4 (9.09 AE 0.50) ×10 -4 (1.32 AE 0.01) ×10 -3 (1.40 AE 0.04) ×10 -2 Fine-tuned (3M, floating, seed þ FSR) ð7.87 AE 0.32Þ × 10 -4 ð9.80 AE 0.27Þ × 10 -4 ð1.22 AE 0.02Þ × 10 -3 ð3.58 AE 0.04Þ × 10 -2 Fine-tuned (3M, floating, seed) ð8.51 AE 0.19Þ × 10 -4 ð9.50 AE 0.18Þ × 10 -4 ð1.36 AE 0.03Þ × 10 -3 ð3.35 AE 0.09Þ × 10 -2 Fully supervised (1M, all) (9.55 AE 0.89) ×10 -4 (1.33 AE 0.04) ×10 -3 (1.45 AE 0.05) ×10 -3 (2.08 AE 0.05) ×10 -2 Fully supervised (3M, all) (8.87 AE 0.20) ×10 -4 (1.14 AE 0.06) ×10 -3 (1.48 AE 0.09) ×10 -3 (2.18 AE 0.15) ×10 -2 Fully supervised (8M, all) ð8.20 AE 0.51Þ × 10 -4 ð9.91 AE 0.33Þ × 10 -4 ð1.29 AE 0.03Þ × 10 -3 ð1.32 AE 0.06Þ × 10 -2 Fully supervised (8M, seed þ FSR) ð8.68 AE 0.56Þ × 10 -4 ð9.95 AE 0.06Þ × 10 -4 ð1.18 AE 0.02Þ × 10 -3 ð3.87 AE 0.04Þ × 10 -2 Fully supervised (8M, seed) ð8.44 AE 0.15Þ × 10 -4 ð1.15 AE 0.01Þ × 10 -3 ð1.32 AE 0.03Þ × 10 -3 ð4.47 AE 0.32Þ × 10 -2
TABLE IV. QCD rejection rates for various training configurations and W efficiencies. The bolded numbers indicate the network with the best performance at a given W efficiency. W efficiency 0.3 0.5 0.7 Training setup 1=ðQCD efficiencyÞ Fine-tuned (1M, floating, all) 1589 AE 88 438 AE 9 134 AE 2 Fine-tuned (3M, floating, all) 1928 AE 31 504 AE 3 147 AE 1 Fully supervised (1M, all) 1288 AE 102 357 AE 16 114 AE 2 Fully supervised (3M, all) 1763 AE 18 459 AE 1 137 AE 2
In addition to background rejection, it is necessary to examine the classifier under other metrics. In particular, we consider how the network response varies under different parton shower configurations used in the simulator (which we call the network's "robustness"). The robustness is quantified by calculating the Wasserstein distance [53][54][55] between the one-dimensional distributions of the classifier output of nominal and, respectively, augmented jets as a measure of the similarity between the responses of the RS3L network to different inputs. The distances are presented in Table III. Each column shows the distance between the distribution for augmented jets (indicated by column title) and nominal jets. A lower Wasserstein distance indicates a network that is less sensitive to the showering configurations. As expected, the distance between nominal jets and jets reshowered with just a different numerical seed is smaller than the distance between nominal jets and any other augmentation ("FSR up," "FSR down," or "Herwig"). Especially the distance between nominal jets and HERWIG7-showered jets is significant.
Comparing the different training strategies, the most robust networks are trained using the RS3L method. With a fixed-weight approach, the Wasserstein distance is significantly smaller than for all other networks. This is a direct benefit of the RS3L pretraining strategy of minimizing the distance between positive pairs. We note the significant trade-off in terms of worse QCD rejection rates (Table II) for the fixed backbone. Including HERWIG7 augmentations in RS3L or, respectively, full supervision greatly reduces the distances compared to setups without HERWIG7 augmentations. Interestingly, we observe that already including seed and/or FSR augmentations in RS3L helps fine-tuned networks to be more robust against HERWIG7-induced domain shift compared to fully supervised strategies. Whether this indicates that a gain in robustness is potentially attainable simply by training the backbone with reseeded jets merits further investigation that we leave to future studies.
The least robust network is a fully supervised network trained only on jets with reseeded configurations (i.e., no systematic variations are seen during training). The deterioration is particularly significant in HERWIG7. More robust networks are derived by including all variations in the finetuning step. As expected, fully supervised networks are more robust when the training dataset is increased.
We now shift to out-of-distribution tasks to address the degree to which a general representation of showering jets was achieved through the RS3L pretraining. To investigate this, we fine-tune RS3L for a classification task that falls outside of the distributions used for pretraining, namely the task of discriminating between QCD jets and jets from "hadronic" W boson decays. Contrary to the Higgs boson decays, jets from W boson decays do not feature B hadrons. 2 The experimental signature therefore falls out of the distribution of the Higgs jets used to train the RS3L backbone.
We follow the same procedure as for the in-distribution classification studies by considering the same pretrained RS3L backbone but now fine-tuned for W vs. QCD classification. We compare with networks fully supervised from scratch for W vs. QCD classification. The results of this study are shown in Table IV. Networks fine-tuned show significant improvement (as much as 10%) in background rejection over fully supervised networks when trained on the same number of examples.
In terms of robustness, the Wasserstein distances between nominal and augmented jets are shown in terms of QCD background rejection in Table V and as a ROC in Fig. 8. The training setups with RS3L are found to be significantly more robust.
RS3L is a novel strategy that combines the concept of resimulation with a contrastive loss function to drive selfsupervised representation learning. The framework provides a natural way for creating a powerful foundation model which can be used for various downstream tasks such as classification and uncertainty mitigation. By encapsulating systematic uncertainties as well as the stochastic variability of the simulation into data augmentations, a latent space with improved robustness compared to other learning strategies can be obtained.
RS3L has been successfully applied to the canonical task of jet tagging, which is the ability to identify the type of TABLE V. Wasserstein distances between distributions of tagger output for nominal and augmented jets in the out-of-domain case. Each column indicates the distances between the nominal and the respective augmentation. The distance is calculated over a sample comprising equal amounts of W and QCD jets. The number of training events is included for each training setup. Finally, uncertainties are derived through measuring the Wasserstein distance for repeated trainings using the same configuration with random initialization.
Training setup Seed FSR up FSR down Herwig Fine-tuned W vs QCD (1M, floating, all) ð2.48 AE 0.95Þ × 10 -4 ð5.18 AE 0.03Þ × 10 -3 ð6.53 AE 0.01Þ × 10 -3 ð4.07 AE 0.05Þ × 10 -2 Fine-tuned W vs QCD (3M, floating, all) ð3.26 AE 0.24Þ × 10 -4 ð4.79 AE 0.09Þ × 10 -3 ð6.08 AE 0.01Þ × 10 -3 ð3.58 AE 0.06Þ × 10 -2 Fully supervised W vs QCD (1M, all) ð3.09 AE 0.39Þ × 10 -4 ð5.02 AE 0.19Þ × 10 -3 ð6.28 AE 0.32Þ × 10 -3 ð4.51 AE 0.39Þ × 10 -2 Fully supervised W vs QCD (3M, all) ð4.08 AE 0.24Þ × 10 -4 ð5.22 AE 0.01Þ × 10 -3 ð6.70 AE 0.13Þ × 10 -3 ð4.33 AE 0.19Þ × 10 -2 2
The authors note that the absolute tagging performance on the W vs. QCD and Higgs vs. QCD datasets cannot be compared directly because of the composition of the sample, see Sec. II. elementary particle at the origin of the evolution of a particle shower in detectors at the Large Hadron Collider. Here, the algorithm revealed identical performance in the limit of large-enough training statistics to state-of-the-art deep learning-based jet taggers that separate jets from Higgs bosons and from QCD partons, improved performance at smaller dataset sizes, and improved robustness against detrimental effects arising from systematic uncertainties. Additionally, it shows excellent transferrability to out-of-distribution tasks, thus having large potential to increase the efficiency of deep learning trainings in high energy physics when employed as a common pretraining.
Enforcing domain completeness through a strategy to embed the known modeling uncertainties within the selfsupervised space directly maps this work to the quality of the simulation used at training. Improved simulators can lead to improved embeddings, which, in turn, lead to a better understanding of the underlying physics. As a result, RS3L can adapt to the next generation of simulation modeling that stands to be substantially better than the current simulation. Through self-supervised approaches, we can continually improve the quality of the physics that is extracted, provided we continue to improve the physics model in the samples used for the self-supervision.
In future work, alternative self-supervised learning approaches to the employed SimCLR strategy, such as VICReg [56], SimSiam [57], and BarlowTwins [58] should be explored to study the practical impact of precise components of their loss functions. Additionally, the dataset size for the pretraining should be varied to systematically study the regimes in which the performances of SSL, fine-tuning, and fully supervised strategies may saturate. Finally, detailed comparisons should be made with other pretraining strategies proposed for HEP, such as masked particle modeling.
We observed that networks trained with 128 and 512 output dimensions performed similarly in the downstream tasks without notable improvement over N ¼ 8.