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

TL;DR

This paper introduces Lorentz Local Canonicalization (LLoCa), a framework that makes any neural network Lorentz-equivariant, improving efficiency and accuracy in high-energy physics applications.

Contribution

LLoCa provides a general method to achieve Lorentz-equivariance in any backbone network using equivariant local reference frames.

Findings

Achieves state-of-the-art accuracy on particle physics tasks.

Runs 4 times faster and uses 10 times fewer FLOPs.

Enables flexible architecture design for Lorentz-equivariant networks.

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.