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

2025-05-26

AI Summary

Existing Lorentz-equivariant neural networks rely on custom-designed layers, hindering integration with general-purpose architectures. Method: We propose the Lorentz Local Calibration (LLoCa) framework—the first method enabling exact Lorentz equivariance for arbitrary backbone networks without structural modification. LLoCa establishes a generic equivariant paradigm grounded in locally calibrated reference frames, extends geometric message passing to the non-compact Lorentz group, and unifies data augmentation as a special case of reference-frame selection. Integrating Lorentz representation theory, spacetime-propagated tensor features, and equivariant graph networks, we design the LLoCa-Transformer architecture. Contribution/Results: On particle physics benchmarks, LLoCa achieves state-of-the-art accuracy while accelerating inference by 4× and reducing computational cost by 5–100× in FLOPs compared to prior Lorentz-equivariant models.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Low Level & Physics-based Vision
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search

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 to 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 surpass state-of-the-art accuracy on relevant particle physics tasks, while being $4 imes$ faster and using $5$-$100 imes$ fewer FLOPs.

Problem

• Enabling Lorentz-equivariance in any neural network architecture
• Overcoming limitations of specialized layers in current implementations
• Improving efficiency and accuracy in particle physics tasks

Innovation

• General framework for Lorentz-equivariant networks
• Uses local reference frames for transformations
• Adapts geometric message passing to Lorentz group