Does equivariance matter at scale?

## Abstract

Abstract Given large data sets and sufficient compute, is it beneficial to design neural architectures for the structure and symmetries of each problem? Or is it more efficient to learn them from data? We study empirically how equivariant and non-equivariant networks scale with compute and training samples. Focusing on a benchmark problem of rigid-body interactions and on general-purpose transformer architectures, we perform a series of experiments, varying the model size, training steps, and dataset size. We find evidence for three conclusions. First, equivariance improves data efficiency, but training non-equivariant models with data augmentation can close this gap given sufficient epochs. Second, scaling with compute follows a power law, with equivariant models outperforming non-equivariant ones at each tested compute budget. Finally, the optimal allocation of a compute budget onto model size and training duration differs between equivariant and non-equivariant models.

## 1 Introduction

In a time of big data and abundant compute, how important are strong inductive biases?
Consider problems governed by known symmetries: should one take these into account by designing and using equivariant neural network architectures (Bronstein et al., ), or is it better to learn them implicitly from data?

A common intuition is that strong inductive biases bring the biggest benefits when little training data is available, and that symmetry properties can just as well be learned from data given sufficient samples and compute.
Recently, high-profile models of protein folding (Abramson et al., ) and conformer generation (Wang et al., ) have received considerable attention for their choice of non-equivariant architectures for geometric problems.

At the same time, there is reason to expect that equivariance is still beneficial in the large-data limit.
Learning means successively narrowing down a hypothesis class based on evidence. From this perspective one can explain (Bahri et al., ) the empirical observation that test losses often scale as a power law with the training compute (Kaplan et al., ; Hoffmann et al., ).
Whereas non-equivariant methods start from the space of virtually all functions, equivariant models start from the subspace of all functions that abide by the symmetries of the problem.
The learning process may benefit from that by focusing solely on further refining this smaller hypothesis class, narrowing down to the correct solution with fewer training steps.

Until the theory of scaling laws is fully understood, the effects of equivariance on scaling is an empirical question, and in this work we study it empirically.
We focus on a benchmark problem of modelling the physical interactions between rigid three-dimensional objects described by meshes. This task is known to be challenging (Allen et al., ). It is manifestly equivariant under $\mathrm{E}(3)$, the symmetry group of rotations, translations, and reflections.
We compare a standard transformer architecture (Vaswani et al., ) to an $\mathrm{E}(3)$-equivariant transformer (Brehmer et al., ).

In this setup we ask three questions:

- 1.
*How do equivariant and non-equivariant models scale as a function of the available data?*
Does data augmentation affect this?
- 2.
*How do equivariant and non-equivariant models scale as a function of training compute?*
Does this scaling follow power laws? Are their coefficients affected by equivariance?
- 3.
*Given a compute budget, how should one allocate it to the model size and the number of training iterations?*
Is this trade-off different for equivariant and non-equivariant models?

In our attempt to answer these questions, we train equivariant and non-equivariant models for different training compute budgets, trade-offs between model size and training steps, and dataset sizes.
We then analyze these results both qualitatively as well as quantitatively by fitting empirical scaling laws.

Our experiments provide evidence for three conclusions.
As expected, equivariance improves *data* efficiency. However, data augmentation largely closes this gap.
Second, equivariant transformers are also more *compute*-efficient, and this advantage persists across all compute budgets studied.
Both model classes exhibit power-law scaling behaviour.
Finally, the optimal allocation of a training compute budget to model size and training steps differs between equivariant and non-equivariant models.
Overall, our findings hint that strong inductive biases may not only yield benefits in the low-data regime, but can also be beneficial with large datasets and large compute budgets.

## 2 Background and related work

#### Neural scaling laws

The scaling of neural network performance as a function of model size or training steps has been studied extensively (Ahmad & Tesauro, ; Hestness et al., ; Rosenfeld et al., ; Henighan et al., ).
Kaplan et al. () first observed that the test loss of autoregressive language models follows a power law over many orders of magnitude.
Hoffmann et al. () improved the methodology further and found the “Chinchilla” scaling laws, which still serve as a reference point for many language models. In our quantitative analysis of compute scaling, we largely follow their approach.

Several works have extended scaling laws from model size and training steps to other dimensions: Muennighoff et al. () studied the effect of the training dataset size, which we also discuss, Alabdulmohsin et al. () analyzed scaling of different architecture hyperparameters separately, and Jones () investigated the scaling with problem complexity.

#### Scaling laws and inductive biases

There has been comparatively little research into the relation between inductive biases and scaling behaviour, perhaps because the transformer architecture (Vaswani et al., ) is so established in language modelling.
Tay et al. () compared the scaling behaviour of different architectures.
Recently, Qiu et al. () investigated how structured linear transformations in transformers affect scaling laws. The authors conclude that imposing structure in them can improve the scaling behaviour.
Our work differs from both of these studies through its focus on symmetric problems and equivariant architectures.

#### Geometric deep learning

Geometric deep learning (Bronstein et al., ) is a paradigm for machine learning in which network architectures are designed to reflect geometric properties of the problem.
One of its core ideas is that of equivariance to symmetry groups (Amari, ; Wood & Shawe-Taylor, ; Makadia et al., ; Cohen & Welling, ): roughly, a network $f$ is said to be equivariant to a symmetry group $G$ if $f(g\cdot x)=g\cdot f(x)$ for all elements $g\in G$ and all inputs $x$, where $\cdot$ is the group action. This means that when you transform the inputs into an equivariant network, its outputs transform consistently.
An equivariant network thus does not have to learn the symmetry structure from data, like a non-equivariant network does.

Equivariance has been found to improve performance, data efficiency, and robustness to out-of-domain generalization in fields as diverse as
quantum mechanics and quantum field theory (Pfau et al., ; Hermann et al., ; Boyda et al., ; Gerdes et al., ),
molecular force fields (Batatia et al., ; Batzner et al., ; Liao & Smidt, ; Musaelian et al., ; Batatia et al., ),
generative models of molecules (Zeni et al., ; Igashov et al., ),
particle physics (Bogatskiy et al., ; Gong et al., ; Spinner et al., ),
biological and medical imaging (Veeling et al., ; Bekkers et al., ; Winkels & Cohen, ; Winkens et al., ; Mohamed et al., ; de Ruijter & Cesa, ; Suk et al., ),
wireless communication (Hehn et al., ),
and robotics (Wang et al., ; ; ; Brehmer et al., ).
The potential of equivariance to improve generalization has also been shown theoretically (Sokolic et al., ; Lyle et al., ; Elesedy & Zaidi, ; Sannai et al., ; Behboodi et al., ; Petrache & Trivedi, ).

At the same time, equivariant architectures are often more complex than non-equivariant architectures. Some researchers believe that equivariant architectures are more difficult to scale up, but to the best of our knowledge there has been little systematic study into this.
However, recent impactful works on protein folding (Abramson et al., ) and conformer generation (Wang et al., ) found that equivariant architectures did not offer any benefits and opted for non-equivariant models and data augmentation instead.

#### E ⁢ ( 3 ) E 3 \mathrm{E}(3) roman_E ( 3 ) equivariance

One symmetry that is important in many scientific and industrial applications is the group $\mathrm{E}(3)$ of isometries of Euclidean space. It consists of translations, rotations, and reflections. This group is the focus of our investigation.

As an $\mathrm{E}(3)$-equivariant architecture, we use the Geometric Algebra Transformer (GATr) (Brehmer et al., ). It has two defining features.
First, GATr uses multivectors from projective geometric algebra as representations, in addition to the usual unstructured representations. These multivectors are 16-dimensional objects that can represent various geometric primitives, including absolute positions in space, directions, as well as translations and rotations. Geometric algebra representations power a number of recent architectures (Brandstetter et al., ; Ruhe et al., ; ; Brehmer et al., ; de Haan et al., ; Spinner et al., ; Zhdanov et al., ; Liu et al., ; ).
Second, GATr is a transformer. It processes inputs in the form of a set of tokens. Pairwise interactions are not computed through local message passing, as in many other equivariant architectures, but through an equivariant dot-product attention mechanism that is compatible with efficient implementations like FlashAttention (Dao et al., ).
We choose GATr as the equivariant model for our scaling investigation because of this similarity to the standard transformer.

## 3 Problem setup

### 3.1 Benchmark problem

#### Desiderata

A benchmark task for this empirical scaling study should be characterized by a low floor and a high ceiling: a small model trained on few samples should perform poorly, while a large model trained on many samples should score orders of magnitude better.
To study data scaling, we need a large number of training samples.
To study equivariance, we look for a geometric problem in which the symmetries and representations are known and exact.

#### Rigid-body modelling problem

We choose a rigid-body modelling problem as our benchmark. Three-dimensional meshes are initialized at some position, orientation, and velocity; they then interact with each other under gravity and collision forces.
Concretely, the inputs to the network consist of a set of triangular meshes for two time points $t=t_{0},t_{0}+\Delta t$, and the task is to predict all mesh vertices at time $t=t_{0}+2\Delta t$.
As a loss function and evaluation metric, we use the mean squared error of the predicted mesh vertex positions.

This problem satisfies all desiderata for our study. Rigid-body interactions are known to be challenging to model: collisions are difficult to detect, since they do not usually occur at or near vertices; the forces acting during a collision are nearly discontinuous (Bauza & Rodriguez, ; Pfrommer et al., ; Allen et al., ).
Synthetic data can be generated cheaply with physics simulators.
Finally, the physics of the process is clearly equivariant under $\mathrm{E}(3)$, provided that the direction of gravity is treated as a feature and rotated along with the scene.

#### Dataset

We construct a dataset of rigid-body interactions following a proposal by Allen et al. (). We use the Kubric simulator (Greff et al., ), which is based on the PyBullet physics engine (Coumans & Bai, ). We recreate the MOVi-B dataset used by Allen et al. () as best as we can, using parameters from their paper and private communication; see Appendix A for details.
Our dataset consists of $4\cdot 10^{5}$ trajectories, each consisting of 96 time steps. Each trajectory includes between 3 and 10 objects, each consisting of between 98 and 2160 mesh faces. The average number of total mesh faces in a scene is 5470.

### 3.2 Models

In selecting architectures, our main objective is not to achieve state-of-the-art results on the particular rigid-body benchmark problem we chose. That would lead us to highly problem-specific architectures (Allen et al., ; Rubanova et al., ). Instead, we aim for general-purpose architectures that are applicable to broad classes of problems.

#### Baseline architecture

The transformer architecture (Vaswani et al., ) has become the de-facto standard across a wide range of machine learning tasks. It is versatile with respect to the input data, propagates gradients effectively, and scales well to large model sizes and input tokens. Most scaling studies have focused on transformers as well.
We therefore use a standard pre-LN (Baevski & Auli, ) transformer with multi-query attention (Shazeer, ) as our non-equivariant architecture.

We represent each mesh face as a token and the positions and velocities of vertices with random Fourier features (Tancik et al., ), which improved performance in initial tests.

Even this baseline architecture is hardly “free from inductive biases”. Because the tokens form not a sequence, but an unordered set, we do not use positional encoding. Therefore, the model is equivariant with respect to one of the symmetries of our problem: that of permutations of the input tokens. In this respect, there is no difference between the two architectures, and we do not compare to any models that are not permutation-equivariant.

#### Equivariant architecture

For the $\mathrm{E}(3)$-equivariant architecture, we again look for broad applicability (at least within the class of $\mathrm{E}(3)$-symmetric problems).
In addition, we would like the architecture to be as structurally similar to the transformer, to isolate the effects of equivariance on scaling as well as possible.
We therefore opt for the (to the best of our knowledge) only $\mathrm{E}(3)$-equivariant architecture that is based on dot-product attention with unlimited receptive fields, and which also otherwise follows the transformer blueprint closely: the Geometric Algebra Transformer (GATr) (Brehmer et al., ).

Again, we represent each mesh face as a token. GATr uses geometric algebra representations in addition to the usual scalar channels, and we can represent the geometric properties of a mesh face in these geometric representations. We describe this embedding in more detail in Appendix B.

#### Hierarchical attention

While we focus on general-purpose architectures, we find that both models benefit from two minor modifications to the transformer blueprint.
First, we use a novel *hierarchical attention* mechanism, in which multiple attention heads use different attention masks: half of the heads are restricted to attend only to mesh faces in the same object, while the other half attends to all tokens (mesh faces).
This allows us to embed awareness of the mesh structure into the transformer architecture, while preserving the efficiency of dot-product attention.

#### Enforcing object rigidity

Second, we enforce *object coherence and rigidity* when computing the outputs. Either transformer model first outputs a translation vector and a rotation quaternion for each mesh face. These are averaged over each object, resulting in a translation vector and a rotation for each rigid object. These $\mathrm{E}(3)$ operations are then applied to the input meshes. In this way, the networks by design translate and rotate rigid objects consistently.
We describe this procedure in more detail in Appendix B.
In preliminary experiments, enforcing object rigidity in this way improved performance substantially compared to directly predicting the positions or velocities of mesh vertices. We also experimented with outputting and exponentiating elements of the Lie algebra for each object, but found that that worked marginally worse.

**Table 1: Architecture hyperparameters as a function of a model size parameter $n$. The equivariant architecture is less wide, but part of their channels are 16-dimensional multivector (MV) channels, which can express a variety of geometric primitives (Brandstetter et al., ; Ruhe et al., ; Brehmer et al., ; Ruhe et al., ; de Haan et al., ).**
| Hyperparameter | Baseline | Equiv. |
| --- | --- | --- |
| Attention blocks | $2n$ | $2n$ |
| Scalar channels | $64n$ | $4n$ |
| MV channels | – | $n$ |
| Attention heads | $2n$ | $2n$ |
| Scalars per key, query, value | $64$ | $8$ |
| MV per key, query, value | – | $2$ |
| Hidden scalar channels in MLP | $128n$ | $8n$ |
| Hidden MV channels in MLP | – | $2n$ |

#### Hyperparameters

We tune the hyperparameters of both models manually. For both the baseline and equivariant transformer, we define a one-parameter family of hyperparameters, fixing the relation between the number of layers, attention heads, and channels to be linear.
Our architectures are shown in Tbl. 1.
Notably, we find that the equivariant transformer benefits from a more narrow architecture, which may be evidence of the expressivity of its multivector channels.

#### Optimization

We train all models with the Adam optimizer (Kingma, ), annealing the learning rate over the course of training from an initial value of $5\cdot 10^{-4}$ on a cosine schedule.
For experiments with small FLOP budgets of less than $10^{18}$ nominal FLOPs, we find that this learning rate can be too small. This is in line with other works that find larger learning rates beneficial for smaller compute budgets (e. g.  Dubey et al., ). We therefore repeat these experiments with a higher learning rate of $10^{-3}$ or $2\cdot 10^{-3}$, depending on the compute budget, and report the better result.
For simplicity, we use the same batch size of 64 samples (or on average $3.5\cdot 10^{5}$ tokens) for all experiments, even though this does not maximize GPU utilization and thus FLOP throughput.
Early stopping is used in all experiments.

### 3.3 Scaling-law analysis

#### Experiments

We perform two series of experiments. First, we study the scaling with compute, in the (practically) infinite-data setting.
We vary a training compute budget over three orders of magnitude, between $10^{16}$ and $10^{19}$ FLOPs. For each FLOP budget, for both the baseline and the equivariant transformer, we perform multiple experiments: each with a different trade-off between model size $N$ and training length $D$.
This requires understanding the relation between $N$, $D$, and the total training FLOPs; we discuss that later in this section.

Second, we study the scaling with training data, fixing the training compute budget, the model size, and the number of training tokens.
For both models we choose settings that performed compute-optimally in the first series of experiments for a compute budget of $10^{18}$ nominal FLOPs.
The number of unique samples in the dataset is varied over five orders of magnitude, from $2\cdot 10^{6}$ tokens to $2\cdot 10^{11}$.
The lower end of this scan corresponds to training for $6\cdot 10^{5}$ epochs, while every sample is seen only once on the upper end of this scan.
For each of these settings, we train a baseline transformer, an equivariant transformer, and a baseline transformer trained with data augmentation, in which symmetry transformations are applied to the samples, independently for each epoch.

#### Counting FLOPs

Setting up our experiments (see above) and analyzing the scaling with compute both require knowing the relation of the total number of training FLOPs $C(N,D)$ and the model size $N$ as well as training tokens $D$.
This relation is different for the baseline and equivariant transformer.

Following Kaplan et al. () and Hoffmann et al. (), we perform this FLOP counting in the limit where the number of model parameters is much larger than the sequence length, which in turn is much larger than 1. The training compute is then dominated by the linear layers.
For both of our models, we find

$$ $C(N,D)\approx\xi ND\,,$ (1) $$

where $\xi$ is an architecture-dependent constant.

For the baseline transformer, famously $\xi=6$ (Kaplan et al., ).
For the equivariant transformer, the value of $\xi$ depends on the ratio of scalar and multivector channels: a model with only scalar channels would also have $\xi=6$, while a pure-multivector model would have more weight sharing and thus a higher FLOPs-per-parameter ratio $\xi=6\cdot 16^{2}/9\approx 171$. For the hyperparameters we use during our scaling study, we find $\xi\approx 61.2$.

Note that these *nominal FLOPs* do not necessarily correspond to the actual compute required to train the model.
For one, the assumed hierarchy between the model parameters and the sequence length is not always satisfied.
Second, our implementations of the models may not be able to fully utilize the GPUs. We observe this in particular for small models and for the implementation of the equivariant transformer, which involves many smaller operations and faces CPU bottlenecks.
Additional overhead comes from inter-GPU communication, data loading, logging, checkpoint saving, validating, and so on.
In our experiments, two models with the same nominal FLOP count would differ by as much as an order of magnitude in real training duration.

So why do we still analyze models in terms of the nominal FLOPs?
While they are an imperfect measure, they do not depend on the implementation and hardware environment, and we believe they are still the best predictor of the theoretically achievable compute cost after sufficient optimization and at scale.

#### Scaling-law ansatz

We model the scaling with compute quantitatively by fitting a scaling law to all of our experiments.
Following Kaplan et al. (), we model the test loss $L$ as a power law in the model parameters $N$ and the training duration $D$, measured in tokens:

$$ $\hat{L}(N,D)=\frac{A}{N^{\alpha}}+\frac{B}{D^{\beta}}+E\,.$ (2) $$

Here $A,B,E,\alpha,\beta$ are fit parameters.

The parameter $E$ represents the irreducible loss that even a perfect model cannot eliminate. Unlike in language or image modelling tasks, there is no clear reason to expect such an irreducible error of practically relevant size for the deterministic physics task we use as a benchmark.
We treat the choice of whether to include $E$ as a fit parameter or fix it to zero as a hyperparameter and choose it through cross validation, as we will describe below.

For the scaling with the size of the training data set, we do not find a scaling law that convincingly describes our experiments.
Our attempts at fitting ’s data-constrained scaling law () to our data did not result in a good agreement.
We therefore refrain from discussing the functional form for this direction of scaling, and will focus on scaling with compute for the remainder of this section.

#### Scaling-law fit

Following Hoffmann et al. (), we fit the scaling-law parameters $(A,B,E,\alpha,\beta)$ separately for each architecture by minimizing the Huber loss (Huber, ) between the predicted and observed log loss values,

$$ $\sum_{\text{experiments}\ i}\mathrm{Huber}_{\delta}\Bigl{(}\log\hat{L}(N_{i},D _{i})-\log L_{i}\Bigr{)}\,.$ (3) $$

Here $\delta$ is a hyperparameter, we choose it based on cross-validation, as we describe in a bit.
We minimize this loss with the L-BFGS optimizer (Liu & Nocedal, ), starting multiple fits from a grid of initializations to avoid getting stuck in local minima.

#### Scaling-law hyperparameters

The scaling-law fit depends on two hyperparameters: whether we include the offset $E$ as a fit parameter and the value of $\delta$.
We determine both through leave-one-out cross-validation, performing scaling-law fits on all but one experiment and evaluating the error $\lvert\log\hat{L}(N_{i},D_{i})-\log L_{i}\rvert$ on the left-out experiment.
In this way, we choose fixing $E=0$ and $\delta=0.001$, though the qualitative fit results are not sensitive to these choices.

#### Compute-optimal performance

From a scaling law as in Eq. (2) and a FLOP function as in Eq. (1), we can derive the compute-optimal model size $N^{*}(C)$ and the compute-optimal training duration $D^{*}(C)$ as a function of the FLOP budget $C$ as

$$ $N^{*}(C)=\frac{G}{\xi^{a}}C^{a}\quad\text{and}\quad D^{*}(C)=\frac{1}{G\,\xi^{ b}}C^{b}\,,$ (4) $$

where $G=(\frac{\alpha A}{\beta B})^{1/(\alpha+\beta)}$, $a=\beta/(\alpha+\beta)$, and $b=\alpha/(\alpha+\beta)$ (Hoffmann et al., ).

The optimal loss achievable for a given FLOP budget is then

$$ $L^{*}(C)=\hat{L}(N^{*}(C),D^{*}(C))=E+\frac{F}{C^{\gamma}}$ (5) $$

with $F=AG^{-\alpha}\xi^{\gamma}+BG^{\beta}\xi^{\gamma}$ and $\gamma=\frac{\alpha\beta}{\alpha+\beta}$.

**Table 2: Scaling-law coefficients. In addition to the central values, we show the 95% confidence intervals from a nonparametric bootstrap.**
| Scaling law | Param. |  | Baseline |  | Equivariant |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Central | Lower | Upper |  | Central | Lower | Upper |  |  |
| Eq. (2): $\hat{L}(N,D)=A/{N^{\alpha}}+B/{D^{\beta}}$ | $A$ |  | $\mathbf{1.27}$ | $0.484$ | $5.07$ |  | $\mathbf{0.000282}$ | $0.000162$ | $0.000607$ |
|  | $B$ |  | $\mathbf{0.202}$ | $0.0108$ | $0.361$ |  | $\mathbf{469}$ | $159$ | $592$ |
|  | $\alpha$ |  | $\mathbf{0.909}$ | $0.832$ | $1.03$ |  | $\mathbf{0.348}$ | $0.293$ | $0.417$ |
|  | $\beta$ |  | $\mathbf{0.379}$ | $0.256$ | $0.404$ |  | $\mathbf{0.734}$ | $0.689$ | $0.747$ |
| Eq. (4): $N^{*}(C)\propto C^{a}$ | $a$ |  | $\mathbf{0.294}$ | $0.215$ | $0.307$ |  | $\mathbf{0.678}$ | $0.619$ | $0.711$ |
|  | $b$ |  | $\mathbf{0.706}$ | $0.693$ | $0.785$ |  | $\mathbf{0.322}$ | $0.289$ | $0.381$ |
| Eq. (5): $L^{*}(C)=F/{C^{\gamma}}$ | $F$ |  | $\mathbf{1.03}$ | $0.124$ | $1.89$ |  | $\mathbf{0.14}$ | $0.0524$ | $0.517$ |
|  | $\gamma$ |  | $\mathbf{0.268}$ | $0.213$ | $0.284$ |  | $\mathbf{0.236}$ | $0.212$ | $0.267$ |

#### Uncertainties

No realistic scaling study directly measures the *optimal* model performance as a function of some parameters.
Reasons for sub-optimality include the choice of hyperparameters, stochasticity in initialization and training, choosing a scaling-law ansatz that does not include the true functional form, and finite sampling of the space of model capacities and training tokens.
We estimate the effect of the latter with a nonparametric bootstrap, similar to Hoffmann et al. (). From $10^{4}$ bootstraps, we construct 95 % confidence intervals on the scaling law coefficients as well as on any derived predictions, using the empirical (or basic) bootstrap method.

## 4 Results

Figure: Figure 3: Test loss (dotted circles) and scaling-law predictions (background colour) as a function of model size and training tokens. Left: non-equivariant transformer. Right: equivariant transformer. In both cases, we observe good agreement of model performance and scaling-law fit.
Refer to caption: x3.png

Figure: Figure 4: Model performance at different training compute budgets (panels) as a function of the model size. We show our experiments (dots) and the predictions of our scaling-law fit (lines). The scaling-law fit describes the measurements well.
Refer to caption: x5.png

### 4.1 Scaling with compute

We first focus on the limit of (essentially) infinite training data and study the model performance as a function of model size $N$ and training tokens $D$.

#### Scaling laws

We fit the scaling law of Eq. (2) with $E=0$ to these experiments. For the baseline transformer, we find coefficients

$$ $\hat{L}_{\text{baseline}}(N,D)=\frac{1.27}{N^{0.909}}+\frac{0.202}{D^{0.379}}\,.$ (6) $$

The equivariant model yields

$$ $\hat{L}_{\text{equivariant}}(N,D)=\frac{2.82\cdot 10^{-4}}{N^{0.348}}+\frac{46 9}{D^{0.734}}\,.$ (7) $$

Confidence intervals are provided in Tbl. 2.

These two models scale quite differently with model size and training length, which has implications for the optimal allocation of a compute budget. We will discuss this later.

#### Fit quality

First, we show how well these fitted scaling laws agree with the data in Figs. 3 and 4. Comparing the observed values of the test loss to the predictions from the scaling laws, we overall find good agreement. There are no glaring deviations, although the power law underestimates the loss for the largest equivariant models and for one baseline outlier.
Most measurements fall within the uncertainty bands, but less than the 95% one would expect if the bootstrap would cover all relevant sources of error. This is evidence that the ansatz of Eq. (2) does not describe the data perfectly.

#### Scaling with compute

Next, we analyze the model performance and its scaling with compute.
From the training laws in Eqs. (6) and (7), we compute best achievable test loss $L^{*}$ as a function of the training compute budget $C$, as given by Eq. (5).
We find

$$ ${L}^{*}_{\text{baseline}}(C)=\frac{1.03}{C^{0.268}}\quad\text{and}\quad{L}^{*} _{\text{equivariant}}(C)=\frac{0.14}{C^{0.236}}\,,$ (8) $$

and the exponents are compatible with each other within the confidence intervals shown in Tbl. 2.
We visualize the empirical compute-loss measurements and the derived optimal compute-loss relationship in Fig. 2.

For any given compute budget, the equivariant transformer significantly outperforms the baseline.
Over the range of compute budgets we tested, the equivariant model achieves a loss that is lower by approximately a factor of 2.

#### Optimal allocation of compute

From the scaling laws we can also derive the optimal allocation of a given computational budget to the parameter count and training duration, see Eq. (4).
We show our results for both models in Fig. 5.

Figure: Figure 5: Optimal parameter allocation. We show the compute-optimal model size as a function of the training compute budget for the equivariant transformer ( ) and the non-equivariant transformer ( ). The equivariant architecture requires smaller models to achieve a compute-optimal performance, but this gap closes for bigger compute budgets.
Refer to caption: x6.png

We find that a compute-optimal equivariant transformer has less parameters than a compute-optimal baseline transformer. This is expected because the equivariant transformer performs more compute per parameter.

Perhaps more surprising is that the optimal trade-off depends on the compute in a different way for the two models. We find that for a regular transformer, one should scale training tokens more steeply than model size. For the equivariant model, we find the opposite trend: one should put additional compute more in the model size than the training tokens. The compute-optimal model sizes thus become more similar for larger compute budgets.

### 4.2 Scaling with data

Next, we turn to the scaling with training data for a fixed training compute budget.
In Fig. 2 we show the test loss as a function of the number of unnique training tokens.
We compare baseline and equivariant transformers, each using a compute-optimal model size and training tokens for a training compute budget of $10^{18}$ nominal FLOPs.

The right end of these curves corresponds to the infinite-data, single-epoch limit considered in the previous section. Here we again see that the equivariant transformer outperforms the baseline model when compared at the same training compute budget.
Moving to smaller training sets, this gap widens substantially, confirming the expectation that equivariance improves data efficiency.

In Fig. 2 we also show results for a baseline transformer model trained with data augmentation. As expected, data augmentation does not make a difference when training for a single epoch. However, it drastically improves the performance in the small-data regime: when training for thousands of epochs, data augmentation makes a baseline transformer as data-efficient as an equivariant model.

## 5 Discussion

Our empirical results provide evidence for the following three conclusions.

#### 1. Equivariant transformers are more data-efficient, but data augmentation largely closes this gap.

The first (and expected) benefit for the equivariant architecture is that it performs better than a non-equivariant architecture when only little training data is available, as we show in Fig. 2.
However, we find a non-equivariant model trained with data augmentation performs just as well as the equivariant architecture, at least when the number of epochs (i. e. repeated uses of the same training sample) is sufficiently large.

#### 2. The scaling with compute follows power laws, and equivariant models outperform non-equivariant ones at each tested compute budget.

Both for non-equivariant and equivariant models, the test loss is well described by the power-law ansatz of Eq. (1), with parameters given in Tbl. 2.
The best achievable model performance for a given training compute budget therefore also scales as a power law, as given in Eq. (8).
We find consistent exponents for the two models, but a substantially smaller prefactor for the equivariant architecture.

This shows a second (and perhaps less expected) benefit for the equivariant architecture: for any fixed compute budget, even in the infinite-data limit, it clearly outperforms the baseline method.
As we show in Fig. 2, this benefit is approximately constant over the range of compute budgets we study.

Under the assumption that the implementations of equivariant and baseline architectures are similarly efficient and one can achieve the same FLOP throughput, this implies that equivariant models can outperform the non-equivariant counterparts even in the large-data, large-compute regime.
In practice, non-equivariant architectures may be easier to optimize for high FLOP throughput, in which case it remains to be seen which architecture is more efficient.

#### 3. Equivariant and non-equivariant models require different trade-offs between model size and training duration.

Our power laws indicate that the optimal allocation of a given compute budget onto the model size and training steps is different for equivariant and non-equivariant transformers, as shown in Fig. 5.
For small compute budget, a compute-optimal equivariant transformer is significantly smaller than a compute-optimal baseline transformer. This gap becomes smaller for larger compute budgets.

We hypothesize three possible explanations for this observation.
First, the baseline transformer, the more mature architecture, may have a better initialization scheme and thus require less training steps to reach a good performance.
Second, the different trade-offs may be related due to the different choice of width and depth between the architectures.
A third possible explanation is linked to the internals of the equivariant transformer architecture, which can express certain primitives particularly efficiently: the free movement and gravitational acceleration of rigid bodies can be represented with few multivector channels, thanks to the geometric product operation integrated into the architecture. This explains why the architecture can achieve a good performance with very few parameters. However, lowering the loss further requires precise collision detection and modelling. These need substantially more computational operations and a substantial amount of scalar channels, similar to the non-equivariant transformer. This offers a possible explanation for why at a larger compute budget, a model size closer to that of the baseline transformer is compute-optimal.

#### Limitations and open questions

As much as we would like to, we cannot conclusively settle the question raised in the title of this paper.
Our work is limited in several ways.
First, we only analyzed a single benchmark problem and two model families. We chose a task with a common symmetry group and general-purpose architectures that are frequently applied to a wide range of problems. We believe it is important to study to what extent our findings generalize to other problems or to other architectures, for instance those based on message-passing over graphs.
Moreover, on the problem we studied, we did not set a new state of the art: we deliberately focused on general-purpose models, which do not achieve the same level of performance as highly problem-specific architectures (Allen et al., ).

Another limitation of our work is that our analysis measures compute with an idealized FLOP counting procedure, as is common practice (Hoffmann et al., ). As we discussed in Sec. 3.3, this does not map one-to-one to real-world run time, at least not before further optimization of the implementation. In Appendix C we show the relation between wall time and nominal FLOPs in our experiments.

Finally, we are only able to study training compute budgets of up to $10^{19}$ FLOPs per model—this does not come close to the approximately $10^{25}$ FLOPs that the currently largest language models are trained for (Dubey et al., ). We did not see power-law scaling break down in the range we studied, but we cannot make claims about the extrapolation beyond it.

Keeping these limitations in mind, we believe that our findings provide some evidence that symmetry-aware modelling can be a sensible choice even for large compute and data budgets.
The benefits and disadvantages of strong inductive biases at scale are important for problems spanning several fields of science and engineering.
We hope that our study can encourage further investigations into this question.

### Acknowledgements

We thank Gabriele Cesa for great discussions as well as useful comments on a paper draft and Kelsey Allen for helping us understand the dataset used in FIGNet.

## Appendix A Dataset

We generate a benchmark dataset of rigid-body interactions with the Kubric simulator (Greff et al., ), which is based on the PyBullet physics engine (Coumans & Bai, ). We follow the MOVi-B configuration as used by Allen et al. (): we generate trajectories trajectories of 96 frames at 48 frames per seconds.
Our training set consists of $4\cdot 10^{5}$ such trajectories, while we use $1000$ trajectories each for the validation and test set.

Our data can be generated with the openly available repository at [https://github.com/google-research/kubric](https://github.com/google-research/kubric). That requires modifying the code to save object meshes, positions, and orientations for each time step, rather than the vision data that kubric stores by default, and running python3 challenges/movi/movi_ab_worker.py --objects_set=kubasic --frame_rate=48.

## Appendix B Models

#### Input representations

For both the baseline and equivariant transformer, we tokenize the problem by assigning one token to each mesh face.
In addition to the vertex positions, we compute the central position of each mesh face, the relative vector from the center to each vertex, the surface normal on the mesh face, and the linearly interpolated velocity between $t_{0}$ and $t_{1}$ for each vertex and the center of each mesh face. Together these form the input features.

For the baseline transformer, these features are embedded using random Fourier features (Tancik et al., ) with 128 frequencies sampled from a Gaussian with standard deviation $0.1$.

For the equivariant transformer, these features are embedded in the projective geometric algebra (PGA) described in Dorst (); Ruhe et al. (); Brehmer et al. (). Specifically, vertex and center positions are represented as PGA trivectors, the relative vector from the center to each vertex as PGA vectors, the mesh face surfaces with the associated normals as PGA vectors, and all velocities as PGA bivectors.

Figure: Figure 6: Throughput of nominal FLOPs per training wall time. The FLOPs we count do not directly correspond to training wall time: larger models and non-equivariant transformers lead to a better GPU utilization, increasing the FLOP throughput.
Refer to caption: x7.png

#### Enforcing object rigidity

The transformer networks output eight features $h$ for each token (mesh face) that represent transformations like translations and rotations.
The final predictions for future vertex positions $\hat{x}(t_{2})$ are then computed by applying these transformations to the current vertex positions $x(t_{1})$.

This is most easily expressed in projective geometric algebra, the representation naturally used by our equivariant transformer models. Here $h$ are the even-grade components of the output PGA multivectors. The predictions are computed as

$$ $\displaystyle h_{\text{agg}}$ $\displaystyle=\mathrm{mean}_{\text{objects}}h\,,$ $\displaystyle\hat{x}(t_{2})$ $\displaystyle=h_{\text{agg}}x(t_{1})\tilde{h}_{\text{agg}}\,.$ (9) $$

In the first line, the mesh-face-level predictions are averaged within each rigid object.
In the second line, the previous position $x(t_{1})$ is translated and rotated with the $\mathrm{E}(3)$ element represented by the network outputs; $\tilde{h}$ is the PGA reverse and $hx\tilde{h}$ (for properly normalized $h$) the sandwich product used to apply transformations to objects (Dorst, ; Ruhe et al., ; Brehmer et al., ).

We also experimented with predicting all vertex positions directly, without enforcing object rigidiy, as well as with parametrizing elements of the Lie algebra of $\mathrm{E}(3)$, which would then be exponentiated to construct transformations $h$. Both approaches performed worse in initial tests.

## Appendix C Scaling-law analysis

In Sec. 3.3 we argue that the nominal FLOPs we count are not fully indicative of real-world run time. We illustrate this in Fig. 6, where we show relation between these FLOPs and wall time in our experiments.
The throughput of nominal FLOPs varies by two orders of magnitude. Larger models as well as non-equivariant transformers lead to a better GPU utilization and thus increase the FLOP throughput.
We expect that the FLOP throughput could be improved by optimizing the batch size or (especially in the case of the equivariant transformer) the model implementation.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Can the authors explain the hierarchical attention part better? Do you assume that you know which tokens belong to which object to design the mask? If so, then isn’t that adding another dimension of problem-specific inductive bias to the architecture? I think this makes things more complicated to generalize the conclusions to any equivariant vs. non-equivariant model.
2. Can the authors explain line 262: “The lower end of this scan corresponds to training for …. the upper end of this scan.”? How many tokens were there for each experiment, and how many times each token was shown—that is, epoch? Maybe making a table of all these experiments in the Appendix would make things easier to understand for readers.
soundness: 3
strengths: 1. The paper is generally well written and tries to address a timely problem. Whether equivariance matters at scale has been under a lot of attention in the community, with arguments for and against equivariant models.
2. This paper is technically sound, with the authors explicitly making their assumptions clear wherever necessary.
3. The experiments with the 3D rigid n-body dynamics dataset are well designed and thoughtful.
4. I appreciate the authors pointing out the limitations of this work in Section 5, in particular, providing the throughput of nominal FLOPs per training wall clock time of different equivariant and non-equivariant models in Figure 6.
summary: It is well known in the literature that equivariance improves sample efficiency and generalization by adding geometric inductive bias to models and encouraging parameter sharing. This makes equivariant models appealing for settings with low data and where there is a constraint on the number of parameters. However, equivariant networks are more expensive to scale, and it is not clear whether there is any advantage to having this inductive bias when a lot of data is available. Recent papers show that non-equivariant models might be a better choice at scale. This paper provides a systematic approach to understanding the scaling behavior of equivariant models and compares them with non-equivariant ones. The authors conclude that non-equivariant models with data augmentation can close the performance gap given enough data and training. However, they also found that, for a fixed compute budget (in nominal FLOPs), equivariant models are better than non-equivariant ones.
weaknesses: 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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: - Are there comments on the quality of data? For instance, if the data is noisy (as a random example, say, images with random noise), is constraining models to the space of equivariant models a good idea?
- Can the authors reference papers (or perform experiments) that support the statement, "... a small model trained on a few samples should perform poorly, while a large model trained on many samples should score orders of magnitude better."?
- What differences would you see in the conclusions if you consider "not-permutation equivariant" models?
- The authors decide on some modifications to the architectures, namely, hierarchical attention and object rigidity. Are these modifications equivalently beneficial for both models (please provide some empirical evidence)?
- Can the optimization techniques affect the outcomes? For instance, how susceptible are the results to different learning rate scheduling and optimizers? Can you also detail the early stopping criterion?
soundness: 2
strengths: - The paper is well-motivated and clearly written. Conducting such a study is crucial, so the paper is very timely, and running the experiments on a large scale to draw insights is appreciated.
- The ablations over different aspects of scaling are exhaustive, as the authors consider scaling model size, training tokens, and compute.
- Several (design) choices are described in detail to help the readers understand the decisions behind the experiments.
summary: The paper discusses scaling laws for equivariant architecture, particularly $E(3)$-equivariant transformer on rigid-body modelling problem. The experiments resulted in three empirical findings - data efficiency of equivariant models, compute scaling follows power-law with equivariant model favoured over non-equivariant one at each compute budget and different trade-offs between model size and training tokens (duration) for each model.
weaknesses: - 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?

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. In lines 256 and 295, $D$ is said to be the training length or duration. However, in line 268, $D$ is the number of training tokens. I'd appreciate a clarification on this matter.

2. I am not sure what the takeaway from some of the presented results in section 4.1 is. More specifically, how should I interpret the coefficients listed in Table 2?
soundness: 3
strengths: - With the tendency towards the use of large and pre-trained models for various tasks, the importance of problem-specific information, such as symmetry and equivariance, has been losing attention. A structured and well-setup experimental setting that sheds light on how incorporating symmetries into the architecture or training process can improve the generalization, performance, etc in various scenarios is much needed. The authors implement a carefully designed setup that partially reveals the characteristics of such architectures. 

- The experimental setting covers a reasonable range of different extremes in terms of data availability, compute budget and parameter size.
summary: The paper studies the scaling properties of an $E(3)$-equivariant transformer architecture with respect to dataset size, parameter size, and computation budget for a rigid body problem. Laying the experimental foundation on the scaling laws developed for LLMs, authors find that the equivariant architecture is both data- and compute-efficient, and the parameter-size requirements for equivariant and non-equivariant architectures are different.
weaknesses: 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.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - How were the hyperparameters for the respective models obtained?
- I could not find what distribution of augmentations was used. In particular for augmentations over translations, it is not obvious what augmentations were used.
- Was the test data augmented or kept in a canonical reference frame?
- It would be interesting to see isoflop lines on figure 3 for the reader to see visually what is the optimal allocation of FLOPs
soundness: 2
strengths: - The paper tackles as important problem for the geometric deep learning. There are not a lot of works performing systematic comparisons of the performance of equivariant models vs non-equivariant ones. The investigation of scaling laws for equivariant models is also of interest
- The findings on optimal allocation of compute for equivariant models and non-equivariant models has value for practitioners, as they are helpful to guide design decisions and hyperparameter searches.
- The experiments have been done in a systematic way, as far as I can tell
summary: This paper proposes to study experimentally the potential benefits of equivariance in different dataset size and compute availability regimes. The study compares an equivariant transformer-type architecture to non-equivariant models with and without data augmentation on a synthetic task of rigid body dynamics. The highlighted findings are an improvement in data efficiency for equivariant models with data augmentation allowing non-equivariant models to come close in performance. It is also argued that given a fix compute budget, equivariant models outperform non-equivariant ones. Finally, power law models for optimal allocation of FLOPs (with model size and training time as variables) are fit, and it is found that the optimal allocation is different for the equivariant model compared to non-equivariant one. The equivariant model tends to achieve better performance by putting more compute in training time, whereas the non-equivariant one performs better when it is made larger.
weaknesses: 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.

Meta-review: The submission studies the scaling laws of equivariant models and compares them with non-equivariant ones.

+ The paper is well written.
+ The experiments are thoughtfully designed.
- The use of certain metrics such as FLOPs is not clearly motivated.
- The experiments are somewhat limited to support the broad claims of the submission.

Decision: Reject