Zero-Permutation Jet-Parton Assignment using a Self-Attention Network

In high-energy particle physics events, it can be advantageous to find the jets associated with the decays of intermediate states, for example, the three jets produced by the hadronic decay of the top quark. Typically, a goodness-of-association measure, such as a $\chi^2$ related to the mass of the associated jets, is constructed, and the best jet combination is found by optimizing this measure. As this process suffers from a combinatorial explosion with the number of jets, the number of permutations is limited by using only the $n$ highest $p_T$ jets. The self-attention block is a neural network unit used for the neural machine translation problem, which can highlight relationships between any number of inputs in a single iteration without permutations. In this paper, we introduce the Self-Attention for Jet Assignment (SaJa) network. SaJa can take any number of jets for input and outputs probabilities of jet-parton assignment for all jets in a single step. We apply SaJa to find jet-parton assignments of fully-hadronic $t\bar{t}$ events to evaluate the performance. We show that SaJa achieves better performance than a likelihood-based approach.