Streamlined jet tagging network assisted by jet prong structure

Abstract: Attention-based transformer models have become increasingly prevalent in collider analysis, offering enhanced performance for tasks such as jet tagging. However, they are computationally intensive and require substantial data for training. In this paper, we introduce a new jet classification network using an MLP mixer, where two subsequent MLP operations serve to transform particle and feature tokens over the jet constituents. The transformed particles are combined with subjet information using multi-head cross-attention so that the network is invariant under the permutation of the jet constituents. We utilize two clustering algorithms to identify subjets: the standard sequential recombination algorithms with fixed radius parameters and a new IRC-safe, density-based algorithm of dynamic radii based on HDBSCAN. The proposed network demonstrates comparable classification performance to state-of-the-art models while boosting computational efficiency drastically. Finally, we evaluate the network performance using various interpretable methods, including centred kernel alignment and attention maps, to highlight network efficacy in collider analysis tasks.