PaperlayerView exact revision on arXiv

Source-linked AI summary

An Efficient Lorentz Equivariant Graph Neural Network for Jet Tagging

Shiqi Gong, Qi Meng, Jue Zhang, Huilin Qu, Congqiao Li, Sitian Qian, Weitao Du, Zhi-Ming Ma, Tie-Yan Liu

arXiv:2201.08187v6hep-phhep-ex

TL;DR

Jet-tagging models often neglect Lorentz symmetry, while tensor-based equivariant designs are computationally costly. LorentzNet uses Minkowski dot-product attention in a symmetry-preserving graph network. Across two representative benchmarks, it substantially improves tagging performance over existing methods.

  • Problem

    Many jet-tagging models neglect Lorentz symmetry, while LGN’s explicit high-order tensor calculations are computationally costly.

    Source: §1 · p. 3, §2.3 · p. 6

  • Method

    LorentzNet uses Minkowski dot-product attention to aggregate particle four-vectors with weights learned from Lorentz-invariant geometric quantities.

    Source: §1 · p. 3

  • Results

    On two representative jet-tagging datasets, LorentzNet achieves substantial performance improvement over existing methods.

    Source: §1 · p. 4, §5 · p. 15

  • Takeaways & Limitations

    LorentzNet reaches highly competitive performance with only a few thousand training jets and reduces training and inference time compared with LGN.

    Source: §5 · p. 15

    The main computational cost comes from fully connected message passing and grows quadratically with the number of particles in a jet.

    Source: §4.6 · p. 15

Abstract

from arXiv · show

Deep learning methods have been increasingly adopted to study jets in particle physics. Since symmetry-preserving behavior has been shown to be an important factor for improving the performance of deep learning in many applications, Lorentz group equivariance - a fundamental spacetime symmetry for elementary particles - has recently been incorporated into a deep learning model for jet tagging. However, the design is computationally costly due to the analytic construction of high-order tensors. In this article, we introduce LorentzNet, a new symmetry-preserving deep learning model for jet tagging. The message passing of LorentzNet relies on an efficient Minkowski dot product attention. Experiments on two representative jet tagging benchmarks show that LorentzNet achieves the best tagging performance and improves significantly over existing state-of-the-art algorithms. The preservation of Lorentz symmetry also greatly improves the efficiency and generalization power of the model, allowing LorentzNet to reach highly competitive performance when trained on only a few thousand jets. Code and models are available at \url{https://github.com/sdogsq/LorentzNet-release}.

1 Introduction

Jet tagging is a major collider-physics application of deep learning, but most models neglect the required spacetime symmetry. LorentzNet targets this gap with a Lorentz-equivariant design and is reported to outperform state-of-the-art methods on two benchmark datasets.

1 Introduction

  • Deep-learning approaches use jet-level observables, images, particle sequences, particle clouds, and clustering trees to improve jet tagging over classic substructure methods.The approaches can use fine detector resolution and learn patterns from low-level inputs.Source: §1 · pp. 2–3
  • Most deep-learning models neglect Lorentz symmetry, although jet-tagging predictions should be independent of a jet’s spatial orientation and Lorentz boost.Preprocessing and existing Lorentz-inspired layers do not guarantee invariance under arbitrary Lorentz transformations for the entire model.Source: §1 · p. 3
  • LorentzNet uses Minkowski dot-product attention over Lorentz-invariant geometric quantities, requiring only two-vector inner products rather than high-order tensor products.The paper presents this design as more efficient in both training and inference than LGN.Source: §1 · pp. 3–4
  • LorentzNet achieves the best tagging performance on two representative datasets, outperforming existing state-of-the-art algorithms such as ParticleNet.With 5%, 1%, and 0.5% of training samples, its performance degradation is smaller than ParticleNet’s.Source: §1 · p. 4

2 Preliminary

The preliminaries set up particle clouds as jet representations and define the Lorentz transformations and equivariance used in the model. They then describe graph-based particle interactions and motivate a construction that avoids explicit high-order tensors.

2.1 Notations

  • A jet is represented as an unordered, permutation-invariant point cloud of constituent particles, with N denoting the particle count, which can vary across jets.The particles’ four-vectors live in R4.Source: §2.1 · p. 4

2.2 Theory of the Lorentz Group

  • Lorentz transformations preserve the Minkowski bilinear form, and equivariance requires outputs to transform consistently with inputs; for scalar outputs, equivariance is invariance.The paper considers scalar or vector outputs.Source: §2.2 · pp. 4–5

2.3 Graph Neural Network for Particle Cloud

  • The graph uses each particle’s four-momentum as its coordinate and scalar attributes such as mass, charge, or identity; all particle pairs are connected for message passing.The fully connected choice avoids assuming prior knowledge of particle interactions.Source: §2.3 · p. 5
  • Ordinary message passing does not ensure Lorentz equivariance, while LGN’s high-order geometric tensors incur high computational cost; LorentzNet is introduced to avoid explicitly calculating them.The next section’s architecture is built on a universal approximation theorem for Lorentz-equivariant functions.Source: §2.3 · p. 6

3 Network Architecture

LorentzNet builds on a characterization of Lorentz-equivariant mappings to replace explicit high-order tensor calculations with attention based on Minkowski dot products. Its stacked equivariant blocks produce Lorentz-equivariant coordinate embeddings and Lorentz-invariant scalar embeddings.

3 Network Architecture

  • Lorentz-equivariant continuous mappings can be constructed without high-order tensors by attending to vectors using their Minkowski dot products.The characterization expresses the mapping as a sum of input vectors weighted by continuous Lorentz-invariant scalar functions.Source: §3 · p. 6

3.1 LorentzNet

  • The input layer uses particle 4-momenta and may include scalar features, while the decoder averages final node embeddings before dropout and binary classification layers.The described scalar inputs include particle mass or particle-identification information when available.Source: §3.1 · pp. 6, 8
  • Each Lorentz Group Equivariant Block forms pairwise edge messages from particle features, including squared distance and Minkowski dot product, then updates coordinate and scalar embeddings.The edge-message network applies ψ(x) = sgn(x) log(|x| + 1) to normalize large values for optimization.Source: §3.1 · p. 7
  • Minkowski dot-product attention aggregates particle 4-momenta with learned scalar weights, while a scale parameter c controls coordinate updates to avoid scale explosion.The attention uses the Minkowski dot product as an input feature; the model also includes squared spatial distance as a prior feature.Source: §3.1 · p. 7

3.2 Theoretical Analysis

  • Theoretical analysis establishes that coordinate embeddings are Lorentz-equivariant and node scalar embeddings are Lorentz-invariant.The attention weights can contain all pairwise Minkowski dot products, which the paper identifies as complete input information for expressiveness.Source: §3.2 · pp. 8–9

3.3 Implementation Details

  • The reported implementation uses six blocks of width 72, dropout 0.2, and task-specific c values of 5 × 10^-3 for top tagging and 1 × 10^-3 for quark-gluon tagging.Training uses AdamW for 35 epochs, with warm-up, cosine-restart, and exponential-decay learning-rate stages.Source: §3.3 · p. 9

4 Experiment

LorentzNet is evaluated on top-tagging and quark-gluon-tagging benchmarks, with comparisons covering performance, sample efficiency, Lorentz-boost robustness, equivariance, and computational cost. It achieves state-of-the-art tagging performance on both tasks.

4.1 Datasets

  • Top tagging uses 1.2M/400k/400k training, validation, and test jets; quark-gluon tagging uses 2M jets and adds particle-type one-hot inputs.Top-tagging jets have up to 200 constituents; the quark-gluon dataset has no detector simulation.Source: §4.1 · pp. 9–10

4.2 Baselines and Tagging Performance

  • LorentzNet achieves state-of-the-art accuracy, AUC, and background rejection at εS = 0.3 and 0.5 on both tagging datasets.The reported results cover top tagging and quark-gluon tagging.Source: Table 1 · p. 11, Table 2 · p. 11, §4.2 · p. 12

4.3 Sample Efficiency

  • As the top-tagging training fraction shrinks, the accuracy and AUC gap between LorentzNet and ParticleNet grows.The comparison uses 5%, 1%, and 0.5% of the training data and the same 400k-jet test set.Source: §4.3 · pp. 12–13, Table 3 · p. 12

4.4 Equivariance test

  • Under Lorentz boosts of increasing β along the x-axis, LorentzNet and LGN remain robust across a large range, while non-equivariant models’ accuracy drops.LorentzNet is more stable than LGN as β approaches 1.Source: §4.4 · p. 13, Figure 3 · p. 13

4.5 Ablation study

  • The non-equivariant LorentzNet variant performs worse than LorentzNet on top tagging.The authors report that this comparison shows the necessity of Lorentz group equivariance for tagging performance.Source: Table 4 · p. 14, §4.5 · p. 14

4.6 Computational cost

  • LorentzNet is almost 5 times faster than LGN on GPU, while its fully connected message-passing cost grows quadratically with particles per jet.Inference timings were collected with batch size 100; the authors note that clustering nodes could reduce the quadratic cost.Source: §4.6 · pp. 14–15, Table 5 · p. 14

5 Conclusion

LorentzNet is a Lorentz group equivariant graph neural network for jet tagging that improves benchmark performance and efficiency, with competitive results from only a few thousand training jets.

5 Conclusion

  • LorentzNet achieves substantial performance improvement over existing methods on two representative jet-tagging datasets.The conclusion identifies the tasks as top tagging and quark-gluon tagging.Source: §5 · p. 15
  • Its message passing preserves Lorentz symmetry while reducing training and inference time compared with LGN.Source: §5 · p. 15
  • LorentzNet reaches highly competitive performance when trained on only a few thousand jets.Source: §5 · p. 15
  • The demonstrated evidence is limited to two jet-tagging tasks, while the architecture may apply to other particle-physics tasks.Examples proposed include jet-mass or four-momentum regression, jet generation, event classification, and pileup mitigation.Source: §5 · p. 15

A Details about LGN

The LGN appendix describes a tensor-product architecture and its computational and stability trade-offs, including the practical difficulty of selecting tensor order and input scaling.

A Details about LGN

  • LGN combines shortcut, self-interaction, and inter-particle interaction terms, using Clebsch-Gordan tensor products to project tensors into Lorentz-group irreducible representations.Source: A Details about LGN · p. 19
  • Universal approximation ideally requires sufficiently high tensor order, but computation requires a cutoff with no direct prior for choosing it in many applications.Source: A Details about LGN · p. 20
  • LGN tensor-product sums can be scale-sensitive and unstable; top-tagging inputs were scaled by 0.005, a dataset-dependent hyperparameter without a principled stabilizing rule that preserves Lorentz symmetry.Source: A Details about LGN · p. 20
  • LGN's quark-gluon-tagging epoch took more than 10 times as long as other baseline models; its cited top-tagging example took 7.5 hours per epoch at batch size 8.Source: A Details about LGN · p. 20

B Relation with EGNN

The appendix contrasts EGNN's Euclidean-distance scalar information with LorentzNet's Minkowski inner-product attention and describes how the comparison was configured.

B Relation with EGNN

  • EGNN uses relative Euclidean distance as its only vector-derived scalar, which cannot recover angle information; LorentzNet uses the Minkowski inner product for information completeness.Source: B Relation with EGNN · p. 20
  • LorentzNet includes Minkowski dot-product attention for scalar prediction, while EGNN lacks a coordinate-update component; the experiments add that update to EGNN for a fair comparison.Source: B Relation with EGNN · pp. 20–21
  • EGNN experiments used batch size 128 for 35 epochs, AdamW with weight decay 0.01, and a warm-up, cosine-restart, then exponential learning-rate schedule.Source: B Relation with EGNN · p. 21

C Computational Cost in Comparison to LGN

The comparison reports faster LorentzNet inference than LGN and competitive results from a smaller LorentzNet model, while LGN large was not evaluated for performance because training was slow.

C Computational Cost in Comparison to LGN

Related papers

Loading 2201.08187v6…