lloca 1.3.6

Lorentz Local Canonicalization

Lorentz Local Canonicalization

This repository contains a standalone implementation of Lorentz Local Canonicalization (LLoCa) by Jonas Spinner, Luigi Favaro, Peter Lippmann, Sebastian Pitz, Gerrit Gerhartz, Huilin Qu, Tilman Plehn, and Fred A. Hamprecht. LLoCa uses equivariantly predicted local reference frames and geometric message passing between these frames to make any architecture Lorentz-equivariant. You can read more about LLoCa in the following two papers and in the LLoCa documentation:

• Lorentz Local Canonicalization: How to make any Network Lorentz-Equivariant (ML audience)
• Lorentz-Equivariance without Limitations (HEP audience)

Installation

You can either install the latest release using pip

pip install lloca

or clone the repository and install the package in dev mode

git clone https://github.com/heidelberg-hepml/lloca.git
cd lloca
pip install -e ".[dev]"
pre-commit install

How to use LLoCa

Please have a look at the LLoCa documentation (WIP) and our example notebook for the LLoCa-Transformer.

Features

• Backbone architectures in lloca/backbone: Transformer, ParticleTransformer, ParticleNet, GraphNet, MLP
• The Transformer backbone supports several attention kernels that can be installed optionally with e.g. pip install lloca[varlen-attention], pip install lloca[xformers-attention], pip install lloca[flex-attention], pip install lloca[flash-attention].
• LLoCaMessagePassing as blueprint for generic LLoCa graph network backbones
• Equivariant vector predictors in lloca/equivectors: MLPVectors, LGATrVectors, PELICANVectors
• Local frames for equivariant architectures on several symmetry groups: SO(1,3) (LearnedPDFrames, LearnedSO13Frames, LearnedRestFrames), SO(3) (LearnedSO3Frames), SO(1,1)xSO(2) (LearnedZFrames) and SO(2) (LearnedSO2Frames); as well as the corresponding random global frames for data augmentation
• Support for arbitrary higher-order representations with the TensorReps class

Coming soon:

• Parity-odd representations
• Support for cross-attention
• More backbone architectures

Examples

• https://github.com/heidelberg-hepml/lloca-experiments: Codebase for the original papers.
• https://github.com/heidelberg-hepml/tagger-quantization: Quantized LLoCa taggers

Let us know if you use lloca, so we can add your repo to the list!

Citation

If you find this code useful in your research, please cite our papers

@article{Favaro:2025pgz,
    author = "Favaro, Luigi and Gerhartz, Gerrit and Hamprecht, Fred A. and Lippmann, Peter and Pitz, Sebastian and Plehn, Tilman and Qu, Huilin and Spinner, Jonas",
    title = "{Lorentz-Equivariance without Limitations}",
    eprint = "2508.14898",
    archivePrefix = "arXiv",
    primaryClass = "hep-ph",
    month = "8",
    year = "2025"
}
@article{Spinner:2025prg,
    author = "Spinner, Jonas and Favaro, Luigi and Lippmann, Peter and Pitz, Sebastian and Gerhartz, Gerrit and Plehn, Tilman and Hamprecht, Fred A.",
    title = "{Lorentz Local Canonicalization: How to Make Any Network Lorentz-Equivariant}",
    eprint = "2505.20280",
    archivePrefix = "arXiv",
    primaryClass = "stat.ML",
    month = "5",
    year = "2025"
}