Lorentz Local Canonicalization: How to Make Any Network Lorentz-Equivariant

TL;DR

Lorentz-equivariant neural networks enable data-efficient modeling in high-energy physics but often rely on specialized layers or incur high costs. Lorentz Local Canonicalization (LLoCa) introduces a general framework that makes any backbone exactly Lorentz-equivariant by learning local reference frames per particle and performing tensorial messages between frames, with Minkowski-product attention and a polar-decomposition-based frame construction. The authors demonstrate state-of-the-art or competitive results on jet tagging and QFT amplitude regression while achieving substantially lower FLOPs and faster training than prior SOTA Lorentz-equivariant models, and they provide a nuanced comparison between exact equivariance and data augmentation. Overall, LLoCa broadens the practical applicability of Lorentz-equivariant learning, enabling efficient deployment of physics-informed architectures across domains that involve non-compact symmetry groups and space-time tensor features.

Abstract

Lorentz-equivariant neural networks are becoming the leading architectures for high-energy physics. Current implementations rely on specialized layers, limiting architectural choices. We introduce Lorentz Local Canonicalization (LLoCa), a general framework that renders any backbone network exactly Lorentz-equivariant. Using equivariantly predicted local reference frames, we construct LLoCa-transformers and graph networks. We adapt a recent approach for geometric message passing to the non-compact Lorentz group, allowing propagation of space-time tensorial features. Data augmentation emerges from LLoCa as a special choice of reference frame. Our models achieve competitive and state-of-the-art accuracy on relevant particle physics tasks, while being $4\times$ faster and using $10\times$ fewer FLOPs.

Figures (7)

• Figure 1: Lorentz Local Canonicalization (LLoCa) for making any architecture Lorentz-equivariant. The input consists of a set of particles each associated with an energy and momentum (and possibly other particle features). The particle features are transformed into learned local reference frames, turning them into Lorentz-invariant local features. The local features can then be processed by any backbone architecture to produce exactly Lorentz-equivariant outputs for a variety of possible tasks. Without including additional domain-specific priors, our approach elevates the performance of domain-agnostic models, such as a vanilla transformer, to the SOTA in the field (see e.g. Fig. \ref{['fig:amp_multiplicity']}).
• Figure 2: Tensorial message passing in LLoCa. Lorentz-invariant local messages are transformed between local reference frames.
• Figure 3: LLoCa surpasses SOTA performance while being $\boldsymbol{4\times}$ faster. The collisions of two particles produces a single $Z$ boson and $n=1,2,3,4$ gluons $g$ (x-ticks $Z+ng$). LLoCa significantly improves the prediction of interaction amplitudes of the non-equivariant GNN and transformer. Left: Our LLoCa-Transformer achieves state-of-the-art results over all four multiplicities. Right: Accuracy and compute for $Z+4g$. The LLoCa-Transformer uses a tenth of the FLOPs and a fourth of the training time relative to the second most accurate model. Uncertainties are standard deviations over three runs. See App. \ref{['app:exp_details']} for more details.
• Figure 4: Data efficiency of LLoCa networks. Lorentz-equivariant amplitude regression for $Z+4g$ production is more data efficient in the large data regime for both graph networks (left) and transformers (right). Perhaps surprisingly, data augmentation (DA) outperforms equivariance for small datasets.
• Figure 5: Comparison between LLoCa models and their non-equivariant counterparts. We consider a simple MLP (left) and a graph network that does not use additional Lorentz invariant edge features (right). The global frames for the data augmentation (DA) results are sampled as described in App. \ref{['app:exp_details']}.