Lorentz Local Canonicalization: How to make any Network Lorentz-Equivariant

Jonas Spinner, Luigi Favaro, Peter Lippmann, Sebastian Pitz, Gerrit Gerhartz, Tilman Plehn, Fred A. Hamprecht

Keywords: Geometric deep learning, equivariance, Lorentz symmetry, Transformer, high-energy physics, particle physics

TL;DR: Make any backbone architecture Lorentz-equivariant with minimal computational cost, with applications in high-energy physics

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.

Primary Area: Machine learning for sciences (e.g. climate, health, life sciences, physics, social sciences)

Submission Number: 7206