Lorentz-Equivariance without Limitations

TL;DR

This work presents Lorentz Local Canonicalization (LLoCa), a framework that imposes exact Lorentz-equivariance on arbitrary neural networks by predicting local frames per particle and performing tensorial message passing across frames. It supports arbitrary representations and provides mechanisms to break or adjust Lorentz symmetry to match realistic detector and event conditions, enabling fair backbone comparisons. Across amplitude regression, end-to-end event generation, and jet tagging, LLoCa-based transformers and GNNs achieve state-of-the-art or competitive performance with favorable compute characteristics, aided by Frames-Net designs and symmetry-breaking strategies. The results demonstrate the practical viability of exact symmetry incorporation in large-scale HEP ML tasks and offer data/code resources to facilitate adoption and further research.

Abstract

Lorentz Local Canonicalization (LLoCa) ensures exact Lorentz-equivariance for arbitrary neural networks with minimal computational overhead. For the LHC, it equivariantly predicts local reference frames for each particle and propagates any-order tensorial information between them. We apply it to graph networks and transformers. We showcase its cutting-edge performance on amplitude regression, end-to-end event generation, and jet tagging. For jet tagging, we introduce a large top tagging dataset to benchmark LLoCa versions of a range of established benchmark architectures and highlight the importance of symmetry breaking.

Figures (9)

• Figure 1: Tensorial message passing in LLoCa. All symbols are described in the text. This figure is also included in Ref. Spinner:2025prg.
• Figure 2: Left: scaling of the prediction error with the multiplicity for all amplitude surrogates. Right: training cost on $Z+4g$ amplitudes. Networks are trained on $10^7$ events for all multiplicities. Error bands are based on the mean and standard deviation of three different random seeds. These results are also included in Ref. Spinner:2025prg.
• Figure 3: Exact Lorentz-equivariance vs. data augmentation for GNNs (left) and transformers (right). Error bands are based on the mean and standard deviation of three different random seeds. These figures are also included in Ref. Spinner:2025prg.
• Figure 4: Effect of the choice of symmetry group (left) and options for the Frames-Net (right) on the LLoCa-Transformer. Error bands are based on the mean and standard deviation of three different random seeds. These figures are also included in Ref. Spinner:2025prg.
• Figure 5: Overview of marginal distributions for $t\bar{t}+1,2,3,4$ jets (top to bottom).