Smooth, exact rotational symmetrization for deep learning on point clouds

Point clouds are versatile representations of 3D objects and have found widespread application in science and engineering. Many successful deep-learning models have been proposed that use them as input. The domain of chemical and materials modeling is especially challenging because exact compliance with physical constraints is highly desirable for a model to be usable in practice. These constraints include smoothness and invariance with respect to translations, rotations, and permutations of identical atoms. If these requirements are not rigorously fulfilled, atomistic simulations might lead to absurd outcomes even if the model has excellent accuracy. Consequently, dedicated architectures, which achieve invariance by restricting their design space, have been developed. General-purpose point-cloud models are more varied but often disregard rotational symmetry. We propose a general symmetrization method that adds rotational equivariance to any given model while preserving all the other requirements. Our approach simplifies the development of better atomic-scale machine-learning schemes by relaxing the constraints on the design space and making it possible to incorporate ideas that proved effective in other domains. We demonstrate this idea by introducing the Point Edge Transformer (PET) architecture, which is not intrinsically equivariant but achieves state-of-the-art performance on several benchmark datasets of molecules and solids. A-posteriori application of our general protocol makes PET exactly equivariant, with minimal changes to its accuracy.

Paper

Similar papers

Peer review

Reviewer oCKF4/10 · confidence 3/52023-07-02

Summary

The authors propose a method for uniting rotational symmetries with translational and permutation symmetries in neural networks for point cloud data. Hereby, the authors particularly focus on molecule data. This is achieved by adding rotational augmentations at training time and evaluating the resulting NN on an ensemble of coordinate systems at test time.

Strengths

- The incorporation of rotational symmetries into NNs applied to point data seems very important to the application of DL to atomic data as it may reduce the data hunger of NNs and improve generalization. - The approach using an ensemble of coordinate systems seems innovative and makes sense.

Weaknesses

- What is the typical number of coordinate systems in eq 3? Since each coordinate system is defined by three atoms, I guess it becomes very large very quickly. - I wonder whether using y_s is even necessary. It seems to greatly increase the runtime at inference while not really adding anything in terms of performance. I.e. Simple training with rotational augmentation seems to yield all of the performance gains. Could the authors elaborate on this? I think it would be necessary to clearly demonstrate the usefulness of this addition - Are competing methods also using rotational augmentations? - -Do the authors have any idea why the SOTA in tab 1 is so much worse on MnO & HEA?

Questions

see weaknesses

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

N/A

Reviewer QskA6/10 · confidence 3/52023-07-03

Summary

This paper presents a novel posteriori rotational equivariant protocol that can be applied to any existing 3D models, irrespective of their initial equivariance to rotations. This protocol ensures not only rotational equivariance but also the preservation of other important model properties, such as permutation equivariance and expressiveness. Additionally, the paper introduces the Point Edge Transformer architecture, which achieves state-of-the-art performance on several benchmarks. The general equivariant design results in a modest performance improvement, highlighting the fact that exact equivariance is not always essential in certain learning tasks.

Strengths

1. This paper alleviates the need to incorporate the rotational equivariance into the model but designs a general method on the top of the trained $3D$ model. The method is proven to be strictly equivariant to 3D rotations, allowing for the utilization of different non-equivariant 3D models and enabling broad applications across various 3D tasks. 2. Through experimental analysis, this paper demonstrates that exact equivariance of the model is not always necessary for certain atomic applications. They show that by employing a powerful model and rotation augmentation techniques, non-intrinsically equivariant models outperform their counterparts. This finding provides valuable insights into the essential factors to consider when designing models for atomic applications. 3. This paper is well-written, and the experimental results effectively support the authors' argument.

Weaknesses

1. The proposed equivariant protocol achieves equivariance but comes with increased computational requirements, particularly in scenarios with large sample sizes. Moreover, exact rotation equivariance may be affected by the presence of noise in the data, thus introducing potential sensitivity issues. 2. The proposed model may not applied to the 3D models like pointnet and pointnet++, which doesn't include the interaction of information between the points. 3. As mentioned in the paper, for GNN network, their method either requires more computation to go over the whole neighboring points, or involves the redesign of the message-passing mechanism.

Questions

I'm not familiar with the topic of atomic-scale applications. Therefore I will ask more questions about the equivariance part. 1. As stated in the limitation, I think the equivariant design might be sensitive the noise in the sample, please correct me if I'm wrong. Could the authors provide some experiments to validate its robustness? 2. Could the model be applied in large point cloud application such as shape segmentation or scene prediction in computer vision? I'm curious the scalabilty of the model and more applications of the model. 3. Since the authors show that in some tasks the equivariance is not the essential part. Could the authors compare their equivariant model with the other equivariant models in some tasks where the equivariance ideed is important, for example, missing atom prediction in tensor field networks?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors adequately addressed the limitations.

Reviewer mhUe5/10 · confidence 3/52023-07-05

Summary

This paper presents a novel protocol for general symmetrization that adds rotational equivariance to any given 3D point-cloud processing model while preserving symmetries with respect to translations and permutations. Applying the protocol to a model a-posteriori (i.e., at inference) makes the model rigorously rotation-equivariant. To demonstrate the potential use of the protocol, the authors present Point Edge Transformer (PET). Albeit not intrinsically equivariant, it achieves state-of-the-art performance on several benchmarks of molecules and solids. --------------------------------------- POST REBUTTAL Given the clarification the authors provided by the rebuttal, as well as the additional demonstration of the robustness of their proposed method wrt. noise, I have raised my score to 5.

Strengths

S1. The proposed idea of training a model that is intrinsically non-SO(3)-equivariant and making it equivariant during inference is novel, and its mathematical presentation appears correct. S2. The proposed experimental evaluation is extensive and substantial.

Weaknesses

W1. Contributions: The main contribution the authors claim is the equivariance scheme. Unfortunately, there is not much in the article to support this. Instead, the bulk of the improvement appears to come from the PET model, which is already sufficiently equivariant across multiple tasks. It is unclear what benefits smoothness gave in practice (I might have missed something in the experiment part). W2. Presentation: The presentation is a bit convoluted and seems too verbose; the provided Appendix contains many important details, e.g., about the PET architecture, that should instead be in the main paper. Furthermore, I had a hard time finding a precise definition of the problem in the paper. The paper addresses an important question of rotational equi- and invariance, but its vague presentation makes it difficult to appreciate the contributions. W3. Related work: Somewhat lacking related work on geometric deep learning. There is a detailed overview of related work on GNNs, but almost no discussion about group equivariant CNNs (G-CNNs). It is unclear how different the proposed model is and if the proposed protocol could be applied to such models.

Questions

Q1. In the abstract, it says (line 11), "We propose a general symmetrization protocol that adds rotational equivariance to any given model". However, the symmetrization seems to assume a base model based on the atom-centric local coordinate system. How common is this? Does this assumption cover the related work, or are there methods that do not use this? Q2. In the analysis (F.8), it is said that one can improve performance further by introducing extra rotations at inference (not on input but on extracted coordinate systems, so the model is still technically invariant/equivariant). I would like to know why this improves performance. If I understood correctly from the training description, it is done in a non-equivariant way with rotational augmentation. Is this the reason for the improvement? Q3. Given that the most common metric is MAE, have the authors tried using MAE as training loss? Q4. Figure 2: $\bar{x}$ seems to refer to edges between atoms in $A_i$, but as far as I can see, $x$ has not been introduced before. The same applies to $E$ and $M$ and $H$. These terms should be clearly introduced in the text before the figure, not in the Appendix. Q5. It appears that **r**$_i$ is the 3D position per atom, and $A_i$ is a group of atoms centered around **r**$_i$. However, it was unclear how features per atom are extracted from the beginning. Are different atoms treated as one-hot embeddings?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

The authors adequately address the limitations and broader impact.

Reviewer r8ta5/10 · confidence 2/52023-07-07

Summary

The paper suggests a network design tackling chemical and materials modeling. The paper identifies rotational symmetries as lacking in existing solutions. The suggested approach is based on symmetrization, i.e., projecting a general function to be equivariant by averaging it over aligning rotations. In addition, a transformer-based architecture is proposed.

Strengths

The idea to enforce smoothness with other symmetries such as rotational symmetries seems to be valuable, and indeed missing from existing solutions.

Weaknesses

The main weaknesses of the paper are related to its level of readability, the quality of its exposition, and its mathematical formulation. In what follows, I will share some relevant examples. It is understood that the selected approach to incorporate rotation equivariance is symmetrization. However, the text does not make it clear what are the challenges in doing so and why the starting point chosen is using frames from pairs. If that is the starting point, it is not entirely clear what is the next point. In addition, it is not shown why this symmetrization is indeed equivariant. In addition, there so no clear and simple mathematical formulation of how the aligning rotations for the symmetrization are selected. Furthermore, all the input and output dimensions are missing from all the functions stated in the paper. Figure 2 is hard to decipher. It is hard to follow the experiment section. I would expect to see some ablation regarding the symmetrization technique. I would expect to see the symmetrization tested with other architectures than the transformer as well. Furthermore, it seems that symmetrization has a minor effect on the model’s generalization capacity, making the benefit from it, given its computational budget increase, questionable. Lastly, the contribution of the proposed transformer architecture seems to be orthogonal to the challenge the paper set to tackle.

Questions

I would appreciate any response from the authors regarding the weakness stated above.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Yes

Reviewer mhUe2023-08-12

I thank the authors for their detailed rebuttal! If I understand it correctly, the main argument the authors make is that equivariance per se is important in many application areas, especially atomistic modeling. In this case, the experiments should have a different focus at the beginning: namely, demonstrating how exact the equivariance attained with the proposed protocol is. Given the clarification the authors provided by the rebuttal, as well as the additional demonstration of the robustness of their proposed method wrt. noise, I'm willing to raise my score to 5.

Authorsrebuttal2023-08-14

Exact equivariance and choice of benchmarks

We thank the reviewer for reading our response and for the further comments! **Exact equivariance** The reviewer is correct; equivariance *per se* is considered essential and is almost universally enforced in the field. We did not include numerical tests of equivariance because the algorithm is designed to make predictions *exactly* equivariant, up to machine precision. However, following our discussion with the Reviewers, we agree it would be worthwhile to add a numerical demonstration similar to the one we did in our responses, as a validation that our implementation is correct. We will do this upon revision. **Choice of benchmarks** The reason we dedicated our experiment section to assessing the performance of our proposed model, PET, on various datasets pertains to a common (mis)belief in the community of atomistic machine learning that exact intrinsic equivariance is not only required for physical consistency but is also necessary to achieve good accuracy. Thus, without such an exhaustive demonstration of PET's excellent and often superior performance, many researchers from this community would doubt the usefulness of our symmetrization method, arguing that there is no reason to fit and next symmetrize non-equivariant models, given that intrinsically equivariant frameworks are already available. With such a demonstration, we have provided strong indications of the opposite, that it is possible to construct better models within an enlarged design space without the burden of incorporating rotational equivariance intrinsically.

Reviewer oCKF2023-08-13

I thank the authors for addressing my concerns sufficiently. I am leaning towards increasing my score. This will also depend on discussions with other reviewers.

Authorsrebuttal2023-08-14

We thank the reviewer oCKF for reading our responses. We remain available for further clarifications, if needed.

Reviewer r8ta2023-08-21

rebuttal

I thank the authors for a detailed rebuttal. Some of my concerns were addressed, specifically, the one clarifying that the application in hand requires equivariant/consistency in predictions. I will raise my score to 5. Thank you.

Authorsrebuttal2023-08-21

We thank the reviewer for reading our response!

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC