Lorentz-Equivariant Geometric Algebra Transformers for High-Energy Physics

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.

The paper is not self-contained.

Thank 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.

The implemented Lorentz equivariant model (what are the inputs?, are they a set of particles? how are they matched to queries, values and keys?)

The 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.

This 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.

The flow-matching approach. Again, what are the inputs, outputs? How's the model defined? What does it produce?

The 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.

To 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.

To 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].

What are the scientific goals of the experiments defined in the paper?

On 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.

As 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.

Now, 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.

We hope that we were able to answer your questions and look forward to discussing further.

References:

• [1] Y. Lipman et al, "Flow matching for generative modelling", ICLR 2023

I appreciate the explanations. I suggest adding them to the paper or appendix. I increased the score.

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.

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].

Indeed, the form of these equations and generally the overall architecture design are similar. There are several differences though:

• The multivectors $x$ in our Eq. (1) and Ref. [1]'s Eq. (4) are elements of different geometric algebras with different metrics.
• The inner product in our Eq. (2) is therefore also different from the one in Ref. [1].
• 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$.
• The geometric product used in the MLP is different between both architectures. This is another consequence of using a different metric.
• 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].
• 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.

The current presentation of the experiments can be a bit hard to understand for people without a physics background.

Thank 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.

In the top tagging experiment, the performance of the proposed method is marginally worse than the baseline method.

It'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.

Although the proposed method is claimed to support symmetry-breaking data, its effect is not well studied in the experiments.

Actually, 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.)

To 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:

The beam direction is thus important to analyze data that involve detector effects.

What are $y_m$, $y_p$, $\eta$, and $\phi$ in (4)?

These 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.

the proposed L-GATr has additional computational overhead compared to traditional transformers.

It'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.

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 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.

We hope that we were able to answer your questions and look forward to discussing further.

References:

• [1] J. Brehmer et al, "Geometric Algebra Transformer", NeurIPS 2023

Thank you for the response. We are glad we were able to explain the novelties in our architecture better.

As for the effect of symmetry breaking, we tend to disagree that the difference between an accuracy of 0.9403 and 0.9417 is insignificant on this benchmark. This difference is similarly big to the gap between equivariant networks and the top-performing non-equivariant networks (ParticleNet, ParT) outlined in Table 1.

To make the relevance of symmetry-breaking inputs clearer, we also re-ran the generative modelling experiment without them. Because of limited time, we ran shorter experiments (50% of the training time quoted in the paper). This is what we find:

So symmetry-breaking inputs are crucial for this experiment. This stark difference is not surprising. Consider the AUC metric, in which we train a classifier to distinguish between samples generated by the model and test data. For a generative model with unbroken symmetry, the distribution is invariant under rotations: it generates samples with particles moving in any direction with equal likelihood. This is very different from the data distribution, which is not invariant under rotations: the beam direction leads to a strongly preferred direction for the produced particles. It is easy for a classifier to spot this difference, leading to the high AUC. A similar argument can be made for the log likelihood.

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.

The model's performance in Table 1 is not optimal

It'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.

The experiments on generative modeling lack comparisons with other equivariant models or SOTA models.

We 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.

The 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.

That 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!

The scalability of the model remains limited. [...] The efficacy and scalability of the proposed framework may not yet suffice to replace traditional methods.

Could you elaborate what you mean here? In terms of scaling the model size, we used L-GATr models with millions of parameters without issues.

In terms of scaling to large training data sets, we see in Fig. 3 that L-GATr benefits from that more than all baselines.

The 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.)

evaluate the proposed model against traditional sampling methods

If 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.)

To 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.

We hope that we were able to answer your questions and look forward to discussing further.

References

• [1] E. Buhmann et al, "EPiC-ly Fast Particle Cloud Generation with Flow-Matching and Diffusion", arXiv:2310.00049
• [2] A. Butter et al, "Jet Diffusion versus JetGPT - Modern Networks for the LHC", arXiv:2305.10475

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.

significant computational overhead

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 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 [1, see Fig. 5 there] from which L-GATr is derived, which shares performance characteristics with L-GATr.

The comparison of generative modeling capabilities is limited to flow models

True. 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.

Thank 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.

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.

In 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.

Is it really necessary to employ transformers since the input dimension is not as high?

It 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.

Why is a smaller model used for the top tagging task?

We 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.

To 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:

This 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).

We hope we were able to address your questions and look forward to discussing further.

References

• [1] J. Brehmer et al, "Geometric Algebra Transformer", NeurIPS 2023
• [2] E. Buhmann et al, "EPiC-ly Fast Particle Cloud Generation with Flow-Matching and Diffusion", arXiv2310.00049
• [3] A. Butter et al, "Jet Diffusion versus JetGPT - Modern Networks for the LHC", arXiv:2305.10475
• [4] S. Badger and J. Bullock, "Using neural networks for efficient evaluation of high multiplicity scattering amplitudes", Journal of High Energy Physics 2020
• [5] S. Badger et al, "Loop amplitudes from precision networks", SciPost Physics 2023