Lorentz-invariant augmentation for high-energy physics deep learning models

In recent years, machine learning models for jet tagging in high-energy physics have gained considerable attention. However, many existing approaches overlook the physical invariants that jets must adhere to, particularly the fundamental spacetime symmetry governed by Lorentz transformations. Setting this statement as the starting point of this work, it is proposed a model-agnostic training strategy that incorporates theory-guided data augmentation to simulate the effects of Lorentz transformations on jet data. The study starts with focusing on the state-of-the-art baseline ParticleNet, a neural network architecture designed for the direct processing of particle clouds for jet tagging. To evaluate the effectiveness of the proposed approach, several experiments are conducted with different augmentation strategies and assess the performance of the augmented models on the widely used top-tagging and quark-gluon reference datasets. The results show that even a small application of the data augmentation strategy increases the robustness of the model to Lorentz boost attacks, i.e., high transformation ß. While the accuracy of the baseline model decreases rapidly with increasing intensity of the transformation ß, the augmented models exhibit more stable performance. Remarkably, models that underwent a moderate level of augmentation demonstrated a statistically significant performance boost on transformations beyond the ones seen at train time. Then the same experimental setup is applied to a second state-of-the-art baseline LorentzNet, a neural network architecture developed to be invariant to Lorentz transformations by design. The performance of the model are also evaluated both on the top-tagging and quark-gluon reference dataset, making possible a full comparison between each experimental setup applied to the two chosen models. The results shows that LorentzNet is more robust to Lorentz boost attacks than ParticleNet, as it is expected to be. Nevertheless the application of the data augmentation strategy to an already invariant architecture, tends to further increase the robustness of the model. This finding highlights the potential of the model-agnostic data augmentation strategy in enhancing model accuracy and robustness while preserving the essential physical properties of the jets.