Learning Symmetries through Loss Landscape

We thank the reviewer for their time and all the comments, which have helped us improve our paper and indicate our contribution. We are also open to answer any further questions.

Limited group transformations:

• We would like to highlight to the reviewer that our experiments primarily focus on point clouds and molecules across three key tasks: 3D Dynamical Systems, Motion Capture, and Molecular Dynamics. While we acknowledge that fully equivariant tasks, such as rotation, can be computationally intensive when sampling across all angles, our empirical results show that applying a single random rotation during training is sufficient to achieve performance comparable to equivariant baselines.

• We also conduct additional experiments with different numbers of samples from the symmetry group during training comparing our method and data augmentation. Our new results confirm that we can achieve reasonable performance using fewer samples from the symmetry group. We included the results in Appendix C.2 in the updated version.

Symmetry group:

• We thank the reviewer for highlighting the need to specify the equivariant tasks our method focuses on. In this work, we consider SE(3) symmetry group (i.e. group of rotations and translations). We have clarified this in the Introduction section of the updated manuscript.

Relaxed equivariance and unconstrained models

• We have updated our manuscript to consider the relevant works on relaxed equivariance in the Related Work section.

• We have added new results on the Motion Capture task. We compare our method against Residual Pathway Priors (RPP) [1], Projection-Based Equivariance Regularizer (PER) [2], and Equivariant MLP (EMLP) [3]. As these architectures are designed based on linear layers and MLP, we apply the augmented loss to standard MLP with a similar number of layers and parameters.

• We also include a new task on the Jet Flow dataset used by [4] (a two-dimensional benchmark that captures turbulent velocity fields): We apply our method to a Convolutional neural network (CNN) and compare it with Relaxed Steerable Convolution (RSteer) [4] and E2CNN [5] (more details in Appendix A in the paper).

• We clarified in the updated version that the high computational cost comes particularly with those relying on spherical harmonics and irreducible representations (Introduction Section).

• We thank the reviewer for the point on the limited expressive power and included it in the updated version. We also added some examples of tasks that exhibit approximate equivariance (Introduction Section).

1. Residual Pathway Priors for Soft Equivariance Constraints. Finzi et al., NeurIPS 2021.

2. Regularizing Towards Soft Equivariance Under Mixed Symmetries. Kim et al., ICML 2023.

3. A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups. Finzi et al., 2021.

4. Approximately Equivariant Networks for Imperfectly Symmetric Dynamics. Wang et al., ICML 2022.

5. General E(2)-Equivariant Steerable CNNs. Weiler et al., NeurIPS 2019.

Thank you for your response. I am pleased that you have carefully read my questions and taken my suggestions. Now, it seems that the paper has indeed improved. However, I still cannot change my score to a positive opinion for several reasons:

1. We often consider equivariance when focus on some complex systems. In this case, adding a random rotation to the loss once is not enough. The model might be able to witness all possible SO(3) transformations by increasing the number of epochs, but this would also greatly increase the training cost. I recommend that the authors test their algorithms on some complex systems, such as QM9, GEOM-Drug, OC20, which are more concerned 3D systems in the field of computational chemistry.

2. Equivariant models are not actually expensive. For example, the TorchMD-Net model only uses low-degree equivariance $l=1$, and it can achieve good results on both QM9 and MD17. Some equivariant models are expensive since the tensor product on higher-degree spherical harmonic representations is very time-consuming.

3. The comparative experiments in the paper still contain irrationalities: MD17 includes both energy and forces properties data, but I am not sure which one the authors are comparing in Table 2 (if the authors have stated this, I apologize). Why not compare both energy and forces? Additionally, the comparison in this table is also unreasonable: why is the error for EGNN so large (MSE for Benzene is 62.40)? The MD17 dataset is a simple molecular dataset, and it should not perform so poorly. You can look at models like TorchMD-Net, Equiformer, etc., where the results on MD17 are almost saturated. Furthermore, if the authors want to prove that your loss helps reduce the computational cost of equivariant models, you should list some detail results of computational metrics (In Figure 5, we only observe that your method outperforms GATr).

In conclusion, I am sorry that I cannot change my score because I believe this paper does not address the key issues of 3D equivariance (performance and computational cost). Adding a random rotation to the loss is a natural idea, and perhaps the authors can look into whether there are similar works that have tried this approach. However, I acknowledge the author's research direction. I believe that approximate or learnable equivariance is valuable because strict equivariance can lead to a loss of expressive power, and unconstrained neural networks are more likely to learn some complex 3D features. The authors could also design algorithms from the perspective of reducing the training cost of equivariant models. Perhaps modifying the loss could solve the bottleneck in equivariance. I encourage the authors to delve deeper into their research.

• "Furthermore, if the authors want to prove that your loss helps reduce the computational cost of equivariant models, you should list some detail results of computational metrics (In Figure 5, we only observe that your method outperforms GATr)" We thank the reviewer for this point. In Figure 5, we report the wall-clock time for the Geometric Algebra Transformer (GATr) and Transformer architectures. We selected models with an equivalent number of blocks and parameters for a fair comparison (please see Appendix B for details). We measured the computational efficiency of each model by recording the time taken for both forward and backward passes during training, as well as inference time. We show that Transformer can achieve up to $10$ times faster inference speed and $2.5$ times faster training speed compared to GATr. However, although we have included the GATr architecture in our discussion, GATr itself is faster than many equivariant architectures, such as SEGNN and SE(3)-Transformer (please see [14]).

• "Adding a random rotation to the loss is a natural idea, and perhaps the authors can look into whether there are similar works that have tried this approach." We thank the reviewer for this point. However, we already included a related work section with a detailed discussion of the existing literature on approximate equivariance, with all the differences between these approaches and our proposed method. If there are any further questions or specific aspects that we have not addressed, or the reviewer would like to point out, we are happy to provide additional clarifications.

• "The authors could also design algorithms from the perspective of reducing the training cost of equivariant models. Perhaps modifying the loss could solve the bottleneck in equivariance. I encourage the authors to delve deeper into their research." We thank the reviewer for their suggestions. However, we believe that both directions are valuable and have a potential impact. In this work, we focus on comparing existing equivariant models with their unconstrained counterparts across a diverse set of benchmarks. Specifically, we evaluate Transformers, Graph Neural Networks (GNNs), Convolutional Neural Networks (CNNs), and Multi-Layer Perceptrons (MLPs) on four distinct tasks: Dynamical Systems, Motion Capture, Molecular Dynamics, and Jet Flow benchmarks. We aim to provide valuable insights into the performance, scalability, and applicability of unconstrained vs equivariant models across various domains.

1. SchNet: A continuous-filter convolutional neural network for modeling quantum interactions. Schütt et al., NeurIPS 2017.

2. Rotation Invariant Graph Neural Networks using Spin Convolutions. Shuaibi1 et al., 2021.

3. Spherical Message Passing for 3D Graph Networks. Liu et al., ICLR 2022.

4. Symmetry-Informed Geometric Representation for Molecules, Proteins, and Crystalline Materials. Liu et al., NeurIPS 2023.

5. Equivariant graph mechanics networks with constraints. Huang et al., ICLR 2022.

6. EqMotion: Equivariant Multi-Agent Motion Prediction With Invariant Interaction Reasoning. Xu et al., CVPR 2023.

7. Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics. Wu et al., NeurIPS 2023.

8. Equivariant Graph Neural Operator for Modeling 3D Dynamics. Xu et al., ICML 2024.

9. Equivariant Graph Hierarchy-Based Neural Networks. Han et al., NeurIPS 20222.

10. Clifford Group Equivariant Simplicial Message Passing Networks. Liu et al., ICLR 2024.

11. Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs Liao et al., ICLR 2023.

12. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations Liao et al., ICLR 2024.

13. Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products. Luo et al., ICLR 2024.

14. Geometric Algebra Transformer. Brehmer et al., NeurIPS 2023.

We thank the reviewer for their time and for engaging with us in the discussion. We appreciate the reviewer’s acknowledgment that we have taken their suggestions and the improvement of the paper.

• "We often consider equivariance when focus on some complex systems. In this case, adding a random rotation to the loss once is not enough". As we pointed out before, we conduct additional experiments with different numbers of samples from the $SO(3)$ group during training comparing our method and data augmentation. In practice, our new results confirm that we can achieve reasonable performance using fewer samples from the symmetry group. Please see the results in Appendix C.2 in the updated version.

• "The model might be able to witness all possible SO(3) transformations by increasing the number of epochs, but this would also greatly increase the training cost." We would like to point out to the reviewer that we trained all the models, unconstrained and equivariant models, in all the experiments and benchmarks using the same number of epochs. Please see the details in Appendix B.

• "I recommend that the authors test their algorithms on some complex systems, such as QM9, GEOM-Drug, OC20, which are more concerned 3D systems in the field of computational chemistry." We want to clarify to the reviewer that both Motion Capture and Molecular Dynamics (MD17) are widely used 3D benchmarks in the literature. Please see for example [1, 2, 3, 4, 5, 6, 7, 8, 9, 10].

• "Equivariant models are not actually expensive. For example, the TorchMD-Net model only uses low-degree equivariance, and it can achieve good results on both QM9 and MD17". We thank the reviewer for pointing out the example of TorchMD-Net. However, there is no contradiction between that and what we pointed out in our paper. Methods that use spherical harmonics and higher-order tensor products (which are more expressive) are computationally expensive. Using a lower degree could reduce the computational cost but it comes at the trade-off of performance. Please see for example [11, 12, 13].

• "Some equivariant models are expensive since the tensor product on higher-degree spherical harmonic representations is very time-consuming". We appropriate the reviewer's acknowledgment that 'tensor product on higher-degree spherical harmonic representations is very time-consuming'. However, this is what we indicated in our paper. Please see the introduction section.

• "The comparative experiments in the paper still contain irrationalities: MD17 includes both energy and forces properties data, but I am not sure which one the authors are comparing in Table 2 (if the authors have stated this, I apologize). Why not compare both energy and forces?" We would like to clarify to the reviewer that MD17 benchmark has two common tasks in the literature:

  • Invariant Task: Predicting energies given molecular states/ positions (Please see for eg [1, 2, 3, 4]).
  • Equivariant Task: Predicting molecular states/ positions after a specific number of time steps given initial states/ positions (Please see for eg [5, 6, 7, 8]).

In this work, we focus on the equivariant task, following the same previous work on this task. Our primary objective is to compare unconstrained models with their corresponding equivariant versions (GNN vs EGNN).

• "Additionally, the comparison in this table is also unreasonable: why is the error for EGNN so large (MSE for Benzene is 62.40)? The MD17 dataset is a simple molecular dataset, and it should not perform so poorly. You can look at models like TorchMD-Net, Equiformer". We would like to point out to the reviewer that we follow the literature on the equivariant task where the models they suggested focus on the invariant task. However, we are not the first to report the results for EGNN on this task and we follow the same setup in the literature. Please see for example [5, 6, 7, 8]. Additionally, all the training details are reported in Appendix B.3 and we also plan to publish the full code if it has been accepted. In our work, we observe that the optimal performance of each molecule is attained at different values of the penalty parameter $\beta$. For instance, Malonaldehyde exhibits a direct correlation between model performance and equivariance, where a higher penalty parameter $\beta$ yields better performance. For this, EGNN achieved good performance. Conversely, for some molecules, there is a trade-off where the best performance is achieved at a lower value of $\beta$. Please see Section 6.3 in the paper.

We thank the reviewer for their time and all the comments, which have helped us improve our paper and indicate our contribution. We are also open to answer any further questions.

Related Work:

• We thank the reviewer for these points. We have added a new section discussing learning symmetries' approximate equivariance and included comparisons in the related work section.

• We have included all the citations and comparisons pointed out by the reviewer.

• The method in [1] used a canonicalization approach (i.e. additional trainable canonicalization network) to orient the input before the main network. For stable canonicalization, the authors used prior distribution over the symmetry group. However, in our work, we didn't use any additional canonicalization network, and we directly learned equivariance through the unconstrained model.

1. Equivariant adaptation of large pretrained models. Mondal et al., NeurIPS 2023.

$\alpha$ and $\beta$ as hyperparameters:

• We agreed with the reviewer that higher $\frac{\beta}{\alpha}$ leads to more equivariant function. However, this will not indicate if the function can approximate equivariance or not, or to what extent, as this will be specific to the parameters $\alpha$ and $\beta$. In other words, we can define the relative equivariance as a fraction of $\beta$ and $\alpha$, but we also want to measure the level of approximate equivariance the model learns, for this we introduced the equivariance metric.

• $\alpha$ and $\beta$ will be treated as additional parameters that are determined from the validation set. Similar to other parameters in the training procedure (eg learning rate, batch size, number of layers in architecture, etc).

Methodology:

Symmetry group: In this work we consider SE(3) symmetry group (i.e. the group of rotations and translations) and SE(3) equivariant architectures. We have clarified this in the Introduction section of the updated manuscript. We are happy to extend this in future for unknown group (eg Lie algebra convolutional by [1]).

if $G$ is very large. As we focus on SE(3) symmetry group, this already considered a very large/ infinite group. We noticed that empirically our algorithm is sufficient to achieve competitive performance with equivariant baselines using very small number of samples. For example, in the N-Body dynamical system, we train all the models using $100$ samples.

The highest level of equivariance is when $\alpha = 0$ and $\beta = 1$. It is important to note that $\beta$ is not limited to $1$. In our formulation, $\beta$ can take any positive real value. By increasing $\beta$ we can place greater emphasis on enforcing equivariance in the model and reduce the equivariance error (Figure 2 in the paper). We observed that this approach is empirically sufficient to achieve competitive performance with advanced equivariant architectures, such as the Geometric Algebra Transformer, when applied to fully equivariant tasks like the N-body dynamical system.

Which metric is better suited for evaluation: We think both metrics are applicable following our observation in the paper, even for nonscalar functions (eg predicting trajectories of the dynamical system). However, the first measure introduced in Equation 12 theoritaclly could have more stability due to the average over the function $f$.

Motion Capture task: For our algortihm, we sample from the SE(3) symmetry group in all the exeriments and not the subgroup for a fair comparison against the equivariant baselines. Furthermore, we have added new results on the Motion Capture task with a comparison to approximate equivariance baselines. We compare our method against Residual Pathway Priors (RPP) [2], Projection-Based Equivariance Regularizer (PER) [3], and equivariant MLP (EMLP) [4]. As these architectures are designed based on linear layers and MLP, we apply the augmented loss to standard MLP with a similar number of layers and parameters. Our new results confirm the applicability of our method to different architectures, including MLPs, CNNs, GNNs, and Transformers, across a wide range of benchmarks.

Minor errors: We thank the reviewer for these points, we have corrected the points in the updated version.

1. Automatic Symmetry Discovery with Lie Algebra Convolutional Network Dehmamy et al., NeurIPS 2021

2. Residual Pathway Priors for Soft Equivariance Constraints. Finzi et al., NeurIPS 2021.

3. Regularizing Towards Soft Equivariance Under Mixed Symmetries. Kim et al., ICML 2023.

4. A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups. Finzi et al., 2021.

We thank the reviewer for their time and all the comments, which have helped us improve our paper and indicate our contribution. We are also open to answer any further questions.

Novelty, and comparison to approximate equivariance:

• Our method is a general training procedure that could be applied to any unconstrained models, without specific assumption of the design (eg GNN/ CNN/ MLP or Transformers) while most of the work on approximate equivariance focus on relaxing equivariant architectures or regularizing specific layers within such architectures. We have added a discussion on this in the Related Work section with red colors in the updated manuscript.

• The method defined in [1] is a regularization for the linear layers of Equivariant MLP [2], and [3] regularized the CNN architecture. However, in our work, we didn't assume a specific class of models, which increases the applicability of our method to various domains. We have clarified this in the updated version.

Empirical evaluation:

• We have added a new comparison on Motion Capture task: We compare our method against Projection-Based Equivariance Regularizer (PER) [1], equivariant MLP (EMLP) [2], and Residual Pathway Priors (RPP) [4]. As these architectures are designed based on linear layers and MLP, we apply the augmented loss to standard MLP with a similar number of layers and parameters.

• We have included a new task on the Jet Flow dataset used by [5] (a two-dimensional benchmark that captures turbulent velocity fields): We apply our method to a Convolutional neural network (CNN) and compare it with Relaxed Steerable Convolution (RSteer) [5] and E2CNN [6] (more details in Appendix A in the paper).

• Our new results confirm the applicability of our method to different architectures, including MLPs, CNNs, in addition to GNNs, and Transformers, across a wide range of benchmarks.

1. Regularizing Towards Soft Equivariance Under Mixed Symmetries. Kim et al., ICML 2023.

2. A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups. Finzi et al., 2021.

3. A simple rotational equivariance loss for generic convolutional segmentation networks: preliminary results.

4. Residual Pathway Priors for Soft Equivariance Constraints. Finzi et al., NeurIPS 2021.

5. Approximately Equivariant Networks for Imperfectly Symmetric Dynamics. Wang et al., ICML 2022.

6. General E(2)-Equivariant Steerable CNNs. Weiler et al., NeurIPS 2019.

We thank the reviewer for their time and comments, which have helped us improve our paper. We are also open to answer any further questions.

Theoretical motivations and differences from data augmentation:

• Our theoretical motivation is that by having an adaptive parameter $\beta$ on the equivariance loss, we can modulate the extent to which a model exhibits equivariance, depending on the requirements of the task. In Section 4, we introduce a measure that quantifies the level of learned equivariance in the model which we use to analyze our results.

• The original difference from data augmentation is that we utilize an additional controlled equivariance loss together with the objective loss that both minimized during training. We consider two distinct approaches to regulate the penalty parameters $\alpha$ and $\beta$: constant penalty and gradual penalty. For constant penalty, we assign a fixed weight to each task’s loss throughout the training process. In contrast, the gradual penalty dynamically adjusts the weights of each task’s loss during training. For gradual penalty, we use GradNorm algorithm [1], which is particularly suited for tasks that involve simultaneous optimization of multiple loss components, as it dynamically adjusts the weight of each loss during training. We clarified this in Section 3.2 of the paper.

• By minimizing the equivariance loss term simultaneously with the objective function, we can control the equivariance objective depending on the parameter $\beta$. This allows us to systematically adjust the degree of equivariance the model learns (Section 6.1 in the paper).

• We noticed this decomposition is important to control the trade-off between equivariance and performance for multiple tasks (Motion Capture in Sections 6.2 and Molecular Dynamics in Section 6.3 of the paper)

• Data augmentation can also be viewed as a special case of our method with $\alpha = 0$ and $\beta = 1$. We clarified this in Section 3.3 of the paper.

Comparing to state of the art:

In this work, we consider Transformers and Graph Neural Networks (GNNs), with their equivariant versions, as our main baselines.

• We compared Transformers against SE(3)-Transformer [2], and Geometric Algebra Transformer [3] which is a recent equivariant architecture for geometric data (Sections 6.2 and 6.3 in the paper).

• We consider another comparison between GNN and EGNN for molecular dynamics tasks (Section 6.3).

• While we have added new results on MLPs and CNNs architectures, we think this will be a future direction to apply our augmented loss to a broader range of unconstrained models.

We sincerely hope that we have addressed the concerns of the reviewer satisfactorily in the revised version and would kindly ask the reviewer to update their score accordingly.

1. GradNorm: Gradient Normalization for Adaptive Loss Balancing in Deep Multitask Networks. Chen et al., ICML 2018.

2. SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks. Fuchs et al., NeurIPS 2020.

3. Geometric Algebra Transformer. Brehmer et al., NeurIPS 2023.