LLoCa vs L-GATr

Most previously published Lorentz-equivariant networks rely on specialized layers to achieve Lorentz-equivariance. Popular examples are

  • LorentzNet, a Lorentz-equivariant graph network using scalar and vector representations

  • PELICAN, a Lorentz-invariant graph network that uses general permutation-equivariant layers operating on Lorentz invariants

  • CGENN, a Lorentz-equivariant graph network using geometric algebra representations

  • L-GATr, a Lorentz-equivariant transformer using geometric algebra representations

Canonicalization, or specifically Lorentz local canonicalization (LLoCa), is an alternative approach that extends an existing backbone architecture with a Lorentz-equivariant canonicalization procedure to achieve Lorentz-equivariance. This section discusses the benefits of both approaches and helps to decide which one to use for your specific application. We focus on L-GATr as a representative of the specialized-layer approach because we compared it carefully with LLoCa-Transformers, but most arguments apply to other specialized-layer networks as well.

Disclaimer: Lorentz-equivariant networks are great, but when starting on a task from scratch we recommend to start with a non-equivariant baseline network first, for instance a vanilla transformer, and move to a Lorentz-equivariant network later on for the extra performance and robustness.

Benefits of L-GATr / specialized layers

  • Simpler, e.g. LGATrSlim fits into a single file. Canonicalization modifies the backbone architecture only slightly, but the Frames-Net and subsequent orthonormalization requires significant extra code.

  • Training dynamics are typically easier, because the interaction between the Frames-Net and backbone in LLoCa can lead to more complex training dynamics. For instance, we found in some cases that the Frames-Net in LLoCa overfits before the backbone does, even though the Frames-Net has far fewer parameters.

Benefits of LLoCa / canonicalization

  • Make any backbone architecture Lorentz-equivariant: We constructed LLoCa-Transformer, LLoCa-GNN, LLoCa-ParticleNet and LLoCa-ParT networks in our publications. LLoCa is particularly useful when your non-equivariant architecture already includes many tricks that would be time-consuming to re-implement in a Lorentz-equivariance-by-layers architecture like L-GATr, such as dynamic graph convolutions, attention mechanisms, or U-Net layouts. See Make any network Lorentz-equivariant for instructions and code examples for how to extend your favorite architecture with LLoCa.

  • Higher-order representations are not only straight-forward to include in LLoCa but also easily mixed in equal or unequal fractions. LLoCa builds on generic tensor representations of the Lorentz group through the TensorReps class, e.g. 16x0n+8x1n+2x2n+1x3n for a direct product of 16 scalar, 8 vector, 2 second-rank tensor, and 1 third-rank tensor representations. Arbitrary higher-order representations are implemented, but using them comes at the cost of slower inference because of the additional matrix multiplications.

  • Reduced resource constraints: For similar network size, we find that LLoCa networks typically require less GPU memory and FLOPs than comparable specialized-layer networks. This is because the backbone processes Lorentz-invariant features only, which are typically lower-dimensional and cheaper to process. The overhead from the Frames-Net and frame-to-frame conversions in message passing increases inference time by a factor of 1.5-2x in our current implementation, but based on our FLOPs analyses we believe that there is significant room for optimization that we want to explore in the future.

  • Flexible for studies on subgroup equivariance and data augmentation: Subgroup-equivariant networks such as SO(3)-equivariant or SO(2)-equivariant networks as well as data augmentation can be implemented easily as modifications in the Frames-Net. This allows for systematic studies of the impact of symmetry-aware designs, which would require entire new network designs when using the specialized-layer approach.