Lorentz Group Equivariant Neural Network for Particle Physics

Alexander Bogatskiy; Brandon Anderson; Jan T. Offermann; Marwah Roussi; David W. Miller; Risi Kondor

Source-linked AI summary

Lorentz Group Equivariant Neural Network for Particle Physics

Alexander Bogatskiy, Brandon Anderson, Jan T. Offermann, Marwah Roussi, David W. Miller, Risi Kondor

arXiv:2006.04780v1Jun 8, 2020hep-phcs.LGhep-exphysics.comp-phstat.ML

TL;DR

The paper addresses machine-learning problems with complex symmetries by developing a neural network architecture equivariant under the Lorentz group. It uses finite-dimensional Lorentz-group representations and Fourier-space equivariance, achieving competitive particle-physics classification with especially strong parameter efficiency and potentially interpretable weights.

Problem
Problems with complex symmetries may benefit from latent representations connected to the underlying symmetry group, beyond translationally equivariant architectures developed for images.
Method
The Lorentz Group Network uses Fourier-space group equivariance and finite-dimensional Lorentz-group representations for its activations.
Results
The architecture achieves competitive particle-physics classification performance, with especially strong performance in the number of learnable parameters and weights that may be interpretable as physical quantities.
Takeaways & Limitations
The work provides an early physics application of fully Fourier-space equivariance and supports developing symmetry- and geometry-based machine-learning methods for scientific applications.
Takeaways & Limitations
The architecture takes more time to develop and its evaluation is more computationally intensive, although it can later apply broadly to problems sharing the same symmetry.

Abstract

from arXiv · show

We present a neural network architecture that is fully equivariant with respect to transformations under the Lorentz group, a fundamental symmetry of space and time in physics. The architecture is based on the theory of the finite-dimensional representations of the Lorentz group and the equivariant nonlinearity involves the tensor product. For classification tasks in particle physics, we demonstrate that such an equivariant architecture leads to drastically simpler models that have relatively few learnable parameters and are much more physically interpretable than leading approaches that use CNNs and point cloud approaches. The competitive performance of the network is demonstrated on a public classification dataset [27] for tagging top quark decays given energy-momenta of jet constituents produced in proton-proton collisions.

1. Introduction

The paper motivates symmetry-aware neural networks for physics data whose transformations are governed by complex groups such as the Lorentz group. It introduces a Fourier-space, Lorentz-equivariant architecture for particle-physics energy-momentum data.

Motivation: Explicit equivariance can reduce parameters and make neural-network structures more interpretable.The motivation draws on convolutional networks, whose kernels provide interpretable structures.
Motivation: Physics problems often require latent representations tied to the underlying symmetry group.The paper argues that symmetry-constrained building blocks may improve generality, interpretability, and uncertainty quantification while simplifying models.
Approach: The work adopts the Fourier-space implementation of group equivariance because it applies directly to vector inputs.Harmonic analysis presents Fourier-space and real-space approaches as theoretically equivalent.
Application: The paper presents a group-equivariant architecture for particle physics, where data typically include energy-momentum 4-vectors from high-energy collisions.The application targets data generated in particle collisions at facilities such as the LHC or by simulation.
Application: The Lorentz group describes transformations between inertial frames, and the architecture works within finite-dimensional representations to avoid computational difficulties from non-compactness.The paper describes this choice as computationally efficient and physically sensible.

2. Related Work

Related work develops equivariant neural networks for groups, geometric spaces, and scientific domains. This paper distinguishes its contribution as a first equivariant neural-network architecture for fundamental physics applications.

Existing equivariant architectures: Prior work spans equivariance for finite groups, permutations, translations, rotations, spheres, the Euclidean group, and Lie-group homogeneous spaces.The cited progression includes graph isomorphisms, spherical harmonics, SE(3), compact Lie groups, and gauge fields.
Fourier-space methods: Fourier-space nonlinear equivariance was developed for SO(3) and successfully applied in computer vision and chemistry.The approach was independently proposed in two works and later used in scientific applications.
Contribution: Although Lorentz-invariant quantities and Lorentz transformations had appeared in earlier networks, this work presents the first equivariant neural-network architecture for fundamental physics applications.The novelty concerns a fully equivariant architecture rather than merely using Lorentz-invariant quantities or transformations.

3. Theory of the Lorentz group

The theory defines Lorentz transformations through Minkowski spacetime and develops their finite-dimensional representation structure. The paper uses these representations to organize equivariant neural-network activations.

Lorentz transformations: The Lorentz group preserves the Minkowski metric and contains spatial rotations together with boosts between inertial frames.The proper orthochronous subgroup is taken as the paper’s Lorentz group and is the fundamental symmetry of relativistic physics.
Lorentz transformations: Lorentz boosts involve the factor γ = (1 − β^2)^−1/2 and approach the identity in the low-speed limit.Boosts are also described as hyperbolic rotations parameterized by rapidity.
Representations: Finite-dimensional representations are vector spaces with smooth group actions, and completely reducible representations decompose into direct sums of irreducible representations.The resulting isotypic decomposition provides a basis for storing representation vectors.
Representations: The finite-dimensional irreducible representations are tensor products of SU(2) representations, with dimensions dim T(k,n) = (k + 1)(n + 1).The Lorentz-group irreducibles are those T(k,n) for which k + n is even.

4. Principles of Equivariant Networks

The paper constructs equivariant networks from representation-valued activations and operations that preserve Lorentz symmetry. Tensor products, Clebsch–Gordan decompositions, equivariant linear maps, and invariant subnetworks provide the approximation framework and motivate the LGN architecture.

Equivariance: An equivariant map intertwines the group actions on its input and output representations.It satisfies F(ρ(g) · v) = ρ′(g) · F(v).
Universal approximation: The universal-approximation construction iteratively takes tensor products, decomposes them into isotypic components, applies equivariant linear maps, and processes invariants with arbitrary subnetworks.These operations approximate equivariant maps between completely reducible representations.
Equivariant linear maps: Equivariant linear maps are parameterized independently within isotypic components after decomposing representations into irreducible multiplicities.This characterization follows from an extension of Schur’s Lemma.
Equivariant linear maps: Equivariant linear layers use significantly fewer learnable parameters than general fully connected layers, depending on the representations.The loss is invariant, so transforming training samples by group elements leaves trained weights unchanged.
Particle interactions: Lorentz-invariant dot products can be efficiently learned because they occur as invariant components in tensor products of 4-vectors.More complicated interactions require higher-order tensor products.
LGN architecture: LGN combines iterated Clebsch–Gordan layers with invariant perceptrons and aggregates particlewise invariant outputs for permutation invariance.The input, invariant perceptrons, and invariant projection share parameters across particles.
LGN architecture: The architecture is designed to procedurally learn relevant Lorentz invariants and equivariants when events contain hundreds of observed particles.Generating all such quantities explicitly to a fixed polynomial degree would quickly become intimidating.

5. Clebsch-Gordan product

The architecture uses tensor products followed by Clebsch-Gordan decomposition as its main equivariant nonlinearity. Lorentz-group representations are organized through finite-dimensional irreducible components and their associated canonical bases and coefficients.

Clebsch-Gordan decomposition: The main nonlinearity forms tensor products and decomposes them into irreducible representations through the Clebsch-Gordan map.The decomposition coefficients are defined in canonical bases, with details of the derivation deferred to supplementary material.
Rotation group: SU(2) irreducible products decompose into representations whose weights range from |l1−l2| to l1+l2 in unit steps.The resulting components have unit multiplicities.
Canonical bases: The product basis is formed from tensor products of canonical basis vectors indexed by the input representation weights and component labels.The decomposed space receives a corresponding canonical basis indexed by the output weight and label.
Lorentz group: For the Lorentz group, Clebsch-Gordan outputs are indexed by k and n ranges determined independently by the two input representations.The allowed ranges run from the absolute difference of each pair to their sum in steps of two.
Lorentz group: Lorentz-group irreducible representations can be viewed through SU(2) decompositions, providing a canonical basis indexed by l and m.The basis construction uses l from |k−n|/2 through (k+n)/2 and m across the corresponding component range.
Clebsch-Gordan coefficients: The Clebsch-Gordan coefficients are expressed using SU(2) coefficients and vanish unless the component labels satisfy m1 + m2 = m.The coefficient formula sums over auxiliary indices within specified ranges.

6. Equivariant Architecture (LGN)

The Lorentz Group Network processes particle 4-momenta and associated scalars with shared, equivariant layers that preserve permutation invariance. Its CG layers combine self-features and pairwise interactions, while invariant outputs support classification.

Permutation invariance: The architecture enforces permutation invariance by sharing learnable weights across particle indices and summing over those indices at appropriate stages.This reflects the physical independence of the output from particle labeling.
Input layer: LGN takes particle 4-momenta and optional scalar attributes as inputs, representing each particle in a Lorentz-group representation.The input layer applies shared weights independently to particles and produces channelized vectors in irreducible components.
CG layers: Each CG layer performs tensor products, Clebsch-Gordan decompositions, and a learnable linear operation on activations living in Lorentz representations.The layers are iterated a specified number of times while maintaining a representation independent of the particle index.
CG layers: Pairwise differences model two-particle interactions, while shared operators and scalar interaction weights preserve permutation invariance and Lorentz equivariance.The architecture also retains previous activations and models self-interactions within each layer.
CG layers: Only early isotypic components may be retained after Clebsch-Gordan decomposition to control memory usage, without affecting the stated equivariant construction.Tensor products are performed channel-wise to reduce computation, with learnable operators mixing channels.
Output layer: The output sums particle activations, extracts the Lorentz-invariant spin-0 component, and applies a final linear layer for binary classification.The output layer produces two scalar weights.

7. Experiments

LGN is evaluated for Lorentz-invariant top-tagging on simulated jet data, using particle 4-momenta and proton-beam inputs. Although its results do not match state-of-the-art performance, it uses 10−1000× fewer parameters.

Task and dataset: The experiment classifies top-quark jets against jets from lighter quarks or gluons, whose classifier outputs should be Lorentz invariant.The study uses a public dataset with 1.2M training, 400k validation, and 400k testing entries.
Task and dataset: Jets arise as collimated sprays of energetic hadrons formed by parton showers and observed in a Lorentz-boosted laboratory frame.The measured jet is therefore represented through particle-level energy-momentum information.
Preprocessing: The preprocessing adds two labeled proton-beam 4-momenta to fix an axis in each event and establish a symmetry-breaking relationship between samples.The added beam vectors are (2, 0, 0, ±1) GeV, and their energy is described as somewhat arbitrary.
Configuration: The top-tagging configuration uses three CG layers and retains irreducible representations with k,n ≤ 2, up to the 9-dimensional T(2,2).The retained set includes T(0,0), T(2,0), T(0,2), and T(1,1).
Evaluation: The comparison reports accuracy, AUC, background rejection at signal efficiency 0.3, and trainable parameter count, with LGN statistics averaged over four trained instances.Higher accuracy, AUC, and background rejection are treated as better.
Results: 10−1000× fewer parameters are used by LGN than competing models, although the reported results do not match the state of the art.The paper provides further training and performance analysis in supplementary material.

8. Conclusion

The paper develops and applies a Lorentz-equivariant architecture for particle-physics classification, emphasizing few learnable parameters and broad reuse across problems with the same symmetry. Its construction relies on finite-dimensional representation theory, Clebsch–Gordan machinery, and Lorentz-aware tensor structures.

The authors developed and successfully applied a Lorentz-equivariant architecture to top-tagging classification in particle physics.
The architecture’s performance is especially strong in the number of learnable parameters.
Evaluation is more computationally intensive, and developing an equivariant architecture takes more time than simpler models.
Once developed for a symmetry group such as the Lorentz group or SL(2, C), the architecture is broadly applicable to many problems sharing that symmetry.
The construction uses finite-dimensional Lorentz representations, Clebsch–Gordan decompositions, and tensor products to represent and test equivariant transformations.

D. Equivariant Universal Approximation

The appendix establishes that equivariant neural networks with tensor-product nonlinearities can approximate continuous equivariant maps. It reduces polynomial equivariant construction to tensor-power decompositions and applies this framework to Lorentz representations.

Polynomial equivariant maps can be represented using tensor products and equivariant linear maps, with a quadratic example v ⊕(v ⊗v).
Constructing all polynomial equivariants reduces to computing isotypic decompositions of tensor powers of the input representation.
Any continuous equivariant map between completely reducible finite-dimensional representations of a classical Lie group can be uniformly approximated by tensor-product equivariant networks.
Nonlinear functions may be applied component-wise to G-invariants while preserving equivariance.
For the Lorentz group, tensor products of input 4-vectors can generate tensors transforming under arbitrary irreducible representations, so the architecture has no restriction on equivariant output types.

E. Equivariance Tests

The network’s Lorentz and rotational equivariance were tested numerically, with errors remaining small across the physically relevant boost range and under internal activation transformations.

Lorentz boosts: Relative output error stayed below 0.1% in double precision up to γ ≈ 5000, covering the physically relevant range [10, 200].With 32-bit precision, the error stayed this low only to γ ∼70 and exceeded 10% after γ ∼200.
Lorentz boosts: Output deviations were evaluated as a function of the boost factor γ applied to network inputs.The experiment used random 4-momenta, multiple boosts, and averages over random inputs and model initializations.
Rotations: Rotational equivariance errors remained of order 10^-16 in double precision across Euler angles from 0 to 10.The authors attribute the residual error to floating-point rounding.
Internal equivariance: Internal equivariance errors were similarly small when transformed inputs were compared with corresponding Lorentz D-matrix transformations of activations.The test applied Lorentz matrices to inputs and compared resulting activations with transformed activations.

F. Computational Cost

The architecture’s computational cost depends strongly on channel and particle counts: parameters grow quadratically with channels, while tensor-product evaluations become slower as particle counts increase.

Memory scaling: GPU memory usage was measured as a function of channels per layer for a three-layer network.Memory use is more efficient when tensor dimensions are multiples of 32.
Parameter scaling: The number of model parameters grows roughly quadratically with the number of channels.The reported scaling is for the three-layer model family varying a uniform channel count.
Parameter scaling: The number of network parameters was evaluated as a function of channels per layer for a three-layer model.These measurements vary the channel count while keeping the experimental model configuration otherwise fixed.
Evaluation cost: Some tensor sizes grow quadratically with the number of particles, and tensor products make evaluations much slower than simpler models.The authors identify optimized tensor-product operations and a custom CUDA kernel as mitigation strategies under development.

G. Network Metrics

Across training, accuracy and AUC approach rough ceilings partway through training, whereas background rejection and loss continue improving; the implementation uses GPU-oriented software dependencies.

Training dynamics: Accuracy and AUC appear to reach rough ceilings partway through training, while background rejection and loss continue improving throughout.These metrics were tracked across epochs using the ensemble from the main experiment.
Implementation: Training on a GPU requires PyTorch and CUDA, while the code is not yet parallelized across multiple GPUs.
Accuracy: Average network accuracy was measured across four independent trained instances on training and validation subsets.Error bars represent standard deviations.
AUC: Average AUC was tracked as a function of epoch number with error bars representing standard deviation.
Background rejection: Average background rejection at 30% signal efficiency was tracked as a function of epoch number.The figure reports epoch-wise evolution with error bars representing standard deviation.
Loss: Average loss was tracked as a function of epoch number, with error bars given by standard deviation.