Lorentz-Equivariant Geometric Algebra Transformers for High-Energy Physics

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.

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. - 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 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. - What are the scientific goals of the experiments defined in the paper? (In the context of fundamental research in high energy physics).

Questions

Please address the questions in the weaknesses section. Having these clearly answered would improve the paper's clarity, reproducibility, and usability. I'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.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Not applicable.

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.

Strengths

1. The paper is well-written and easy to follow. In addition, the problem is well-motivated. 2. The proposed L-GATr generalizes GATr from $E(3)$ to the Lorentz group. 3. As stated by the authors, the Lorentz-equivariant flow matching proposed in section 3.2 is the first generative model proposed for particle physics. 4. Compared to graph-based Lorentz equivariant networks, the transformer architecture is more efficient and scalable. 5. The proposed method is shown to be more data efficient than the baselines in both amplitude surrogates and generative modeling experiments. 6. The ability to scale with the data is also verified in several experiments. 7. 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. 2. 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. 3. In the top tagging experiment, the performance of the proposed method is marginally worse than the baseline method. 4. Although the proposed method is claimed to support symmetry-breaking data, its effect is not well studied in the experiments.

Questions

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

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

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.

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.

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. - The paper clearly distinguishes their work from existing research, emphasizing the significance of their contributions. - 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. - The experiments on generative modeling lack comparisons with other equivariant models or SOTA models. - The scalability of the model remains limited, which reduces its suitability for high-energy physics applications.

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.

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

4

Limitations

The efficacy and scalability of the proposed framework may not yet suffice to replace traditional methods.

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.

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. * The comparison of generative modeling capabilities is limited to flow models, which may not represent the full spectrum of possible approaches. * 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.

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

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors briefly discussed the limitations in the Discussion section.