Pairton: Iterative Reconstruction of Short-Lived Particles

Pairton: Iterative Reconstruction of Short-Lived Particles Andreas Hermansen ,∗ Chris Scheulen ,† and Tobias Golling Département de physique nucléaire et corpusculaire, University of Geneva (Dated: August 17, 2026) We present Pairton, an iterative framework for reconstructing short-lived particles in high-energy collision events. By formulating particle reconstruction as a masked prediction process over graph structures, Pairton learns conditional distributions consistent with a factorised decomposition of decay products and iteratively predicts edges in the adjacency matrix representing particle decay relationships. Leveraging a pairformer-based architecture with dynamically updated pairwise representations, our method incorporates global event consistency. We demonstrate state-of-the-art performance on fully hadronic tt̄ decays. Pairton provides a general, flexible paradigm for particle reconstruction and can be readily extended to other topologies, bridging ideas from modern generative modelling and high-energy physics.

arXiv:2608.14278v1 [hep-ph] 14 Aug 2026

I.

INTRODUCTION

At the Large Hadron Collider (LHC) [1], short-lived particles, such as top quarks and Higgs bosons, are produced in abundance by colliding beams of protons at unprecedented energies. These particles decay almost instantaneously to form leptons and hadronic jets, collimated showers of hadrons, that can be detected by the various sub-detectors surrounding the collision point. Reconstructing these short-lived particles from their decay products is a crucial step in studying their properties and interactions. In processes with high particle multiplicities, such as those involving fully hadronic decay of a top quark pair, the combinatorial background can be significant. As a result, it is challenging to accurately identify the decay products of the short-lived particles. In this paper, we propose Pairton, a generative perspective on particle reconstruction. We formulate the problem as structured probabilistic inference over partially observed graphs, where nodes represent final state objects and edges encode their decay relationships. Pairton reconstructs missing edges iteratively by predicting them conditioned on previously inferred structure. We apply Pairton to the common benchmark of reconstructing the all-hadronic decay mode of top quark pairs (tt̄), where the full decay chain is tt̄ → W + b W − b̄ → q1 q̄1′ b q2 q̄2′ b̄. All six quarks hadronise and can be detected as jets in the detector, whereby the two jets originating from b quarks are referred to as b-jets, which can be distinguished from the jets originating from lighter quarks. To fully reconstruct the tt̄ system, one therefore needs to correctly identify which jets originate from the W bosons and which jets originate from the b quarks. Additional jets from gluon radiation and pile-up interactions are possible in

∗ andreas.hermansen@unige.ch † chris.scheulen@cern.ch

the collision events and complicate this task. Methods developed and tested on the tt̄ system can additionally be extended to other final states involving short-lived particles. The implementation of Pairton and the workflows for training, inference, and ablation studies are available at: https://github.com/AMHermansen/pairton.

II.

CURRENT APPROACHES

Traditional methods for short-lived particle reconstruction, such as the likelihood-based KLFitter algorithm [2], often rely on a combinatorial approach, where all possible combinations of detected particles are considered as potential decay products. In these algorithms, the combination which achieves the lowest value of a predefined cost function is selected as the best candidate. Alternatively, multiple machine learning methods have been developed in recent years to address this task: SPANet [3–6] uses a transformer-based architecture with an attention mechanism to capture the underlying symmetries of the system. Topograph [7] views the problem as a graph reconstruction task, wherein top quark and W boson candidates are injected as additional nodes, and the model classifies edges between jets and these candidates. The idea of framing the problem as a graph reconstruction task without the need for additional top and W boson nodes has already been discussed by the Topograph authors, but was not explored further there. HyPER [8] uses a hypergraph representation, where hyperedges of degree two represent the decays of W bosons, and hyperedges of degree three represent the decays of top quarks. TIGER [9] creates a topology agnostic, hierarchical graph neural network to infer particle decay assignments without assuming a fixed event topology. As a result, the model allows for general-purpose reconstruction and classification across varied processes. All of these machine learning methods achieve enhanced performance compared to traditional methods [2, 3, 7–9].

2 III.

METHOD

The reconstruction of the all-hadronic tt̄ topology involves assigning the correct final-state objects of two t → W b decays when neglecting the differentiation into particles and anti-particles. As a result, one has to find the two jets belonging to the W boson produced in each top quark decay, and assign a third jet for the b quark produced in the same decay. Generally, the precise assignment of the jets initiated by the quark and the anti-quark in the W boson decay is not performed in most physics analyses, such that permutations among these are permitted. Similarly, permuting the top quark and the anti-top quark is equally allowed. In summary, we are therefore interested in assigning a pair of jets for each W boson decay, and a third jet belonging to the corresponding b quark, resulting in a total of six assigned jets. Mathematically, the optimal reconstruction can be expressed as the most probable reconstruction, Ŵ1 , Ŵ2 , b̂1 , b̂2 = arg

max

W1 ,W2 ,b1 ,b2

P (W1 , W2 , b1 , b2 | x) ,

(1) where x represents the data associated with the event, and W1 , W2 , b1 , b2 represent the reconstructed W bosons and b-jets, respectively. We note that, while bi represents a single index, Wi represents an unordered set of two indices, when assigning a unique index to each jet in the collision event under consideration. By the chain rule of probability, the joint distribution factorises as the product of conditional probabilities. In the case of reconstructing an all-hadronic tt̄ decay, this results in

P (W1 , W2 , b1 , b2 | x) = P(W1 | x) × P(W2 | W1 , x) × P(b1 , b2 | W1 , W2 , x),

feed back into the model in subsequent particle assignments. In Appendix A, we establish the formal connection between the factorised probability distribution and our masked denoising procedure, which we present in the following.

Graph Representation

3

2

4

1 2 3 4 5 6 7 8 1 A A 2 A A 3 A A A 4A 5 A A 6 A A 7 A A A 8 A t,W

1

t

t

t,W

t

W

t

t,W

t

t,W

5

8 6

t

7

t,W

t

t

W

t,W

FIG. 1. Representation of an all-hadronic tt̄ event with eight jets as (left) a graph of the decay products of the individual tops and (right) the combined view of the corresponding top and W boson adjacency matrices At and AW , respectively. In the chosen example, one top quark decays into the jets 1, 4, and 7, where jet 7 originates from the b quark while jets 1 and 4 originate from the subsequent W boson decay. Similarly, the other top quark decays into jets 2, 3, and 6 with jet 3 originating from the b quark. Jets 5 and 8 are additional jets not originating from the top quark decays. In the adjacency matrix view, grey entries correspond to zeros, while the coloured entries correspond to graph edges (including self-loops) corresponding to the graph assignment on the left: Blue entries and solid edges are ones in both At and AW , green entries and dashed edges are ones in At only, and purple entries are ones in AW only.

(2)

where the indexing of the W bosons and b quarks is arbitrary. Each of the conditional probability terms can be understood as a partial event reconstruction step of the corresponding particle if the reconstruction is performed iteratively. Ablation experiments included in Appendix C show no measurable benefit from further factorising the distribution of the b-jets. Therefore, we keep the b-jets together in the last term to reduce the number of function evaluations on our model. For completeness, we list the validation performance of fully factorising the joint distribution in Appendix C. To learn the conditional distributions, we reframe the problem using a graph representation of the particle assignments in order to deploy a masked denoising algorithm in our approach. As a result, the conditional probabilities introduced in Equation (2) can naturally be represented in iteratively updated graphs, which we

We represent each event as a collection of simple graphs Gα = (V, Eα ) with a shared set of nodes. Each graph corresponds to a particle type α ∈ {W, t} for the allhadronic tt̄ decays that we are considering. The node set V corresponds to detected jets. An edge (i, j) ∈ Eα indicates that jets i and j are decay products of particle α. Self-loops encode that a jet is not associated with the corresponding particle type. This representation allows us to encode particle assignments as adjacency matrices Aα = aα ij defined through aα ij =

( 1 0

if (i, j) ∈ Eα , otherwise.

(3)

This is illustrated in Figure 1. Beyond providing a natural representation of the physical decay topology, formulating the problem via adjacency matrices offers an architectural advantage. It allows the partially observed adjacency matrices, Ãα , to be directly embedded and incorporated into the model as

3 dynamic pairwise edge features. By treating these structural assignments as input features, a graph neural network can seamlessly condition its subsequent predictions on the explicitly represented relational state between all jets without requiring complex message-passing operations over higher-order tensors.

3

2

4

t,W

5

8 6

7

Inspired by discrete diffusion [10] we formulate the reconstruction problem as a discrete denoising problem over adjacency matrices defined above. Let A denote the ground-truth adjacency matrices. We then define a partially-deterministic Markov chain responsible for the forward masking process, which produces partially observed adjacency matrices Ã,

3

2

t ∈ {0, 1, 2},

(4)

t,W

1

Masked Denoising Formulation

qt (Ã | A),

1 2 3 4 5 6 7 8 1 A 2 3 4A 5 6 7 8 1 2 3 4 5 6 7 8 1 A 2 A 3 A 4A 5 A 6 A 7 A 8 A t,W

t,W

4

1

W

t,W

5

8

W

t,W

where the masking level t determines which particles are visible:

6

W

7

W

• both W bosons revealed, • one W boson revealed, 3

• no particles revealed. Entries logically inferable from revealed particles are filled with their true values in Ã, while remaining entries are replaced with a mask token M. This is a special case of the absorbing state discrete diffusion processes used in language modelling adapted to the physical system. Instead of relying on fully random masking directly applied to the adjacency matrices, we opt for this particular Markov chain to accommodate the hierarchical nature of the tt̄ decay, but note that any Markov chain that eventually ends up with no revealed particles will work. This particular Markov chain effectively exploits the physical nature of the tt̄ decay to define an efficient corruption scheme allowing us to reverse the process using few steps. The hierarchical corruption scheme directly mirrors the physical decay topology of the tt̄ system, allowing the model to unambiguously condition on previously predicted particles at each step of the reverse process. Conversely, an unstructured approach that randomly masks individual edges instead would frequently yield intermediate configurations lacking a clear physical interpretation. Furthermore, such random masking risks creating states where missing edges could be deterministically inferred due to the strict structural constraints of the adjacency matrices. The simplicity of the above process is a consequence of the hierarchical and symmetric nature of the tt̄ process, which limits the number of meaningful intermediate masking states. More complicated topologies like the associated production of a Higgs boson with a top quark pair, tt̄H, would have a greater number of semantically meaningful masking states, such as one visible

2

4

1 2 3 4 5 6 7 8 1 A A 2 A A 3 A A A 4A 5 A A 6 A A 7 A A A 8 A t,W

1

t

t

t,W

t

W

t

t,W

t

t,W

5

8 6

7

t

t,W

t

t

W

t,W

FIG. 2. Sequential reconstruction of W bosons and top quarks in three steps for the event depicted in Figure 1 following the same visual style. Additionally, white entries indicate unknown adjacency, while the scratched entries indicate that no edge can exist in the W boson edge-set, so that the associated jets cannot be decay products of the same W boson, but might still be part of the top quark edge-set.

W boson or one visible Higgs boson, which in turn would lead to richer and more expressive Markov processes. We train a model fθ (Ã, x) to predict masked entries conditioned on the observed structure and jet features. The training objective in this case is " # X X
α α α L(θ) = E δM (ãij ) ℓbce aij , (fθ (x, Ã))ij , (5) α i,j∈V

where δM is 1 if the entry is masked and 0 otherwise, and the expectation runs over both the training data and over the masking level t. This objective corresponds to masked conditional likelihood training, analogous to the discrete diffusion models introduced by LLaDA [11]

4 for natural language processing. More broadly, this paradigm of absorbing state discrete diffusion is a compelling alternative generative framework for large language models compared to the more traditional nature of left-to-right auto-regressive decoding [10, 12–14]. While masked and predictive methods have gained traction in high-energy physics, they have largely been used in BERT-style models with a fixed masking rate (e. g. Masked Particle Modeling on Sets [15]) or standard GPT-style auto-regressive masking (e. g. OmniJetα [16]). Our work introduces discrete diffusion to highenergy physics. By formulating event reconstruction as a discrete diffusion process over the adjacency matrices, our method learns to reverse the masking process, which enables it to iteratively predict the structural topology of an event. At inference time, reconstruction proceeds by iteratively predicting particle assignments, thus corresponding to the reverse of the masking process. Starting from a fully masked adjacency matrix, the model predicts particles and refines new predictions with previously selected structures, reflecting the factorisation as outlined in Equation (2). A visual illustration of this approach is provided in Figure 2. We approximate the maximum a posteriori (MAP) assignment by greedily selecting edges that maximise the sum of predicted logits, X Zij logit pij , (6) i≤j

where Zij indicates selected edges, and pij = fθ (x, Ã)ij is the corresponding model output. We note that it would be possible to explore more sophisticated sampling algorithms like beam searches [17], but leave the exploration of such methods for future work. The pseudo-code for our full reconstruction algorithm is included in Appendix D. It is important to note that while our algorithm used at inference is deterministic, our model is learning the conditional distributions outlined in Equation (2) as shown for completeness in Appendix A.

Model Architecture

To model pθ (aij | Ã, x), we use a pairformer-based architecture initially developed as the primary component of the AlphaFold3 architecture for protein structure prediction [18]. Conceptually, the architecture is a modified version of a transformer, where the model also uses pairwise features to augment the attention scores, Sij , with an additional bias term qi · k j Sij = √ + bij . d

(7)

Here, bij is a learnable function of the pairwise features. This is conceptually very similar to the particle attention mechanism used in the Particle Transformer (ParT),

Loss/ Predict

W Adjacency Prediction (N x N)

Top Adjacency Prediction (N x N)

Linear Outer sum

Linear Outer sum Latent jet features (N x Ls)

Pairformer Stack Mask Level ∅ → W1 → W1W2

Pairwise Features (N x N x Lp)

Jet features (N x Fs)

Pairwise jet features (N x N x Fp)

Visible Adjacency

True Adjacencies (N x N x 2)

FIG. 3. Visual illustration of the model during training and inference. During training, the mask level is chosen at random and used together with the true adjacency matrices to construct the visible adjacency matrices given to the model. During inference, the visible adjacency matrices are constructed iteratively. The true adjacency matrices are also used as targets for the loss calculation during training.

which has achieved remarkable performance in jet tagging [19]. However, where ParT uses a fixed pairwise representation, the pairformer architecture dynamically updates the pairwise features. Edge predictions are computed as
α âα (8) ij = gϕ W x̃i + W x̃j , where W is a learnable linear projection, gϕα is an MLP, and x̃i are latent node embeddings. A visual model representation during training and inference is shown in Figure 3. As the input jet features, we use the jet energy E, transverse momentum pT , pseudo-rapidity η, azimuthal angle φ, and whether or not the jet was b-tagged. Additionally, we use the pairwise jet features of the di-jet invariant mass mij , the difference in pseudo-rapidity ∆ηij and azimuthal angle ∆φ p ij between the jets, and the angular distance ∆Rij = (∆φij )2 + (∆ηij )2 of the jets from each other. During the model pre-processing, some of these input features are transformed as listed in Table I. Additionally, the visible adjacency matrices determined from the true adjacency matrices and the mask step during training or from the previous inference steps during inference are embedded through an embedding layer and

5 TABLE I. Feature transformation applied for both single- and pairwise jet features. Only features for which transformations are applied are listed, while the other features mentioned in the text are propagated to Pairton as is. Feature

Feature(s) after Transformation

E pT

log(E/GeV) − 4.5 log(pT /GeV) − 4.5 cos φ sin φ

φ

∆φij mij

∆φij ∆(sin φ)ij ∆(cos φ)ij log(mij /GeV) − 4.5

added to the embedded pairwise jet features before passing the features into the pairformer stack as shown in Figure 3.

IV.

RESULTS Dataset

For model training and evaluation, the dataset originally published in conjunction with HyPER [20] was used. It contains a total of 6 × 107 all-hadronic √ tt̄ events simulated at a centre-of-mass energy of s = 13 TeV using MadGraph5 aMC@Nlo 2.9.16 [21] at next-toleading order in QCD for the matrix element calculation and Pythia 8.306 [22] for the parton shower. For the detector response, Delphes 3.5.0 [23] was used in a configuration simulating the ATLAS detector [24]. Jets were reconstructed using the anti-kt algorithm [25] implemented in FastJet [26] with a radius parameter of R = 0.4. The resulting jets are required to pass a transverse momentum threshold of pT = 25 GeV and a cut in the absolute pseudo-rapidity of |η| < 2.5. The tagging of b-jets was simulated by applying a pT -dependent tagging efficiency based on Reference [27] to jets containing b-hadrons. Reconstructed objects were matched to finalstate particles of the matrix element generator with a maximal angular distance of ∆R < 0.4. All events are required to contain at least six jets, of which at least two are b-tagged. The dataset containing events passing the selection criteria was randomly partitioned into training, validation, and testing sub-sets containing 90 %, 5 %, and 5 % of events, respectively. After applying the event selection criteria described above and requiring at least one fully reconstructible W boson in each event, roughly 8.8 × 106 events remain in total for training and validating the models. Out of these, roughly 2.4 × 106 are fully reconstructible. For the exact numbers and a complete description of the event generation setup we refer to Reference [8].

TABLE II. Comparison of model reconstruction efficiencies for full-event (εtt̄ ), top quark (εt ), and W boson (εW ) reconstruction. The values for SPANet, Topograph, and HyPER are taken from Reference [8]. The uncertainties of Pairton were calculated by varying the initial random initialisation and shuffling the training batches five times. Since only statistical uncertainties of the test dataset size were considered for the other models, we only quote the central value for these, and note that the statistical uncertainties are generally around 0.001. While only fully reconstructible events were considered for the full-event reconstruction efficiencies, reconstructible W bosons and top quarks from partially reconstructible events were additionally considered for their respective particle reconstruction efficiencies. Jets

SPANet Topograph HyPER

Pairton

εtt̄

6 7 ≥8 Inclusive

0.801 0.654 0.487 0.651

0.805 0.649 0.472 0.645

0.816 0.672 0.513 0.670

0.8298 ± 0.0004 0.6988 ± 0.0011 0.5459 ± 0.0004 0.6944 ± 0.0003

εt

6 7 ≥8 Inclusive

0.667 0.649 0.586 0.640

0.658 0.638 0.573 0.629

0.663 0.654 0.601 0.644

0.6760 ± 0.0008 0.6693 ± 0.0005 0.6195 ± 0.0002 0.6589 ± 0.0005

6 7 εW ≥8 Inclusive

0.688 0.671 0.608 0.662

0.683 0.662 0.594 0.653

0.697 0.683 0.628 0.675

0.7010 ± 0.0018 0.6925 ± 0.0012 0.6409 ± 0.0008 0.6826 ± 0.0013

Reconstruction Performance and Statistical Stability

As the central performance metric, we record the efficiencies, also referred to as true positive rates, of reconstructing full events (εtt̄ ), top quarks (εt ), and W bosons (εW ), representing a common way of determining the performance of reconstruction algorithms [3, 7, 8]. Due to the high degree of symmetry of the fully hadronic tt̄ decay, we record a W boson as correctly reconstructed if the two associated jets are correctly identified, although they may be permuted. For a top quark to be correctly reconstructed, its daughter W boson and the remaining b quark must be correctly predicted. For the full tt̄ system to be correctly reconstructed, both the top and antitop quark have to be correctly reconstructed. However, the prediction of the top and anti-top quarks may be permuted, since jet charge measurements of the decay products of fully hadronic tt̄ decays are challenging and commonly not attempted in analyses of the process. Using these metrics, Pairton significantly outperforms current approaches when compared to SPANet, Topograph, and HyPER as shown in Table II, where the efficiency figures for the reference methods are taken from Reference [8]. The largest improvements can be observed at high jet multiplicities, where reconstructing the full tt̄ topology is especially difficult due to large jet combinatorics and a crowded detector environment. Since TIGER [9] aims to be topology agnostic, only the efficiency where the model had been trained on all events

Event Efficiency ( tt)

6

0.80 0.75 0.70 0.65 0.60 0.55 0.50 0.45 0.40

find that Pairton outperforms earlier methods across the board, with the largest gains observed for low mtt̄ , where the reconstruction efficiencies are lowest.

Pairton HyPER Topograph SPANet

Ablation Studies

) ) ) ) [0 , 3 5 0 [3 5 0 , 4 5 0 [4 5 0 , 5 5 0 [5 5 0 , 6 5 0

650

Mass Interval (mtt) [GeV]

FIG. 4. Event reconstruction efficiencies of the different models in dependence on mtt̄ , the invariant mass of the tt̄ system as determined from truth-matched final-state reconstructed jets. While Pairton consistently outperforms the other methods with HyPER showing the second-highest efficiencies, Topograph outperforms SPANet in the low-mass region while their ordering is reversed for high mtt̄ . The performance metrics of HyPER, SPANet, and Topograph are taken from Reference [8].

is included in Reference [9]. However, while Pairton requires a complete ground-truth assignment and likewise predicts full assignments, TIGER is trained on the entire dataset and can predict partial reconstructions. As a result, a fair direct comparison of efficiencies is not possible, since both the training statistics and the models differ between the two. We note, however, that the efficiencies achieved by TIGER are comparable to those of HyPER. For Pairton, uncertainties were estimated by retraining the model five times with different random initialisations and shuffled training batches. The resulting variations are small, indicating stable optimisation and low sensitivity to initialisation. Equivalent retraining statistics were not available for the comparison methods. Instead, their uncertainties reported in Reference [8] correspond to statistical uncertainties of the test dataset only and are therefore not included in Table II. In addition to the phase-space inclusive reconstruction efficiencies, full-event reconstruction efficiencies were also calculated in dependence on the mass of the tt̄ system as determined from the truth-matched reconstructed jets in fully reconstructible events. This differential efficiency is especially interesting to analyses interested in only measuring parts of the invariant mass spectrum, such as quantum entanglement measurements [28, 29] or some analyses investigating toponium states at the tt̄ production threshold [30]. The corresponding comparison of Pairton with existing methods is shown in Figure 4. Similar to the comparisons split by jet multiplicity, we

To evaluate the individual contributions of our architectural choices and the iterative reconstruction procedure, we performed a series of ablation studies, whose full details are provided in Appendix C. First, we compared the core pairformer architecture against a standard transformer and a bias-augmented transformer baseline. The pairformer yielded substantial improvements, achieving roughly 5 % absolute increase in full-event reconstruction efficiency over the standard transformer. Second, we tested the impact of our iterative prediction formulation. We found that our 3-step generative process outperforms both single-step prediction and a hyperedge based prediction modelled after the prediction head used in Reference [8]. Furthermore, simply increasing the model depth to match the computational cost of the iterative process did not yield any increased performance. This confirms that the observed state-of-the-art gains originate fundamentally from the iterative nature of our reconstruction algorithm in addition to the performance gains due to the pairformer architecture.

V.

CONCLUSION

We have presented Pairton, a new approach to the reconstruction of short-lived particles by reformulating the problem as a discrete diffusion process over graphs. By factorising the joint distribution over reconstructed particles and training a masked edge predictor to learn the corresponding conditional distributions, we obtain a probabilistic formulation for event reconstruction. Conceptually, Pairton differs in two important ways from previous approaches. First, inspired by absorbing state discrete diffusion language modelling, we perform reconstruction iteratively by conditioning on previously partially reconstructed structures. This enables the model to incorporate global consistency constraints in a natural and probabilistically sound manner. Second, to facilitate this conditioning, our model takes as additional inputs previously predicted structures. In order to achieve this more smoothly we opt for a prediction method, that purely relies on the partially predicted adjacency matrices, since these can be seamlessly used as pairwise edge features in a graph neural network. Using a pairformer-based architecture with dynamically updated pairwise representations, we achieve stateof-the-art performance on the fully hadronic tt̄ reconstruction benchmark. Across all jet multiplicities, Pairton consistently outperforms current methods in fullevent, top quark, and W boson reconstruction efficiencies. The largest improvements can be observed at high

7 jet multiplicities. The presented framework is general and can be extended beyond the all-hadronic tt̄ system. While our masking process follows a relatively straightforward hierarchical strategy, the discrete diffusion formulation enables the definition of arbitrary noising processes provided they transform the fully visible adjacency matrices to fully masked ones. In particular, more complex topologies, such as tt̄H reconstruction, would inherently be less linear, since the ordering of predicting W boson edges and Higgs boson edges could be made interchangeable. We believe that iterative graph reconstruction provides a powerful and flexible paradigm for particle reconstruction, bridging ideas from modern generative modelling and high-energy physics event interpretation.

[1] L. Evans and P. Bryant, LHC Machine, JINST 3 (08), S08001. [2] J. Erdmann, S. Guindon, K. Kröninger, B. Lemmer, O. Nackenhorst, A. Quadt, and P. Stolte, A likelihoodbased reconstruction algorithm for top-quark pairs and the KLFitter framework, Nucl. Instrum. Meth. A 748, 18 (2014). [3] M. J. Fenton, A. Shmakov, T.-W. Ho, S.-C. Hsu, D. Whiteson, and P. Baldi, Permutationless many-jet event reconstruction with symmetry preserving attention networks, Phys. Rev. D 105, 112008 (2022). [4] A. Shmakov, M. J. Fenton, T.-W. Ho, S.-C. Hsu, D. Whiteson, and P. Baldi, SPANet: Generalized permutationless set assignment for particle physics using symmetry preserving attention, SciPost Phys. 12, 178 (2022). [5] M. J. Fenton, A. Shmakov, H. Okawa, Y. Li, K.-Y. Hsiao, S.-C. Hsu, D. Whiteson, and P. Baldi, Reconstruction of unstable heavy particles using deep symmetry-preserving attention networks, Commun. Phys. 7, 139 (2024). [6] H. Li et al., Reconstruction of boosted and resolved multi-Higgs-boson events with symmetry-preserving attention networks, JHEP 2025, 119. [7] L. Ehrke, J. A. Raine, K. Zoch, M. Guth, and T. Golling, Topological reconstruction of particle physics processes using graph neural networks, Phys. Rev. D 107, 116019 (2023). [8] C. Birch-Sykes, B. Le, Y. Peters, E. Simpson, and Z. Zhang, Reconstructing short-lived particles using hypergraph representation learning, Phys. Rev. D 111, 032004 (2025). [9] N. Soybelman, N. Kakati, E. Gross, and F. A. Di Bello, Topology-agnostic, hierarchical graph network for event reconstruction, Phys. Rev. D 113, 012014 (2026). [10] J. Austin, D. D. Johnson, J. Ho, D. Tarlow, and R. van den Berg, Structured Denoising Diffusion Models in Discrete State-Spaces (2023), arXiv:2107.03006 [cs.LG]. [11] S. Nie, F. Zhu, Z. You, X. Zhang, J. Ou, J. Hu, J. Zhou, Y. Lin, J.-R. Wen, and C. Li, Large Language Diffusion Models (2025), arXiv:2502.09992 [cs.CL].

VI.

ACKNOWLEDGMENTS

We would like to thank Ethan Simpson and Zihan Zhang for stimulating discussions providing insights into HyPER and for their kind support in making available to us additional reference model performance metrics. This work is supported by the European Union’s Horizon Europe research and innovation program under the Marie Sklodowska-Curie grant agreement No 101168829 called “Challenging AI with Challenges from Physics: How to solve fundamental problems in Physics by AI and vice versa (AIPHY)”. We would also like to acknowledge funding through the SNSF project grant 200020 212127 called “At the two upgrade frontiers: machine learning and the ITk Pixel detector”. The training and evaluation of the models developed in this work were performed at the University of Geneva using the Baobab high-performance computing cluster.

[12] S. S. Sahoo, M. Arriola, Y. Schiff, A. Gokaslan, E. Marroquin, J. T. Chiu, A. Rush, and V. Kuleshov, Simple and Effective Masked Diffusion Language Models, in Advances in Neural Information Processing Systems, Vol. 37, edited by A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (Curran Associates, Inc., 2024) pp. 130136–130184. [13] J. Ou, S. Nie, K. Xue, F. Zhu, J. Sun, Z. Li, and C. Li, Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data, in International Conference on Learning Representations, Vol. 2025, edited by Y. Yue, A. Garg, N. Peng, F. Sha, and R. Yu (2025) pp. 64972–65009. [14] A. Lou, C. Meng, and S. Ermon, Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution, in Proceedings of the 41st International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 235, edited by R. Salakhutdinov, Z. Kolter, K. Heller, A. Weller, N. Oliver, J. Scarlett, and F. Berkenkamp (PMLR, 2024) pp. 32819–32848. [15] T. Golling, L. Heinrich, M. Kagan, S. Klein, M. Leigh, M. Osadchy, and J. A. Raine, Masked particle modeling on sets: towards self-supervised high energy physics foundation models, Mach. Learn. Sci. Tech. 5, 035074 (2024). [16] J. Birk, A. Hallin, and G. Kasieczka, OmniJet-α: the first cross-task foundation model for particle physics, Mach. Learn. Sci. Tech. 5, 035031 (2024). [17] P. S. Ow and T. E. Morton, Filtered beam search in scheduling, Int. J. Prod. Res. 26, 35 (1988). [18] J. Abramson et al., Accurate structure prediction of biomolecular interactions with AlphaFold 3, Nature 630, 493 (2024). [19] H. Qu, C. Li, and S. Qian, Particle Transformer for Jet Tagging, in Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 162, edited by K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato (PMLR, 2022) pp. 18281–18292. [20] Z. Zhang, C. Birch-Sykes, B. Le, Y. Peters, and E. Simp-

8 son, Top quark pair production with fully hadronic final states at the LHC, 10.5281/zenodo.10653837 (2024). [21] J. Alwall, R. Frederix, S. Frixione, V. Hirschi, F. Maltoni, O. Mattelaer, H. S. Shao, T. Stelzer, P. Torrielli, and M. Zaro, The automated computation of tree-level and next-to-leading order differential cross sections, and their matching to parton shower simulations, JHEP 2014, 079. [22] C. Bierlich et al., A comprehensive guide to the physics and usage of PYTHIA 8.3, SciPost Phys. Codeb. 2022, 8 (2022). [23] J. de Favereau, C. Delaere, P. Demin, A. Giammanco, V. Lemaı̂tre, A. Mertens, and M. Selvaggi (DELPHES 3 Collaboration), DELPHES 3: a modular framework for fast simulation of a generic collider experiment, JHEP 2014, 057. [24] ATLAS Collaboration, The ATLAS Experiment at the CERN Large Hadron Collider, JINST 3, S08003. [25] M. Cacciari, G. P. Salam, and G. Soyez, The anti-kt jet clustering algorithm, JHEP 2008, 063. [26] M. Cacciari, G. P. Salam, and G. Soyez, FastJet User Manual, Eur. Phys. J. C 72, 1896 (2012). [27] ATLAS Collaboration, Expected performance of the ATLAS b-tagging algorithms in Run-2 , Tech. Rep. ATLPHYS-PUB-2015-022 (CERN, 2015). [28] ATLAS Collaboration, Observation of quantum entanglement with top quarks at the ATLAS detector, Nature 633, 542 (2024). [29] CMS Collaboration, Observation of quantum entanglement in top √quark pair production in proton–proton collisions at s = 13 TeV, Rept. Prog. Phys. 87, 117801 (2024). [30] ATLAS Collaboration, Observation of a cross-section enhancement near the tt̄ production threshold in √ s = 13 TeV pp collisions with the ATLAS detector, Rept. Prog. Phys. 89, 057801 (2026). [31] A. Paszke et al., Pytorch: An imperative style, highperformance deep learning library (2019). [32] W. Falcon et al., PyTorch Lightning (2020). [33] F. Mölder et al., Sustainable data analysis with Snakemake, F1000Research 10, 33 (2025).

Appendix A: Connection between Masked Denoising and Auto-regressive Factorization

In this section, we clarify the relationship between the training objective and the learned conditional distributions. The training objective is given by     X X α α  L(θ) = E  δM (ãα . ij ) ℓbce aij , (fθ (x, Ã))ij α i,j∈V

(A1) It is well known that, for a Bernoulli variable y, the minimiser of the binary cross-entropy satisfies f ∗ (x) = P(y = 1 | x).

qt (Ã | A),

(A4)

which determines which variables are observed and which are predicted. In our setting, the masking process is structured such that: • when no particles are revealed, all entries are masked, • when one W boson is revealed, all entries implied by that W boson are observed, • when both W bosons are revealed, all entries implied by them are observed. Therefore, the learned conditional distributions take the form

P(aαij | x, Ã),

(A5)

where à encodes the currently revealed particle assignments. Because the masking process reveals particles in stages, the model predictions can be associated explicitly to the conditional probabilities: • The edges of a W boson correspond to modelling P ({W1 = (i, j)} ∪ {W2 = (i, j)} | x), and since the events W1 = (i, j) and W2 = (i, j) are mutually exclusive, this becomes P (W1 | x) + P (W2 | x). Since the ordering of the W bosons is arbitrary this effectively models P (W | x). • The probabilities of the edges corresponding to the second W boson map directly to P (W2 = (i, j) | W1 , x). • The top quark edges connecting a node k of a b quark to an already existing W = (i, j) and a second visible W correspond to P({b1 = k} | {W1 = (i, j)} , W2 , x). Thus, the training procedure is consistent with learning the conditional distributions appearing in an autoregressive factorisation of the joint distribution over particle assignments. By reversing this specific masking chain, the inference procedure exactly mirrors the chain rule factorisation of the joint distribution presented in Equation (2).

(A2)

Applying this result entry-wise, we obtain that the optimal predictor satisfies α (f ∗ (x, Ã))α ij = P(aij = 1 | x, Ã).

Thus, the model learns the conditional distribution of each adjacency entry given the observed structure and input features. The masking process defines a distribution over partially observed adjacency matrices,

(A3)

Appendix B: Model Hyperparameters

For the hyperparameters of the pairformer we generally follow the original hyperparameters from Reference [18].

9 TABLE III. Summary of the most important hyperparameters. Name

Value

Number of pairformer blocks 4 Sequence embedding dimension 128 Pairwise embedding dimension 64 Optimiser AdamW Learning rate 0.001 Learning rate scheduler Cosine annealing Optimiser steps 826 k (100 epochs) Batch Size 256

However, we use a much smaller architecture and replace LayerNorm with the now more common RMSNorm. A summary of hyperparameters used is included in Table III. Overall, the parameters chosen resulted in a model with approximately 1.2 × 106 learnable parameters. All models were implemented in python v3.12.3 (https://www.python.org) using pytorch v2.8.0 [31] and trained using pytorch-lightning v2.6.1 [32]. For workflow management we used snakemake [33]. Our models were trained on a single RTX 3080 unless otherwise specified. Despite the iterative nature of our model, it achieves a reconstruction throughput of roughly 2000 events per second during inference.

Appendix C: Model Ablation Studies

We have performed two ablation studies on details of the neural network architecture and the reconstruction method. In both studies, efficiencies were evaluated on the validation set in order to not bias the performance estimates provided on the test set. As a result, the listed performance of our final model differs slightly between the ablation studies and the efficiency comparison listed in Table II, which was evaluated on the test set. In the first study on the choice of neural network architecture, we compared the transformer architecture, which has been widely adopted in many physics applications, a transformer with a bias term in the attention matrix similar to the ParT architecture used for jet-tagging, and the pairformer architecture. To keep the compute budgets roughly similar for a fair model comparison, we kept the dimension of the neural networks fixed at 128 for the sequence embeddings, while a dimension of 64 was used for the pairwise embedding. To compensate for the fact these architectures do not have the same number of parameters per layer, we decided to increase the number of layers to eight for the transformer (six for the transformer with bias), while the pairformer maintained four layers. We found that this gave roughly the same throughput for each model during training. We note that this is not the only way to increase complexity of the alternative models, but a full hyperparameter search for all models was outside the

TABLE IV. Comparison of the reconstruction efficiencies from the ablation study comparing different neural network architectures. The efficiencies shown are the maximum efficiencies achieved on the validation set during model training. Model

εtt̄

εt

εW

Transformer 0.635 0.623 0.655 Transformer with bias 0.668 0.649 0.678 Pairformer 0.688 0.656 0.682 Iterative pairformer 0.692 0.661 0.684

scope of this ablation study. The remaining hyperparameters were held constant between the models. Since the standard transformer is unable to use pairwise features, we decided to not make information from the partially masked adjacency matrices available to any model during the ablation study, as opposed to the architecture of the final model described in Section III. Stated mathematically, for this ablation study we define our models such that fθ (x, ã) = fθ (x, ∅).

(C1)

To capture the gain achieved from access to the partially masked adjacency matrices, we also trained an “Iterative pairformer”, which is conditioned on these matrices. Table IV shows convincingly that the pairformer architecture achieves superior performance compared to the more traditional transformer-based architectures. The conditioning of the iterative pairformer leads to a further performance improvement. This observation is also exemplified in our second ablation study, which specifically compares the effects of our iterative prediction method against other reconstruction methods. As a baseline, we included the unconditioned pairformer from the previous ablation study and added the iterative pairformer running on three steps: the reconstruction of the first W boson, the reconstruction of the second W boson, and the simultaneous reconstruction of the two b quarks. In addition to this 3-step iterative model, we included a 4-step generative procedure where the b quark prediction was factorised instead of the simultaneous prediction of both b quarks. We also compare to the hyperedge-based reconstruction method utilised by HyPER, where we use a pairformer architecture as the embedding graph neural network instead of the graph neural network trained in the original HyPER paper. Since the HyPER method needs to materialise the hyperedges, we found that the resulting model required slightly more memory than our reconstruction method, which meant it just barely did not fit on an RTX 3080 GPU. As a consequence, it was trained on an RTX 3090 instead. In order to investigate the effect of a larger test-time compute budget, we included a deep pairformer architecture as a last reference model. This model consists of twelve pairformer blocks as opposed to the four blocks

10 b quarks

correct: 80 967 (0.953) partial: 1533 (0.018) incorr.: 2496 (0.029)

second-highest scoring W boson

first W boson one-shot model

second W boson

b quarks

correct: 85 272 (0.820)

correct: 81 449 (0.955) partial: 1379 (0.016) incorr.: 2444 (0.029)

iterative model

correct: 84 996 (0.818) correct: 103 971 (0.885)

correct: partial: incorr.:

8208 (0.461) 9450 (0.498) 1317 (0.069)

1528 (0.371) 2094 (0.509) 493 (0.120)

correct: partial: incorr.:

4731 (0.503) 3916 (0.417) 754 (0.080)

8107 (0.434) 9266 (0.496) 1326 (0.071)

correct: partial: incorr.:

1471 (0.376) 1997 (0.510) 445 (0.114)

correct: partial: incorr.:

4831 (0.503) 4038 (0.420) 734 (0.076)

incorr.: 18 699 (0.180)

incorr.: 18 975 (0.182) correct: partial: incorr.:

correct: partial: incorr.:

correct: 4115 (0.304)

correct: 3913 (0.290)

incorr.: 13 516 (0.115)

incorr.: 9401 (0.696)

incorr.: 9603 (0.710)

FIG. 5. Comparison of the reconstruction performance of the 3-step iterative pairformer model (towards the right) and a oneshot pairformer model (towards the left) on a total of 117 487 fully reconstructible all-hadronic tt̄ events of the validation set. To directly compare the model behaviour, the reconstruction of the first W boson is common between the two methods. Afterwards, the iterative model reconstructs further particles in additional model passes using partial adjacency matrices from previous iterations as described in Section III. Meanwhile, all particles are reconstructed from the first model pass in a greedy approach for the one-shot model, whereby the first W boson selected is the one maximising the edge assignment heuristic detailed in Appendix D. In each reconstruction step, the number of correct and incorrect assignments, as well as their proportion with respect to the previous assignment outcome, are shown. For the b quark assignment, the number and proportion of partial correct assignments are additionally shown while respecting consistency in the top quark matching between possible correctly reconstructed W bosons and b quarks.

used in the baseline model. Similar to the model with the hyperedge-based reconstruction head, this larger pairformer architecture was trained on an RTX 3090 using the same batch size as the other models. TABLE V. Comparison of the reconstruction efficiencies from the ablation study comparing different reconstruction methods. The efficiencies shown are the maximum efficiencies achieved on the validation set during model training. Model

εtt̄

εt

εW

Pairformer 0.688 0.656 0.682 Pairformer HyPER 0.688 0.657 0.681 Iterative pairformer 3-step 0.692 0.661 0.684 Iterative pairformer 4-step 0.691 0.660 0.684 Pairformer deep 0.687 0.658 0.684

The iterative predictive model outperforms the noniterative baseline and the hyperedge-based reconstruction, as shown in Table V. Additionally, the iterative models also outperform the deeper alternative, showcasing that our performance gain cannot exclusively be attributed to a higher test-time compute budget. From a step-by-step comparison of the reconstruction performance between jointly trained one-shot pairformer and 3-step iterative models shown in Figure 5, it can be concluded that the improvement of the iterative model mainly stems from a better prediction performance after successful previous predictions. At the same time, the

iterative model shows a poorer ability to recover from incorrect assignments, which may be due to the current training procedure, in which the iterative model only trains on correct partial adjacency matrices for the assignments of the second W boson and the b quarks. Augmenting the training to make use of both correct and incorrect previous assignments might therefore lead to further performance improvements. Since we observe similar performance between the 3step and the 4-step models, we opted for a 3-step iterative reconstruction in the final model to increase event throughput.

Appendix D: Algorithm for Graph Reconstruction

Here we provide pseudo-code for how we reconstruct edges given the model outputs. The decoding step uses a greedy approximation, motivated by the optimal assignment in the case of independent predictions which becomes a tractable problem. This means that we select the edge which locally maximises Y Z X 1−Zi,j pi,ji,j (1 − pi,j ) ∝ Zi,j logit pi,j (D1) i≤j

i≤j

in a greedy manner at each prediction step in the notation of Equation (6). We note that other reconstruction algorithms could be explored, and that attempting to

11 Algorithm 1: Algorithm for choosing W boson assignments from W adjacency logits. W Input: σij : Logits for the W adjacency matrix Input: (W1∗ , W2∗ ): Optional W boson indices W W W W gij ← σij − σii − σjj W gii ← −∞ // W boson assignments only if W ∗ exists then gW1∗ k , gk W1∗ ← −∞ ∀k // Force new W boson gW2∗ k , gk W2∗ ← −∞ ∀k W W = (W1 , W2 ) ← arg maxi,j gij return W

Algorithm 2: Algorithm for choosing b quark assignments from top logits and W boson assignments.

k W1

b(2)

← σt

W/t

Input: fθ : Model to predict logits from data and partial adjacency matrices σ W ← fθW (x, ∅) W (1) ← ChooseWAssignments(σ W )   σ W ← fθW x, Ã(W (1) ) // Update W logits

b(1) , b(2) ← ChooseBAssignments(σ t , W (1) , W (2) ) return W (1) , W (2) , b(1) , b(2)

k W2

t (2) − σll l W1 l W2 b(1) b(2) ckl ← gk + gl ckk ← −∞ ∀k // Choose (b(1) , b(2) ) ← arg maxk,l ckl (1) (2)

return b

Algorithm 3: Overall heuristic algorithm for choosing top quark and W boson assignments.

W (2) ← ChooseWAssignments(σ W , W (1) )   σ t ← fθt x, Ã(W (1) , W (2) ) // Compute t logits

t Input: σij : Logits for top adjacency matrix Input: W (1) , W (2) : Pairs of chosen W boson jets b(1) t gk ← σ t (1) + σ t (1) − σkk

gl

Therefore, the assignment takes into account both the possible edge between the two nodes i and j as well as the possible self-edges of each node. The exact procedure is described in Algorithm 1.

t (2) + σ

different b quarks

,b

maximise the above quantity might not yield the best reconstruction algorithm in general. For the W boson assignment we simply select the edge indices î, ĵ which locally maximise î, ĵ = arg max σi,j − σi,i − σj,j . i,j

(D2)

For selecting the b quarks for the already chosen W boson, we again aim to maximise Equation (D1). However, since we already have an edge, we will instead add the contributions from the two edges connecting to the new node, while subtracting the self-contribution of the node. We do this for both b quarks and select the pair of b quarks which maximises the sum of their contributions. The pseudo-code for this is shown in Algorithm 2. For the full reconstruction algorithm we interchangeably predict the logits of the edges and then apply the appropriate selection algorithm until we have selected all particles as shown in Algorithm 3.

Related documents

• Pairton: Iterative Reconstruction of Short-Lived Particles#620472
• Pairton: Iterative Reconstruction of Short-Lived Particles#618793
• DeSyR: A Decoupled Symbolic Recovery Framework with PINN-Guided Structure Search and Physics-Informed Coefficient Refinement#819971
• Precise QCD Predictions for the Scotogenic Model at Colliders#1063087
• Measurement of Lepton-Jet Correlation in Deep-Inelastic Scattering with the H1 Detector Using Machine Learning for Unfolding#1119942