Equivariant Neural Operator Learning with Graphon Convolution

Summary

This paper proposes an architecture for learning mappings between signals/functions, which is equivariant to rigid body transformations. The architecture is based on a spectral graph convolutional network based on a novel design of a graphon/graph shift operator (or sequence of graphons/graph shift operators).

Strengths

The construction of the graphon/graph shift operator underlying the spectral graph convolutional network is motivated by geometric and physical considerations. The method seems to achieve good results in experiments.

Weaknesses

The construction of spectral convolution for graphons is well known, and not a contribution of this paper. See for example works by Alejandro Ribeiro, Luana Ruiz, Ron Levie, Gitta Kutyniok, and Nicolas Keriven. The exposition and mathematical formulations are often not clear and inaccurate. It is hence hard to follow the paper. For example, the description about the construction of the graphon is not clear. Is there a single graphon, or is there a different graphon for each pair of basis elements? Moreover, the discussion around bases lacks mathematical rigor. There seems to be a confusion between SE(3) and SO(3). Unless I am missing something, the theorems show equivariance with respect to SO(3), but the abstract and introduction talk about SE(3). See more details in “questions.” The exposition of the method and its motivation could be explained better. What the paper does is define a good graphon (or sequence of graphons) for problems (primarily) in quantum chemistry, and uses this graphon (or this sequence of graphons) in spectral graph convolutional networks. It seems like this paper has potential, but it is not ready for publication at its current state.

Questions

Paragraph of line 88: Note that SE(3) also includes translations, not only rotations. Perhaps you mean to develop the theory for SO(3)? Line 99 : “the message passing must be expressed as an integral”. You define the aggregation here, not a full message passing layer. There are already papers that defined such a graphon aggregation, as I wrote above. Line 104: “approximate the integral with Monte Carlo estimation over all grid points” - Monte Carlo methods are based on random points, not grid points, and are not affected by the curse of dimensionality directly. Line 105: “Instead, we follow a similar idea in the coefficient learning … to define a set of complete basis functions” – basis expansion does suffer from the curse of dimensionality. So it seems like the motivating text is misleading. Line 115: “To address the limitation mentioned in the previous subsection, we leverage the discrete graph structure to build…” Which graph? You did not discuss a graph in connection to the domain D until now. You did not discuss the setting properly. What is the data structure/discretization that you work with? Equation (4): In what sense is this a basis of the Hilbert space $L^2(R^3)$? It’s not an orthogonal basis or a frame of $L^2(R^3)$. It is also not a Hamel basis. It should be written more rigorously what type of expansions you consider. Perhaps instead of the term “basis” you should say a complete system (if you use a complete system), a representation system, or series expansion. Moreover, in line 132 you write “equivariant basis”. In what sense is it equivariant and with respect to which group representation? In the text before, you discussed SE(3), but spherical harmonics have nice equivariance properties with respect to SO(3), not SE(3). You should define what you mean by equivariant basis, probably meaning that the change of basis operator is an intertwining operator of two representations of SO(3). Line 142: What is S? Did you forget the pair of indices? In what sense can S be written in the representation system of (4)? This is not trivial as (4) is not an orthonormal basis or a frame. Perhaps you could show that (4) is complete (?), and then you can approximate S using a finite sum of elements of (4), but can you get an exact equation with a series expansion? This should be shown or explained. Line 146: spherical harmonics are not irreducible representations. A representation is a mapping from a group to operators. The spherical harmonics provide a basis that can span the invariant subspaces of the rotation representation of SO(3). Formula (5): you wrote an implicit equation, which is probably not what you meant. The left-hand-side should be the output feature, which is different from the input feature. Formula (5): you are using notations that you have not defined. What is $\phi_J^{lk}$? Note that before, you missed the indices of S, so $\phi$ should have more indices. Where is the learnable part of this message passing scheme? This seems like a definition of aggregation or shift operator. Section 4: spectral graphon GNN is a well-known construction. Moreover, the model is not clear. Is there a single graphon underlying the method, or is there a different graphon for each pair of basis elements?

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

2 fair

Presentation

1 poor

Contribution

3 good

Limitations

Limitations are discussed.

Summary

This work presents a new model architecture that combines the coefficient learning scheme with coordinate-based residual operator layer for learning mapping between continuous function in 3D Euclidean space. The authors further show the proposed model can be interpreted as graph convolution on a graphon. Experiments on various datasets show the effectiveness of the proposed model.

Strengths

1. Generalizing neural operator learning to 3D Euclidean space with equivariance guarantee is an interesting and underexplored research question. 2. The authors provide theoretical explanation to the model by generalizing graph convolution to graphons. 3. The proposed method performs better than other neural operators in considered datasets.

Weaknesses

1. The paper is generally hard to follow as many parts are not clearly stated or explained. For example: - Why the target feature function $\rho$ can be expressed as the equation in line 117? What is the rationale behind the equation? - Why the basis function is constructed as in equation 4? How it is connected with atomic orbitals? Is such a construction necessary for achieving equivariance and good empirical performance? - How is the spherical tensor in section 3.3 connected with the basis function in the construction of $S$? - What is the motivation of using residual operator layer? Why it can mitigate the expressivity issue? 2. Following above, the presentation could be significantly improved by giving some preliminary introductions of the key technical parts used in the method, e.g. coefficient learning, concepts and background in quantum chemistry, TFN. 3. Since the method appears to be a mixture of existing approaches from different areas, it is unclear what is the novelty of this work compared with prior art on methodology aspect. 4. The authors claim the proposed neural operator to be a general framework for learning mapping between functions but only test on electron density datasets. The chosen baselines are somewhat outdated, most mainstream equivariant GNNs e.g. [1] are not considered. [1] E (n) equivariant graph neural networks

Questions

See the weakness section.

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

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

There is no negative societal impact as far as I can tell.

Summary

This paper presents a novel framework for learning mappings between continuous functions in the 3D Euclidean space. First, the framework combines the coefficient learning scheme and the residual operator layer to maintain the sensitivity of the model to equivariance. In addition, the proposed approach utilizes both continuous and discrete graph structures of the input data to efficiently capture geometric information. Experimental results on a large-scale electron density of datasets show that the proposed model outperforms the current state-of-the-art architectures.

Strengths

This paper proposes a method to ensure that the model preserves equivariance by introducing a residual operator layer in the coefficient learning framework. The paper discusses related work, outlines the approach, presents experimental results, and concludes with the potential implications of the proposed approach. The strengths of this paper are the focus on the 3D Euclidean space problem, the innovation of the equivariance techniques, and the emphasis on discrete structural information. In terms of contributions, this paper introduces a new equivariant neural operator learning framework for the 3D Euclidean space problem. The framework effectively combines coefficient learning and residual arithmetic to leverage the strengths and mitigate the weaknesses of existing approaches. In addition, this paper provides a detailed theoretical explanation of the proposed neural operator learning scheme from a graph spectrum view, treating the proposed model as applying the transformation to a spectrum of continuous feature function. Experimental results conducted on a widely used large-scale electron density of datasets validate the superior performance of the proposed approach.

Weaknesses

Some issues are as follows: 1. Although the proposed approach performs well in this paper, further baseline methods from recent years need to be added for comparison to more fully validate the effectiveness of the proposed approach. 2. It is worth noting that the proposed method might demand substantial computational resources, and its computational complexity is not explicitly discussed. As a result, its scalability to handle very large graphs might be a concern. 3. It is also necessary to pay attention to the writing standard of the paper format. For example, the title names of sections 2.1 and 3.3 are the same and need to be corrected.

Questions

see weakness

Rating

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

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

While the paper does not explicitly state the limitations of the proposed method, it is important to consider certain factors that may affect its effectiveness. Specifically, the performance of the method may vary depending on the characteristics of the particular input graph. Further research is necessary to assess the generalizability and scalability of the proposed method.

Summary

This paper introduces a method for learning SE(3) equivariant continuous functions in R^3 given an approximation of a continuous input and structural information in the form of a graph. The authors represent continuous inputs using ideas from coefficient learning (i.e., learning functions as the sum of learned coefficients multiplied by elements of a predefined basis), and incorporate structural information from discrete graphs by centering the basis functions at the node coordinates. To improve the quality of the learned function a residual term inspired by coordinate-based interpolation methods is used. The authors also provide a graph spectral interpretation of their method showing that the approach resembles graphon convolution. The proposed method outperforms baselines.

Strengths

* Originality: The work appears to be original. Several components of the proposed method come from existing efforts, however, the method to integrate structural information and the graph spectral interpretation appear to be new. * Quality: The work appears to be of very good quality. * Clarity: The paper is well written and the ideas are communicated in an accessible way. * Significance: This paper proposes a learning based approach for electron density estimation. Traditional approaches are computationally expensive, and therefore do not scale well. The proposed approach addresses this challenge by allowing for estimation in a feedforward deep learning framework.

Weaknesses

* Quality: Although the model is constructed to be SE(3) equivariant, this doesn’t seem to be reflected very well in the analysis. Can the authors explain why the error for rotated QM9 is so high relative to the unrotated version in table 1 and 2?

Questions

* Questions * As I understand S_{ij} is given by a neural network; is there a mechanism to make sure it is close to the integral form? * Possible Typos * Equation 3, change f_{i,u} to f_{u,i} * Equation 5, and 7, consider adding \qquad after the comma for readability * Fig 3. Consider using the same axis upper bound for both plots to improve readability

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors have communicated limitations of their work

Summary

The paper introduces an rchitecture combining a coefficient learning scheme with a residual operator layer for learning mappings between continuous functions in 3D Euclidean space designed to achieve SE(3)-equivariance. The proposed method can be seen as convolution on graphons, which are dense graphs with infinitely many nodes. They claim that their approach, called InfGCN, captures geometric information while preserving equivariance by leveraging both the continuous graphon structure and the discrete graph structure of the input data. The authors perform experiments on large-scale electron density datasets, and they claim that the proposed model demonstrates significant improvement over current state-of-the-art architectures.

Strengths

The paper is well written and the topic is very relevant. The notation is clear despite the technical nature of the work.

Weaknesses

A bit more clarity in the definition of equivariant is desired. Even though SO(3) irrep representation are mentioned, it should be clearer which group is referred to, e.g., SO(3), O(3), E3, etc?

Questions

- Are there any advantages for explainability from the framework?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed 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 discuss the limitations of their work.

Increase rating

Thank you for the detailed response. Conditioned on all the discussed changes being implemented in the camera ready paper, I increase my score.