Does equivariance matter at scale?

I think the study is limited in significant ways that prevent me from saying that the question in the title was satisfyingly answered. I also think that the result analysis and the design of some elements of the study lead to conclusions that are potentially misleading.

Choice of task

• First, I don't think rigid body dynamics is an appropriate choice of benchmark task, since there is almost no overlap between the rigid body dynamics and equivariance communities. To the best of my knowledge, none of the state-of-the-art models on this problem (which the models implemented in this paper do not match) are rotation equivariant. It seems like it has been found by the rigid body dynamics community that equivariance is not worth it in this setting. If the authors want to suggest that equivariance might be of importance for this problem, they should compare an equivariant version of one of the state-of-the-art models like FIGNet.
• Second, the necessity of E(3)-equivariance for this task is questionable. This is because a natural reference frame is given by the direction of gravity and the position of the center of mass. Allen et al. 2022 use this reference frame in addition with O(2) augmentations around the Z axis with great success. A realistic comparison should use this canonical reference frame with O(2) augmentations for the non-equivariant model. I expect that this will make the task easier for the non-equivariant model and data augmentation more computationally efficient.
• Third, there are problems for which equivariance is thought to be more important, and the study does not give evidence that the findings would extend on these problems, despite making general claims. Examples include potential energy/force field prediction, protein folding, molecular generative modelling and medical image segmentation. I think the conclusions of the study would be much more valuable if the authors had chosen to evaluate on one or ideally some of these tasks.
• Fourth, the only type of equivariance investigated is continuous E(3) equivariance. I therefore think it is an overstatement to extend the conclusions to any type of equivariance. Notably, no discrete groups are considered. Both equivariant models and non-equivariant models have been argued to be successful for many discrete groups such as image rotations (ViT not being equivariant for example).

Summary: The scope of the paper is ambitious, as from the title and abstract, it proposes to provide insights on the necessity of equivariance in general. However, the only experimental setup investigated is rigid body dynamics with E(3) equivariance. Even for this problem I find the experimental setup too artificial to give practical insights. Given this, I think the scope is overstated and that the experiments are lackluster.

Choice of architectures and training setup

• Another important limitation is that only one non-equivariant and equivariant architectures are compared. I think the diversity of architectures used in geometric deep learning does not compares with NLP where Transformers are dominant. Notably, I think a comparison using GNN type architectures would be valuable. These architectures are as far as I know more prevalent in geometric deep learning (at least definitely for rigid body dynamics) than Transformers on problems with E(3) equivariance, due to locality being a relevant inductive bias and offering computational advantages. I also think at least one spherical harmonics based architectures like Equiformer should be investigated since these are more widely used than geometric algebras.
• It is not clear how the family of hyperparameters for the two models is obtained. An important aspect of an experimental study like this one is to specify a consistent procedure to come up with hyperparameters since these are important for model performance. I encourage the authors to give more details on this and if possible to use a budget method for hyperparameter tuning.
• I don't see a reason why the learning rate should be the same for the two models. It is mentioned in the paper that different learning rates were experimented with, so I think the authors should tune such parameters separately for each model. The learning rate is known to be important and the optimal one is not expected to the same for every model.
• I do not think that the same batch size should be used for both models. I understand that there is an issue of fair comparison here, since we expect that bigger batch size should lead to better performance. However, it would be more useful to know if the memory utilization of the equivariant model is significantly larger in a way that would make larger batches usable for the non-equivariant one. This is necessary to answer the question asked by the paper since in practice, people look to maximally use the available resources. Being able to do this efficiently should therefore be taken into account when choosing which model to use. I would be more interested in seeing experiments for which resources are used maximally for both models.

Summary: I find the study to be too narrow in terms of architectures investigated. The method used for hyperparameter tuning should be reported and made consistent. Finally, the experiment should be designed such that the different models maximally use the available ressources.

Result analysis

• The number of FLOPs is not really a metric of computational cost that matters in practice. While I understand that analyzing models in terms of FLOPs has the advantage of being hardware agnostic and to facilitate analysis, I think FLOPs should not be used to draw general conclusions. Therefore, I think the conclusion in the abstract "equivariant models outperform non-equivariant ones at each tested compute budget" is misleading. I think comparing models in terms of inference/training time on given hardware would be significantly more insightful. This is the result that should be put at the forefront of the paper. In particular, for rigid body dynamics forward time is much more important in practice than any measure based on training FLOPs, since the main goal of machine learning simulators is speed.
• From the FLOP throughput given in appendix, it seems like the conclusions obtained from figure 1 would be different if we looked at training time. The gap between equivariant and non-equivariant models would maybe become smaller. I think a version of figure 1 with forward time should be provided in the main paper and the conclusions nuanced appropriately if needed.
• I find the results provided in Figure 1 (and Figure 2 to some extent) to be quite confusing. The authors mention that the results are obtained in the infinite-data setting. If I understand correctly, this is training on a single epoch of data. As mentioned in the paper, this has the effect of making data augmentation irrelevant, since data augmentation will not be able to increase the effective amount of data seen by the non-equivariant model. I think this scenario is however unrealistic. I cannot think of real applications of equivariant models for which data is so abundant that one could only afford to train for a single epoch. If these applications exist, I encourage the authors to highlight them. Otherwise, I think it is misleading to consider that limit over the relevant setting where available compute can handle the available amount of data. I think this limit should be discussed separately from the case N_epochs >> 1 that is more relevant in practice.

Summary: I think conclusions on computational efficiency/budget should not be based on number of FLOPs, but on a practically relevant metric like inference time. It also seems like the paper overemphasize the one epoch limit (where data augmentation is useless), which is not relevant in practice.

1. While the experimental setup is carefully designed, the experiments are limited to only one architecture and one problem. Equivariances and symmetries are present in a wide range of problems and go beyond the studied $E(3)$ group. For instance, Lie symmetries have recently been utilized for training Neural Operators; See [1].

2. Also, in many domains, transformers are still not the staple architecture. Focusing on only one architecture class while trying to answer the broader question of the "effectiveness of equivariant networks" is limiting. For instance, with PDEs again, Fourier Neural Operators and Convolutional networks are still preferred over transformers (and there exist many equivariant convolutional architectures that can be studied, e.g. [2]).

3. Some parts of the paper regarding the scaling laws (section 3.3) are not very clear and hard to follow (See question 1). Also, some of the results in Section 4.1 are hard to interpret (see question 2)

4. In the end, while I appreciate the authors' efforts in designing a well-structured experimental setup, I believe the paper is severely limited in terms of experiments and generalizing to other problems.

[1] Brandstetter, Johannes, Max Welling, and Daniel E. Worrall. "Lie point symmetry data augmentation for neural PDE solvers." International Conference on Machine Learning. PMLR, 2022.

[2] Finzi, Marc, et al. "Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data." International Conference on Machine Learning. PMLR, 2020.

• The primary concern is the limited scope of tasks (benchmarks) and equivariant models:

  • Are these insights going to generalize across all tasks? For instance, the authors point out that one of the motivations for the study is the recent high-profile models' (for protein folding and conformer generation) choice of non-equivariant models - why didn't the authors consider equivariant counterparts of these models and these tasks?
  • Will these insights generalize to other kinds of equivariant models instead of the $E(3)$ equivariant transformer designed as in GATr, such as a stack of $E(3)$ equivariant layers? Multiple experiments have shown that the choice of equivariant architecture can change the training time.
  • Will there be any differences in the observations for discrete groups, such as $C_n$ or $D_n$, instead of continuous groups, especially for test performance with training data?

• Availability of data/compute budgets:

  • The paper isn't conclusive for seemingly infinite data and large compute (most present-day experiments); instead, it states that for finite compute budgets up to a specific limit ($10^{19}$ FLOPs), the equivariant model outperforms the non-equivariant model (Fig 1). Can the authors share if the two lines are coming close in Fig 1? And given the simplicity of non-equivariant models (and, generally, ease of training), what choice is favourable in the above-described setting?

• Unknown symmetries:

  • Consider that the exact symmetry (or underlying group of transformations) is unknown for a particular task, or some invariant transformations and their effects on the task are known. Is it better to use an equivariant model with a larger group (encompassing these transformations) or a non-equivariant model with these transformation augmentation?

1. One major concern with this paper is that the authors use nominal FLOPs as the training compute budget to draw a conclusion in line 85/86: “Equivariant transformers are also more compute-efficient, and this advantage persists across all compute budgets studied.” However, the authors themselves point out later that nominal FLOPs do not reflect the real-world run-time and training cost, as equivariant models have poor GPU utilization. I think it’ll be fair to provide an equivalent of Fig. 1 that uses the real training cost in GPU-hours. This can be better for practitioners to draw conclusions on which model to go for with limited compute budget. I understand the authors' point of using nominal FLOPs when analyzing each model separately. However, while recommending which model leads to better loss given a fixed compute budget, the authors should consider GPU hours/wall clock time, which is more practical than nominal FLOPs.
2. When drawing conclusions on whether to use equivariant vs. non-equivariant models on a fixed compute budget, the paper doesn’t discuss the inference speed benefits of non-equivariant models over equivariant ones. Any inference wall clock comparisons of equivariant models with non-equivariant ones would shed more insights on which model to go for given a fixed latency budget. For example, if a bigger non-equivariant model with lower loss takes the same wall clock time for inference (with better GPU utilization), then it might make more sense to go for that model.
3. I think “Overall, our findings hint that strong inductive biases …… large datasets and large compute budgets.” is a very strong claim given that the authors only tested their hypothesis in one synthetic 3D rigid body interaction dataset. I understand that the reason for choosing the dataset was to be able to generate as much data as you want cheaply. I wonder why the authors didn’t consider other domains like images, where a lot of data, models (both equivariant and non-equivariant), and self-supervised learning techniques are available. It would have made the paper stronger to validate their point for both E2 and E3 equivariant models. I think the current experimental setup is too simplistic to make general conclusions on whether equivariance matters at scale or not. The authors should be more transparent about the limitations of this study early on in the paper. The hypothesis the authors tested in this paper is more like “Does E3 equivariance matter at scale for 3D n-body dynamics?”
4. I am not sure if the authors are familiar with the body of work that tries to solve the inefficiencies of equivariant networks by frame averaging or learning canonicalizations of the data. As the purpose of that line of work is mainly to scale equivariance to large models, I think a discussion of that in the related work would be nice.