Jet Reconstruction with Mamba Networks in Collider Events

TL;DR

This paper presents a new end-to-end deep learning framework using Mamba networks for detailed jet reconstruction in collider events, integrating multiple tasks into one system for improved accuracy and robustness.

Contribution

The authors develop a multi-task learning model that combines instance segmentation, classification, and kinematic regression for jet reconstruction, utilizing a novel MILP-based method for label generation.

Findings

Achieves high classification accuracy with AP scores around 0.57 for key jet types.

Maintains stable performance under high pileup conditions.

Successfully reconstructs mass peaks of beyond standard model particles.

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.