Jet Reconstruction with Mamba Networks in Collider Events

Abstract

We introduce a novel end-to-end framework for jet reconstruction in high-energy collider events, leveraging the efficiency and long-range modeling capabilities of the Mamba architecture. Our model unifies instance segmentation, classification, and kinematic regression into a single multi-task learning system, enabling a sophisticated multi-level reconstruction that simultaneously identifies primary heavy jets ($t$, $H$, $W/Z$) and their constituent sub-jets. To facilitate supervised learning for this complex task, we develop a novel method for assigning final-state hadrons to their ancestor colored partons using a Mixed-Integer Linear Programming solver, which generates high-fidelity ground-truth labels. The model achieves high classification accuracy, with an Average Precision score of 0.569 for $W/Z$-jets and 0.568 for $b$-jets, and shows exceptional precision in kinematic reconstruction. Furthermore, we show that the model not only maintains stable performance in high-pileup environments but also successfully reconstructs the mass peaks of beyond the standard model particles. This work presents a powerful and versatile new tool for comprehensive event reconstruction at the LHC.