lgatr 1.4.4

Lorentz-Equivariant Geometric Algebra Transformer for High-Energy Physics

Lorentz-Equivariant Geometric Algebra Transformer

This repository contains a standalone implementation of the Lorentz-Equivariant Geometric Algebra Transformer (L-GATr) by Jonas Spinner, Víctor Bresó, Pim de Haan, Tilman Plehn, Huilin Qu, Jesse Thaler, and Johann Brehmer. L-GATr uses spacetime geometric algebra representations to construct Lorentz-equivariant layers and combines them into a transformer architecture. You can read more about L-GATr as well as the more efficient L-GATr-slim in the following three papers and in the L-GATr documentation:

• Lorentz-Equivariant Geometric Algebra Transformers for High-Energy Physics (L-GATr; for ML audience)
• A Lorentz-Equivariant Transformer for All of the LHC (L-GATr; for HEP audience)
• Economical Jet Taggers - Equivariant, Slim, and Quantized (L-GATr-slim)

Installation

You can either install the latest release using pip

pip install lgatr

or clone the repository and install the package in dev mode

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

How to use L-GATr

Please have a look at the L-GATr documentation and our example notebooks for LGATr and ConditionalLGATr.

Overview of features in L-GATr:

• L-GATr encoder and decoder as LGATr and ConditionalLGATr
• Additional attention backends, installation via pip install lgatr[varlen-attention], pip install lgatr[xformers-attention], pip install lgatr[flex-attention], pip install lgatr[flash-attention] or any combination. You might have to run python -m pip install --upgrade pip setuptools wheel because extra imports require modern versions of pip, setuptools, wheel.
• Support for torch's automatic mixed precision; critical operations are performed in float32
• Interface to the geometric algebra: Embedding and extracting multivectors; spurions for symmetry breaking at the input level
• Many hyperparameters to play with, organized via the SelfAttentionConfig, CrossAttentionConfig, MLPConfig and LGATRConfig objects
• LGATrSlim and ConditionalLGATrSlim as more efficient variants that use only scalar and vector representations

Examples

• https://github.com/heidelberg-hepml/lorentz-gatr: Original LGATr implementation used for the papers. This repo doesn't import the lgatr package, but has its own (outdated) lgatr/ folder. (paper1 paper2)
• https://github.com/heidelberg-hepml/lloca-experiments: Code for the LLoCa project, including L-GATr as a baseline. The main results from https://github.com/heidelberg-hepml/lorentz-gatr can be reproduced here using the lgatr package. (paper1 paper2)
• https://github.com/spinjo/weaver-core/blob/lgatr/weaver/nn/model/LGATr.py: L-GATr in the CMS boosted object tagging library weaver. Includes examples for how to use L-GATr without the xformers package.ing
• https://github.com/heidelberg-hepml/high-dim-unfolding: Generative jet substructure unfolding with L-GATr, uses the ConditionalLGATr. (paper)
• https://github.com/gregorkrz/jetcluster: IRC-safe jet clustering with L-GATr, starting from the https://github.com/heidelberg-hepml/lorentz-gatr repo. (paper)
• https://github.com/heidelberg-hepml/tagger-quantization: Quantized jet taggers, including float8+ternary weight implementations of L-GATr and L-GATr-slim. (paper)

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

Contributing

Contributions are welcome! To get started:

1. Fork this repository and create a new branch for your feature or fix.
2. Make your changes, following the existing code style (and using pre-commit).
3. Add or update tests where appropriate.
4. Open a pull request with a clear description of your changes.

If you’re not sure where to begin, feel free to open an issue to discuss your idea first.

Citation

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

@article{Brehmer:2024yqw,
    author = "Brehmer, Johann and Bres{\'o}, V{\'\i}ctor and de Haan, Pim and Plehn, Tilman and Qu, Huilin and Spinner, Jonas and Thaler, Jesse",
    title = "{A Lorentz-equivariant transformer for all of the LHC}",
    eprint = "2411.00446",
    archivePrefix = "arXiv",
    primaryClass = "hep-ph",
    reportNumber = "MIT-CTP/5802",
    doi = "10.21468/SciPostPhys.19.4.108",
    journal = "SciPost Phys.",
    volume = "19",
    number = "4",
    pages = "108",
    year = "2025"
}
@article{Petitjean:2025zjf,
    author = {Petitjean, Antoine and Plehn, Tilman and Spinner, Jonas and K{\"o}the, Ullrich},
    title = "{Economical Jet Taggers -- Equivariant, Slim, and Quantized}",
    eprint = "2512.17011",
    archivePrefix = "arXiv",
    primaryClass = "hep-ph",
    reportNumber = "IPPP/25/93",
    month = "12",
    year = "2025"
}
@inproceedings{spinner2025lorentz,
  title={Lorentz-Equivariant Geometric Algebra Transformers for High-Energy Physics},
  author={Spinner, Jonas and Bres{\'o}, Victor and De Haan, Pim and Plehn, Tilman and Thaler, Jesse and Brehmer, Johann},
  booktitle={Advances in Neural Information Processing Systems},
  year={2024},
  volume={37},
  eprint = {2405.14806},
  url = {https://arxiv.org/abs/2405.14806}
}
@inproceedings{brehmer2023geometric,
  title = {Geometric Algebra Transformer},
  author = {Brehmer, Johann and de Haan, Pim and Behrends, S{\"o}nke and Cohen, Taco},
  booktitle = {Advances in Neural Information Processing Systems},
  year = {2023},
  volume = {36},
  eprint = {2305.18415},
  url = {https://arxiv.org/abs/2305.18415},
}