Learning Symmetries through Loss Landscape

Authors: Ahmed A. A. Elhag, T. Konstantin Rusch, Francesco Di Giovanni, Michael M. Bronstein
Primary area: learning on graphs and other geometries & topologies
Keywords: Unconstrained models, equivariant models, symmetries.

Abstract:
Incorporating equivariance as an inductive bias into deep learning architectures, to take advantage of the data symmetry, has been successful in multiple applications such as chemistry and dynamical systems. The build of equivariance architecture, particularly w.r.t. roto-translations, is crucial for effectively modeling geometric graphs and molecules, where the understanding of 3D structures enhances generalization. However, despite their potential, equivariant models often pose challenges due to their high computational complexity. In this paper, we study the capabilities of unconstrained models (which do not build equivariance into the architecture) and how they generalize compared to equivariant models. We show that unconstrained models can learn approximate symmetries by minimizing additional simple equivariance loss. By formulating equivariance as a new learning objective, we can control the level of approximate equivariance in the model. Our method achieves competitive performance compared to equivariant baselines while being 10x faster at inference and 2.5x at training.

Decision: reject
Reviewer scores: 5, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730120237756

thread_type:
official_review

<thread_object>
creator: reviewer_jqbk

<content>
summary:
In this paper, the authors analyze the application of unconstrained models in equivariant tasks by conducting a comprehensive analysis of unconstrained models, comparing their performance and computational efficiency against equivariant models. Besided, the authors introduce a novel, simple loss function that enables these models to approximate symmetries, which can be optimized during training.

strengths:
Relaxing equivariance is a valuable research direction that can break through the constraints on generalization or expressive power caused by strictly equivariant operations.

weaknesses:
1. The primary concern is the authors' motivation. The idea of using group transformations for data augmentation is native, but for many equivariant tasks, it is challenging to obtain a general model through data sampling due to the bias introduced by limited sampling. For instance, for point clouds or molecules, sampling across all angles would expand the dataset by hundreds of times and still struggle to enable the model to effectively learn fine-grained rotation equivariance. I suggest the authors validate their approach on common 3D datasets such as QM9 or ModelNet40.

2. The authors base their introduction in the first three sections on general equivariance, yet the impact of different equivariance groups on algorithms varies. For example, permutation equivariance and translation equivariance can be directly covered by simple operations, making the paper's method inapplicable. The authors should specify which equivariant tasks their method focus on.

3. Relaxing equivariance is a widely discussed topic, and the authors lack relevant citations and analysis [1] [2] [3] [4]. Moreover, the main advantage of unconstrained models is their ability to learn more complex features. It is worth noting that strictly equivariant operations can limit the expressive power of GNNs [5] [6], but unconstrained models may surpass these limitations. Additionally, some tasks are not strictly equivariant, allowing unconstrained models to be applicable. The authors' emphasis on the lower computational complexity of unconstrained operations is incorrect. In the 3D domain, models like torchmd are strictly equivariant yet have low complexity.

    [1] Residual pathway priors for soft equivariance constraints, Finzi, et al.

    [2] Approximately equivariant networks for imperfectly symmetric dynamics. Wang, et al.

    [3] Relaxing equivariance constraints with non-stationary continuous filters. van der Ouderaa, et al.

    [4] Learning Partial Equivariances from Data. Romero, et al.

    [5] On the Universality of Rotation Equivariant Point Cloud Networks. Nadav Dym, Haggai Maron. 

    [6] On the Expressive Power of Geometric Graph Neural Networks. Chaitanya K. Joshi, Cristian Bodnar, Simon V. Mathis, Taco Cohen, Pietro Liò.

4. I do not understand how the loss surface in Figure 1 was created and why it demonstrates the advantages of unconstrained models.

5. There are numerous issues with the paper's presentation:

    (a) The equations in lines 218, 224, and 227 lack numbering.

    (b) In line 215, the definition of G is finite, which is problematic for integrals where the group size can be infinite. Most groups mentioned in the paper are infinite, and I do not understand why the authors restrict groups to being finite in their initial definition.

    (c) All the references have formatting issues because none of them specify the source of the papers. For instance, "Equivariant Graph Hierarchy-Based Neural Networks" in your paper is accepted in NeurIPS 2022, not arxiv.

    (d) Appendix B is incomplete; several titles are clustered together without any explanatory text.

questions:
See weekness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730648015113

thread_type:
official_review

<thread_object>
creator: reviewer_eqn6

<content>
summary:
The paper attempts to build equivariance into unconstrained models by posing equivariance as a constrained optimization problem, which can, in turn, also control the level of approximate equivariance in the models. The authors demonstrate results in N-body dynamical systems, motion capture, and molecule dynamics, and they analyze the effect of the level of approximate equivariance on task performance.

strengths:
- The paper is well-motivated and clearly written (particularly the sections on background and methods).
- The limitation section discusses an important limitation of the interplay between optimization paths and loss landscape.
- The experiments are conducted in different domains and examine several essential aspects of the algorithm, giving more insights into the method and how levels of equivariance can affect downstream task performance.

weaknesses:
- **Related work**:
   - Although the paper uses equivariance as a constrained optimization problem and discusses it in the context of unconstrained models, it misses several crucial relevant works. Discussion of these works would help to place the submission in the literature and give a view of how this work differs from and compares to existing works.
  - Learning equivariance from data [1, 2], approximate/soft equivariance [3, 4, 5, 6], equivariance as a constrained optimization problem [7, 8], and equivariance with unconstrained models [9, 10, 11, 12].
  - Can the authors highlight the differences from Sec 3.1 and Sec. 3.2 of [10]?

- **$\beta$ and $\alpha$ as hyperparameters**:
  - The authors suggest that the level or extent of equivariance can be controlled with $\beta$ and $\alpha$ - is there a formal way to define this "level" of equivariance or is it an intuition tied to the loss itself, i.e., higher $\frac{\beta}{\alpha}$ indicates more equivariant? 
  - Next, how would someone know the optimal level of equivariance while using your proposed algorithm - $\beta$ is not learned, and the results indicate that the optimal $\beta$ can be identified from the test data results, which is not ideal. Rephrasing this, how do you know how much equivariance is required for the task, and thus what values of $\alpha$ and $\beta$ to set?


- **Methodology**:
  - How will your algorithm work if group $G$ is unknown?
  - How can your method reasonably approximate equivariance if $G$ is very large and the duration of training is not enough?
  - The highest level of equivariance is when $\alpha = 0$ and $\beta=1$. However, this is equivalent to data augmentation, which does not guarantee exact equivariance. Can your algorithm guarantee exact equivariance?
  - While the trends are consistent for both metrics, as reported in the paper, it might be helpful to discuss which metric - Eq. 9 or Eq. 10 is better suited for evaluation. How does Equation 9 work (or make sense) when $f(x)$ is non-scalar?
  - For Motion Capture, if the symmetry constraints are already known, instead of complete SE(3) equivariant baselines, why didn't the authors select appropriate equivariant models that are equivariant to the required SE(3) subgroup or consider symmetry breaking [14, 15]? What $G$ did your algorithm use? If it is the subgroup of SE(3), then it is an unfair comparison.

- **Minor spelling errors**:
  - L156 "requiring equivariant into" should be "requiring equivariance in"
  - L396 "it is" should be "its"


**References**:
1. Equivariance Discovery by Learned Parameter-Sharing. Yeh et al., AISTATS 2022.
2. Learning Equivariances and Partial Equivariances from Data. Romero et al., NeurIPS 2022.
3. Learning Layer-wise Equivariances Automatically using Gradients. Ouderaa et al., 2023.
4. Residual Pathway Priors for Soft Equivariance Constraints. Finzi et al., NeurIPS 2021.
5. Almost Equivariance via Lie Algebra Convolutions. McNeela et al., 2024.
6. Regularizing Towards Soft Equivariance Under Mixed Symmetries. Kim et al., ICML 2023.
7. Improved Canonicalization for Model Agnostic Equivariance. Panigrahi et al., 2024.
8. Structuring Representations Using Group Invariants. Shakerinava et al., NeurIPS 2022.
9. Equivariance with Learned Canonicalization Functions. Kaba et al., ICML 2023.
10. Equivariant adaptation of large pretrained models. Mondal et al., NeurIPS 2023.
11. Equi-Tuning: Group Equivariant Fine-Tuning of Pretrained Models. Basu et al., AAAI 2023.
12. Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance. Kim et al., NeurIPS 2023.
13. Steerable Equivariant Representation Learning. Bhardwaj et al., 2023
14. Symmetry breaking and equivariant neural networks. Kaba et al., NeurIPS NeuReps workshop 2023
15. Improving Equivariant Networks with Probabilistic Symmetry Breaking. Lawrence et al., ICML GRaM workshop 2024.

questions:
- How will your algorithm fare when there are data constraints? Equivariant models are inherently data efficient, but your algorithm does not seem to be.
- The loss landscape plots depend on the selected directions - so how can we infer from just two random directions that the loss landscape is better for Transformers or GATr? The optimization paths should affect, and although it is discussed to some extent in limitations, it would be better if there is more discussion on this.

Most of the other questions I had are listed in the Weaknesses section. I will be happy to improve the score if the authors address the questions and weaknesses with supportive evidence during the discussion phase.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730673148347

thread_type:
official_review

<thread_object>
creator: reviewer_jguc

<content>
summary:
The paper introduces an augmented loss function that incorporates a measure of average equivariance, aiming to enhance the prediction of approximately equivariant information. The authors validate this approach on both equivariant and non-equivariant tasks, utilizing transformer and graph neural network architectures. Furthermore, they conduct a visual analysis of the loss landscape, comparing two different architectures: Transformer with the augmented loss and GATr without.

strengths:
1. The augmented loss function is generalizable across various architectures.

2. The augmented loss function requires relatively few samples to work effectively making it computationally efficient.

weaknesses:
The proposed methodology can be described as approximate equivariance, but the paper lacks an adequate background and comparative analysis against existing works on approximate equivariance. This raises concerns about both the novelty and the empirical validation of the approach.

1.  Novelty: Augmented loss functions enforcing approximate equivariance have been studied (e.g. [1]) including an average measure (e.g. [2]).

2. Empirical Support:  The paper does not benchmark against other methods that address approximate equivariance (e.g., [1]), nor does it consider theoretically grounded approaches to symmetry breaking (e.g., [3], [4]) or simpler strategies like combining SE3Transformer with MLPs.

[1] Kim, Hyunsu, Hyungi Lee, Hongseok Yang, and Juho Lee. "Regularizing towards Soft Equivariance under Mixed Symmetries." Proceedings of the 40th International Conference on Machine Learning, ICML'23, 2023, pp. 686, JMLR.org.

[2]  K. Lin, B. Huang, L. M. Collins, K. Bradbury and J. M. Malof, "A simple rotational equivariance loss for generic convolutional segmentation networks: preliminary results," IGARSS 2019 - 2019 IEEE International Geoscience and Remote Sensing Symposium, Yokohama, Japan, 2019.

[3] Wang, Rui, Robin Walters, and Tess Smidt. "Relaxed Octahedral Group Convolution for Learning Symmetry Breaking in 3D Physical Systems." NeurIPS 2023 AI for Science Workshop, 2023, https://openreview.net/forum?id=B8EpSHEp9j.

[4] Lawrence, Hannah, Vasco Portilheiro, Yan Zhang, and Sékou-Oumar Kaba. "Improving Equivariant Networks with Probabilistic Symmetry Breaking." ICML 2024 Workshop on Geometry-grounded Representation Learning and Generative Modeling, 2024, https://openreview.net/forum?id=1VlRaXNMWO.

questions:
1. To address fundamental concerns, I recommend a more comprehensive background and analysis, alongside broader empirical comparisons. Specifically, the study should include a wider range of architectures that go beyond strictly equivariant and non-equivariant models, particularly for the motion capture task, which is inherently non-equivariant.

2. In practice, if a single randomly sampled group element is used per sample in each training step (as mentioned at the end of Section 3.1), this should be explicitly stated in Section 3.2 where the sampling procedure is discussed and the number of samples M is introduced.

3.  The paper lacks heuristic, theoretical and/or empirical justification for the choice of one group element per sample per training epoch sufficient. As a result the particular choice of measure is understudied.  Moreover, it remains unclear how the equivariance error is measured in Section 6. How many samples are used for this computation?

4.  It's entirely unclear if the difference in the loss landscape is a result of the augmented loss function or a result of the architectural difference. I would strongly recommend performing comparisons with a fixed architecture.

5. The MD17 dataset includes two regression targets: energy and forces. Please clarify in the text which target is being used (likely force regression) and how the results are generated.

6. Near the end of Section 6.2, the statement "Best performance is observed at an intermediate level of equivariance..." is confusing. Since the paper modifies the loss function, not the architecture, this needs further explanation for supporting the proposed methodology. Otherwise, the conclusion is simply to not utilize strict equivariant architectures for non-equivariant tasks.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730706900420

thread_type:
official_review

<thread_object>
creator: reviewer_eyak

<content>
summary:
The paper investigates unconstrained models for handling data symmetry. The authors demonstrate that by designing a loss function specifically tailored for learning equivariance, unconstrained models can approximate data symmetry by minimizing this equivariant loss. This approach allows the models to efficiently control the level of equivariance while maintaining flexibility.

strengths:
- The paper is well-written and presents the core ideas in a clear and accessible manner.
- Using a "landscape" to describe the benefits of unconstrained models is particularly novel and insightful.

weaknesses:
- The study introduces this equivariant loss without providing a strong theoretical foundation for the proposed approach. Also, the proposed method is quite straightforward, and its distinction from data augmentation is unclear. It essentially computes the loss on a larger augmented dataset by sampling transformed data.   I suspect this method is already well-known within the community, which limits the novelty of the contribution.

- The experimental comparisons are performed on a limited set of classic models rather than state-of-the-art models, raising concerns about the practical applicability of the method to more advanced techniques.

questions:
Please address my concerns in the weakness part, especially the novelty of the proposed method and the theoretical foundations.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732237306480

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers and ACs

comment:
We thank all the reviewers for their time and valuable feedback, which have helped improve our paper and confirm our contributions.

The reviewers have highlighted several strengths of our work:

- eYaK: "The paper is well-written and presents the core ideas in a clear and
    accessible manner. Using a "landscape" to describe the benefits of unconstrained models is particularly novel and insightful."
- JGuC: "The augmented loss function is generalizable across various architectures.
    The augmented loss function requires relatively few samples to work effectively making it computationally efficient."
- EQn6: "The experiments are conducted in different domains and examine several
    essential aspects of the algorithm, giving more insights into the method
    and how levels of equivariance can affect downstream task performance."
- JQbk: "Relaxing equivariance is a valuable research direction that can break
    through the constraints on generalization or expressive power caused by
    strictly equivariant operations."

Our work considers the current active topic on learning symmetries in unconstrained models versus constrained equivariant models. We consider two general architectures, Graph Neural Networks (GNNs) and Transformers, along with their equivariant versions. We test our approach on three different tasks: N-body dynamical system, Motion Capture, and Molecular Dynamics.

However, we found that most of the concerns mentioned by reviewers focus on additional evaluation and analysis compared to the prior work on approximate equivariance. Specifically, there are concerns about the Motion Capture task not being fully E(3) equivariance. We have updated the manuscript (changes highlighted in red color) following their feedback. The updated version includes:

- A new section discussing approximate equivariance and included comparisons to relevant prior work.

- A new comparison on 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. 

**Table: Performance on Motion Capture dataset (MSE × 10⁻²)**

|            | EMLP | RPP | PER | MLP  | Data Augment. | Ours  
|----------------------|------------|--------------|------------|--------------|------------|--------------|
| **Walking (Subject #35)**          | 7.01 ± 0.46      | 6.99 ± 0.21      | 7.48 ± 0.39        | 6.80 ± 0.18    | 6.37 ± 0.04         | **6.04 ± 0.09**     |
| **Running (Subject #9)**           | 57.38 ± 8.39      | 34.18 ± 2.00      | 33.03 ± 0.37       | 39.56 ± 2.25    | 40.23 ± 0.94       | **32.57 ± 1.47**   |

- 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 Convolutional neural network (CNN) and compare it with Relaxed Steerable Convolution (RSteer) [4] and E2CNN [5] (more details in Appendix A in the updated version).


**Table: Performance on Jet Flow dataset (RMSE)**

| Model   | Future               | Domain               |
|---------|----------------------|----------------------|
| E2CNN   | 0.21 ± 0.02          | 0.27 ± 0.03          |
| RSteer  | *0.17 ± 0.01*        | **0.16 ± 0.01**      |
| Ours    | **0.16 ± 0.003**     | *0.18 ± 0.003*       |


Our new results confirm the applicability of our method to different architectures, including MLPs, CNNs, GNNs, and Transformers, across a wide range of benchmarks. We thank the reviewers again for their valuable feedback and kindly ask them to consider increasing their scores if that addresses their concerns. 



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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732511325789

thread_type:
official_comment

<thread_object>
creator: area_chair_3gmh

<content>
comment:
Dear reviewers,

If you haven’t done so already, please engage in the discussion as soon as possible. Specifically, please acknowledge that you have thoroughly reviewed the authors' rebuttal and indicate whether your concerns have been adequately addressed. Your input during this critical phase is essential—not only for the authors but also for your fellow reviewers and the Area Chair—to ensure a fair evaluation.
Best wishes,
AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a framework for learning approximate symmetries from data by starting with a neural network that is not inherently equivariant and training it with an equivariance regularizer. The authors made commendable efforts to address several concerns raised by the reviewers, such as the lack of comparisons with key existing works and the need to demonstrate the framework's effectiveness on large-scale, complex systems. However, despite these improvements, the reviewers remain unenthusiastic about the paper. Their main reservations lie in its limited novelty (similar approaches, such as Kim et al., 2023, have already been explored, albeit with a focus on linear layers) and its restricted applicability (e.g., requiring multiple training runs with varying $\alpha$ and $\beta$ values). The AC encourages the authors to further refine this research direction and address the outstanding issues regarding the framework's applicability and practicality.