Using XGBoost to correctly pair muons in events with two dimuons from off-shell boson decays

Machine learning applications for tasks in high energy physics have proliferated in recent years—applications can be found in the realm of jet and particle classification, event selection, and triggering to name a few. One area that has received scant attention is the reconstruction of decay products from off-shell parent particles. For instance, in a physics analysis where the final state contains two pairs of muons (known as dimuons), one can reconstruct the final-state muon pairs by calculating the invariant masses of the dimuon pairs and then finding the correct combination by minimizing the difference in the invariant masses of the dimuon pairs if the parent bosons of the dimuons are on their mass shells. This method becomes invalid when these muons decay from off-shell bosons, necessitating a different method to determine the correct muon pairing. To tackle this problem, we present an eXtreme Gradient Boosted (XGBoost) decision tree model trained on Monte-Carlo simulated data that classifies the correct and incorrect pairings of the muons in the final-state dimuons. We also present our results of using data augmentation, feature engineering, and hyperparameter tuning to maximize the performance metrics of our model. Preliminary results indicate a maximum accuracy of 0.972, a maximum area-under-the-curve of 0.997, and a maximum Matthews correlation coefficient of 0.936. The aim of these studies is to eventually create an extensible model which can be employed in searches for bosons in dark matter models.