Lorentz-Equivariant Geometric Algebra Transformers for High-Energy Physics
Authors: Jonas Spinner
Victor Breso Pla
Pim De Haan
Tilman Plehn
Jesse Thaler
Johann Brehmer
Venue: NeurIPS 2024 poster (2024)
Keywords: Geometric deep learning
equivariance
Lorentz symmetry
Transformer
flow matching
high-energy physics
particle physics
Primary area: machine_learning_for_physical_sciences
TL;DR: A Lorentz-equivariant Transformer architecture plus Lorentz-equivariant flow matching for high-energy physics

Abstract: Extracting scientific understanding from particle-physics experiments requires solving diverse learning problems with high precision and good data efficiency. We propose the Lorentz Geometric Algebra Transformer (L-GATr), a new multi-purpose architecture for high-energy physics. L-GATr represents high-energy data in a geometric algebra over four-dimensional space-time and is equivariant under Lorentz transformations, the symmetry group of relativistic kinematics. At the same time, the architecture is a Transformer, which makes it versatile and scalable to large systems. L-GATr is first demonstrated on regression and classification tasks from particle physics. We then construct the first Lorentz-equivariant generative model: a continuous normalizing flow based on an L-GATr network, trained with Riemannian flow matching. Across our experiments, L-GATr is on par with or outperforms strong domain-specific baselines.

Lorentz-Equivarian Geometric Algebra Transformers for High-Energy Physics
Jonas Spinner"
Victor Bresó"
Heidelberg University
Heidelberg University
j spinner@thphys..... uni hew.elberg.e...ee.eeee.bbeewieeh de
v. breso@thphys wew.eeleeeee.deee de
Pim de Haan
Tilman Plehn Heidelberg University
Jesse Thaler MIT LAIFI
Johann Brehmer Qualcomm AI Research!
Qualcomm AI Research
Abstract
Extracting scientific understanding from particle-physics experiments requires solving diverse learning problems with high precision and good data efficiency. We propose the Lorentz Geometric Algebra Transformer (L--GATr), new multi- purpose architecture for high-energy physies. L--ATT represents high-energy data in a geometric algebra over four-dimensoonal space-time and is equivariant under Lorentz transformations, the symmetry group of relativistic kinematics. At the same time, the architecture is Transformer, which makes it versatile and scalable to large systems. L-GATr first demonstrated on regression and classification tasks from particle physics. We then construct the first Lorentz-equivaaiann generative model: a continuous normalizing flow based on an L-GATr network, trained with Riemannian flow matching. Across experiments, L-GATr is par with outperforms strong domain-specifie baselines.
Introduction
In the quest to understand nature the most fundamental level, machine learning is omnipresent [23]. Take the most complex machine ever built: at CERN's Large Hadron Collider (LHC), protons are accelerated close the speed of light and interact; their remnants recorded by various detector components, totalling around 10 bytes of data per second [24]. These data are filtered, processed, and compared theory predictions, as we sketch in Fig. Each step of this pipeline requires making decisions about high-dimensional data. More often than not, these decisions are rooted in machine learning, increasingly often deep neural networks [12, 13, 20, 32, 44, 49, 52, 66]. This approach powers most measurements in high-energy physics, culminating in the Higgs boson discovery in 2012 [5,
High-energy physics analyses put stringent requirements on network architectures. They need to be able represent particle data and have be expressive enough learn complex relations in high-dimensional spaces precisely. Moreover, training data often come from precise theory computations and complex detector simulations, both of which require a considerable computational cost; architectures therefore need to be data efficient. Generative models of particle-physice data face additional challenges: because of detector boundaries and selection cuts, densities frequently feature sharp edges; at the same time, it is often important to model low-density tails of distributions precisely over multiple orders of magnitude of probability densities.
Off-the-shed architectures originally developed for vision or language are popular starting points for
"quual contribution
Qualcomm AI Research is an initiative of Qualcomm Technologies, Inc.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

high-energy physics applications [19, 37], but do not satisfy these goals reliably. We argue that this is because they do not make systematic use of the rich structure of the data. Particle interactions are governed by quantum field theories and respect their symmetries, notably the Lorentz symmetry of special relativity [39, 63]. First Lorentz-equivariant architectures have recently been proposed [10, 42, 68], but they are limited to specific applications and not designed with a focus on scalability.
In this work, we introduce the Lorentz Geometric Algebra Transformer (L-GATr), new general- purpose network architecture for high-energy physics. It is based on three design choices. First, L-GATr is equivariant with respect to the Lorentz symmetry." supports partial and approximate symmetries as found in some high-energy physics applications through symmetry- -ereekkng inputs. Second, as representations. L-GATr uses the geometric (or Clifford) algebra over the voctorectors of special relativity. This algebra is based on the scalar and four-vector properties that LHC data are naturally parameterized and extends them to higher orders, increasing the network capacity. Finally, L-GATTr is a Transformer. It supports variable-length inputs, as found in many LHC problems, and even large models can be trained efficiently. Because it computes pairwise interactions through scaled dot-product attention, for which there are highly optimized backends like Flash Attention [33], the architecture scales particularly well to problems with many tokens or particles.
L-GATr is based on the Geometric Algebra Transformer architecture [14, 36), which was designed for non-relativistic problems governed by the Euclidean symmetry E(3) of translations, rotations, and reflections. Our L-GATr architecture generalizes that to relativistic scenarios and the Lorentz symmetry. To this end, we develop several new network layers, including maximally expressive Lorentz-equivariant linear map, Lorentz-equivarian attention mechanism, and Lorentz-equivariant layer normalization.
In addition to the general-purpose architecture, we develop the first (to the best of our knowledge) Lorentz-equivariaat generative model. We construct a continuous normalizing flow with an L-GATr denoising network and propose training it with a Riemannian flow matching approach [25]. This not only lets us train the model in a scalable way, but also allows us encode more aspects of the problem geometry into the model: we can even hard-code phase-space boundaries, which are commonplace in high-energy physics.
We demonstrate I.-GATr on three particle-physics applications. We first train neural surrogates for quantum field theoretic amplitudes, regression problem with high demands on precision. Next, train classification models and evaluate L-GATr on the popular benchmark problem of top tagging. Finally, we turn to the generative modelling of reconstructed particles, which can make the entire analysis pipeline substantially more efficient. The three applications differ in the role they play in the LHC analysis pipeline (see Fig. 1), data, and learning objective, highlighting the versatility of L-GATr. We find that L-GATr is on par with or outperforms strong domain-specific baselines across all problems, both in terms of performance and data efficiency.
Our implementation of L-GATr is available at http://ggitnbbbc.cc//eeeeeeeeeee lorentz- gatr.
2 Background and related work
High-energy physics In Fig. I we sketch the typical data-analysis pipeline in particle physics. Its central idea is to take the data collected in the detectors as well as the predictions from different theories of physics, process both in parallel, and ultimately compare their predictions. The pipeline includes various steps, including the computation of scattering probabilities or amplitudes in mathe- matical frameworks called quantum field theories, the Monte-Carlo sampling from the theory, the simulated interaction of particles with the detector, the dimensionality reduction of the raw detector output to small number of observables, the data filtering extract only collisions of interest, and the statistical analysis of whether two predictions are consistent.
What most steps in this pipeline have in common is the notion of particles, the main representation of data in high-energy physics. A particle is characterized by a discrete type label, an energy E' e I, and spatial momentum E 3 Types include fundamental particles like electrons, photons, quarks, and
*One could extend L- GAT to the full Poincaré symmetry, which additionally includes space- space-time translations. However, this is not necessary for most particle-physie physics applications. as usually only the momentum, and not the absolute position, of particles is of interest. A key exception is the study of lived particles.

---

Recon- struction
Reconstr. particles
Event se selection
Nature
Detector
Inference
X
Scattering amplitudes
MC sampler
Shower hadron. sim
Detector sim
Recon-
Reconstr, particles
Event selection
Discoveries &
Quantum Theory
4.1 Amplitude surrogates
4.3 Generative modelling
tagging
Figure 1: Schematic view of the data-an Sata-analysis workflow in high-energy physics. Measurements (top) are processed in parallel with simulated data (bottom); their comparison is ultimately the basis for most scientific conclusions. In orange, show how the three applications of L-GATr experiment with in this paper fit into this workflow. The architecture is also applicable several other stages, including reconstruction and inference.
gluons, composite particles like protons, as well reconstructed objects like "jets" [70] or "particle- flow candidates" [71], which are the outputs complex reconstruction algorithms. The energy and spatial momentum of particle are conveniently combined into a four-mommntum p (E. p).
The laws of fundamental physics [41, 75] are invariant with respect to the choice of an inertial reference frame [39, 63]: they do not change under rotations and boosts from one un-accelerated reference frame into another.* Together, these transformations form the special orthochronous Loreniz group SO+ (1, (1,3) 3). This group is the connected component of the orthogonal group on the four-vector space 1.3 with Minkowski metric diag(+1, -1, -1, -1) [61]. Lorentz transformations mix temporal and spatial components. Space and time should therefore not be considered as separate concepts, but rather as components of four-dimensional space-time. Particle foue-momenta another instance of this: they transform in the vector representation of the Lorentz group as p" + p' = Σ, Agp" for E SO (1, with the Lorentz transformation mixing energy and spatial momentum.
Geometric deep learning The central tenet geometric deep learning [15, 31] to embed known structure of problem into the architecture used to solve it, instead of having learn it completely from data. The key idea that of equivariance to symmetry groups: when the inputs network are transformed with symmetry transformation g. the outputs should transform under the same element of the symmetry group. f(g x) f()) where denotes the group action. What is known "equivariance" in machine learning often called "covariance" in physics [28].
GATr Our work is rooted in the Geometric Algebra Transformer (GATr) [14, 36), a network architecture that is equivariant E(3). the group of non-relativistic translations, rotations, and reflections. GATr represents inputs, hidden states, and outputs in the geometrie (or Clifford) algebra G3,0011 [29, 43, 67]. A geometric algebra extends base space like RR3 to higher orders and adds bilinear map known as the geometric product. We provide a formal introduction in Appendix A. What matters in practice is that this vector space can represent various 3D geometric objects. Brehmer et al. [14] develop different layers for this representation and combine them in a Transformer architecture [73]. For L-GATr, build on the GATr blueprint, but re-design all components such that they can represent four-momenta and are equivariant with respect to Lorentz transformations.
Lorentz-equivariant architectures Recently, some Lorentz-equivariant architectures have been proposed. Most closely related to this work is the Clifford Group Equivariant Neural Networks (CGENN) Ruhe al. [68] and their extensions to simplical complexes [54] and stecrable convolu- tions [78]. Like us, they use the geometric algebra over four-vectors. While they also use Lorentz- equivariant linear maps and geometric products, our architectures differ number of ways. In particular, they propose message-passing graph neural network, while we build Transformer archi-
"Allowing for accelerating reference frames would bring us to the general theory of relativity. which is irreevvnt for particle physics experiment long as they "Speecia" and "orthochronous" here mean that spatial and temporal reflections are not considered as symmetries. In fact, the fundamental laws of nature are not invariant under those transformations, effect known P-violation and T-vioaation.
performed close to black hole.

---

tecture based on dot-prot pot-prruuct attention.
Other Lonentz-equivariant architectures include LorentzNet [42] and the Permutation Equivariant and Lorentz. Invariant or Covariant Aggregator Network (PELICAN) [10]. Both are message-passinn graph neural network well. Given set of four-vectors, PELICAN computes all pairwise inner products, which are Lorentz invariants, and then processes them with a permutation-equivarini architecture. LorentzNet maintains scalar and four-vector representations and updates them with graph attention mechanism similar to the one proposed by Villar et al. [74].
Flow matching Continuous normalizing flows [26] are a class of generative models that push a sample from a base density through a transformation defined by an ordinary differential equation.
Specifically, the evolution of point x E RRd is modelled as a time-dependent flow t/z : IRd Rd with (xx) u, (x)), (z) z, where ue is a time-dependent vector field.
Conditional flow matching [3, 53] a simple and scalable training algorithm for continuous nor- malizing flows that does not require the simulation of trajectories during training. Instead, the ob- jective is to match a vector field ey (z), parametrized by a neural network, onto a conditional target vector field un (z(21) along conditional probability path pt (2|21), minimizing the loss LCFM = Ex-erjo.11, (a (1x)) - ur (x(rr))) )|l where x ~ g(x1) are samples from the base dis- tribution.
Choosing a well-suited probability path and corresponding target vector field substantially improve the data efficiency and sampling quality. A principled approach to this choice is Riemannian flow matching (RFM) [25]. Instead of connecting target and latent space points by straight lines in Euclidean space. RFM proposes to choose probability paths based on the metric of the manifold structure of the data space. If available in closed form, they propose to use geodesies as probability paths, which corresponds to optimal transport between base and data density.
3 The Lorentz Geometric Algebra Transformer (L-GATr)
3.1 Lorentz-equivariant architecture
The inputs, hidden states, and outputs of L-GATr are variable- size sets of tokens. Each token consists of n copies of the geometric algebra G1,3 and m additional scalar channels.
Geometric algebra representations
The geometric algebra G1,3 is defined formally in Appendix A. In practice, G1,3 is a 16-dimensional vector space that consists of multiple subspaces (or grades). The 0-th grade consists of scalars that do not transform under Lorentz transformations, for instance embeddings of particle types or regression amplitudes. The first grade contains space-timg time four-vectors such the four-momenta = (E. p). The remaining grades extend these objects to higher orders (i. e. antisymmetric tensors), increasing expressivity. In addition, the geometric algebra defines a bilinear map, the geometric product G1,3 x G1,a + G1,a, which contains both the space-time inner product and a generalization of the Euclidean cross product.
This representation naturally fits most LHC problems, which are canonically represented as sets of particles, each parameterized with type information and a token, store the particle type as a one-hot embedding in the scalar channels and the four-momentum in the first grade of the geometric algebra.
We represent each particle as
Lorentz-equivariant linear layers We define several new layers that have both G1,3 and additional scalar representations inputs and outputs. For readability, will suppress the scalar channels in the following. We require each layer ((x) to be equivariant with respect to Lorentz transformations E SO* (1,3): f(A-z) = A f()), where denotes the action of the Lorentz group on the geometric algebra (see Appendix A). Lorentz equivariance strongly constrains linear maps between geometric algebra representations:"
0 3) equuivaaia linear maps are restricted to the first sum; the additional equivariance under reflections forbids the multiplication with the pseudoscalar.

---

Proposition 1. Any linear map Linear G1. - G1,3 that is equivariant to SO* (1,3) ) is of the form
Linear(x)) Σ (1(
2012 (2)
(I)
for parameters v, E T Here 60123 is the pseudoscalar, the unique highest-grade basis element in G1,3: (x)x is the blade projection of a multivector, which sets all non-grade-k elements to zero.
We show this in Appendix A. In our architecture, linear layers map between multiple input and output channels. There are then ten learnable weights Uk, Wi for each pair of input and output G1,3 channels (plus the usual weights for linear maps between the additional scalar channels).
Lorentz-equivariint non-linear layers We define four additional layers, all which are manifestly Lorentz-equivariant. The first is the scaled dot- dot-product attention
(gv kio)
Attention(g,k, v)ve
Σ Softmax,
Vie'
VI6ng
where the indices i, label tokens, e label channels, ne is the number of channels, and (s is the G1.3 inner product. This inner product can be rewritten as a pre-computed list of signs and Euclidean inner product, which why we can compute the attention mechanism with efficient backends developed for the original Transformer architecture, for instance Flash Attention [33]. This key to the good scalability of L-GATr, which we will demonstrate later.
When defining a normalization layer, we have to be careful: in the GLS inner product, cancellations between positive-norm directions and negative-norm directions can lead to norm values much smaller than the scale of the individual components; dividing by the norm then risks blowing up the data. These cancellations are unavoidable consequence of the geometry of space-time. We mitigate this issue by using the grade-wise absolute value of the inner product in the norm
LayerNorm(()) =
1 The
(xe))e (ee))) te,
k=0
applying absolute value around each grade of each multivector channel (xc) Here e > 0 is a constant that further numerically stabilizes the operation. This normalization was proposed by De Haan et al. [36] for -invariiari architectures, we adapt it to the Lorentz-equivariant setting.
We also the geometric product GP(x. y) zy defined by the geometric algebra G1,3- Finally, we use the scalar-gated GELU [46] nonlinearities GatedGELU(z) = GE)U(()))),, as proposed by Brehmer al. [14].
Transformer architecture
We combine these layers into a Transformer architecture [73, 76]:
F LayerN)
AttentionBlock(a) Lincar Attention Linear()), Linearzz)
MTPPBlok(22) = Linear o GatedGELU o Linear o GP(Linear(i), Linear(T)) + Block(æ) MLPBlock AttentiionBlck(())
L-GATr(z) = Lincar o Block e Block o ... Block e Lincar(r).
This L-GATr architecture is structurally similar to the original GATr architecture [14], but the representations, linear layers, attention mechanism, geometric product, and normalization layer are different accommodate the different nature of the data and different symmetry group.
Lorentz symmetry breaking While fundamental physics is (to the best of our knowledge) symmet- ric under Lorentz transformations, the LHC measurement process is not. The direction of the proton beams presents the most obvious violation of this symmetry. Smaller violations are due to the detec- tor resolution: particles hitting the central part of the detector (orthogonal to the beam the detector rest frame) are typically reconstructed with higher precision than those emerging narrow angle

---

to the beam. Even smaller violations come, for instance, from individual defunct detector elements. Solving some tasks may therefore benefit from a network that can break Lorentz equivariance.
L--GATr supports such broken or approximate symmetries by including the symmetry-breaking effects additional inputs into the network. Concretely, whenever analyze reconstruction-level data, include the beam directions; see Appendix C. This approach combines the strong inductive biases of Lorentz-equivariant architecture with the ability learn to break the symmetry when required.
3.2 Lorentz-equivariant flow matching
In addition to regression and classification models, we construct a generative model for particle data. Besides the strict requirements precision, flexibility, and data efficiency, generative models of LHC data need to able to address sharp edges and long tails in high-dimensional distributions.
We develop a continuous normalizing flow based on an L-GATr vector field and train it with Riemannian flow matching (RFM) [25]. This approach has several compelling properties: training simulation-free and scalable and the generative model is corentz-equivaricn approach allows us to deal with sharp edges and long tails in a geometric way: we parameterize the reachable nuurmmoentumm space for each particle manifold and use geodesics on this manifold as probability paths from base samples data points.
In addition, the RFM
Probability paths perfect for particles Concretely, reconstructed particles p p (E, p) are often required to satisfy constraints of the form p + p 2 Pr min and p² > O. Following Ref's [17, 18, 45, 48], we parameterize this manifold with physically motivated coordinates y = (ymm Vpp 77, 6). These variables form alternative basis for the particle nou-momenta and defined through the map
p = (E, L.. Po. Po. P2) = f(y)= = (Vmi² p cosh? n,, pr cos Φ, pr sin Φ, pr pish (4)
where m² exp(ym) and pT PT min + exp( This basis is better aligned with the physically relevant proper-
60
ties of particles the context of collider experiment: n and ø represent the angle in which particle is moving, Up a measure the momentum with which moves away from the collision, and ym is related to its mass.
20-
Py [Gev]
Pr < Pimmn
We define constant diagonal metric in the coordinates y and use the corresponding geodesics as probability paths. This Riemannian manifold is geodesically convex, mean- ing any two points are connected by unique geodesic, and geodesically complete, meaning that paths thus never en- ter Nour-momentum regions forbidden by the phase-space cuts. By also running the ordinary differential equation (ODE) solver in these coordinates, we guarantee that each sample satisfies the four-mon momentum constraints. As an added benefit, this choice of metric compresses the high- energy tails of typical particle distributions and thus sim- plifies learning them correctly.
20
40
60
-25
0
25
50
75
Px [GeV]
Figure Target vector field for Riemannian flow matching. Our choice of metric space guarantees that the generative model respects phase-spass boundaries (red circle).
In Fig. 2, we show target probability paths generated in this way. Our approach ensures that none of the trajectories pass through the phase-space region pr < Pr. where the target density does not have support; instead, the geodesics lead around this problematic region.
"Strictly speaking, only the map from the base density to data space is equivariant with respect to the full Lorentz group. The base density and thus also the density of the generative model are only invariant with respect to rotations. This is because the group of boosts is not compact: it is impossible to define properly normalized density that assigns the same probability corentz-invarie
every boosted data variation. In theory, could define fully measure: then the flow would define Lorentz-invarian measure normalizable- good luck with that. In practice, compact subsets of the orbits, for instance characterized by a limited range of the center-of -aasss mass momentum, suffice. All of this is in analogy to "E(3 (3)-iiivrraan" generative models [47], which strictly only invariant rotations, but not to (pon-compact) translations.
would not be

---

MLP
Z 244g
10"1
Transformer
10
DSI
MSE normalized amplitudes
10"2
MSE on normalized amplitudes
CGENN
10"2
10-1
L-GATr
10"4
10-3
MLP
10"5
10-4
Transformer
DSE
10
10
CGENN
1077
L-GATr
10
7+1g
7+22g
7.43g
7.44g
10
Number of training samples
Figure 3: Amplitude surrogates. Left: Surrogate error for processes of increasing particle multiplicity and complexity, training on the full dataset of 4 10 samples. L-GATr outperforms the baselines, especially at more complex processes. Right: Surrogate error function of the training dataset size.
4
Experiments
We now demonstrate L-GATr in three applications. Each addresses different problem in the data- analysis pipeline sketched in Fig.
4.1 Surrogates for QFT amplitudes
Problem We first demonstrate L-GATr as neural surrogate for quantum field theoretical ampli- tudes [6-8, 59, 60), the core of the theory predictions that L.HC measurements are compared to. These amplitudes describe the (un-normalized) probability of interactions of fundamental particles as function their four-mon momenta. As this is a fundamental interaction and does not include the measurement process, is exactly because it requires solving complex integrals, on the other hand because the number relevant terms combinatorially grows with the number of particles. Neural surrogates can greatly speed up this process and thus enable better theory predictions, but accurately modelling the amplitudes of high-multiphl processes has been challenging.
Evaluating them is expensive, the one hand
As example processes, we study qğ + ng, the production of Z boson with n 1..... additional gluons from a quark-antiquark pair. For each gluon multiplicity, we train a L-GATr model to predict the amplitude
function of the nur-momenta of the initial and final particles." The generation of the training data and the precise setup of the learning problem are described in Appendix C. We compare L-GATr to various baselines, including the Lorentz-equivarian message- passing architecture CGENN [68], a Transformer [73], and DSI, baseline based on the Deep Sets framework [77] that we develop ourselves; we describe in detail Appendix
Surrogate quality L-GATr consistently approximates the amplitudes with high precision, as we show in the left panel of Fig. For a small number of particles, it is slightly worse than our own baseline DSI, but scales much better to large number of particles, where outperforms all other methods. This is exactly the region in which neural surrogates could have the highest impact.
Data efficiency In the right panel of Fig. we study the data efficiency of the different architectures. We find that L-GATr is competitive at any training data size, combining the small-data ad advantages of its strong inductive biases and the big-data advantages its Transformer architecture.
'We also experimented with training a single L-GATr model to learn the amplitudes of all processes jointly, finding similar performance.

---

Model
Accuracy
AUC
L/en (es 0.5)
1/eB (es 0.3)
TopoDNN [49]
0,916
0.972
295 + 5
LoLa [16]
0.929
0.980
722 17
P-CNN [1]
0,930
0,9803
201 + 4
759 +
N-subjettiness [62]
0,929
0,981
867 1 15
PFN [51]
0,932
0.9819
247 1 3
888 1 17
TreeNiN [58]
0,933
0.982
1025 1 11
ParticleNet [64]
0,940
0,9858
397 +
1615 =
ParT [65]
0,940
0,9858
413 + 16
1602 - 81
LorentzNet [42]
0,942
0.9868
498 + 18
2195 - 173
CGENN [68]
0.942
0.9869
500
2172
PELICAN* [10]
0.9426
0.0002
0.9870
0.0001
2250
75
I.-ATT (ours)*
0.9423
0.0002
0.9870
0.0001
540 + 20
2240 +
Table 1:
Top tagging. We compare accuracy, area under the ROC curve (AUC), and inverse background acceptance rate 1/eg at two different signal acceptance rates (or recall) es E (0.3, 0.5) for the top tagging dataset from Kasieczka al. [50]. Lorentz-equivariant methods indicated with asterisk ; the hest results for each metric are bold. For L-GATr, we show the mean and standard deviation five random seeds. Baseline results are taken from the literature.
4.2 Top tagging
Problem Next, we turn to the problem of classifying whether a spray of reconstructed hadrons originated from the decay of a top quark any other process. This problem of top tagging is an important filtering step in any analysis that targets the physics of top quarks, the heaviest elementary particle in the Standard Model. Particle collisions involving these particles to physicists because the production and decay probabilities of top quarks are sensitive to several proposed theories of new physics, including for instance the existence of "supersymmetrie" particles. We use the established top tagging dataset by Kasieczka et al. [49, 50] as benchmark and compare the published results for many algorithms and architectures.
of particular interest
Results As shown in Tbl. 1, L-GATr on par with or better than even the strongest baselines on this well-studied benchmark.
4.3 Generative modelling
Problem Finally, we study the generative modelling of reconstructed events an end-to-end generation task [17, 18], bypassing the whole simulation chain visualized in Fig. 1. Such generative models can obliterate the computational cost of both the theory computations and the detector simulation once. However, the high-dimensional distributions of reconstructed particles often have non-trivial kinematic features that are challenging for generative models to learn, for instance the properties of unstable resonances and angular correlations. We focus on the processes pp - tī + jets, the generation of top pairs with 0... 4 additional jets, where the top quarks decay hadronically, - by'd""
We train continuous normalizing flows based on an L-GATr network with the Riemannian flow matching objective deseribed in Sec. 3. As baselines, we consider similar flow matching models, but use MLP and Transformer networks as score models, as proposed by Refs. [18, 45). We also construct a flow matching model using the E(3)-equivariant GATr from Ref. [14]. Finally, we also train JetGPT [18] model, an autoregressive transformer architecture developed for particle physics that is not equivariant to the Lorentz symmetry.
Kinematic distributions We begin with qualitative analysis of the samples from the generative models. In Fig. 4 we show example marginal distributions from the different models and compare them to the ground-truie -tuthh distribution in the test set. We select three marginals that are notoriously difficult to model correctly for generative models. While the differences are subtle and only visible in tails and edges of the distributions, L-GATr matches the true distribution better than the baselines. However, none of the models are able capture the kinematics of the top mass peak at percent-level precision yet.
8

---

10-2
0.015
0.3
Marginal density
Marginal density
Marginal density
Ground truth MLP
Transformer L-GATr
10
Model
Model Gmundmma
0.0
Medel Gmerdde
0.0
1.1
1.0
1.0
0.9
100
300
250
Pry [GeV]
AR
m, [GeV]
Figure 4: Generative modelling: Marginal distributions of reconstructed particles in the pp ti 4 jets process. We compare the ground-truth distribution (black) to three generative models: continuous normalizing flows based on a Transformer. MLP, or our L-GATr network. The three marginals shown represent kinematic features that are known to be challenging. The L-GATr flow describes them most accurately.
Log likelihood Next, we evaluate the generative models quantitatively through the log likelihood of data samples under the trained models; see Appendix C for details. The left panel of Fig. 5 shows that the L-GATr models outperform all baselines across all different jet multiplicities. They maintain this performance advantage also for smaller training data size, as shown the right panel. The flow models, including L-GATr, are more data-efficient than the autoregressive transformer JetGPT.
Classifier two-sample test How close to the ground-truth distribution are these generative models really? Neither marginal distributions nor log likelihood scores fully answer this question, as the former neglect most of the sigh-dimensional information and the latter do not have a known ground- truth value compare to. We therefore perform a classifier sample test [55]. We find that L- GATr samples are difficult to distinguish from the ground-truth distribution: a classifier trained to discriminate them achieves only a ROC AUC of between 0.51 and 0.56, depending on the process. In contrast, Transformer and MLP distributions are more easily discriminated from the background, with ROC AUC results between 0.58 and 0.85. For details, see Appendix C.
Effect of Riemannian flow matching How important was our choice of probability paths through Riemannian flow match- ing for the performance these models? Tbl. we compare the log likelihood of CFM L-GATr models that differ only in the probability paths. Clearly, the Riemannian flow matching approach that allows us encode geometric constraints is cru- cial for a good performance. We find similarly large gains for all architectures.
Probability paths
NLL
Euclidean RFM
-30.11 1 0.98
-32.65 1L 0.01
Table 2: Benefit of Riemannian flow matching generative models. show the negative log likelihood on the tÉ Oj test set (lower is better).
4.4 Computational cost and scalability
Finally, we briefly comment on L-GATr's computational cost. Compared to vanilla Transformer, the architecture has some computational overhead because of the more complex linear maps. However, scales exactly in the same way to large particle multiplicities, where both architectures are bottle- necked by the same dot-pro attention mechanism. At the same time, L-GATr is substantially more efficient than equivariant architectures based message passing, both in terms of compute and memory. This is because high-energy physics problems do not lend themselves to sparse graphs, and for dense graphs, dot-product attention is much more efficient. See Appendix C for measurements.
Discussion
Out of all areas of science, high-energy physics a strong contender for the field which symmetries play the most central role. Surprisingly, while particle physicists were quick to embrace machine learning, architectures tailored to the symmetries inherent in particle physics problems have received comparably little attention.

---

-5.3
JeiGPT
JetGPT
Negative log likelihood per particle
Transformer CPM
Transformer CFM
B(3)-GATY CFM
-26
E(3)-GAIt CFM
Negativee log likellhood
L-GAII CFM
L-GAT: CFM
-28
5.4
-30-
-32
cZ+0f
-5.5
+
rE+1j
t
+
+-
10a
10P
10
Number training samples
Figure 5: Generative modelling: negative log likelihond
test set (lower better). Left: For different processes. Right: As function of the training dataset size. We show the mean and standard deviation of three random seeds. The L-GATr conditional flow matching (CFM) model outperforms all other CFM models well the autoregressive transformer JetGPT, across all processes and all training sizes.
We introduced the Lorentz Geometric Algebra Transformer (L-GATr), versatile architecture with strong inductive biases for high-energy physics: its representations are based on particle four- momenta, extended to higher orders in geometric algebra, and its layers equívariant with respect to the Lorentz symmetry of special relativity. At the same time, L-GATr is Transformer, and scales favorably to large capacity and large numbers of input tokens.
We demonstrated L-GATr's versatility on diverse regression, classification, and generative modelling tasks from the LHC analysis workflow, For the latter, we constructed the first Lorentz-equiivarian generative model based on Riemannian flow matching. Across all experiments, L-GATr performed as well as better than strong baselines.
Still, L-GATr has its limitations. While the architecture scales better than comparable message- passing networks, it has some computational overhead compared to, for instance, efficient Transformer implementations. And while L-GATr should in principle be suitable for pretraining across multiple problems, we have not yet investigated its potential as foundation model.
While the LHC is preparing for the high-luminosity runs and its legacy measurements, the high- energy physics community is optimizing all steps of the analysis pipeline. Deploying performant and data-efficient architectures such L-GATr could improve this pipeline in many places. We hope that this will ultimately contribute to more precise measurements of nature at its most fundamental level.
Acknowledgements
We would like to thank Taco Cohen and Anja Butter for fruitful discussions, Nathan Hütsch for help with conditional flow matching, and David Ruhe for support with the CGENN code.
J.S., V.B. and T.P. are supported by the tionale Spitzenforschung, project Uncertainties - Teaching Al Limits NWWSTIIF222-0011 Deutsche Forschungsgemeeinscaaff (DFG, German Research Foundation) under grant 396021762 TRR 257 Particle Physics Phenomenology after the Higgs Discovery, and through Germany's 's Ex- cellence Strategy EXC 21 1/11 390900948 (the Heidelberg STRUCTURES Excellence Cluster). J.S. is funded by the Sarl-Zeiss-Se StifttSn through the project Model-Basee Al: Physical Models and Deep Learning for Imaging and Cancer Treatment. V.B. is supported by the BMBF Junior Group Generative Precision Networks for Particle Physics (DLR 011S22079). V.B. acknowledges finan- cial support from the Grant No. ASFAE/2022/009 (Generalitat Valenciana and MCIN, NextGienrr- tionEU PRTR-C17.101, J.T. is supported by the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelligence and Fundamental Interac- tions, http::/itaai:i org/), by the U.S. Department of Energy Office of High Energy Physics un- der grant number DE-SC0012567. and by the Simons Foundation through Investigator grant 929241.
Stifftu through h the program Interna-
10

---

References
[1] Boosted jet identification using particle candidates and deep neural networks. 2017. URL https://cds cern. ch/record/2295725. (Cited on page 8)
[2] Georges Aad et al. Measurements of top-quark pair single- and double-differretiall cross- sections the all- hadronic channel in pp collisions at v 13 TeV using the ATLAS detector. JHEP, 01:033, 2021. doi: 00.100//HEPP01(2211000
(Cited on page 19)
[3] Michael S Albergo. Nicholas M Boffi, and Eric Vanden Eijnden. Stochastic interpolants: unifying framework for flows and diffusions. arXiv preprini arXiv:2303.08897, 2023, (Cited on page 4)
[4] Johan Alwall, Michel Herquet, Fabio Maltoni, Olivier Mattelaer, and Tim Stelzer. MadGraph 5 Going Beyond. JHEP, 06: 128, 2011. doi: 10.1007/HHEP062201] 1)128. (Cited on pages 17 and 19)
[5] ATLAS collaboration. Observation of particle in the search for the Standard Model Higgs boson with the ATLAS detector at the LHC. Physics Letters B, 716(1):1-29) 2012. (Cited on page 1)
[6] Joseph Aylett-Bullock, Simon Badger, and Ryan Moodie. Optimising simulations for diphoton production at hadron colliders using amplitude neural networks. Journal of High Energy Physics, 2021(8):1-30, 2021. (Cited on page 7)
[7] Simon Badger and Joseph Bullock. Using neural networks for efficient evaluation of high multiplicity scattering amplitudes. Journal of High Energy Physics, 2020(6): 2020. (Not cited.)
[8] Simon Badger, Anja Butter, Michel Luchmann, Sebastian Pitz, and Tilman Plchn. Loop amplitudes from precision networks. SciPost Physics Core, 6(2):034, 2023. (Cited on page 7)
[9] Eduardo Bayro-Corrochano Sven Buchholz, and Gerald Sommer. A new self-organizing neural network using geometric algebra. In Proceedings of 13th International Conference on Pattern Recognition, volume 4, pages 555-559. IEEE, 1996. (Cited page 16)
[10] Alexander Bogatskiy, Timothy Hoffman, David W. Miller, and Jan T. Offermann. Pelican: Permutation equivariant and lorentz invariant or covariant aggregator network for particle physics, 2022. (Cited pages 2, 4, and 8)
[11] Johannes Brandstetter, Rianne van den Berg, Max Welling, and Jayesh K Gupta. Clifford neural layers for PDE modeling. 2oXiv:2009.04934, 2022. (Cited page 16)
[12] Johann Brehmer and Kyle Cranmer. Simulation- based inference methods for particle physics. In Artificial Intelligence for High Energy Physics, pages 579-611. World Scientific, 2022. (Cited
on page 1)
[13] Johann Brehmer, Kyle Cranmer, Gilles Louppe, and Juan Pavez. Constraining effective field theories with machine learning. Physical Review Letters, 121(11)111101, 2018. (Cited on page I)
[14] Johann Brehmer, Pim de Haan, Sönke Behrends, and Taco Cohen. Geometric Algebra Trans- former. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 37, 2023. (Cited on pages 2, 3, 8, 16, and 17)
[15] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković, Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. 2021. (Cited on page 3)
[16] Anja Butter, Gregor Kasieczka, Tilman Plehn, and Michael Russell. Deep-leamedd top tagging with lorentz layer. SciPost Physics, 5(3):028, 2018. (Cited on page 8)
[17] Anja Butter, Theo Heimel, Sander Hummerich, Tobias Krebs, Tilman Plehn, Armand Rousselot, and Sophia Vent. Generative networks for precision enthusiasts. SciPost Phys., 14(4):078, 2023. doi: 10.2 1 102114////iisstPhyss 14.4.078.
(Cited pages 6 and 8)
11

---

[18] Anja Butter, Nathan Huetsch, Sofia Palacios Schweitzer, Tilman Plehn, Peter Sorrenson, and Jonas Spinner. Jet Diffusion versus JetGPT Modern Networks for the LHC. 5 2023. (Cited on pages 6, 8, and 20)
[19] Anja Butter, Nathan Huetsch, Sofia Palacios Schweitzer, Tilman Plehn, Peter Sorrenson, and Jonas Spinner. Jet diffusion versus jetgpt-modern modern networks for the lhc. arXiv preprint arXiv:2305.10472. 10475, 2023. (Cited page 2)
[20] Anja Butter, Tilman Plehn. Steffen Schumann, Simon Badger, Sascha Caron, Kyle Cranmer, Francesco Armando Di Bello, Etienne Dreyer, Stefano Forte, Sanmay Ganguly, et al. Machine learning and LHC event generation. SciPost Physics, 14(4):079, 2023. (Cited on page 1)
[21] Matteo Cacciari, Gavin Salam, and Gregory Soyez. The anti-k; jet clustering algorithm. JHEP, 04:063, 2008. doi: 10.1088/1126 02108///126-60088/000000066663 (Cited page 19)
[22] Mattco Cacciari, Gavin P. Salam, and Gregory Soyez. FastJet User Manual. Eur. Phys. J. C, 72:
1896, 2012. doi: 10.1 10.11440epic//10002220000 1896-2. (Cited on page 19)
[23] Giuseppe Carlco, Ignacio Cirac, Kyle Cranmer, Laurent Daudet, Maria Schuld, Naftali Tishby, Leslie Vogt-Maranto, and Lenka Zdeborováá Machine learning and the physical sciences. Reviews of Modern Physics, 91(4):04000, 2019. (Cited on page 1)
[24] CERN Data Centre. Key facts and URL cern. ch/site//eeeeeeeeeee stt/sites////e (Cited on page 1)
web. http:////w.......coooo/////////a
[25] Ricky TQ Chen and Yaron Lipman. Riemannian flow matching general geometries. arXiv:2302.nm 03660, 2023. (Cited on pages 2, 4, and 6)
[26] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 31, 2018. (Cited on pages and 20)
[27] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoe V. Le. Symbolic discovery of optimization algorithms, 2023. (Cited on page 19)
[28] M C Cheng, V Anagiannis, M Weiler, and others. Covariance in physics and convolutional neural networks. arXiv preprint arXiv, 2019. URL https://arxiv org/abe///////eoooeo/ooooeee abs/ /196. 02481. (Cited page
[29] William Kingdon Clifford. Applications of Grassmann's 's Extensive Algebra. Amer. J. Math., 1 (4):350-358, 1878. (Cited on page 3)
[30] CMS collaboration. Observation of a new boson at a mass of 125 GeV with the CMS experiment the LHC. Physics Letiers B, 716(1)330-11, 2012. (Cited page 1)
[31] Taco Cohen. Equivariant Convolutional Networks. PhD thesis, University of Amsterdam, 2021.
(Cited on page 3)
[32] Kyle Cranmer, Johann Brehmer, and Gilles Louppe. The frontier of simulation-based inference. Proceedings of the National Academy of Sciences, 117443)30055-3000,,, 2020. (Cited on page 1)
[33] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. Advances in Neural Information Processing Systems, 111166644435599 2022. (Cited on pages and 5)
[34] Ranit Das, Luigi Favaro, Theo Heimel, Claudius Krause, Tilman Plehn, and David Shih. How to understand limitations of generative networks. SciPost Phys., 16(1):031, 2024. doi: 10.2 (0.21468//ciPostPPys. 16. 1.031. (Cited page 21)
[35] J. de Favereau, C. Delaere, P. Demin, A. Giammanco, V. Lemaitre, A. Mertens, and M. Selvaggi. DELPHES 3, modular framework for fast simulation generic collider experiment. JHEP, 02:057, 2014. doi: 10. 10.1007/JHEE02//0440005 (Cited on page 19)
12

---

[36] Pim De Haan, Taco Cohen, and Johann Brehmer. Euclidean, projective, conformal: Choosing a geometric algebra for equivariant transformers. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li, editors, Proceedings of The 27th International Conference Artificial Intelligence and Statisties, volume 238 of Proceedings of Machine Learning Research, pages PMLR, 02-04 May 2024. (Cited on pages 2, 3, 5, and 16)
[37] Luke de Oliveira, Michael Kagan, Lester Mackey, Benjamin Nachman, and Ariel Schwartzman. Jet-images deep learning edition. Journal of High Energy Physics, 2016(7):1-32, 2016. (Cited on page 2)
[38] C Doran and A Lasenby. Geometric algebra for physicists. Cambridge University Press, 2003. (Cited on page 16)
[39] Albert Einstein. Zur Elektrodynamik bewegter Körper. Annalen der Physik, 4, 1905. (Cited on pages and 3)
[40] Mare Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. April 2021. URL http: //arxiv.org/ Frrzzi..org/ab//1/000 09459. (Cited on page 16)
[41] Sheldon Glashow. Partial-symmetries of weak interactions. Nuclear physics, 22(4):579-588,
1961. (Cited on page 3)
[42] Shiqi Gong, Qi Meng. Jue Zhang, Huilin Qu, Congqiao Li, Sitian Qian, Weitao Du Zhi- Ming Ma, and Tie-Yan Liu. An efficient lorentz equivariant graph neural network for jet tagging. Journal of High Energy Physies, 2022(7), July 2022. ISSN 1029-8479. doi: 10. 1007//hep07(2022)0300 URL http:///// dow. http://xx.....oo//////////////// 101000000HHPPO (2022)030. (Cited on pages 2, 4, and 8)
[43] Hermann Grassmann. Die lineale Ausdehnungslehre. Otto Wigand, Leipzig, 1844. (Cited on page 3)
[44] Dan Guest, Kyle Cranmer, and Daniel Whiteson. Deep learning and its application to LHC physics. Annual Review of Nuclear and Particle Science, 68:161-181, 2018. (Cited on page 1)
[45] Theo Heimel, Nathan Huetsch, Ramon Winterhalder, Tilman Plchn, and Anja Butter. Precision- Machine Learning for the Matrix Element Method. 10 2023. (Cited on pages 6 and 8)
[46] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv:1606.08415, 2016. (Cited on page 5)
[47] Emiel Hoogeboom, Víctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant dif- fusion for molecule generation in 3D. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Confer- ence on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8867-8887. PMLR, 2022. URL http:: https://prrcceeedngg mlr wal...prss/vvve wwteesssvvvv2//eee hoogeboom22... html. (Cited on page 6)
[48] Nathan Huetsch et al. The Landscape of Unfolding with Machine Learning. 2024. (Cited on page 6)
[49] Gregor Kasieczka, Tilman Plehn, Anja Butter, Kyle Cranmer, Dipsikha Debnath, Barry M Dillon, Malcolm Fairbairn, Darius A Faroughy, Wojtek Fedorko, Christophe Gay, al. The machine learning landscape of top taggers. SciPost Physics, 7(1)0114, 2019. (Cited on pages 1, 8, and 18)
[50] Gregor Kasieczka, Tilman Plehn, Jennifer Thompson, and Michael Russel. Top quark tagging reference dataset, March 2019. URL https:://oi....oeo org/ 5281/zenodo 2603256. (Cited on pages 8 and 18)
[51] Patrick Komiske, Eric M Metodiev, and Jesse Thaler. Energy flow networks: deep sets for particle jets. Journal of High Energy Physics, 2019(1): 2019. (Cited on page
[52] Andrew J Larkoski, Ian Moult, and Benjamin Nachman. Jet substructure at the Large Hadron Collider: a review of recent advances theory and machine learning. Physics Reports, 841: 1 1-63, 2020. (Cited on page 1)
13

---

[53] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv:2210.02747, 2022. (Cited on page 4)
[54] Cong Liu, David Ruhe, Floor Eijkelboom, and Patrick Forré. Clifford group equivariant simplicial message passing networks. arXiv preprint arXiv:2402.10011, 2024. (Cited on page
[55] David Lapez-Paz and Maxime Oquab. Revisiting classifier two-sample tests, arXiv preprint arXiv:1610.06545 2016. (Cited on page 9)
[56] Uyy Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with restarts. CoRR, hbss/608.00398: 2016. URL http://ttaaatiti org/ aabs/160.. 03983. (Cited on page 19)
[57] Pertti Lounesto. Clifford Algebras and Spinors. London Mathematical Society Lecture Note. Cambridge University Press, 2001. (Cited on page 16)
[58] Sebastian Macaluso and Kyle Cranmer. Tree Network Network (TreeNiN) for jet physics, 2019. URL https: //gps:///iii github. 0do.2582216. (Cited on page 8)
DOI: 10.528 1/zen-
[59] Maître and Truong. One-loop matrix element emulation with factorisation awareness. Journal of High Energy Physics, 2023(5):1-21, 2023. (Cited on page
[60] Daniel Maître and Henry Truong. factorisation-awar matrix element emulator. Journal of High Energy Physics, 2021(11):1-244 2021. (Cited on page 7)
[61] Hermann Minkowski. Die Grundgleichungen für die elektromagnetischen Vorgänge in bewegten Körpern. Nachrichten von der Gesellschaft der Wissenschaften zu Göttingen, Mathematisch- Physikalische Klasse, 1908:53-111, 1908. (Cited page 3)
[62] Liam Moore, Karl Nordström, Sreedevi Varma, and Malcolm Fairbaim. Reports of my demise are greatly exaggerated: n-subjettiness taggers take on jet images. SciPost physics, 7(3):036, 2019. (Cited on page 8)
[63] Henri Poincaré. Sur la dynamique de l'électron. Circolo Matematico di Palermo, 1906. (Cited on pages 2 and 3)
[64] Huilin Qu and Loukas Gouskos. Jet tagging via particle clouds. Physical Review D, 101(5), March 2020. ISSN 2470-0029. doi: 10.1103/phy 101.00000.0000001.... URL http://dz. doi.
org/10 11/1110//PPP 11100/hyseevDD..... 056019. (Cited on page 8)
[65] Huilin Qu, Congqiao Li, and Sitian Qian. Particle transformer for jet tagging. 2024. (Cited on page 8)
[66] Alexander Radovic, Mike Williams, David Rousseau, Michael Kagan, Daniele Bonacorsi, Alexander Himmel, Adam Aurisano, Kazuhiro Terao, and Taritree Wongjirad. Machine learning at the energy and intensity frontiers of particle physics. Nature, 560(7716):411455 2018. (Cited
on page 1)
[67] Martin Roelfs and Steven De Keninck. Graded symmetry groups: plane and simple. arXiv:2107.03771, 2021. (Cited on page 3)
[68] David Ruhe, Johannes Brandstetter, and Patrick Forré. Clifford group equivariant neural networks. In Advances in Neural Information Processing Systems, volume 37, 2023. (Cited on pages 2, 3, 8, 16, 17, 18, 21, and 22)
[69] David Ruhe, Jayesh K Gupta, Steven de Keninck, Max Welling, and Johannes Brandstetter. Geometric clifford algebra networks. In International Conference on Machine Learning, 2023. (Cited on page 16)
[70] Gavin P. Salam. Towards Jetography, Eur. Phys. J. C, 67:637-666, 2010. doi: 10.1 140/ep// s 1000520010-1144-2 (Cited on page 3)
[71] A. M. Sirunyan et al. Particle-flow reconstruction and global event deseription with the CMS detector. JINS,, 12(10):P 2017. doi:
(Cited on page
14

---

[72] Torbjörn Sjöstrand, Stefan Ask, Jesper R. Christiansen, Richard Corke, Nishita Desai, Philip Ilten, Stephen Mrenna, Stefan Prestel, Christine O. Rasmussen, and Peter Z. Skands. An introduction to PYTHIA 8.2. Comput. Phys. Commun., 191:159-177, 2015. doi: 10.1016/j.cpp
2015.01.024. (Cited on page 19)
[73] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017. (Cited on pages 3, 5, 7, and 19)
[74] Soledad Villar, David W Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith. Scalars universal: Equivariant machine learning, structured like classical physies. Advances in Neural Information Processing Systems, 14:28848-28886 2021. (Cited on page 4)
[75] Steven Weinberg. A model of leptons. Physical Review Letters, 19(21):1264, 1967. (Cited on page 3)
[76] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. International Conference on Machine Learning, pages 10524-10533, PMLR, 2020. (Cited on page 5)
[77] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. Advances in neural information processing systems, 30, 2017. (Cited pages? and 17)
[78] Maksim Zhdanov, David Ruhe, Maurice Weiler, Ana Lucic, Johannes Brandstetter, and Patrick Forré. Clifford-steerabe convolutional neural networks. arXiv preprint 2402:22402.17730 2024. (Cited on page 3)
15

---

A Geometric algebra
Geometric algebras are mathematical objects that were initially used for physics. Although they have been used in machine learning for decades [9], they have seen a recent uptick in popularity [11, 14, 38, 57, 68, 69]. In this section, will introduce geometric algebras and the relevant concepts.
An algebra is vector space that is equipped with an associative bilinear product. Given a vector space V with a symmetric bilinear inner product, we can construct an algebra G(V), called the geometric or Clifford algebra, in the following way: choose an orthogonal basis e; of the original d-dimensional vector space V. Then, the algebra has 2d dimensions with basis given by elements en "hee....... =: da. with S j1 < J2 < ... < jk S d, 0 S S d. For example, for R3, with orthonormal basis 61, e2, C3, a basis for the algebra G(R²) is
(5) An algebra element spanned by basis elements with k indices is called k-vector or vector of grade k. generic element whose basis elements can have varying grades is called multivector. multivector x can be projected to k-vector with the grade projection (r)).
1, e1 G22 c3, e12, c13: G23, 6123-
The product on the algebra, called the geometric product, is defined to satisfy e;e; = -yee; if + j and e,e; = (ex, ex), which by bilinearity and associativity fully specifies the algebra. Given an algebra GV). there is a group Pin(V) that is generated by the 1-vectors in the algebra with norm +1, and whose group product is the geometric product. This group has linear action p: Pin(V) x G(V) + (V) on the algebra defined such that for any unit 1-vector u Pin(V) and 1-vector I E G(V)
p(u, x)) = -wr."
The action defined be an algebra homomorphism, meaning that for any Pin(V), x. G(V), plu = p(u, p),, y). Also, is group action, meaning that for any two group elements u,u Pin(V), p(uu,x)) = p(u, 2(), As the group Pin(V) is generated by products of I-vectors, and the algebra G (V) is generated by linear combinations and geometric products, this fully specifies the action p.
Space-time time geometric algebra In this paper, we use the geometric algebra G1.3 = G based on cour-dimensional Minkowski space R. 1.3 which has an orthogonal basis with one basis vector eo satisfying (Cu, ea) +1 and for i = 1, 2, 3 a basis vector e, satisfying (e,ec) 4 = 11 The Pin group Pin(RR.3) 1.3) a double cover of the Lorentz group 311,3). As we do not require equivariance to time reversals or spatial mirrorings, we are only interested in equivariance to the connected subgroup SO* (1,3).
Equivariance The fact that is an algebra homomorphism is equivalent to saying that the geomet- ric product is equivariant Pin(V). Furthermore, the grades in geometric algebra form subrep- resentations [14, Prop. 2]. Thus, the grade projections are equivariant. The pseudoscalar is one- dimensional real representation, and thus must be invariant to any connected subgroup of Pin(V). Therefore, multiplying by the pseudoscalar is equivariant the connected group SO*(1, (1,3). Hence, the linear layer in Eq. (1) is equivariant, for any value of the parameters v, w.
To show that this forms a complete basis of all equivariant linear maps, use the numerical approach of De Haan et al. [36], based on Finzi et al. [40]. Numerically, we find a 10-dimensional space of equivariant maps, indicating the basis in Eq. (1) is complete, which proves Prop. 1.
Expressivity A fortiori, De Haan al. [36] showed that for several geometric algebras, any equivariant map G(V) x x G(V) - G(V) that is a polynomial function of the coefficients of G(V), can be expressed as linear combinations of grade projections, geometric products and invariant multivectors, and is thus expressible by GATr. This argument holds for any geometric algebra that is based on vector space V with a non-degenerate metric, which is the case for the Minkowski space 1,3 Hence, this expressivity argument can be extended to L-GATr: the operations in its MLPs are able to express any polyomial map of G1,3 mutlivectors.
B Architectures
Lorentz Geometric Algebra Transformer (L--GATr) Our main contribution is the L-GATr archi- tecture, described in detail in Sec. 3. Our implementation is in part based the Geometric Algebra
16

---

Transformer [14] code in version 1.0.0.' Unlike [14], we use multi-head attention, not multi-query attention.
Clifford group equivariant neural network (CGENN) We use the CGENN architecture [68] as a baseline. We use the official implementation" 10 and adapt their top-tagging code to our amplitude regression experiments.
Deep Sets with invariants (DSI) We build new architecture based on the Deep Sets frame- work [77] that uses momentum invariants as part of the input for tackling the amplitude surrogate task. Deep Sets is a permutation-invariant architecture that applies the same function to each ele- ment of an input set, aggregates the results with pernutation-inwe invariian operation like a sum, and processes the outputs with another function.
Our adaptation for the amplitude regression tasks applies the Deep Sets approach to each subset of identical particles in the inputs, as the amplitudes are manifestly invariant under permutations of the four-momenta of particles of the same type. We thus apply a different preprocessing to each particle type and aggregate them separately. In addition to the particle-specific latent space samples, the input to our main network also includes the momentum invariants for all the particles involved in the process. DSI thus combines Lorentz-squivariant, permutation-cquivarian path with non- Lorentz-equivariann non-permutann poro-prruuaatooeeuuvvrrraaa path, allowing the network to learn whether to rely on equivariant or non-equivariant features.
Both preprocessing units and the main network are implemented MLPs with GELU nonlinearities. Using this model, we are able to obtain optimal performance for simple interactions, but we observe a poor scaling behavior for prediction quality as we increase particle multiplicity.
Geometric Algebra Perceptron (GAP) To ablate what extent L-GATr's performance is due the geometric algebra representations and its equivariance properties and to what extent due to the Transformer architecture, we L--GATr's MLP block as standalone network. We call it "Geometric Algebra Perceptron" (GAP). Instead of structuring our data as a set of tokens, we format the particle data as a single list of channels. This implies that interactions between particles within the model will be carried out in the linear layers, as opposed to the attention mechanism we use in L-GATr.
Transformer Orthogonally to GAP, we also use a vanilla Transformer as a baseline in our experi- ments. We use pre-LayerNorm Transformer with multi-head attention and GELU nonlincarities. This setup mirrors that of L-GATr as closely as possible.
Multilayer perceptron (MLP) The MIP represents our simplest baseline, formulated as a stack of linear layers with GELU nonlinearities.
C Experiment details
C.1 Surrogates for QFT amplitudes
Dataset We generate training and evaluation data consisting of phase space inputs and their corresponding interaction amplitudes for processes q7 Z ng. {1, 2, 4). where initial quark-antiquark pair interact to produce a Z boson and a variable number of gluons. The amplitudes invariant under the permutation of identical particles and under Lorentz transformations. These datasets are generated by the MadGraph Monte Carlo event generator [4] in two steps." " First, we use standard run to generate the phase space distributions. This standard run applies importance sampling to produce unweighted samples, that is, events that are distributed according to the probability distribution that describes the physical interactions. Second, we re-compute the amplitude values corresponding to these phase-space samples with MadGraph's 's standalone module.
Availablee at https: h/tps:///gttubb http:/gggtub.cooooeee www.///aalooom....oooooo wwtp://eaal.oomm--eeeeeeea- Aww.eeAerrrrrrreeeeeeeeeee-ee wwtp://eemttttgggeeeeee -e--
transformer
under BSSD33CClause-CLL license.
19 Available at https: thub.
under MIT license.
Available https: / wwwchpad..nt//ngoaaee. under Uol-NCSA open source license.
17

---

We produce four datasets, each with different number of gluons. Each dataset consists of 4 x 10a samples for training, 102 for validation, and x 10 for testing. Our event sets feature kinematic cuts in the transverse momentum of the outgoing partieles (pr > 20 GeV) and on the angular distance between the gluons (AR- V + Ad² > 0.4).
For the learning problem, we affinely normalize the amplitudes y to zero mean and variance one:
J log (y4 - log(
O1gg(())
Models For the L-GATr model, we embed each particle token. is characterized with its four-momentum, embedded as a grade-1 multivector in the geometric algebra, as well as a one-hot embedding of the particle type, embedded unit variance, using the same normalization for each component to preserve Lorentz equivariance. In addition to the particle tokens use one "global" token, initialized to zero. After processing these inputs with an L-GATr network, we select the scalar component of the global token and identify it as the amplitude output. We use 8 attention blocks, 32 multivector and 32 scalar channels, and 8 attention heads, resulting in 1.8 x 10 learnable parameters.
scalar. We standardize the four-momentum inputs to
For the CGENN, we minimally alter the graph neural network version of the model built for the top tagging classification task so that it is able to perform amplitude regression. We keep the perparaaeess proposed by [68] and use 72 hidden node features, 8 hidden edge features, and 4 blocks. This model features around 3.2 x 10 trainable parameters.
For the GAP we use the same procedure as with L-GATr to embed (and preprocess) the inputs and extract the outputs, the only difference is that the different particles in given event are distributed as individual channels in the input. This model consists of 8 blocks, 96 multivector and 96 scalar channels, resulting 2.5 x 10 learnable parameters.
For the Transformer baseline, we again include particle tokens through a one-hot embedding. In this case the inputs a are preprocessed by performing standarization, defined
(8)
where the mean and the standard deviation are computed over each particle input separately. As for the network structure, we 8 attention blocks, 128 hidden channels and 8 attention heads, resulting in 1.3 x 10 learnable parameters.
For the DSI, we implement input standarization in the same way we do with the Transformer, but also apply the same transformation to each of the momentum invariant inputs separately. As for the layer structure, all MLP modules have 4 layers with 128 hidden channels each, and we set up the preprocessing units so that they output 64-dimensional latent space samples. All in all, we end up with 2.6 x 10 parameters in total.
For the MLP, we once again apply standarization, this time over the whole input. The network consists of 5 layers and 128 hidden channels amounting to 7 x 10a learnable parameters.
Training All models are trained by minimizing a mean squared error (MSE) loss on the prepro- cessed amplitude targets and by making use of the Adam optimizer. We use a batch size of 256 and a fixed learning rate of 10-4 for all baselines. As for the number of training steps, MLP and DSI are trained for around 2.5 x 10 iterations, the Transformer for around 10 iterations and GAP, CGENN and GATr for 2.5 x 10 iterations. We use no regularization method for any of our baselines. We use early stopping across all training runs.
C.2 Top tagging
Dataset We use the reference top quark tagging dataset by Kasieczka al. [49, 50]. The data samples are structured as point clouds, with each event simulating a measurement by the ATLAS experiment at detector level. Signal samples originate from the decay of a top quark, while the rest of the events are generated by standard background processes. The dataset consists of 1.2 x 106 events for training and 4 x 100 each for validation and testing.
Available https: //t::://///ooo//e http://eennnoooereere/r/////////
under CC-BY 4.0 license.
18

---

Models As the top-tagging datasets operate with reconstructed particles measured by detector, we include the proton beam direction as an extra particle imput the network, which partially breaks the Lorentz equivariance of the process. We encode it as a rank-2 multivector representing the plane orthogonal to the beam direction. We also add another token containing the time direction (1, 0, O, 0) as rank-1 I multivector. The time direction is required to break the special orthochronous Lorentz group SO+ (1, down to the subgroup SO(3).
Otherwise, we use the same setup as in the amplitude regression task. We use 12 attention blocks, 16 multivector and 32 scalar channels and 8 attention heads, resulting in 1.1 10 learnable parameters.
Training L-GATr is trained by minimizing a binary cross entropy (BCE) loss on the top quark labels. We train it for 2 x 102 steps using the EvoLved Sign Momentum (LION) [27] optimizer with a weight decay of 0.2 and a batch size of 128. We use a Cosine Annealing scheduler [56] with the maximum learning rate set at 10-1
C.3 Generative modelling
Dataset The tř + n jets, = 0...4 dataset is simulated with the MadGraph 3.5.1 event generation toolchain, consisting of MadEvent [4] for the underlying hard process, Pythia 8 [72] for the parton shower, Delphes 3 [35] for a fast detector simulation, and the anti-kr jet reconstruction algorithm [21] with R 0.4 as implemented in FASTJET [22]. The Pythia simulation does not include multi-parton interactions. We use the ATLAS detector card for the Delphes detector simulation, apply the phase space cuts pr > 22 Ve In| < 5, AR A + Anj < 0.5 and require 2 b-tagged jets. The events are reconstructed with a ---be reconstruction algorithm [2], and identical particles are ordered by Pr.
The sizes of the tf + n jets. n = 0...4 datasets reflect the frequency of the respective processes. resulting in 9.8 x 10 (n 0), 7.2 x 10 (n 1), 3.7 10 (n 2). 1.5 10 (n = 3) and 4.8 x 10 (n = 4) events. On each dataset, 1 1% of the events are set aside as validation and test split. We rescale the four-momenta p by the standard deviation of all five datasets 206.6 GeV for processing with neural networks.
Models The L-GATr score network operates in Minkowski space p (E. Py: Pz)) whereas flow matching happens in the physically motivated coordinates y = (ym yp. 17, o) defined in Eq. (4), After transforming y into p. embed each particle p into geometric algebra representations. We use scalar channels for the one-hot-encoded particle type and the flow time, for which we use Gaussian Fourier Projection with 8 channels [73]. We add the same symmetry breaking inputs that we used for the top tagging task, this time as extra multivector channels. The output of the L-GATr network is a vector field in Minkowski space (UÆ, p/ Pp,, vp..). To obtain vector fields in the manifold coordinates y. multiply with the Jacobians of the transformation p y to obtain (U ymm uy Un. vg). Because of the logarithm in the change of variables, the Jacobians for ym Up can take on large values, which makes training unstable. To avoid this complication, we extract Tym Spp directly from the scalar output channels of L-GATr and use these values as vector field Uym. Py..- Like the -breakgg inputs, this procedure breaks the Lorentz symmetry down the residual symmetry group of the measurement process. We find it beneficial to perform the transformation y ++ at 64-bit floating- point precision, which has noticable effect the computational cost. We an L-GATr network with 16 multivector channels, 32 scalar channels, 6 L-GATr blocks, and 8 attention heads, totalling 5.4 10s learnable parameters.
We implement the E(3)-GATr score network following the same prescription outlined above for L- GATr. The only difference is that only transform between y and (ym. Pa Py, Pa) to avoid the large jacobians from the yom ++ m transformation. Our task only requires SO0()--quivariiace, therefore we do not use of the translation-equive the internal translation translation-equivariann representations. For each particle, we embed the three-momentum (Pa, Py, p=) into the vector component of the multivector, and ym scalar.
input and output representations of B(3)-GATr, but we keep
The Transformer score network directly operates in the physically motivated y space. Each particle embedded into one token, with 4 channels for the components of y, 8 channels for the flow time in Gaussian Fourier Projection, and the one-hot-encoded particle type. The network has 108 channels, Transformer blocks, and 8 attention heads, totalling to 5.7 x 10th learnable parameters.
19

---

The MLP score network also operates in y space. Each event is embedded as a list of the components of y, together with the time embedded with Gaussian Fourier Projection using channels. The network has 336 channels and 6 blocks, with 5.9 x 10 learnable parameters.
We implement the JetGPT model closely following Ref. [18]. JetGPT is autoregressive mixture model, i.c. its parameters are predicted autoregressively for each component of the high-dimensional distribution. We use the same transformer as for the Transformer-CFM to predict the mixture parameters, but using a triangular attention matrix to achieve the autoregressive structure. The components are the same y coordinates that we use for the flow matching models. amounting to four tokens for each particle. We use gaussian mixture models for the non non-periodic coordinates (ym: Up. 1i). and von Mises mixture models for the periodic coordinate ø. Due to the autoregressive structure. the ordering of components affects the performance. As more conditions are added, the components will be harder to learn using a fixed amount of training data. We order the particles as (41. q1, 92, 22, 93, J1..... j4), We train the model to minimize the joint log-likelihood and sample the components sequentially. Each component is embedded into one token, with I channel for the value, and 24 to 40 channels for the mocoode component type. We use the same transformer as for the Transformer-CFM, with 108 channels, 6 Transformer blocks, and 8 attention heads, but the last layer maps onto the 108 parameters of a mixture model of 36 gaussian or von Mises distributions.
and within each particle we order the components as (ø, Up, 7) V.m)
Training All networks are trained for 2 x 10 iterations with batchsize 2048 using the Adam optimizer with default settings and an initial learning rate of 0.001. We evaluate the validation loss every 102 iterations and decrease the learning rate by a factor of 10 after no improvements for 20 validation steps. We stop training after the validation loss has not improved after 50 validation steps, and pick the network with the best validation loss.
Base distribution The base distribution is defined in the rescaled Minkowski space discussed above. We use standardized Gaussians for the spatial momentum Pæ,y. ~ NN 1) and the log-transformed squared mass ym log m² NN, 1). We ensure the constraints pr > 22 GeV, AR 0.5 through rejection sampling. We have experimented with other base distributions and find similar performance.
Probability paths The target probability paths for RCFM linearly change the physically motivated coordinates y (um, Vp: 77. Φ) defined in Eq. (4). We use pseudorapidity 7 instead of true rapidity, since this is easier to implement given the cut on transverse momentum pr. We use a constant diagonal metric, with the squared inverse standard deviation of these coordinates in the training dataset on the diagonal. This is equivalent to a standardization step. We construct periodic target vector fields for angular coordinates @ by adding factors of 2z until angular coordinates and angular velocities end up in the interval n. T].
Negative log-likelihood (NLL) metric For any sample p in Minkowski space, we evaluate the log- density of the transformed sample y using the instantaneous change of variables [26]. In other words, we solve the ODE
( It f((()))
(
ve(( (-div(v))(2))
dt
with the initial conditions I1 y, (z1) O. We use the Hutchinson trace estimator to evaluate the divergence of the vector field. Using the base density Po. we then evaluate the density P Pmodel of CFM-generated samples in Minkowski space
log P (p) (p) = - log Po( + fo(zo) + log det Opo Oyo
0y1 0p1
+ log det
(10)
The last two terms are the logarithms of the jacobian determinants for the transformations between Minkowski space and the physically motivated space.
Classifier two-sample test We train a MLP classifier to distinguish generated events from the ground truth. The classifier inputs are full events in the representation, together with challenging correlations, in particular all pairwise A R values, and the y representations of the reconstructed particles I, T, W+ W- The classifier network has 256 channels and layers. is trained for 500 epochs with batchsize 1024, dropout rate of 0.1, and the Adam optimizer with default settings.
20

---

1.0
1.0
JeGGPT
JerGPT
MLP CFM
MIP CFM
0.9
Transformer CFM
0.9
Transformer CFM
B(3)-GATr CFM
E(3)-GATY CFM
Neural classifier ADC
CFM
Neural classitier AUC
L-GATr CRM
0.8
0.8
0.7
0.7
0.6
11 + 1E+1J
1E+ 20
ti+3
ti 4f
10a
10a
Number of training samples
Figure 6: Generative modelling: classifier two-sample tests. We show how well classifier can discriminate model samples from test samples, measured through the under the ROC curve (lower is better, 0.5 is ideal). Left: For different processes. Right: As a function of training dataset size. show mean and standard deviation of three random seeds. The L-GATr flow outperforms the baselines in all processes and training set sizes.
We start with an initial learning rate of 0.0003 and decrease the learning rate by a factor of 10 after no improvements in the validation loss for epochs. We stop training after 10 epochs without improvements in the validation loss and load the best-validation model afterwards. We use the full ground truth dataset as well as 1M generated events and split into 80% for training and 10% each for testing and validation. We use the AUC of this classifier evaluated on the test dataset as scalar metric, with the value 0.5 for perfect generator. Neural classifiers approximate the event-wise wise likelihood (r)////oeeee ratio Pdata (x)) (modeel (r) of single events, which is the most powerful test statistic according to the Neyman Pearsonn lemma, and opens many ways further study the performance of the generator beyond scalar metrics [34].
Our results are shown in Fig. 6. Across training data sizes and processes, the L-GATr flows are more difficult to distinguish from the test samples than the baselines.
C.4 Computational cost and scalability
In Fig. we compare L-GATr to pessage-passing graph neural network (we use CGENN [68]) and a vamilla Transformer in terms of their test-time computational costs. For this comparison, we use small versions of all architectures, consisting of single model block and around x 10 learnable parameters. In the case of the Transformer and L-GATr, we fix their layer structure so that inputs going into the attention layer consist of 72 channels. Our measurements are performed with datasets made up by a single sample and all models are run on an H100 GPU.
L-GATr in its current implementation is not yet as efficient as Transformer for small systems: for up to hundreds of particles, L-GATr takes an order of magnitude longer to evaluate. This is caused by the linear layers in L-GATr, which are more costly to execute than their non-equivariant counterparts and represent constant computational overhead. However, because L-GATr based on the same efficient backend for dot-product attention, it scales just like Transformer larger systems, and we find the same computational cost for 5000 particles or more.
Compared to an equivariant graph network, L-GATr is clearly more efficient in terms of compute time and memory. Already at small systems, L-GATr can be evaluated an order of magnitude faster. The difference is even more pronounced in terms of memory: the graph network ran out of memory for more than a thousand particles. This largely because graph network implementations are often optimized for sparse computational graphs, but here we use fully connected graphs: LHC problems often benefit from fully connected computational graph, because pairwise interactions not usually decay with the Minkowski norm of the distance between momenta. Transformer-ban approaches
21

---

10a
Transformer
L-GATr
10a
Inference time (ms)
101
10"
10P
10
102
10a
10
10a
Number of items
Figure 7: Inference cost (wall-time per forward pass) function of the number of particles. We compare L-GATr, Transformer, and message-passing graph neural network (we use CGENN [68] but expect similar results other architectures). The latter runs out memory when evaluating more than thousand particles. While we do our best to find comparable settings, such comparisons depend on lot of choices and should be interpreted with care. Nevertheless, we believe they illustrate that L-GATr scales to large systems like Transformer, thanks it being based dot-product attention.
like L-GATr can thus have substantial computational advantages in particle physics problems that involve large number of particles.
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We justify our main claims in Sections 3 to 4. Guidelines:
The answer NA means that the abstract and introduction do not include the claims
made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer:
[Yes]
Justification:
We discuss L-GATr's shortcomings throughout the paper and in Sec. 5.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the
implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms
and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers penalize honesty concerning limitations.
will be specifically instructed
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Proposition 1 is shown in Appendix A.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if
23

---

they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented
by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: See Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our work is open open- open-sourced at wtttppp//ggttu...oo//eeeeeeeeeep Lorentz-gatr,
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.ct
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to
24

---

reproduce the results. See the NeurIPS code and data submission guidelines (httpss
CooeeSumeesseeeeeeee ssi ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperpa- rameters, how they were chosen, type of optimizer, etc.) necessary to understand the results? Answer: [Yes]
Justification:
See Appendix C and
1 lorentz-gatze gatr.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide an estimate of epistemic uncertainty through the variance between multiple independent results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula,
call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 8. Experiments Compute Resources Question:
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [No]
Justification: Our experiments are of a sufficiently small scale that they should be easily
reproducible in typical academic environment.
Guidelines:
The answer NA means that the paper does not include experiments.
25

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't Itdntt make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //terr///uuisspp cc htte//uulicceeeeeeee http://///etccccccdeeeeeeeeeeeee Answer: [Yes]
Justification: We fully conform to the Code of Ethics. Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there special consid-
eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Since this work specific to high-energy physics, we expect societal impact through the effect that our architecture has on research in that field. We do not anticipate any particular risk of negative societal impact. Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (c.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deplommeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and funetioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA]
Justification: We are not aware of any such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing
26

---

safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors
should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the existing assets we use in the paper. Licenses are mentioned in Appendix C.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes)
Justification: See https / //gi /gtgttttttu..t.////////.g..g.tet www.heideeeeerreeeeeeecceceeeeee httpa///ooalallbrrgg.a/////rette hhapgg/e/lll/leetttgaatr Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human
27

---

Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "hwLbl7RP2A", "reviewer": "Reviewer_PF88", "summary": "The authors propose an architecture for high-energy physics events – the Lorentz Geometric Algebra Transformer, which is equivariant under Lorentz transformations. The architecture is based on the Geometric Algebra Transformer architecture, and generalizes to relativistic scenarios and the Lorentz symmetry. The architecture is demonstrated on regression, classification and generation tasks in particle physics.", "questions": "* The experiments are conducted on four vectors of the hard particles. In this case, is it really necessary to employ transformers since the input dimension is not as high? Especially for the regression task, the training set is a relatively small one, yet the model has around 2 million parameters. \n* Why is a smaller model used for the top tagging task? Obviously, this task has higher dimensionality and larger datasets compared to the regression task.", "limitations": "The authors briefly discussed the limitations in the Discussion section.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This article is well-written and effectively communicates its ideas. While the novelty of the research may not be groundbreaking, it offers valuable contributions to the field. The authors conduct a sufficient number of experiments to test their proposed architecture.", "weaknesses": "* The proposed architecture, while theoretically appealing, suffers from significant computational overhead due to the addition of Lorentz layers to the already resource-intensive Transformer model.\n* The comparison of generative modeling capabilities is limited to flow models, which may not represent the full spectrum of possible approaches.\n* The model's performance, while adequate, does not demonstrate a significant improvement over baseline models, which may be a concern given the added computational cost.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough and constructive review. We are happy to hear that you found our architecture a valuable contribution and the paper well-written. Thanks as well for the questions and criticisms, which we address one by one.\n\n> significant computational overhead\n\nIt's true that L-GATr has some computational overhead over a standard transformer architecture, but we find it sufficiently fast for all particle-physics applications we encountered.\n\nIn Appendix C.4 of our paper we show timing measurements. A forward pass through the tested model takes between 2ms and 100ms when up to 40k particles are processed. This is fast enough for almost all steps of the data-analysis pipeline and in fact negligible compared to the simulators. (There are a few exceptions that require faster inference, most notably the trigger, but these specialized cases go beyond the scope of the paper.)\n\nCompared to a standard transformer, L-GATr is slower for small numbers of particles. However, it is comparably fast when the number of particles is large, since then the attention operation is the bottleneck for both. Compared to a message-passing network, a popular type of equivariant architecture, L-GATr is consistently faster and scales much better to large numbers of particles. This is consistent with the prior work GATr [1, see Fig. 5 there] from which L-GATr is derived, which shares performance characteristics with L-GATr.\n\n> The comparison of generative modeling capabilities is limited to flow models\n\nTrue. We initially focused on flow matching baselines because of the stable training, high-quality samples, and the (approximately) tractable likelihood function. Because of these advantages, flow matching is establishing itself as a standard in high-energy physics. For instance, Buhmann et al. [2] compare the performance of GAN, diffusion, and flow-matching models based on the same architectural elements and find flow matching to work best.\n\nThank you for the suggestion to expand this. We are working on generating additional baseline results for the generative modelling task, including an SO(3)-equivariant architecture and a non-equivariant autoregressive transformer model [3]. We will include the results and discuss them in the final version of the paper.\n\n> The model's performance, while adequate, does not demonstrate a significant improvement over baseline models, which may be a concern given the added computational cost.\n\nIn particular in the amplitude regression task, we find L-GATr to consistently outperform the baselines. The advantage over the baselines gets larger for higher-multiplicity final states and more complex interactions. This regime of amplitude modelling is known to be the most challenging for neural surrogates, but is at the same time the most practically important, as theory computations can become prohibitively expensive there [4, 5]. We believe that L-GATr's performance improvements here can have substantial real-world impact.\n\n> Is it really necessary to employ transformers since the input dimension is not as high?\n\nIt is not strictly necessary, but transformers lead to the best performance in our experiments. In Fig. 3, for instance, we show that a transformer outperforms a simple MLP on the amplitude regression task. Similarly, our L-GATr outperforms the GAP model, a Lorentz-equivariant MLP we constructed. This is in line with the larger trend in machine learning, which finds transformers to often outperform other models on a large variety of tasks.\n\n> Why is a smaller model used for the top tagging task?\n\nWe tuned hyperparameters independently for each of the three problems. We interpret the differences in optimal model size as follows: the top-tagging problem takes many particles as inputs, but the task itself is a comparably simple binary classification problem. In contrast, amplitude regression requires learning complex functions of the particle momenta with high precision, especially for the higher multiplicities; these can be more accurately expressed with a bigger network.\n\nTo quantify how important the network capacity is for the amplitude regression problem, we compared our standard L-GATr models with 2 million parameters to a smaller version with 0.7 million parameters:\n\nModel size | MSE ($Z + 1g$) [$10^{-7}$] | MSE ($Z + 4g$) [$10^{-5}$]\n--- | --- | ---\n2M | 1.91 | 1.48\n0.7M | 1.62 | 2.71\n\nThis confirms that the capacity is more important for the more complex $Z + 4g$ process (though the smaller L-GATr model still outperforms all baselines there). \n\nWe hope we were able to address your questions and look forward to discussing further.\n\n**References**\n\n- [1] J. Brehmer et al, \"Geometric Algebra Transformer\", NeurIPS 2023\n- [2] E. Buhmann et al, \"EPiC-ly Fast Particle Cloud Generation with Flow-Matching and Diffusion\", arXiv2310.00049\n- [3] A. Butter et al, \"Jet Diffusion versus JetGPT - Modern Networks for the LHC\", arXiv:2305.10475\n- [4] S. Badger and J. Bullock, \"Using neural networks for efficient evaluation of high multiplicity scattering amplitudes\", Journal of High Energy Physics 2020\n- [5] S. Badger et al, \"Loop amplitudes from precision networks\", SciPost Physics 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9LBwVu1bqd", "reviewer": "Reviewer_srsp", "summary": "The paper proposes the Lorentz Geometric Algebra Transformer (L-GATr) for high-energy physics tasks. This model extends the Geometric Algebra Transformer by incorporating relativistic considerations. Specifically, L-GATr supports partial and approximate symmetry for symmetry-breaking inputs and is applied to generative models.", "questions": "I think another interesting comparison would be to evaluate the proposed model against traditional sampling methods, especially if the goal is to design machine learning models that are more efficient than traditional methods (within an acceptable error tolerance). Could the paper include this comparison? I believe a positive result could attract attention from researchers in the high-energy physics community.", "limitations": "The efficacy and scalability of the proposed framework may not yet suffice to replace traditional methods.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- The motivation of the paper is strong, addressing an important application of equivariance models. While symmetries are prevalent in high-energy physics, there are few applications of equivariance models in this field.\n- The paper clearly distinguishes their work from existing research, emphasizing the significance of their contributions.\n- They connect the proposed model to the generative framework, making it potentially applicable to a broader range of tasks.", "weaknesses": "- The model's performance in Table 1 is not optimal.\n- The experiments on generative modeling lack comparisons with other equivariant models or SOTA models.\n- The scalability of the model remains limited, which reduces its suitability for high-energy physics applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough and constructive review. We are glad that you liked the motivation of our work, found our contributions significant, and appreciated the generative modelling part. Thanks as well for the questions and criticisms, which we address in the following.\n\n> The model's performance in Table 1 is not optimal\n\nIt's true we do not outperform the strongest baselines on this task (unlike on the other two tasks). We consider our performance and that of the strongest models in the literature to be on par for all practical purposes. In terms of accuracy, the strongest model and L-GATr differ only by 0.09 percentage points, and on one of the four established metrics L-GATr achieves the best score to date.\n\n> The experiments on generative modeling lack comparisons with other equivariant models or SOTA models.\n\nWe are not aware of any other Lorentz-equivariant generative models. If you know of any, we would appreciate if you could point us to them; we would be eager to add them to the comparison.\n\nThe transformer-based flow-matching model that we compare to is the strongest baseline we could come up with. Flow matching is establishing itself as a standard in high-energy physics. For instance, Buhmann et al. [1] compare the performance of GAN, diffusion, and flow-matching models based on the same architectural elements and find flow matching to work best.\n\nThat being said, we agree that we should compare to more diverse baselines, especially those with open-source implementations. We started experiments with SO(3)-equivariant architectures and non-equivariant autoregressive transformers [2] and will include them in the final version. Thank you for the suggestion!\n\n> The scalability of the model remains limited. [...] The efficacy and scalability of the proposed framework may not yet suffice to replace traditional methods.\n\nCould you elaborate what you mean here? In terms of scaling the model size, we used L-GATr models with millions of parameters without issues.\n\nIn terms of scaling to large training data sets, we see in Fig. 3 that L-GATr benefits from that more than all baselines.\n\nThe scaling to large number of particles (tokens) is studied in Appendix C.4. A forward pass through the tested L-GATr model takes between 2ms and 100ms when up to 40k particles are processed. That is fast enough for almost all steps of the particle-physics data-analysis pipeline and in fact negligible compared to the simulators. (There are a few exceptions that require faster inference, most notably the trigger, but these specialized cases go beyond the scope of the paper.)\n\n> evaluate the proposed model against traditional sampling methods\n\nIf we understand this suggestion correctly, we are already doing this: we evaluate our generative model by comparing samples from it to samples from the simulator, which are based on traditional Monte-Carlo sampling. We quantify the difference between these two distributions by comparing marginal distributions, evaluating the log likelihood, and through a classifier two-sample test. (Unfortunately, the density of the traditional samples is not tractable, which makes it difficult to define more objective metrics for the high-dimensional samples.)\n\nTo the extent that these metrics can measure it, we find that the L-GATr model gives us samples that are close to the traditional sampling method, while being orders of magnitude faster to generate: sampling from an L-GATr flow-matching model takes milliseconds per sample, from a fast classical sampler seconds per sample, and from the most accurate state-of-the-art samplers minutes per sample.\n\nWe hope that we were able to answer your questions and look forward to discussing further.\n\n**References**\n\n- [1] E. Buhmann et al, \"EPiC-ly Fast Particle Cloud Generation with Flow-Matching and Diffusion\", arXiv:2310.00049\n- [2] A. Butter et al, \"Jet Diffusion versus JetGPT - Modern Networks for the LHC\", arXiv:2305.10475", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LHZuDUgr8k", "reviewer": "Reviewer_w7CH", "summary": "The paper proposes a Lorentz equivariant transformer (L-GATr) based on geometric algebra for high energy physics. It generalizes the Geometric Algebra Transformer (GATr) from $E(3)$ equivariance to the Lorentz group. The proposed transformer is then developed into a generative model based on Riemannian flow matching for particle data. L-GATr is evaluated on several high energy physics tasks, including quantum field theory amplitude surrogates, top tagging, and generative modeling for event reconstruction.", "questions": "1. What are $y_m, y_p, \\eta$, and $\\phi$ in (4)?", "limitations": "As mentioned in the paper, the proposed L-GATr has additional computational overhead compared to traditional transformers. Secondly, even though the framework allows additional inputs to address symmetry breaking issues, the effect of such an approach is not well studied. It is unclear how well the proposed method can handle symmetry breaking inputs.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow. In addition, the problem is well-motivated.\n2. The proposed L-GATr generalizes GATr from $E(3)$ to the Lorentz group.\n3. As stated by the authors, the Lorentz-equivariant flow matching proposed in section 3.2 is the first generative model proposed for particle physics. \n4. Compared to graph-based Lorentz equivariant networks, the transformer architecture is more efficient and scalable. \n5. The proposed method is shown to be more data efficient than the baselines in both amplitude surrogates and generative modeling experiments. \n6. The ability to scale with the data is also verified in several experiments. \n7. The benefits of Riemannian flow matching compared to the Euclidean version are demonstrated in the experiments.", "weaknesses": "1. It is a bit unclear to me what has been modified for Lorentz equivariance in the transformer framework. Specifically, (1) and (2) look the same as (4) and (5) in [13]. It’ll be great if the authors can state the changes explicitly. \n2. The current presentation of the experiments can be a bit hard to understand for people without a physics background. Adding some basic introduction to the problems can strengthen the paper. \n3. In the top tagging experiment, the performance of the proposed method is marginally worse than the baseline method. \n4. Although the proposed method is claimed to support symmetry-breaking data, its effect is not well studied in the experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough and constructive review. We are glad to hear that you appreciated how we generalized the GATr architecture from E(3) to the Lorentz symmetry and the development of the first Lorentz-equivariant architecture. We were particularly happy that you found the paper easy to follow. Thanks as well for the questions and criticisms, which we address one by one.\n\n> It is a bit unclear to me what has been modified for Lorentz equivariance in the transformer framework. Specifically, (1) and (2) look the same as (4) and (5) in [1].\n\nIndeed, the form of these equations and generally the overall architecture design are similar. There are several differences though:\n\n- The multivectors $x$ in our Eq. (1) and Ref. [1]'s Eq. (4) are elements of different geometric algebras with different metrics.\n- The inner product in our Eq. (2) is therefore also different from the one in Ref. [1].\n- The second term in our Eq. (1) and Ref. [1]'s Eq. (4) is different. In both cases, we have a term for the five grade projections, and we can multiply with a multivector that is invariant to the group. For Lorentz equivariance, this invariant multivector is the pseudoscalar $e_{0123}$. In the $E(3)$ case, this invariant multivector is the vector $e_0$.\n- The geometric product used in the MLP is different between both architectures. This is another consequence of using a different metric.\n- L-GATr uses only the geometric product as a bilinear operation in the MLP, while the original GATr had to concatenate that with an equivariant join operation to ensure expressivity. The technical reasons for this are discussed in the appendices of Ref. [1].\n- The layer normalization in Eq. (3) needs to be different from Ref. [1]'s layer normalization. The reason is that the Minkowski norm of special relativity can be negative and the normalization operation has to be robust against that.\n\n> The current presentation of the experiments can be a bit hard to understand for people without a physics background.\n\nThank you for this feedback. It is important to us that both machine learners and particle physicists can understand the paper. We will use the extra page in the final version to substantially expand the introduction to particle physics and to the concrete experiments.\n\n> In the top tagging experiment, the performance of the proposed method is marginally worse than the baseline method.\n\nIt's true we do not outperform the strongest baselines on this task (unlike on the other two tasks). We consider our performance and that of the strongest models in the literature to be on par for all practical purposes. In terms of accuracy, the strongest model and L-GATr differ only by 0.09 percentage points, and on one of the four established metrics L-GATr achieves the best score to date.\n\n> Although the proposed method is claimed to support symmetry-breaking data, its effect is not well studied in the experiments.\n\nActually, in two of our three experiments we use symmetry-breaking data: the direction of the particle beam is provided as an input in the top-tagging and generative modelling problems. This is important, because the detector measurements depend on this direction. (In contrast, the amplitude regression task does not include detector effects, so additional symmetry-breaking inputs are neither needed nor provided there.)\n\nTo test the relevance of these symmetry-breaking inputs, we re-ran the top-tagging experiment for a version of L-GATr *without* the beam direction as an input. We found that this decreased the model's performance as follows:\n\nModel | Accuracy | AUC | $1/\\epsilon_B$ ($\\epsilon_S = 0.5$) | $1/\\epsilon_B$ ($\\epsilon_S = 0.3$)\n--- | --- | --- | --- | ---\nL-GATr (original) | 0.9417 | 0.9868 | 548 | 2148\nL-GATr (no beam direction) | 0.9403 | 0.9860 | 440 | 1840\n\nThe beam direction is thus important to analyze data that involve detector effects.\n\n> What are $y_m$, $y_p$, $\\eta$, and $\\phi$ in (4)?\n\nThese four variables form an alternative basis for the particle four-momenta that is better aligned with the physically relevant properties of a particle in the context of a collider experiment: $\\eta$ and $\\phi$ represent the angle in which a particle is moving, $y_p$ is a measure of the speed with which it moves away from the collision, and $y_m$ is related to its mass. A convenient property for flow matching is that the support of the distribution is convex in this space. The dictionary that maps $(y_m, y_p, \\eta, \\phi)$ to four-momentum $(E, p_x, p_y, p_z)$ is given in Eq. (4) and just after it. We will add the inverse map from four-momenta to this alternative basis and improve the discussion. \n\n> the proposed L-GATr has additional computational overhead compared to traditional transformers.\n\nIt's true that L-GATr has some computational overhead over a standard transformer architecture, but we would like to stress that it is managable for almost all particle-physics applications.\n\nIn Appendix C.4 of our paper we show timing measurements. A forward pass through the tested model takes between 2ms and 100ms when up to 40k particles are processed. This is fast enough for almost all steps of the data-analysis pipeline and in fact negligible compared to the simulators. (There are a few exceptions that require faster inference, most notably the trigger, but these specialized cases go beyond the scope of the paper.)\n\nCompared to a standard transformer, L-GATr is slower for small numbers of particles. However, it is comparably fast when the number of particles is large, since then the attention operation is the bottleneck for both. Compared to a message-passing network, a popular type of equivariant architecture, L-GATr is consistently faster and scales much better to large numbers of particles.\n\nWe hope that we were able to answer your questions and look forward to discussing further.\n\n**References**:\n\n- [1] J. Brehmer et al, \"Geometric Algebra Transformer\", NeurIPS 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DmEi30yI1L", "reviewer": "Reviewer_PeE1", "summary": "This paper implements a Lorentz-equivariant transformer and applies it to several problems in particle physics. The main contribution of the paper is not the definition of a conceptually different model. At least it doesn't claim to define a model significantly different from previously proposed models. But it provides a software library (under BSD-3-Clause-Clear license) and a nice set of numerical experiments inspired by physics problems. This is a nice contribution to the community.", "questions": "Please address the questions in the weaknesses section.\n\nHaving these clearly answered would improve the paper's clarity, reproducibility, and usability. \n\nI'd be willing to raise my score to 7-8 if all these aspects were clearly explained with concrete explanations and equations, since I believe this can turn into a great tool to the AI 4 physics community.", "limitations": "Not applicable.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper combines interesting ideas (imposing symmetries arising from physics using invariant theory and clifford algebras) with state-of-the-art machine learning models (transformers, Riemannian flow maching generative modeling). It comes with a software package and multiple experiments where several methods are compared. I see this as a tool that may be used by several members of the AI4science and AI4physics communities.", "weaknesses": "The paper is not self-contained. In order to be a valuable tool for our community it would benefit from a more comprehensive explanation of several of its technical aspects. \n\n- The implemented Lorentz equivariant model (what are the inputs?, are they a set of particles? how are they matched to queries, values and keys?)\n- The flow-matching approach. Again, what are the inputs, outputs? How's the model defined? What does it produce? The closest thing to an explanation is in page 24, in the context of describing the experiment, but still I don't find this explanation sufficient to understand the model conceptually.\n- What are the scientific goals of the experiments defined in the paper? (In the context of fundamental research in high energy physics).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough and constructive review. We are particularly happy that you appreciate the open-source release of our library. Thanks as well for the questions and criticisms, which we address in the following.\n\n> The paper is not self-contained.\n\nThank you for this feedback. Due to the fact that we had three diverse experiments, we had to put many experimental details in the appendix. We will use the extra page in the final version to expand in the main paper our description of architecture, flow-matching approach, and the experiments.\n\n> The implemented Lorentz equivariant model (what are the inputs?, are they a set of particles? how are they matched to queries, values and keys?)\n\nThe inputs to our architecture are indeed a set of particles. Each particle is characterized by a) an energy and a momentum, which are concatenated into a \"four-momentum\" vector, and b) a discrete type; see the beginning of Sec. 2. As we describe in the beginning of Sec. 3.1, the four-vector is embedded into the vector grade of the multivector, and the other grades are zero-padded. The type is represented with a standard one-hot encoding.\n\nThis data is then processed with the L-GATr architecture defined in line 177. It consists of a sequence of layers of different types: layer normalization, the construction of keys, queries, and values with linear layers, an attention mechanism, geometric products, gated nonlinearities, and skip connections. We define all these maps on pages 4 and 5.\n\n> The flow-matching approach. Again, what are the inputs, outputs? How's the model defined? What does it produce?\n\nThe key component here is a network $v_t(x)$ that takes as inputs a set of particles $x$ as well as a time variable $t$ and outputs one number for each element of the input particle properties. As an example, consider our $t\\bar{t}+0j$ experiment, where we generate 6 particles, each characterized by the 4 properties $y_p, y_m, \\phi, \\eta$ defined in Eq. (4). The L-GATr vector field model embeds them into 6 tokens, each containing a multivector with the 4 properties as well as particle type and symmetry breaking information. The architecture processes this data and outputs a single multivector for each item, and we extract the vector field required for flow matching from the vector component of this multivector. For each of the 6 particles, this vector field contains 4 values corresponding to the 4 properties. For details, see the Models section in Appendix C.3.\n\nTo sample with such a model, one first samples initial particle properties from a base distribution (like a multivariate Gaussian in a suitable basis) and calls this $x(1)$. Then one updates the particle properties $x(t)$ by integrating the differential equation $d/dt x(t) = v_t(x)$ from $t = 1$ to $t = 0$ using an off-the-shelf ODE solver. The final values $x(0)$ are the outputs: a set of particle properties sampled from the model. By repeating this process, we can generate different samples of sets of particles.\n\nTo train the model, we use the conditional flow matching loss defined in line 122. Intuitively, it trains the network $v_t(x)$ such that the vector field it defines moves particles on shortest paths from the base density to the data density. This is explained in much more detail in the flow-matching paper [1]. \n\n> What are the scientific goals of the experiments defined in the paper?\n\nOn a high level, all three experiments are part of the data-analysis pipeline at the LHC. We roughly sketch how they fit into this bigger picture in Fig. 1. The goal of any such analysis is either the measurement of parameters of fundamental theories of nature, such as the mass of a particle or the strength of an interaction between particles, or to discover or exclude hypotheses altogether, like in the discovery of the Higgs boson in 2012. \n\nAs a concrete example, consider the top-tagging problem of Sec. 4.2. This classification task is used to filter data collected at the Large Hadron Collider for those collisions that are likely to be the result of a top quark decay. Top quarks are the heaviest fundamental particle currently known. We are interested in selecting data that involves their decays because the production and decay probabilities of top quarks depend on whether certain \"supersymmetric\" particles exist. Analyzing top-decay events will therefore allow us to discover or exclude the existence of supersymmetric particles, but to do that well, we need to be able to filter them precisely, and here we show that L-GATr is a powerful tool for that task.\n\nNow, why do we care about the existence of supersymmetric particles? If they exist, they would help us answer the \"hierarchy problem\", one of the biggest open questions in fundamental physics. It can be phrased as \"Why is the gravitational force so much weaker than the electromagnetic, weak, and strong forces?\". The existence of supersymmetric particles would also be relevant for a number of other scientific questions, including the nature of Dark Matter. Apart from supersymmetry, there are more hypotheses for new physics that benefit from confident top-taggers. All in all, if L-GATr can improve the performance on the top tagging task, it may contribute to a better, clearer, or sooner answer to such fundamental physics questions.\n\nWe hope that we were able to answer your questions and look forward to discussing further.\n\n**References:**\n\n- [1] Y. Lipman et al, \"Flow matching for generative modelling\", ICLR 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all reviewers for their detailed feedback and questions.

We are excited to read that the reviewers found the work a "valuable contribution to the field" (reviewer **Pf88**), that they appreciated the "significance of [the] contributions" (reviewer **srsp**), and that they found our L-GATr "a tool that may be used by several members of the AI4science and AI4physics communities" (reviewer **PeE1**).

On a technical level, they appreciated that our transformer "generalizes GATr from E(3) to the Lorentz group" (reviewer **w7CH**), that "the Lorentz-equivariant flow matching is the first [such model]" (reviewer **w7CH**), and that we "conduct a sufficient number of experiments to test" it (reviewer **PF88**). 
We are particularly encouraged by the reviewers finding that the paper is "well-written and effectively communicates its ideas" (reviewer **srsp**).

&nbsp; 

Their critique and questions are helping us improve the paper further. Here we want to highlight three points.

First, reviewers asked about **additional baselines** for the generative modelling experiment.

This is a great suggestion. We started experiments with SO(3)-equivariant architectures and non-equivariant autoregressive transformers [1]. We were not able to finish these in time for this rebuttal, but we will add the results and discuss them for the final version of this paper.

&nbsp; 

Second, reviewers asked about the relevance of **symmetry-breaking inputs**.

In two of our three experiments, we break the full Lorentz symmetry by providing the direction of the collider beam pipe as an input to L-GATr. This information is important, both to capture explicit symmetry breaking from the initial beam directions and to model soft breaking of the symmetry from detector effects. In the original paper, we did however not demonstrate this relevance empirically.

We performed a new experiment to measure the importance of symmetry-breaking inputs. We re-ran the top-tagging experiment for a version of L-GATr *without* the beam direction as an input. We found that this decreased the model's performance as follows:

Model | Accuracy | AUC | $1/\epsilon_B$ ($\epsilon_S = 0.5$) | $1/\epsilon_B$ ($\epsilon_S = 0.3$)
--- | --- | --- | --- | ---
L-GATr (original) | 0.9417 | 0.9868 | 548 | 2148
L-GATr (no beam direction) | 0.9403 | 0.9860 | 440 | 1840

L-GATr's ability to break Lorentz symmetry through the inputs is thus relevant in practice.

&nbsp; 

Finally, the reviewers were concerned about the **computational overhead and scalability** of L-GATr.

It's true that L-GATr has some computational overhead over a standard transformer architecture, but we find it sufficiently fast for all particle-physics applications we encountered.

In Appendix C.4 of our paper we show timing measurements. A forward pass through the tested model takes between 2ms and 100ms when up to 40k particles are processed. This is fast enough for almost all steps of the data-analysis pipeline of the Large Hadron Collider and in fact negligible compared to the simulators. (There are a few exceptions that require faster inference, most notably the trigger, but these specialized cases go beyond the scope of the paper.)

Compared to a standard transformer, L-GATr is slower for small numbers of particles. However, it is comparably fast when the number of particles is large, since then the attention operation is the bottleneck for both. Compared to a message-passing network, a popular type of equivariant architecture, L-GATr is consistently faster and scales much better to large numbers of particles. This is consistent with the prior work GATr [2, see Fig. 5 there] from which L-GATr is derived, which shares performance characteristics with L-GATr. Using a message-passing network with sparse connections, which can make it possibly faster than a transformer, is not possible in these particle-physics experiments, because the interaction strength between particles does not decay with the Minkowski distance between their momenta.

In addition to this scaling with the number of particles (or tokens), we also demonstrate in the paper that L-GATr scales well to millions of parameters and to large training datasets (see Fig. 3). Overall, we consider the scalability one of the strongest properties of the model.

&nbsp; 

We would like to thank the reviewers again. We hope that we were able to address their questions adequately and look forward to the discussion period.

&nbsp; 

**References**:

- [1] A. Butter et al, "Jet Diffusion versus JetGPT - Modern Networks for the LHC", arXiv:2305.10475
- [2] J. Brehmer et al, "Geometric Algebra Transformer", NeurIPS 2023

Meta-review: This paper designs the first Geometric Algebra Transformer (GATr) equivariant to the action of Lorentz group, specifically suited for data of observed particles represented by four-momenta. Reviewers have found this paper well written and well motivated, presents a novel  contribution in its construction of the Lorentz-equivariant transformer architecture, demonstrates (this first example of Lorentz equivariant) generative model for particle physics, and demonstrates several real-life physics applications. On the down side, some reviewers felt the paper not sufficiently accessible to the NeurIPS community, and there were some concerns regarding the computation overhead of the Lorentz transformer layers. The authors have addressed this last point in their rebuttal. Overall, this paper seems like a solid contribution to geometric ML.
Decision: Accept (poster)