Equivariant Neural Operator Learning with Graphon Convolution
Authors: Chaoran Cheng
Jian Peng
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Neural Operator Learning
Spectral Graph Theory
Graphon
TL;DR: We propose InfGCN, an equivariant neural operator learning architecture which can be interpreted as graphon convolution. InfGCN achieved SOTA performance across several electron density datasets.

Abstract: We propose a general architecture that combines the coefficient learning scheme with a residual operator layer for learning mappings between continuous functions in the 3D Euclidean space. Our proposed model is guaranteed to achieve SE(3)-equivariance by design. From the graph spectrum view, our method can be interpreted as convolution on graphons (dense graphs with infinitely many nodes), which we term InfGCN. By leveraging both the continuous graphon structure and the discrete graph structure of the input data, our model can effectively capture the geometric information while preserving equivariance. Through extensive experiments on large-scale electron density datasets, we observed that our model significantly outperformed the current state-of-the-art architectures. Multiple ablation studies were also carried out to demonstrate the effectiveness of the proposed architecture.

Equivariant Neural Operator Learning with Graphon Convolution
Chaoran Cheng University of Illinois Chaapa-Champaign Chaoran7@illimg edu
Jian Peng University of Illinois Urbana- -hampaan jianpeng@illinoie edu
Abstract
We propose a general architecture that combines the coefficient learning scheme with residual operator layer for learning mappings between continuous func- tions in the 3D Euclidean space. Our proposed model is guaranteed achieve SE(3)-equivariance by design. From the graph spectrum view, our method can be interpreted as convolution on graphons (dense graphs with infinitely many nodes), which we term InfGCN. By leveraging both the continuous graphon structure and the discrete graph structure of the imput data, our model can effectively capture the geometric information while preserving equivariance. Through extensive ex- periments on large-scale electron density datasets, we observed that our model significantly outperformed the current state-of state-of-the-art the-art architectures. Multiple abla- tion studies were also carried out to demonstrate the effectiveness of the proposed architecture.
1 Introduction
Continuous functions in the 3D Euclidean space are widely encountered in science and engineering domains, and learning the mappings between these functions has potentially amplitude of applica- tions. For example, the Schrödinger equation for the wave-like behavior the electron in a molecule, the Helmholtz equation for the time-independent wave functions, and the Navier-Stoke equation for the dynamics fluids all output a continuous function spanning over TR given the initial input. The discrete structure like the coordinates of the atoms, sources, and sinks also provides crucial information. Several works have demonstrated the rich geometric information of these data to boost the performance of other machine learning models, e.g., incorporating electron density data to better predict the physical properties of molecules [ 22 52].
It is common that these data themselves have inherently complicated 3D geometric structures. Work on directly predicting these structures, however, remains few. The traditional ways of obtaining such continuous data often rely quantum chemical computation the approximation method to solve ODEs and PDEs. For example, the ground truth electron density is often obtained with ab initio methods [88 26 with accurate results but an N computational scaling, making it prohibitive or inefficient for large molecules. Other methods like the Kohn-Sham density functional theory (KS-DFT) 24 has an 3 computational scaling with relatively large error. Therefore, building an efficient and accurate machine learning-based electron density estimator will have a positive impact on this realm.
Similar to the crucial concept of equivariance for discrete 3D scenarios, we can also define equivari- ance for function defined on R as the property that the output transforms in accordance with the transformation on the input data. The equivariance property demonstrates the robustness of the model in the sense that it independent of the poses of the input structure, thus also serving an implicit way of data augmentation such that the model is trained on the whole trajectory of the input sample. Equivariance on point clouds can be obtained with vector neuron-based models 6 16 45 [5] and
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

We notice the close relationship between the tensor field network (TFN) and the equivariance the continuous functions and also propose our equivariant architecture based on the tensor product.
tensor field networks 50 10).
In this way, we define task equivariant neural operator learning. We roughly summarize
previous work on operator learning into the following four classes: 1) voxel-based regression (3D-
CNN)
47 40
coefficient learning with a pre-defined set of basis functions 2 27
49 4
3)
coordinate-based interpolation neural networks 17 [8]; and 4) neural operator learning 228
25
32.. The roxel-based 3D-CNN models are straightforward methods for discretizing the continuous
input and output but are also sensitive to the specific discretization [25]. The coefficient learning models provide an alternative to discretization and are invariant to grid discretization. However, as the dimension of the Hilbert space is infinite, this method will inevitably incur errors with a finite set basis functions. The coordinate- basen networks take the raw coordinates as input and use a learnable model to "interpolate" them to obtain the coordinate-specitn output. They leverage the discrete structure and provide a strong baseline, but hard cut-off distance prevents long-distance interaction. The neural operators (NOs) are the newly-emerged architecture specifically tailored for operator learning with strong theoretical bases [28]. However, current NOs are mostly tested only on ID or 2D data and have difficulty scaling up large 3D voxels. They also ignore the discrete structure which provides crucial information in many scenarios.
To leverage the advantages of these methods while mitigating their drawbacks, we build our model upon the coefficient learning framework with an additional equivariant residual operator layer that finetunes the final prediction with the coordinate-specift information. A graphical overview of our model architecture is shown in Fig T We also provide a theoretical interpretation of the proposed neural operator learning scheme from the graph spectrum view. Similar to its discrete counterpart of graph convolutional network, our proposed model can be viewed as applying the transformation to the spectrum of the continuous feature function, thus can be interpreted as the spectral convolution on a graphon, a dense graph with infinitely many and continuously indexable nodes. In this way, we term our proposed model "InfGCN". Our model is able to achieve state-of-the-arl performance across several large-scale electron density datasets. Ablation studies were also carried out to further demonstrate the effectiveness of the architecture.
To summarize, our contributions are, 1) we proposed novel architecture that combines coefficient learning with the coordinate-based residual operator layer, with our model guaranteed to preserve by design; 2) we provided a theoretical interpretation of our model from the graph spectrum point of view graphon convolution; and 3) carried out extensive experiments and abla- tion studies on large-scale electron density datasets to demonstrate the effectiveness our proposed model. Our code is publicly available at sulhttps:/// tttpt::///thub thub. com/ ccr- cheng/ Biienn/InfGCN-pytorre
x(x)
p(r)
(4x) F20 weils) + =(x)
Figure 1: Overview of the model architecture. (a) The input molecule with node-wise spherical tensor features. (b) The message passing scheme in InfGCN. & denotes the tensor product of two spherical residual operator layer (Sec 3.4p. (d) Spherical harmonics. (e) The final prediction combining the expanded basis functions and the residue.
tensors f, (x) (Sec 3.3). (c) Coordinate-spent
2 Preliminary
We use = (V, E) to denote the (discrete) graph with the corresponding node coordinates {x, [=]- A continuous function over the region IB is also provided as the target feature: p D - R. We also assume that there is an initial feature function fin : - R either obtained from less accurate methods, a random guess, or some learnable initialization. Formally, given Jin E (D), which is square-integrable input feature function over D. and the target feature function p E (D), we want to learn an operator in the Hilbert space T (D) + (D( approximate the target function.

---

Different from common regression tasks over finite-dimension vector spaces, the Hilbert space (D) is infiite-dimensional.
2.1 Equivariance
Equivariance describes the behavior of the model when the input data transformed. Formally, for group G' acting on X' and group H acting on J, for a function f X + V, if there exists a homomorphism S G - H such that [' x) = [Sg) [x) holds for all g G, X, then is equivariant. Specifically, if 9 > maps every group action to the identity action, we have the definition of invariance: [(g x) (x). Vg E
In this work, we will mainly focus on the 3D Euclidean space with the special Euclidean group SE(3), the group of all rigid transformations. As translation equivariance can be trivially achieved by using only the relative displacement vectors, we usually ignore in our discussion and focus on rotation (i.e., SO(3)) equivariance. We first define the rotation of continuous function f E L²R as (x) := f(R "x), where R. is the rotation matrix associated with the rotation operator R. Note that the inverse occurs because In this way, the equivariance condition of an operator T with respect to rotation can be formulated as
rotating the coordinate frame instead of the coordinates.
T(Rs) R(T,) VR
(1)
For clarity, we will distinguish equivariance and invariance, and use equivariance for functions satisfying [I.
Method
3.1 Intuition
Intuitively, we would like to follow the message passing paradigm [12 to aggregate information from every other point x DD In our scenario, however, as the nodes indexed by the coordinate x e are infinite and continuous, the aggregation of the messages must be expressed as an integral:
Tw f(x) = S. W(x,))f(y)dy
(2)
where W x DDD [0, 1] is square-integrable kernel function that parameterizes the source node features f(y). There are two major problems regarding the formulation Eq. 1) unlike the discrete counterpart in which the number of nodes is finite, parameterization of the kernel W in the continuous setting hard; and 2) even W is well-defined, the integral is generally intractable. Some NOs 299 32 directly approximate the integral with Monte Carlo estimation
all grid points, which makes it harder to scale to voxels. Instead, we follow a similar idea in the coefficient learning methods 2 27 define set of complete basis functions {dxx 10c over L (D). In this way, the feature function can be expanded onto the basis as f(x) = k= f where fx are the coefficients. We can then parameterize the message passing in Eq. the coefficient learning with truncation to the N-th basis. We call such an expansion method unicentric as there is only one basis for expansion. In theory, as the size of the basis goes to infinite, the above expansion method can approximate any function y E (D) with a diminishing error. In practice, however, using a very large number of bases is often impractical. The geometric information of the discrete graph is also not leveraged.
3.2 Multicentric Approximation
To address the limitation mentioned in the previous subsection, we leverage the discrete graph structure to build mulficentric expansion scheme. We use the node coordinates in the discrete graph as the centers of basis sets: p(x) = Euev Eis fr,tii(x ra). We demonstrated in
Appendix that with some regularity and locality assumptions, the message passing in Eq 2 can be parameterized as
fra=
Σ vEN(a) 1=1
Σ
wiy Sy
(3)
3

---

where (r)(r) = fo (x)) ridx models the interaction between the two displaced basis at centers i, The outer summation is over ((u) (u) {u), the of neighboring centers of including u, and Wij are learnable parameters. Note that, once the basis functions are assigned, Siz only depends on r, but it is generally hard to obtain the closed-form expressions. We can use neural nets to approximate it and coalesce the weight parameter into the nets.
The integral (ry( is often referred functions can be viewed as the atomic orbitals and, in this way, the integral can therefore be interpreted the overlap between two displaced atom-centered electron clouds. The evaluation of the overlap integral is important in the self-consistent field method (Hartrec-Fock method) [48].
the overlap integral in quantum chemistry. The basis
3.3 Equivariant Message Passing
We will now consider the functions on the 3-dimensional Euclidean space, i.e., D= = R3, as they are widely encountered practical applications and non-trivial to achieve equivariance. It is not easy to find set of equivariant basis that satisfies Eq.(T. Inspired by the atomic orbitals used in quantum chemistry, we construct the basis function with Gaussian- based radial function R, (r) and spherical harmonics Y (f):
(4) where r rrl is the vector length and = r/r the direction vector on the unit sphere. Grd are normalizing constants such that f3a WWem (r)mm()) (rV = 1. The degree of the spherical harmonics l takes values of non-negative integers, and the order m takes integers values between -l and & (inclusive). Therefore, there are 2 + spherical harmonics of degree f. In this way, the basis index i are now triplets of (n, f, m).
Wntmm (r) = RC, (r) Gut
YY
To further incorporate the directional information, we follow the Tensor Field Network (50) to achieve
cquivariance based on the tensor product. Note that for any index pair (n1& (mm,, n22 L2m2)) the overlap integral (r) can also be expanded onto the basis as S(r) Entm 8nêm Unem (r) =; Lmtm innem (r)
For a fixed r and radial index n, the coefficient sequence = {ptm E 0, 0 m < e) can be viewed as a spherical tensor. Notice that the node feature f = (f' 2 0} can also be viewed as a spherical tensor. In the following discussion, we will omit the radial function index n for clarity it is independent of rotation. TFN leverages the fact that the spherical harmonics span the basis for the irreducible representations of SO(3) and the tensor product of them produces equivariant spherical tensors. The message passing scheme in TFN is defined as:
J
P
Σ (EN 420
Σ
W (x. )
( = Σ
(T)
Σ
YJ (T)
(5)
J= 1
where 2% is the Clebsch- Gordan matrix of shape (2f + 2) x (2k + 1) and 4 RR+ - IR are learnable radial nets that constitute part of the edge tensor features. A detailed deduction is provided in Appendix[A] [A] Intuitively, as TFN can achieve equivariance for spherical tensors, the output spherical tensor interpreted the coefficients should also give an equivariant continuous function p(x). Indeed we have
Theorem. Given an equivariant continuous input, the message passing defined Eq. 5 gives equivariant output when interpreted as coefficients of the basis functions.
rigorous proof of rotation equivariance can be found in Appendix A The translation equivariance also trivially holds as we only the relative displacement vector Fue X, X,, The equivariance of this scheme relies on the equivariance of the input feature map fm- Note that the O-degree features that correspond pure Gaussians isotropic, we use these features as the initial input. In practice, we use atom-specific embeddings to allow more flexibility in our model.
Also note that for orthogonal, the overlap integral is non zero only if mi = m2. Therefore,
the message passing can be simplified. As the spherical harmonics
e = wf + Σ Σ W (x, - x
(6)
VEN( R20
'If' we use an infinite number of complete orthonormal basis, the expansion can be achieved without error. However, if we use some none-complete or finite basis, this should be viewed as an approximation in the subspace spanned by the basis.

---

The first term is referred to as self- interratttaan in previous papers [4] but can be naturally inferred from our message passing scheme. For the nonlinearity, we follow 50 use the vector norm of each degree of vector features:
-
f (1
where o% are the activation functions. The equivariance holds the vector norm is invariant to rotation. Also, to avoid over- paramerrizaiiionn and save computational resources, we only consider the interactions within the same radial index: 'n (r) == onn Sum' eer (r). Note that this assumption generally does not hold even for orthogonal radial bases, but in practice, the model was still able to achieve comparable and even better results 5.3)
3.4 Residual Operator Layer
The dimension of the function space is infinite, but in practice, tion. Therefore, the expressiveness of the model will be limited by the number of basis functions used. Also, as the radial part of the basis in Eq. (4. is neither complete nor orthogonal, it can induce loss for the simple coefficient estimation approach. To mitigate this problem, we apply an additional layer to capture the residue given query point p at coordinate x. More specifically, the residual operator layer aggregates the neighboring node features to produce an invariant scala to finetune the final estimation:
only the finite approxima-
=(x) = Σ
Σ
W (x, x)
vEN(p) R20
This scheme resembles the coordinate- interpolation nets and was proved effective in our ablation study (Sec..3.. Therefore, the final output function is
((x)
Σ fnem Wnem (x) + xx)
modm
The equivariance of the residual operator layer as well as in the finite approximation case is also provided in Appendix A The loss function can be naturally defined with respect the norm in (TR) as L = 1 p JE lp(x)
ddx)
4 Graph Spectral View of InfGCN
Just the Graph Convolutional Network (GCN) 21 can be interpreted as the spectral convolution of the discrete graph, we also provide an interpretation of InfGCN as the transformation on the graphon spectra, thus leading to similar concept of graphon convolution. We will first introduce the (slightly generalized) concept of graphon. Defined on region DD a graphon, also known as graph limit or graph function, is symmetric igterabite function:
/ D2
W D [0, 1],
<
(10)
Intuitively, the kernel W(x, y) can be viewed as the probability that an edge forms between the continuously indexable nodes x, E 'D. Now, consider the operator Tw defined in (q. As the integral kernel W is symmetric and square-integrable we can apply the spectral theorem to conclude that the operator Tw it induces is self-adjoint operator whose spectrum consists of a countable number of real-valued eigenvalues { with d + 0. Let [ov be the cigenfunctions such that Tw Φx Ā+ Similarly to the graph convolution for discrete graphs, any transformation on the eigenvalues F { + ] can be viewed as the graphon convolution back in the spatial domain. We note that GCN uses the polynomial approach to filter the graph frequencies as
Hx Ek Wiy L. L where Wk are the parameters. Define the power series of Tw
Th Tw ((x) I W(x (yy)dy, TV I
(11)
where Z is the identity mapping on D. A graphon filter can be then defined as H.f = Exe we T f. We also follow GCN use the Chebyshev polynomials approximate the graphon filter H:
HS 0, Tw
(12)
Note that for scalars, SO(3) equivariance is equivalent to invariance.
5

---

Just as the message passing on discrete graphs can be viewed as graph convolution, we point out here that any model that tries approximate the continuous analog Tw as defined in Eq. can also be viewed as graphon convolution. This includes InfGCN, all NOs, and coefficient learning nets. more formal statement using the graph Fourier transform (GFT) and the discrete graph spectral theory are provided in Appendix B for completeness.
Another related result was demonstrated by Tsubaki et al. [5]] that the discrete graph convolution is equivalent to the linear transformation on a poor basis function set, with the hidden representation being the coefficient vectors and the weight matrix in GCN being the basis functions. As we have shown above, the same argument can be easily adapted for graphon convolution that the message passing in Eq. can be also viewed as the linear combination of atomic orbitals (LCAO) [37] in traditional quantum chemistry.
Furthermore, based on Eq 3 we can now give more intuitive interpretation of the radial network in TEN; it captures the magnitude of the radial part of the overlap integral S(r) of the basis in = From the point convolution aspect, the TFN structure can be also considered special case of our proposed InfGCN model. The discrete input features can be regarded as the summation of Dirac measures over the node coordinates fin [xn Suō(x xx).
Experiments
We carried out extensive experiments on large-scale electron density datasets to illustrate the state-of- the-art performance of our proposed InfGCN model over the current baselines. Multiple ablation studies were also carried out to demonstrate the effectiveness of the proposed architecture.
5.1 Datasets and Baselines
We evaluated our model on three electron density datasets. As computers cannot truly store continuous data, all datasets provide the electron density in volumetric form on a pre-defined grid. Atom types and atom coordinates are also available as discrete features.
QM9. The QM9 dataset 41 39] contains 133,885 species with up to nine heavy atoms (CONF). The density data as well as the data split come from [17 18] [7]88 which gives 123835 training samples, 50 validation samples, and 10000 testing samples.
This large-scale dataset contains electron densities 17,418 cubic inorganic materials (53) In our experiment setting, we first filtered out the noble gas (He, Ne, Ar, Kr, Xe) and kept only the crystal structure whose primitive cell contains less than 64 atoms. This gave 16,421 remaining data points. A data split of 14421, 1000, and 1000 for train/validatiinntess was pre-assigned.
Cubic.
MD. The dataset contains 6 small molecules (ethanol, benzene, phenol, resorcinol, ethane, malonalde- hyde) with different geometries sampled from molecular dynamics (MD). The former molecules from M| with 1000 sampled geometries each. The latter two from geometries each. The models were trained separately for each molecule.
2 with 2000 sampled
To evaluate the models, we followed our evaluation metrics:
to define the normalized mean absolute error (NMAE) as
J3a l0(x) - p(x)
NMAE
(13)
JEa lp(x) dx
To avoid randomness, different from the sampling evaluation scheme in 17 we did the evaluation on the partitioned mini-batches of the full density grid. Also, to demonstrate the equivariance of InfGCN, the density and atom coordinates were randomly rotated during inference for the QM9 dataset. The rotated density was sampled with trilinear interpolation from the original grid. Equivariance is trivial for crystal data, there is a canonical way of assigning the primitive cell. Similarly, for the MD dataset, the authors described canonical ways to align the molecules 2 so we also did not rotate them. More details regarding the datasets can be found in Appendix
We compared our proposed model with wide range different baselines including CNN 40 interpolation networks (DeepDFT [17. DeepDFT2 [8] EGNN 43). DimeNet 23 DimeNet++
[TI, and neural operators (GNO 28]. FNO 29). LNO 32). For InfGCN, we set the maximal
degree of spherical tensors to L 7, with 16 radial basis and 3 convolution layers. For CNN and

---

neural operators, an atom tppe-specific initial density function is constructed. A sampling scheme is used for all models except CNN. All models were trained single NVIDIA A100 GPU. More specifications on the model architecture and the training procedure can be found in Appendix D
5.2 Main Results
The NMAE of different models on various datasets are shown in Table 1 Models with the best performance highlighted in bold. Our model is able achieve state-of-the-aa performance for almost all datasets. For QM9 and Cubic, the performance improvement is about 1% compared to the second-best model both rotated and unrotated data, which is significant considering the small loss. We also noticed that CNN worked well on small molecules MD and QM9, but quickly ran out of memory for larger Cubic data samples with even batch size of (marked "OOM" in the table). This because CNN ran on the full grid with a maximal number of 3843 voxels.
Table 1: NMAE (%) on QM9, Cubic, and MD datasets.
Interpolation Net
Neural Operator
Dataset/Model
IutGCN
CNN
DeepDFT
DeepDFT2
EGNN
DimeNet
DimeNet++
GNO
FNO
LNO
QM9
rotated unrotated
4.73
5.89
4,98
12.13
12.98
12.75
33.25
24.13
0.93
2.01
2.95
1.03
11.92
11.97
11.69
40.86
28.83
26.14
Cubic
8,98
OOM
14.08
10.37
11.74
12.51
12.18
53.55
48.08
46.33
ethanol
8.43
13.97
7.34
8.83
13.90
13.99
14.24
82.35
31.98
43.17
benzene
5.11
11.98
5.49
13.49
14.48
14.34
82,46
20.05
38.82
phenol
5.51
11.52
9.09
7.00
13.59
12.93
12.99
66.69
42.98
60.70
resorcinol
5.95
11.07
8.18
6.95
12.61
12.04
12.01
58.75
26.06
35.07
ethane
7.01
14.72
6,36
15.17
13.11
12.95
26.31
77.14
nakozikkydde
10.34
18.52
9.31
10.68
12.37
18.71
16.79
84,52
34.58
47.22
IWEGUN
LNN
Deppo/
Depli1'12
LINN
LIO
LS00
Daaaky
Grondl Touth
InfCICN
DeppPT
DeeDier2
IGNN
DimeNer
LilmeNet
GNO
FNG
LNO
Pondiction
Density Eiroe
NMAE/%
Figure 2: Visualization of the predicted density, the density error, and the NMAE, Up: Indole (file ID 24492 from QM9). Down: Cr,CuNiSes (mp-122640)) from Cubie). The colors of the points indicate different atom types, and the isosurfaces indicate different density values. The pink and green isosurfaces in the error plots represent the negative and positive errors, respectively.
The visualizations of the predicted densities in Fig 2 provide more insights into the models. On QM9, the error of InfGCN had regular spherical shape, indicating the smoothness property of the coefficient-bas methods. The errors for the interpolation nets and CNN had a more complicated rugged spatial pattern. For Cubic, InfGCN was able to capture the periodicity information whereas almost all other models failed. The neural operators showed distinct spatial pattern on the partition boundaries of the grid, are provided in Appendix E on Cubic and QM9 with a wide range of representative molecules to demonstrate the generalizability of InfGCN.
was demonstrated to be sensitive to the partition. More visualizations
The plots of model sizes versus NMAE on the QM9 dataset are shown in Figure 3 seen from the figure that InfGCN achieved better performance with relatively small model size. The interpolation nets and CNN (in red) provide strong baselines. The neural operators (in orange), on the other hand, fail to scale to 3D data sampling scheme is required.
It can clearly
7

---

5.3 Ablation Study
Rotated QM9
Unrotated QM9
500
GNO
400
mNo
PNNO
UNO
NMAE/ %
200
NMAE/ / %
50
DimeNet.-
Dimehie:
Oimertt
EGNN
v10
mmor DDepDFF
20
DeepDFT
"
of 23
InfEE DeoppFF2
10-
InfGEN
"DeepOFT2 20
11
20
403 1/1
400
Model size /
M
Model size M
Figure 3: Plot of model sizes NMAE for different models and ablation studies on the QM9 dataset.
To further demonstrate the effectiveness of InfGCN, we also carried out extensive ablation studies on various aspects of the proposed architecture on the QM9 dataset. The results are summarized in Table and and are also demonstrated in Figure 3 in blue and green.
Table 2: NMAE (%) and the parameter count of different model settings on the QM9 dataset.
Model
InIGCN(s7)
86
85
84
83
82
81
80
no-res
fc
QM-rot (%)
4.73
4.77
4.76
4.77
4.86
6.95
9.56
12.62
6.14
4.95
QM-umrot (%)
0.93
1.01
1.11
1.08
1.46
4.65
8.07
12.05
3.72
1.36
Parameters (M)
1.20
0.85
0.58
0.39
0.26
0.17
0.13
0.11
1.16
17.42
Number of spherical basis. For coefficient learning models, using more basis functions will naturally lead to a more expressive power of the model. For discrete tasks, [50] [0] used only the degree- spherical tensor which corresponds to vectors. We experiments with the maximal degree of the spherical tensor 0 < L < 7 (s1 columns). Note that so corresponds to atom-centered Gaussian mixtures. It can be shown in Figure B (in blue) that the error smoothly drops as the maximal degree increases. Nonetheless, the performance gain is not significant with L 2 4. This is probably because the residual operator layer can effectively finetune the finite approximation error and therefore allows for the trade-off between performance and efficiency. In this way, our proposed InfGCN can potentially scale up larger datasets with an appropriate choice of the number of spherical basis.
Residue prediction. The residue prediction layer is one of the major contributions of our model that tries to mitigate the finite approximation error. It can be shown (under no- no-res) that this design significantly improves the performance by nearly 2% with negligible increases in the model size and training time. These results justify the effectiveness of the residue prediction scheme.
Fully-connected tensor product. As mentioned Sec we used channel-wie wise tensor product instead fully connected one that allows inter-channel interaction. We also tried the fully-connect tensor product under fc. turns out that the fully-connected model was 15 times larger than the original model and took 2.5 times as long as the latter to train. The results, however, are even worse, probably due to overfitting on the training set.
Related Work
6.1 Neural Operator Learning
We use the term neural operator in a wider sense for any model that outputs continuous data here. For modeling 3D densities, statistical approaches are still widely used in quantum chemistry realms. For example, 2 and [49] used kernel ridge regression to determine the coefficients for atomic
orbitals. [13 used symmetry-adaplee Gaussian process regression for coefficient estimation. These
8

---

traditional methods are able to produce moderate to good results but are also less flexible and difficult to scale. For machine learning-basse methods, 47 utilized a voxel-based 3D convolutional net with a U-Net architecture [40] to predict density at a voxel level. Other works leveraged similar idea of multicentric approximation. [27] and [3] all designed a tensor product-based equivariant GNN to predict the density spectra. These works are more flexible and efficient, but coefficient learning models inevitably have finite approximation errors.
Another stream of work on neural operator learning focused on directly transforming the discretized input. Tasks of these models often involve solving PDE or ODE systems in 1D or 2D scenarios. For example, [30] proposed the infinite- network to approximate the continuous output. Graph Neural Operator (28) approximated the operator with randomly sampled subgraphs and the message passing scheme. (29] and [9] tried to parameterize and learn the operator from the Fourier domain and spectral domain, respectively. (38] proposed an analog to the U-Net structure to achieve memory efficiency. These models are hard to scale to larger 3D data and are also sensitive to the partition of the grid if a sampling scheme is required. They also do not leverage the discrete structure.
6.2 Interpolation Networks
The term interpolation network coined in [25 for models that take raw query coordinates as input. As graph neural networks have achieved tremendous success in discrete tasks, they are usually the base models for interpolation nets. [55] and [4] constructed the molecule graph perform
variant message passing schemes with the final query-specific prediction. [77 proposed the DeepDFT
model which also considered the graph between query coordinates and [18 further extended it to
a locally equivariant GNN [5).. (0) proposed a similar model on crystalline compounds. Besides these specialized models, we also point out that current equivariant models for discrete graphs can all be adapted for continuous tasks in principle, just like DimeNet and DimeNet++ that we used as the baselines. Models that use only the invariant features including distance, angles, and dihedral angles can be trivially equivariant but lacking expressiveness 7 5 23
[16 15 proposed the GVP model in which features are partitioned into scalars and vectors with carefully designed interaction to guarantee equivariance. Other works leveraged the canonical local frame [19] or tried to learn such a local frame (33]. Another line of works, the tensor field network (50 [O], utilized the group theoretical results of the irreducible representations of SO(3) and proposed a tensor product based architecture. We follow the last method as we notice the close relationship between the spherical
tensor and the basis set. Though previous works with similar architecture exist [27] 3
we first give
rigorous proof of the equivariance of the continuous function.
6.3 Downstream Applications
Several previous works tried to leverage the geometric information of the continuous function. [2 utilized the charge density and spin density as both the supervising signal and the additional input for predicting molecular energies, which achieved significant performance improvement compared to traditional DFT- methods. [5I first projected the density onto the pre-defined basis set and then applied different neural nets the coefficients make predictions downstream tasks. [46] used a 3D CNN to encode the electron density of the protein complex to predict the backbone structure. These works have demonstrated the significance of continuous data.
7 Limitation and Conclusion
In this paper, we introduce a novel equivariant neural operator learning architecture with the core com- ponent interpretable as the convolution on graphons. With extensive experiments, we demonstrated the effectiveness and generalizability of our model. We also discuss the limitation and potential improvement of the proposed InfGCN model in future work. As the choice of the radial basis is arbitrary, there is no theory or criterion for better radial basis, and therefore, it leaves space for improvement. For example, we may use Slater-type orbitals (STO) instead of Gaussian- We may further orthogonalize the basis, which leads to the series of solutions to the Schrödinger equation for the hydrogen-like atom with more direct chemical indications. For structures with periodic boundary conditions, Fourier bases may provide a better solution. A learnable radial basis parameterized by neural net is also a feasible option to provide more flexibility.
orbitals.

---

References
[1] M. Bogojeski. L. Vogt-Maranto, M. E. Tuckerman, K.-R. Müller, and K. Burke. Quantum chemical accuracy from density functional approximations via machine learning. Nature Communications, 11(1) Oct 2020. ISSN 2041-1723. doi: 10.10389/41147--00-119009 URL https://///d org/ 10. 1038/ -4114460000-1995-3
[2] F. Brockherde, I.. Li, K. Burke, and K. Müller. By-passing the kohn-sham equations with chine learning. CoRR, htss1160.02215. 2016. URL httpp://///rxx http:////aaw.w..oo/aa/////a/a/// htrggabs/16099 02815
[3] Z. Chen, N. Andrejevic, T. Smidt, Z. Ding, Q. Xu, Y.T.. Chi, Q. T. Nguyen, A. Alatas, J. Kong, and M. Li. Direct prediction of phonon density of states with euclidean neural networks. Advanced Science, 8112)2004214, 2021. doi: https:///oi.org//0.. http:///....1100//ol/20000000000 attps://onliieelbrrayy wiley.. wwwmdoo//ee//ooooooooo.oooooo aa/////ab///// 1atp//11.00022000 12:0000/avvvv 2tw.0002222v..222222222222 202004214
URL
[4] Ö. Çiçek, A. Abdulkadir, S. S. Lienkamp. Brox, and O. Ronneberger. 3d u-net: Learning dense volumetric segmentation from sparse amnotation. In S. Ourselin, L. Joskowicz, M. R. Sabuneu, G. B. Ünal, and W. M. W. III, editors, Medical Image Computing and Computer- Assisted Intervention MICCAI 2016 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part 11, volume 9901 of Lecture Notes Computer Science, pages 424-4322 2016. doi: 778-3-319-4672--
_9, URL wttpt::///oooo org/10.1007
[5] B. Coors, A. Condurache, and A. Geiger. Spherenet: Learning spherical representations for detection and classification in omnidirectional images. In V. Ferrari, M. Hebert, C. Sminchisescu, and Y. Weiss, editors, Computer Vision ECCV 2018 15th European Conference, Munich, Germany, September 8- 14. 2018, Proceedings, Part IX, volume 11213 of Lecture Notes in URL Computer Science pages 525-541. Springer, 2018. doi: 10. (0000/978-33000-0222400 htttps://ooi. org/ 10. 1007 000007/97883-000-12244-00
[6] C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasacchi, and L. Guibas. Vector neurons: A general framework for so(3) -equivvriaai networks. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC. Canada, October 10-17, 2021, pages 12180- 12189. IEEE, 2021. doi: 10. 109/CCC448922.2222.0001
URL https: //doi org/10
1109/ICCV43922. 2021. 01198
[7] H. Deng. T. Birdal, and Ilic. Ppfnet: Global context aware local features for ro- bust 3d point matching. In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, pages 205. Computer Vision Foundation / IEEE Computer Society, 2018. doi: 10. 109//R.. 2018.00028. UtL PPFNet PPFNet_Global PPPlet__looal__CotteetttCCC 2018 20we_paperr_p___.
thecvf. com/ comm/ottenttccpp.c... 2atpttt_cvrr__011111e11 2atprr_221//ttmm/eeeee htmal
[8] A. Edmonds. Angular Momentum in Quantum Mechanics. Investigations in physics. Princeton University Press, 1957.
[9] V. Fanaskov and I. V. Oseledets. Spectral neural operators. CoRR, abs/2205.0002 10573, 2022. doi: 20.48550/a 10573. https:: http://d......r/// 10. 20.5500/rXiv.0000 10573 www.:r/........r.........oo
[10] F. Fuchs, D. E. Worrall, V. Fischer, and M. Welling, Se(3) transffrrnss 3d roto-translation equivariant attention networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020. URL httpps:/pproceddnngg seuripps.cc/ppper/20000
wwtt---bbbtttt.... .hmml
[11] J. Gasteiger, S. Giri, J. T. Margraf, and S. Günnemann. Fast and uncertainty- directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NeurlPS, 2020.
[12] J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In D. Precup and W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017,
10

---

volume 70 of Proceedings of Machine Learning Research, pages 1263-1272. PMLR, 2017. html
URL htsLiittprooceedinngge mlr www.ppress/////.
[13] A. Grisafi, D. M. Wilkins, G. Csányi, and M. Ceriotti. Symmetry-adapted machine learn-
ing for tensorial properties of atomistic systems. Phys. Rev. Leut. 120:036002, Jan 2018 doi: 10.1103/PhysRevLett 12t..ss...e........00000000 120.036002. URL https://link: aps www//oi//////oooo 1000.1100 1103/
PhysRevLett 120. 036002
[14] G. Hegde and R. C. Bowen. Machine-lcameed approximations to density functional theory hamiltonians, 2016. UtL.tttss://aaxx http://awxxv...rrgbaaa//a// http://wwwww..rr////////.././///
[15] B. Jing. S. Hismann, P. N. Soni, and R. O. Dror. Equivariant graph neural networks for 3d macromolecular structure. CoRR, abs/2106.03843. 2021. URL BRLhttts://aare www.v/..rg//b////////
2106. 03843
[16] B. Jing, S. Eismann, P. Suriana, R. J. L. Townshend, and R. O. Dror. Learning from protein structure with geometric vector perceptrons. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
URL httpps://ppenreviek wwwpearrwweeeenetfoouiiiii.ee..e net/forumiid 1
[17] P. B. Jørgensen and A. Bhowmik. Deepdft: Neural message passing network for accurate charge density prediction. CoRR, abs/201 I. .33346, 2020. URL
[18] B. Jørgensen and A. Bhowmik. Equivariant graph neural networks for fast electron density estimation of molecules, liquids, and solids. npj Computational Materials, 8(1):183, Aug 2022. ISSN 2057-3960. doi: 10.1038/s41524-000 022-00262-00
URL https: org/ 1038/
[19] J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneherger, K. Tunyasuvunakool, R. Bates, A. Židek, A. Potapenko, A. Bridgland, C. Meyer, S. A. A. Kohl, J. Ballard, A. Cowie, B. Romera-Paredes S. Nikolov, R. Jain, J. Adler, T. Back, S. Petersen, D. Reiman, E. Clancy, M. Zielinski, M. Steinegger, M. Pacholska, T. Berghammer, S. Bodenstein, D. Silver, O. Vinyals, A. W. Senior, K. Kavukcuoglu, P. Kohli, and D. Hassabis. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-56)
Aug 2021. ISSN 1476-4687, doi: URL URL htttp::////// //doii org 10. 1038/ s41586- 021-5886--21--0899---1
[20] P. R. Kaundinya, K. Choudhary, and S. R. Kalidindi. Prediction of the electron density of states for crystalline compounds with atomistic line graph neural networks (ALIGNN). CoRR, abs/2201. .08348, 2022. URI. ticetps://arxivv org/ abs 2201.22222222222 08348
[21] N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. In 5th International Conference Learning Representations, ICIR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https:// openreview net wett//oorumidd-SSAAaaa
[22] Kirkpatrick, B. McMorrow, D. H. Turban, A. L. Gaunt, J. S. Spencer, A. G. D. G. Matthews, A. Obika, L. Thiry, M. Fortunato, D. Pfau, L. R. Castellanos, S. Petersen, A. W. R. Nelson, P. Kohli, P. Mori-Sánchez, D. Hassabis, and A. Cohen. Pushing the frontiers of density functionals by solving the fractional electron problem. Science, 374 2021, doi: 10. 126/science.abj651111 URL htt.hhttp///// //t:::///wwww www.escceenne www.sccieccc.......o......ee..ee http://///oooooooooee//oo// acw.a/nnc....rg/////a//////a.... 1126/ acience. acience.abj6611
[23] J. Klicpera, J. GroB, and S. Günnemann. Directional message passing for molecular graphs. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia. April 26-30, 2020. OpenReview.net, 2020. URL https:://peereviiiie wetp/formmffi BieWbxStPH
[24] W. Kohn and I. J. Sham. Self-consistent equations including exchange and correlation effects. Phys. Rev, 140:A1133-A113: Nov 1965. doi: 10.1 103/Physsev.. 140.A1133 URL https
//link. aps org/ /ww//0.i101111// 11t.:...11111111111. 140. A1133
11

---

[25] N. Kovachki, Z. Li, B. Liu, K. Azizzadenesheli, K. Bhattacharya, A. Stuart, and A. Anandkumar. Neural operator: Learning maps between function spaces, 2021. URL https://axxxi org/
hbss2200.008481
[26] H. G. KÜMMEL. A biography of the coupled cluster method. International Journal of Modern Physics B, 17(28):5311-3325 2003, doi: 10.1422/921177792000000442 URLhttps https
//doi (ddii.or//ie 10gg/1..i 11/1/11/////A 502179992030224442
[27] A. J. Lee, J. A. Rackers, and W. P. Bricker. Predicting accurate ab initio DNA electron densities with equivariant neural networks. Biophys J, 121(20::3833385, Sept. 2022.
[28] Li, N. B. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. M. Stuart, and A. Anand- kumar. Neural operator: Graph kernel network for partial differential equations. CoRR, 0bs/2003.03485, 2020. WRI.jhttps://arxtt www./agg///////o/oo//...o...eeee abs/2003... 03485
[29] Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. M. Stuart, and A. Anand- kumar. Fourier neural operator for parametric partial differential equations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https:://ppereeiiii http://pwwrrviiiee.eeeerrrrrrrre wttp://orrummiid--80000vaaaaam
[30] R. Livni, D. Carmon, and A. Globerson. Learning infinite layer networks without the kernel trick. In Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 2198-2207. PMLR, 2017. / //proceedings.niu mlr wtwsllpppeee//veeeee grress/v70/liviiin httml html
[31] L. Lovász. Large Networks and Graph Limits. American Mathematical Society colloquium publications. American Mathematical Society, 2012. ISBN 9780821890851. URL https
//books //booss.oogge. google.oooooo.o.ooo..o wom.gle...omeeeeo wwt.com////////// books Fid=FsFFHLiddBSCCC
[32] L. Lu, P. Jin, G. Pang, Z. Zhang, and G. E. Karniadakis. Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators. Nature Machine Intelligence, 3(3):218 229,, mar 2021. doi: 10. 000101142256-020--00000 //dii.org/ 10. 1038/22 021-0060--200-0-0-0--00000
URL https
[33] S. Luo, J. Li, J. Guan, Y. Su, C. Cheng, J. Peng, and J. Ma. Equivariant point cloud analysis via learning orientations for message passing. IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 18910- 18919. IEEE, 2022. doi: 10. 01109/CVPR22688.00222.00 CL09/CVPR52688.. 2022.01836
URL https://dol org/10
[34] S. Maskey, R. Levie, Lee, and G. Kutyniok. Generalization analysis of message passing neural networks on large random graphs In NeurlPS, 2022.
URL http://papers. nips. httpeecccpppppeeepeee filef/ wwwpl///ppp///222221. haahh/ (eeaae//88d884926577776666666633
Conference html
[35] B. Mohar and W. Woess. survey on spectra of infinite graphs. Bulletin the London Mathe- matical Society, 233: 1989. doi: https//4oi.org//..1112/4/mmm
URL /htpps 3......mm........... 209
//oodmathsoc..oo wow.asawoo....mmmmmmmmrmraaa.aee /ooodhaabbacc.omllmmmleeeaeaaay
[36] M. W. Morency and G. Leus. Graphon filters: Graph signal processing in the limit. JEEE
Trans. Signal Process., 69:1740-1754, 2021. doi: 10. 10.1100/TSP222213306115 URL https
//doi oor.oog//100 10. 1109/TSP 2021 3061575
[37] L. Pauling.
The nature of the chemical bond. application of results obtained from the quantum mechanics and from a theory of paramagnetic susceptibility to the structure of molecules. Journal of the American Chemical Society, 53(4):1367-1400 1931. doi: 10 URL httphhttp.... : //oii wwg/oo.///...... 00010/000000001110000333500001
[38] M. A. Rahman, Z. F. Ross, and K. Azizzadeneshelli U-NO: u-shaped neural operators. CoRR, abs/2204.11 127, 2022, doi: 10.44550\axxx.222441117 URI.httss://aat http::///oooooo////////// 48550/
arXiv. 2204. 11127
12

---

[39] R. Ramakrishnan, P. O. Dral, M. Rupp, and O. A. von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 2014.
[40] O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In N. Navab, Hornegger, W. M. III, and F. Frangi, editors, Medical Image Computing and Computer-Assisted Intervention MICCAI 2015 18th International Conference Munich, Germany, October 5 9, 2015, Proceedings, Part III. volume 9351 of Lecture Notes in Computer Science, pages 234-241. Springer, 2015. doi: 10. (0.1107/978-33319-25577 URL htllhttps:// //ooi org/ 10. 0970000//77733 - 3-11----119----4111
[41] L.. Ruddigkeit, R. Deursen, L. C. Blum, and J.-L.. Reymond. Enumeration of 166 billion
organic small molecules the chemical universe database gdb-17. Journal of Chemical Information and Modeling, 52(1 1):2664 2875, 2012. doi: 10.1021/ci300415d. URL https PMID: 23088335.
[42] L. Ruiz, L. O. Chamon, and A. Ribeiro. Graphon neural networks and the transfer- ability of graph neural networks. In H. Larochelle, M. Ranzato, R. Hadsell, M. Bal- can. and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12. 2020, virtual, 2020. URL 0a1GG5880005400bbbcccc222bbbb000 wwwp:///4odd--bbbbacccoo Abstract html
[43] V. G. Satorras, E. Hoogehoom, and M. Welling. E(n) equivariant graph neural networks. In M. Meila and T. Zhang. editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9323-9333 PMLR, 2021. http::/prooceeedinggg
mlr. Sttrs///133/aaaorrreeee
html
[44] K. Schütt, P. Kindermans, H. E. S. Felix, S. Chmiela, A. Tkatchenko, and K. Müller. Schnet: A continuous- convolutional neural network for modeling quantum interactions. In Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Gar- nett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 991-1001, 2017. URL https: //t http://pooeeeeddeeeeeeee neur pss wttppppcppppppppppppooocpccppt ct/ htcppaper//0u/ (amper/20117/aaa
300e1446694411333300080888833383 wwtp://30000--Attatt.tom Abstract.
[45] K. Schütt, O. T. Unke, and M. Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9377-9388 PMLR, 2021. http:: http:///pooceeeneee//ee/ mlr /ttess/v13//// http:////99/ctht22eeee schutt2la html
[46] D. Si, S. A. Moritz, Pfab, J. Hou, R. Cao, L. Wang. T. Wu, and J. Cheng. Deep learning to predict protein backbone structure from high-resolution cryo-em density maps. Scientific Reports, 1021):4282, Mar 2020. ISSN 2045-2322. doi: 10.
URL
https://doi... org/10../........ 000100000000/10000000000000 620-6598--
[47] A. V. Sinitskiy and V. S. Pande. Deep neural network computes electron densities and energies of a large set of organic molecules faster than density functional theory (dft), 2018. URL
www.oogg////ooeeee 02723
[48] C. Slater. The self consistent field and the structure of atoms. Phys. Rev, 32:339-348, Sep 1928. doi: 10. 10.1003/hysseee.33333. URL https://link. aps. org/ org/doi/100 1103/ PhysRev. PhysRev.32.333 33.339
[49] M. Su, J. Yang, H.-J Xiang, and X.-G. Gong. Efficient prediction of density functional theory hamiltonian with graph neural network. 2022. URL http:tttarxxi..oottt org/ abs/ 2205 05475
[50] N. Thomas, T. F. Smidt, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. Riley. Tensor field networks: Rotation- and translation-eqqiivrrann neural networks for 3d point clouds. CoRR, abs/1802.08219, 2018. URL httpht/t//aaaa org/ /tbs/800..0... 08219
13

---

[51] M. Tsubaki and T. Mizoguchi. On the equivalence of molecular graph convolution and molec- ular wave function with poor basis set. Larochelle, M. Ranzato, R. Hadsell, M. Bal- can, and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 12, 2020, virtual, 2020. URL procpedg// wwapiiinnnnnn.uuuiiee cc/ httpp/pppp///////////// 115177766222cff11122200001111113 Awwtract.tttttteem html
[52] M. Tsubaki and T. Mizoguchi. Quantum deep field:
Data-driven wave function, electron density generation, and atomization energy prediction and extrapolation with machine learning. Phys. Rev. Lett. 125:206401, Nov 2020. doi: 10.1 103/PhysRev..cco 125.206401. URL https
//link. wtt..kaap...rg/oooo/.ooo..o
http//ooo/1.../ 11w.:0111111////hhseeeettt. 125 206401
[53] Q. Wang, K. Choudhary, Y. Liu, Hu, and M. Hu. Large scale dataset of real space electronic charge density of cubic inorganic materials from density functional theory (dft) calculations. Scientificc Data, 91):59, Feb 2022, ISSN 2052-4463. doi: 10.1038441197-02--115--- URL
https: //ooi. http///i.o//////////// org/10. 1038/ 0008//411577002--11----1
[54] 7. Wang, I.. Ruiz, and A. Ribeiro. Convolutional neural networks on manifolds: From graphs and back. In 56th Asilomar Conference on Signals, Systems, and Computers, AC- SSC 2022, Pacific Grove, CA, USA, October 31 Nov. 2, 2022, pages 356-360. IEEE, URL httph: http///oi...rg/0.11999999/
2022. doi: 10. (00///EEEEOON566442200000000 TEBECONF56349 2022. 10051964
[55]
L. Zepeda-Núñezz Y. Chen, J. Zhang, W. Jia, L. Zhang, and L. Lin. Deep density: Cireumventing
the kohn-sham equations via symmetry preserving neural networks. J. Comput. Phys., 443:
1 110523, 2021. doi: 10. 10.10166/.ep..221.11555 10523 U https: http://do...org///// 10166/ jep
2021. 110523
14

---

Supplementary Material
A Proof of Rotation Equivariance
In this section, we will give rigid proof of rotation equivariance of our proposed InfGCN model with finite approximation. Just mentioned in the main text, we will ignore the radial index n for clarity. Recall that we want to generate equivariant density functions as shown Figure
R
p(r)
(Rp)(r)
Figure 4: An illustration of rotation equivariance the electron density function benzene on R
Proposition A.1. Rotation of spherical harmonic of degree l and order m (RY
ymm transforms into linear combination of spherical harmonics of the same degree:
(RY
Σ (R))( (r)
(14)
where D Dm (R) an element of the Wigner D-matrix.
The proof of this property of spherical harmonics can be found in books on quantum mechanics, for example, Eq.4. in (8] Therefore, for a qquare-iteeraabl function defined on the unit sphere in TR3. we can also describe the rotation of the function with Wigner D-matrics:
Proposition A.2. Assume f E (S²) and the rotation of f have the (infinite) expansions onto the spherical harmonic basis
f(r) = Σ Yrn(r)
S y
Em
(15)
Rf((r)
I Y (r)
Then, we have
g = ==D"
(16)
where P the coefficient vector of degree with m 28 elements, and D is the corresponding Wigner D-matrix of degree f.
Proof. Notice that for each degree l, the coefficients are transformed linearly according to Eq which concludes the proof.
14)
Define spherical tensor (ff l 2 0}, we can further simplify the notation in Proposition A.2 as
R. (f)
(17)
A pictorial illustration of the rotation of the spherical harmonics is provided in Figure It can be shown that the computational diagram commutes in a sense it is equivalent to applying Wigner D-matrices on the coefficients and then projecting them back as continuous function.
One crucial property of the spherical tensors that the tensor product is equivariant to rotation R. E 50(3): Proposition A.3. The tensor product of two spherical tensors satisfies the following identity:
DR(a & b) DR(a) & DR(b)
(18)
15

---

-0.257
0.350
0.906
p(x)
R
0.874
-0.331
0.356
(x)
00.45
-0.878
0.238
Dx=
D
Dt/
T
13
8.53
9.94 P-
0. 0.54
6.53
8.67 S+ e. e. 0.63 co
Figure 5: An illustration of rotation equivariance of linear combination of spherical harmonics as a continuous function on the unit sphere. For better visualization, the radial value used to plot the spherical harmonics the squared density [p|2. The calculation, however, still done the original (real-valued) spherical harmonics.
The proof of this property can be found in the TFN paper [50] in Appendix C. Essentially, it is natural corollary as the property of irreducible representations. Combining Eq. (1) and 33 we can then design equivariant message
Σ
Σ
WiEk
(((or)) f
(19)
(EN(M)
Jm
where (r) - ("( Here, we use the same technique as TFN to restrict the radial functions to be independent of the order m so that it is isotropic. Win are the learnable weights. The tensor product gives 4-dimensional tensor. The summation indices J, correspond to the two dimensions other than É, k. One way to define such tensor product for two spherical tensors comes from the coupling of angular momentum in physics. The tensor product
@ b defined as
CJm=
aém, bkmn (
(20)
my magm=
where (omn,mmgJee Jm)) are the (lebsch-Gordanl coefficients, and nonzero only when / X:] S < &8+k -J S m < J. Substituting Eq.(20) into Eq. summation of
(9) and ignoring the zero terms, we will get
k+t Σ we)ps(r) NEN(W) k20 J=ik-E
r = Σ ΣE
Σ (T) m=-j
(21)
where Q4h (m)m2) (em kmg Jm). Coalescing the weight into the learnable radial function
E (r))sp(r) have final message passing scheme defined in (q. With Proposition
A.3 we immediately have the following corollary:
Theorem A.4. The message passing defined in Eq. 19 (with infinitely many basis functions) is
equivariant to rotation.
Proof. According to Proposition A.2 Therefore, we have
the two spherical tensors in (9 transform as in Eq. 17
T(Rf.)
Σ
Σ
Wiek Σ
Rup(raw) Rf,
vEA(u)
8k
Jm
Σ
Σ
Wekk Σ
(R(((T @ DRI,
vEA(u) Ek
Jm
Σ
Σ
Wikk Σ
DR(P(TT))) (& f.)
(22)
(EN(u) Ek
Dr Σ
Σ
Wikk Σ ((Tww) f
(EN(()
&k
Jm
= R((F.))
16

---

Wes can be moved inside DR because it is essentially a linear combination of equivariant functions, thus is also equivariant.
Now let's consider finite approximation where 0 s l S L for fixed L. We have the following proposition:
Proposition A.5. Let PL {a. 2 03 + PLa faz 0 l S. L be the projection tensor operator that projects the spherical tensor onto the bases with degree less or equal to L, then PL is rotation equivariant:
Pr( = R(P1.a)
(23)
Proof. It suffices to note that for each degree L, the corresponding component af rotates according to 17], which only depends on the components with the same degree. Therefore, Pr is equivariant are preserved on both sides.
as the components with degree
Proposition A.6. The composition of two equivariant operators Ti To also equivariant.
Proof.
(T1 (Ra) (I(T2(Ra)) (T((((aa)) R2T1
(24)
L
Combining Theorem A.4 and Proposition A.5 A.6 we have the equivariant property with finite approximation:
Corollary A.7. The result in Theorem A.4 also holds for a finite degree of spherical basis 0 S E L.
Notice that for degree O features (scalars), equivariance is equivalent to invariance, and it can be obtained with projection operator Po. we have
Corollary A.8. The residual operator layer defined Eq ( with finite approximation is invariant with respect to the frame thus equivariant to rotation under the global frame.
Combining the results above, we immediately obtain the rotation equivariance of our proposed model. Theorem A.9. The proposed model in Eq. with finite approximation and the residual operator satisfies the equivariance condition defined in Eq. 1
B Graph Spectral Theory and Graphon Convolution
In this section, we will introduce some preliminary for graph spectral theory and demonstrate the deduction for graphon convolution with more details. The basic concept of graphon can be found in various mathematics or signal processing references [36 35 31.
B.1 Graph Spectral Theory
We begin with the graph spectral theory for discrete graphs. For discrete graph g (V,E), the graph Fourier transform (GFT) is defined as
x
(25)
where S = UAUT is the eigenvalue decomposition of the graph shift operator S. A graph shift operator is diagonalizable matrix N satisfying S.j = 0 for i / j. (i, & The graph Laplacian L = I A and the normalized version L = 1-22 112 as was used GCN (21) where A is the adjacency matrix are such operators. As clear this definition of GFT is, it remains computationally prohibitive to implement the filtering operation on large graph in this way. To filter the graph frequencies, polynomial approach is thus adopted on the eigenvalues:
K
Hx
L
(26)
k=
where wi are the learnable parameters. the GCN formulation, the authors used the diagonalized matrix A instead of L, which is essentially equivalent.
17

---

We now switch to the continuous graphon setting. As defined in Eq.(TO). graphon is a symmetric
square-integrable function. The kernel W induces an operator Tw defined in Eq. As W is symmetric and square-integrable, Tw is self-adjoint operator that can be decomposed as
uTw A2
(27)
where 2/ is some unitary operator and A is mulplication operator, i.c., there exists function {(x) such that for all f(x), Af(x) = f(x)f(x), This directly follows the result the spectral theorem for self-adjoint operators. In this way, may similarly define the graphon Fourier transform
uu
(28)
Following the polynomial approach of approximating the graphon filter, we first define the power series of Tw as in (l and use the Chebyshev polynomials of Tw approximate the graphon filter H as H.f = f Tw Either way, parameterization and evaluation Tw f is required. As mentioned in Sec our model essentially operates on the eigenvalues of the operator Tw. In the graph spectral point of view, the eigenvalues are the spectrum of the operator. Therefore, any spectral filtering be effectively viewed as graphon convolution. More details regarding parameterization will be discussed below.
B.2 Approximating Graphon Convolution
We now consider parameterization and evaluation of Tw deduce Eq For any complete orthonor- mal basis { 1 === of I. (D), any square-integrable function f can be expanded as f Σx= f& ths where /& Jo (x) thu (x)dx. We can then arrange the transform = Tw as the following matrix-vector form g Wf, where
We = 1/ (x((x) S. W W(x, (xx
(29)
If (thh 100 coincide with the eigenfunctions (dx J00 which satisfy
Tw = Au
(30)
We have Wy = A fp (x) o, (x)xx. For the unicentric setting, Wij non-zero only when i= (the self-interaction term). For the multicentric setting, however, the computation is different. Recall that in the multicentric setting, we assume the global feature function is the summation of all atom-centered functions
p(x)
Σ REV i=1
fx rw)
(31)
where is the coordinate of center Similarly, considering center at the origin with the other at r, we have
WG = S (xxx) - r)dx = >3 (x(x)),(xx -anxx-rrdc r)dx
(32)
The "overlap integral" Sty (r) arises here and we further parameterize the integral with Wij as the basis transformation also involves index i. Therefore, using the above matrix-vector formulation, we have the following parameterization:
20 Σ
f.
wiy S.ySS,, (T)
(33)
we also assume the locality of the interaction between two expansion centers, we can sum over all neighboring nodes to give the result in (3. The locality assumption often holds as the basis functions decay exponentially. Therefore, ideally, the overlap integral between two far-away centers should be negligible.
18

---

B.3 Graphon Convolution and Continuous MPNN
Previous work regarding continuous pesssage-passingg neural networks is available. We briefly review them here and discuss their relation to our proposed method of graphon convolution. The idea of generalizing the discrete meessage-passing paradigm to continuous cases is essentially the same procedure as we have described in Sec 3.1 and all previous work used Eq. 2 to formulate the continuous message passing. For example, 42] proposed WNN (W refers to the graphon kernel W) as the limiting object GNNs with an increasing number of nodes and explored the transferability of this continuous formulation. (34] proposed cMPNN that explicitly modeled continuous functions and provided theoretical guarantees of the convergence and generalization error under discretization. [54] proposed MNN for modeling mappings between continuous manifolds and also leveraged the graph limit view of large graphs.
Though sharing a similar idea of utilizing continuous geometric structures, our method is funda- mentally different from the above models. Most significantly, in the above work, the authors either explicitly constructed the graphon kernel W (WNN, MNN) or directly estimated the Fredholm integral (cMPNN) in similar fashion various neural operators. Our approach, on the other hand, implicitly constructed the graphon and parameterized them in the spectral domain. We noted in See B.I that the kernel W does not have a canonical form in our setting and Monte Carlo estimation is prohibitive for large voxels. Instead, we defined a basis set and demonstrated in previous subsections that transformation on the coefficients can be also viewed as graphon convolution in the spatial domain. In this way, we implicitly assume that there exists different graphon for each data sample defined by their discrete structure and their categorical information. Nonetheless, the parameterization of graphon was done with the same graphon convolution for the whole dataset, as we expected this information to be generalizable across different samples. In the abovementioned work, however, a different net needs to be trained for different graphon.
In terms of the problem formulation, we further assume there exists an underlying discrete structure that has a significant physical connotation, c.g., atoms in electron density. The datasets on electron density we experimented with are real-world data and are significantly larger than those used in previous graphon-relate work. Based on the above difference in data structure and tasks, we designed a new network architecture that is different from the work before. We approximated the graphon convolution with neural nets on the coefficients instead of the feature functions themselves, and we also proposed the residual operator layer to mitigate the finite approximation error. Also, we extended the definition of rotation equivariance to continuous functions and provided rigid proof that our model achieves such a desirable property by design.
In conclusion, our work should be viewed as parallel to the existing work on graphons We are also aware of the values of previous theoretical work. As our model still followed the framework on estimating and parameterizing the integral in E the theoretical results on convergence and transferability could be adapted for our model to make more concrete and solid.
C
Datasets
In this section, we provide more details about the datasets we used in the experiments.
QM9.
The densities are calculated with VASP using the Baddew-Burk--Ernzeeee (PBE) functional
17 18. The grid coordinates are guaranteed be orthogonal but generally different for different molecules.
Cubic. The densities are caleulated with VASP using the projector augmented wave (PAW) method [53]. As the crystal structure satisfics the periodic boundary condition (pbc), the volumetric data are given for primitive cell with translation periodicity. We only focus the total charge density and ignore the spin density. Note that though all materials belong to the cubic crystal system, some of the face-center cubic (fec) structures given in its not-orthogonal primitive rhombohedral cell.
MD. Densities from I are calculated with the PBE XC functional; densities from [2] are calculated with Quantum ESPRESSO using the PBE functional. Both datasets are simulated in a cubic box with a length of 20 Bohr and a uniform grid size of 50 The result volumetric density is represented in Fourier basis, so we first converted it into the Cartesian form for all models.
19

---

The raw data of these datasets come in different formats. We defined a unified data interface to facilitate experiments and easy extension other datasets, A data point consists of the following data fields for training and inference:
atom_type. Atom types of size N.
atom_coord: Atom coordinates of size (N, 3).
density: Voxelized density value of size 2₂ N₂ x N₂. Note that was flattened have the order of X-Y-Z.
grid_coord Coordinates of the grid points where densities are sampled, with size (N₂ x Ny x 3)..
shape: vector representing the discretization sizes of each of the three dimensions.
cell: 3-by-3 matrix representing the cell vectors.
Other common statistics of the datasets summarized in Table The MD dataset does not have a validation split. The number of grids and grid lengths for a single dimension so the number of voxels scales cubically with respect to it.
Table 3: Dataset details.
Dataset
QM9
Cubic
MD
train/val/test split max/min/mean #grid max/min/mmen #node
123835/50/10000
14421/1000/1000
1000(2000)550(4000)
160/40/87.8 86
448/32/93.97
20/20/20
39/3/17.998
64/1/10.49
14/8/10 1088
max/min/mean length (Bohr)
15.83/4.00/8.65
26.20/1.7855
20.00/20.00/20.00
#node type
5
84
D Model and Training Specification
In this section, we provide our model specifications well as the baseline model specifications. Training- and testing-related hyperparameters used in the experiments are also provided.
D.1 Model Specification
We provide more details about the proposed InfGCN model and the baseline models. The baseline models' architectures are briefly described and major hyperparameters are provided. The model sizes provided here are for QM9.
InGGNN. We used spherical degree 7 and the number of radial bases n 16 with the Gaussian parameters aş starting at 0.5 Bohr and ending at 5.0 Bohr. The distance 64-dimensional vector and went through two fully-connected layers with a hidden size of 128. We used 3 InfGCN layers. This model has 1.20M trainable parameters and was used for all datasets.
first embedded in
We used 3D-CNN with the U-Net architecture which has been successful in biomedical imaging tasks. CNN is generally not rotation equivariant. As the density grids in the datasets are not necessarily the same, we manually injected the grid information by pre-computing the initial feature map on the grid points as:
CNN 40 4
fx(x)
Σ exp NEV
-ag
XX
X
1
(34)
Tu
where r. the covalent radius atom u and {au are pre-defined Gaussian parameters that contribute to the feature channel. The initial feature map was built with 16 Gaussian parameters dę starting 0.5 Bohr and ending at 5.0 Bohr. We used a 3-level U-Net with 32, 64, and 128 feature channels, respectively. The resultant model has 990k trainable parameters.
The following 5 baselines are interpolation nets, as they take query coordinates and try to interpolate them from the node (atom) information.
20

---

DeepDFT [77] and DeepDFT2 DeepDFT is GNN-based network that models the interaction between the atom vertices and the query vertices for which the charge density predicted. As DeepDFT only takes the invariant features of atom types and edge distance as input, it is also globally equivariant. DeepDFT2 uses PaiNN - as the GNN architecture. PaiNN designs equivariant interaction between scalar and vectorial features. Therefore, DeepDF12 is locally equivariant. We simply followed the original model architectures which gave models of 2.04M and 2.93M trainable parameters, respectively.
EGNN EGNN defines an equivariant message passing based on invariant edge features like distance embedding. We used 4 layers of EGNN with an input dimension of 128 and hidden and output dimension of 256, resulting larger model than the original EGNN paper. We also added a query-specific residual GNN similar to InfGCN. The model has 2.27M trainable parameters.
DimeNet 24 and DimeNet++
spherical 21) Fourier- Bessel basis functions (2D) analogs to spherical harmonics) to embed bond angles, hoping to capture geometric information about the interaction between atoms. They have achieved SOTA performance on physical property prediction. We slightly modified the original models to output a 128-dimensional feature for each atom and added query-specifie residual GNN similar to InfGCN. All the other model hyperparameters are the same as the original models. As result, the two models have 2.3 M and 2.02M parameters, respectively.
DimeNet
The following 3 baselines are neural operators. They directly try to parameterize the Fredholm operator in Eq. 2 using various approaches. Same CNN, they cannot automatically capture the grid information, so we also use the feature in Eq. as the initial feature. The initial feature map for these models is built with 32 Gaussian parameters og starting at 0.5 Bohr and ending at 5.0 Bohr. For all NOs, a sampling training and inference scheme is utilized. We will discuss it in detail in the next subsection.
The Graph Neural Operator (referred to as Graph Kernel Network or GKN in the original paper) tries to parameterize the Fredholm operator in Eq. 2 with the message passing on Monte Carlo sampled random subgraphs:
GNO 28
f(x) +
W [S(()+ + (V(I) F(xx
(35)
N(w)
where F is neural net. Note that GKN is neither translation equivariant nor rotation equivariant. We used a feature size of 384 and stacked 4 convolutional layers. The cut-off distance was 5.0 Bohr for all datasets. The resultant model has 1.84M trainable parameters.
FNO 29 The Fourier Neural Operator does the parameterization in the Fourier domain:
f(x) + (W + +F +F ((F)(x))
(36)
where F,F-1 are the Fourier transform and inverse Fourier transform over the sampled data points, and W. R are learnable parameter matrices. For the parameterization in the Fourier domain, only a
fixed number of uwwffreuency Fourier modes are kept for efficiency. We used a feature size of 128 with a number of Fourier modes of 128 and stacked 4 layers. The cut-off distance was 5.0 Bohr for all datasets. The resultant model has 33.63M trainable parameters.
LNO 32 The Linear Neural Operator is based on the low-rank decomposition of the kernel W W(X y) = EE= 9, (x))), similar to the unstacked DeepONet proposed in (32). We used feature size of 384 with rank of 64 and stacked 4 layers. The cut-off distance was 5.0 Bohr for all datasets. The resultant model has 803k trainable parameters.
To facilitate model training and evaluation, we also defined a unified model interface such that each model takes the atom types, atom coordinates as defined above, and the sampled densities and sampled grid coordinates which we will cover below. The model outputs predicted density values at each sampled grid point.
D.2 Training Specification
We followed 17 to use a sampling training scheme, and we also adapted it for every model except for 3D-CNN. During training and validation, only small portion of the grid is randomly sampled
21

---

which with the corresponding density values. This scheme drastically reduces the required GPU memory, as there were cases when the whole voxel could not fit into a single GPU. During inference. however, all voxels were partitioned into mini-batches for full evaluation to avoid randomness for a more convincing result. The 3D-CNN model required the whole voxel information, so the sampling scheme was not used.
As was demonstrated in 288] this sampling scheme can be best understood as Nyström approximation of the integral in Eq. 2 The original FNO and LNO models used the whole grid points for the estimation of the integral (Monte Carlo approximation). This one of the major reasons that these models cannot scale to 3D voxel data. In our experiment, FNO and LNO would cause OOM even for the QM9 dataset with a batch size of 1.
The training and testing specifications are provided in Table The cut-off in Bohr refers to the cut-off distance for building the atom-atom graph and the atom-query graph for InfGCN and interpolation nets. All training was done on single NVIDIA A 100 GPU. For efficiency, the testing for QM9 was done on the last 1600 samples, larger molecules were tested. For Cubic and MD, testing was done on all the test samples.
Table 4: Training specifications.
Model
Dataset
cutoff'
n. jter
Ir
patience
batch_ size
Ir decay
train_ sample inf sample
QM9
3.0
40k
lc-3
10
64
InfGCN
Cubic
5.0
10k
5e-3
5
32
0.5
1024
4096
MD
3.0
5e-3
5
64
QM9
100k
3e-4
10
4
CNN
NA
0.5
NA
NA
MD
4k
lc-3
5
32
QM9
3.0
40k
3e-4
10
64
DeepDFT/
Cubic
5.0
10k
3e-4
10
32
0.5
1024
4096
DeepDFT2
MD
3.0
le-3
5
64
EGNN/
QM9
5.0
40k
3e-4
10
DimeNet/
Cubic
5.0
10k
3e-4
10
64
0.5
1024
4096
DimeNet++
MD
3.0
2k
le-3
5
GNO
QM9
80k
3c-4
10
FNO/
Cubic
NA
10k
3c-4
10
32
0.5
1024
4096
LNO
MD
2k
1c-3
5
D.3
Complexity Analysis
The naive implementation of a mnessage-passing layer with padding and masking scales O (EE 1)") where |E| is the number of edges and C' is the number of channels. This is because the message passing step involves a tensor product of two spherical tensors and the Trebsch-Gordand coefficients have their six indices all summed. However, note that there are only (e + 1)2 spherical harmonics with the degree up to l. If coefficients are compressed into one long vector, the complexity can be reduced to O CEE + 1). During the expansion of the basis functions on the voxels, the time complexity O KC(E + 1)2) where is the number of voxels sampled. In practice, we noticed that a small e suffices that (e+ 1)a can be viewed constant (e = 7 corresponds to 4096). Also, the Ttebsch-Gordae efficiently done by index manipulation. Our implementation was based on the implements efficient spherical vector manipulations.
coefficients can be pre-computed and stored, and the summation can be
In comparison, most GNN-based layers scale (?E where D is the hidden feature size. There- fore, in our standard-setiing setting (C = 16, L 7), the time complexity approximately of the same order (with D = 256). For GNO, FNO, and LNO, one layer scales as O D2), O(K D O(K (())) respectively. The additional coefficients are for the Fourier modes or the rank. For 3D-CNN, the time complexity scales as OCi,n Nr N,,,,, N33) where k is the kernel size. This is significantly larger than any of the GNN-based methods, as the whole voxel needs to be encoded and decoded. In practice, we found the interpolation nets ran slightly quicker than InfGCN and NOs, but our proposed InfGCN was able to achieve better performance.
https://e3nn org/
22

---

E Additional Results
In this section, we provide more visualization results on the QM9 and Cubic datasets to further study the generalizability of InfGCN. For the Cubic dataset, we further provide a sample on the cubic primitive cell Figure
Groand
LOOOK
DepOeT
Dep0F72
ECKS
Dandis
Phodariic
Dauetr
SVAR (3)
Figure 6: Visualization of the predicted density, the density error, and the NMAE for RbgAsO, (mp-1103055 in Cubic).
For the density visualization, we used linearly spaced values from 0.05 to 3.5 with 5 isosurfaces for the ground truth density and -0.03 (deep pink), -0.01, 0.01, and 0.03 (deep green) for the density errors for QM9 in Figure 2 We used values of 0.3, 3, and for the ground truth density and -0.3, -0.1, 0.1, and 0.3 for the density errors for the rhombic primitive cell in Figure 2 We used linearly spaced values from 0.5 to 8.0 with 5 isosurfaces for the ground truth density and -0.3, -0.1, 0.1, and 0.3 for the density errors for Cubic in Figure 6
As QM9 covers broad range of different types of molecules, we manually picked some representative molecules and evaluated all the models. The results are provided in Table In order to provide finer- ther-grained comparison between different models, we used finer- merrggranee isosurfaces with values of -0.03, -0.01, 0.01, and 0.03, respectively. The corresponding NMAE (%) is also provided below the plot. Molecule names, their corresponding file IDs, chemical structures, and ground truth densities are also provided in the table.
The selected molecules cover variety of chemical types, ranging from alkane, alcohol, and ester to aromatic heterocyclic compounds. InfGCN has a significant advantage over other baselines, demonstrating its generalizability across different molccules. We also observe some patterns of the density estimation error:
All models performed better on alkanes or the alkyl part of the molecules, e.g., the lincar nonane, the branched r-butyl group, and the cyclic cyclohexane. One exception is cubane, which adopts an unusually sharp 90º bonding angle that would be highly strained as compared to the 109.45 angle of tetrahedral carbon. In this example, InfGCN was still able to make a relatively accurate prediction.
The predictions were worse for atoms with high electronegativity (F, O, N), even if the errors are normalized by the number of electrons. From chemical aspect, higher density will lead to a better "polarizing" ability to distort the electron cloud of the covalent atom, leading to more complicated higher-order interactions between atomic orbitals [ For example, the carbon atom in CF2 has significant positive partial charge, but DeepDFT overestimated its
density (in pink). Noticeably, InfGCN can estimate the oxygen density with great accuracy, e.g., in glycerol, di-r-butyl ether, and isoamyl acetate.
The predictions were even worse for a conjugated and aromatic system where electrons are highly delocalized Delocalization of electrons allows for long-distance interactions which are harder to capture. The presence of electron-donating groups (EDGs) and electron- withdrawing groups (EWGs) contributes greatly to the conjugated system. For example,
"The electron cloud actually won't undergo any polarizing or hybridization procedure, but it is still beneficial to think of it in this way. as it provides an intuitive interpretation of the linear combination of atomic orbitals (LCAO). In practice, such ideas still widely used in chemistry.
'The term delocalized is also not accurate for describing electron density. However, as mentioned above, it is a convenient way of building molecular orbitals from atomic orbitals, and so are the concepts of EDG and EWG mentioned in the following text.
23

---

Figure 7: Resonance forms of aniline and nitrobenzene with formal charges.
the amino group 2H an EDG when bonding to a conjugated system like benzene, facilitating ortho- and para-electrophilic reactions. In contrast, the nitro group -NO is EWG that facilitates ortho- and vara-nucleophilic reactions (See Fig It can be seen from the visualization that the ortho- and para-positions positions of aniline underestimated (in green) and those of nitrobenzene are overestimated (in pink) with DeepDFT and other models. For cytosine and fluorouracil, the amide (lactam) tautomeric form predominates pH further making the electron structures more complicated to achieve accurate predictions. More examples of the conjugated systems include the nitro group itself where the density of oxygen is overestimated and the amide group in asparagine where the density of the amide oxygen underestimated and that of nitrogen is overestimated.
Table 5: More results on QM9.
File
Ground Truth
Name
Structure
InfGCN
CNN
DecpDFT DepDFT2 EGNN DimeNet DimeNetee
GNO
FNO
LNO
ID
ammonia
2
NH3
3
22.79
5.34
2,22
12.27
11.45
15.22
STAF
137.5
urea
20
NyM
44.72
22
acetone oxime
49
furan
52
tetrafluoro- methane
184
CF4
1.76
4.63
L95
1.00
9.75
0.49
7.75
3A.19
30.31
27.36
glycerol
397
0.79
252
2.77
0:38
11.73
11.36
11.03
20121
23.70
cyclohetane
658
1.55
11:21
Continued next page

---

61439 Table 5: More results QM9. (Continued)
aniline
940
1.22
4.77
2.85
1.30
12:82
12:50
12.16
35.25
51.35
24.00
cytosine
4318
125
5.62
2.55
1.37
16335
J0J18
10
24:38
20
cubane
19116
2.70
11MM
123I
purine
24537
di-z-butyl ether
57520
isoamyl acetate
60424
067
2
2
0:7
12.1
12.13
LL
201
27.16
30
U
asparagine
1/12
371
3,12
1:11
10.0
M2
JAM
20
nonane
114514 865
254
2.31
6A5
12:1
JN32
2010
nirrobnnoeee
131915
22:93
Auorouracil
132701
25


Reviews:
[{"review_id": "npYO8sLqQU", "reviewer": "Reviewer_cHMR", "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.\nThe 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. ", "questions": "- Are there any advantages for explainability from the framework?", "limitations": "The authors discuss the limitations of their work. ", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal to Reviewer cHMR\n\nWe really appreciate your high recognition of the high quality of our work. We will address your concerns and questions as follows.\n\n\n## Q1 SE(3)/SO(3) equivariance\n\nClarification regarding SE(3)/SO(3) is made in the **rebuttal to all reviewers**. We thank you for pointing this out and we will make changes to the revised manuscript to avoid potential ambiguity. \n\n\n\n## Q2 Explainability \n\nThe explainability of our proposed model may be elaborated on in the following aspects:\n\n- The designed message passing can be interpreted as the linear combination of atomic orbitals (**LCAO**), which has a close connection with quantum chemistry (Sec.4).\n- The overall architecture can be viewed as **graphon convolution**. In this sense, every grid point was able to aggregate information from all other grid points in $\\mathbb{R}^3$, making the model more expressive.\n- As a model for the operator learning task, **intermediate results** from different graphon convolutional layers can also be exacted and visualized to provide more insights.\n\nStill, we note that, as a machine learning model, the detailed information flow of our model also relies on the development of the theoretical parts in ML and GNN to provide more insights into the whole model.\n\n\n\nWe appreciate your review to help make our work more comprehensive and solid, and we hope the above clarification will properly address your concerns and problems. Please comment if you have further questions. We are more than happy to address any further concerns."}, {"review_id": "bS21vD6SPq", "reviewer": "Reviewer_Wmsq", "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.", "questions": "* Questions\n  * 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?\n* Possible Typos\n  * Equation 3, change f_{i,u} to f_{u,i}\n  * Equation 5, and 7, consider adding \\qquad after the comma for readability\n  * Fig 3. Consider using the same axis upper bound for both plots to improve readability\n", "limitations": "The authors have communicated limitations of their work", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n* Quality: The work appears to be of very good quality.\n* Clarity: The paper is well written and the ideas are communicated in an accessible way.\n* 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal to Reviewer Wmsq\n\nWe really appreciate your recognition of our **originality** in designing the novel equivariant operator learning scheme and the good **quality** of the overall work. We also thank you for acknowledging the **significance** of our work in scientific domains. We will address your concerns and questions as follows.\n\n\n\n## Q1 SE(3) equivariance\n\nWe also noticed the performance gap between the rotated and unrotated QM9 dataset, and we believe it may be attributed to the following two factors:\n\n- We used **trilinear interpolation** on grid values to obtain the rotated target function instead of recalculating them from quantum chemical methods. Therefore, the interpolation may introduce errors to the rotated ground truth target function.\n- The grids rotated outside the original cell were **omitted** and the new grids that did not have near interpolation anchors were **masked out** (instead of zero-padded). In other words, only the intersected region between the original and the rotated grids was kept. Normally, the masked grids are often close to the cell boundary with near-zero target values that are easy to predict for all models. So, after masking out, the normalized MAE (MAE divided by the summation of all grid target values) will naturally increase.\n\nWe would also like to point out that the equivariance property of some baselines is more obvious to demonstrate than InfGCN. For example, DeepDFT used only the scalar features like the distances between atoms and grids and was trivially equivariant. DeepDFT2 was built upon PaiNN, an equivariant GNN that designs tensorial operations between scalars and vectors, so it was also equivariant. We analyzed the equivariance of all models in Appendix D.1. Nonetheless, there was still a **performance gap in these equivariant baselines**. Therefore, we think this gap should be attributed to the bias when we constructed the rotated dataset instead of the expressiveness of the models.\n\n\n\n## Q2 Estimation of overlap integral $S$\n\nThe learnable radial net $\\varphi$ in Eq.5 is closely related to the overlap integral $S$ but is not identical to the latter (as we mentioned in Sec.4 line 211). As we need to **parameterize** each InfGCN layer with different parameters, $\\varphi$ is actually the combination of $S$ and the learnable weight (see Eq.3 or Eq.21 in Appendix A). Therefore, it is unnecessary and difficult in practice to enforce a strict equation between $S$ and $\\varphi$. On the other hand, our proposed model is quite stable during training. Therefore, we may assume that $\\varphi$ at least captures the magnitude or other features of the overlap integral. \n\n\n\n## Q3 Typos & formatting\n\nWe appreciate you pointing out these typos and formatting problems in our paper, and we will fix them in the revised manuscript.\n\n- **Eq.3**. We will modify the notation to be consistent with line 117.\n- **Eq.5 & 7**. We will add space for readability.\n- **Fig.3**. We originally used a different scale for better visualization, as there existed a performance gap between the rotated and unrotated QM9. We will revise the figure.\n\n\n\nWe value your review to help make our work more comprehensive and solid, and we hope the above clarification will properly address your concerns and problems. Please comment if you have further questions. We are more than happy to address any further concerns."}, {"review_id": "iwZwRBP5jv", "reviewer": "Reviewer_sQju", "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.", "questions": "see weakness", "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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nIn 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.\n", "weaknesses": "Some issues are as follows:\n\n1. 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.\n\n2. 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.\n\n3. 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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal to Reviewer sQju\n\nWe appreciate your recognition of our **novel design of an equivariant operator learning** scheme and **detailed theoretical interpretation**. We also thank you for acknowledging our **extensive experiments** on large-scale datasets and **superior performance** over the baselines. We will address your concerns and questions as follows.\n\n\n## Q1 More baselines \n\nWe understand your concerns regarding the baseline models, and we would like to first elaborate on our experiments and baselines to provide justification for the proposed model's effectiveness.\n\n- We have experimented with a **wide variety of baseline models** ranging from GNNs, CNNs to neural operators. Many of the baselines came from recent work and often achieved SOTA performance in their tasks, as we mentioned in the **rebuttal to all reviewers**.\n- We carried out **extensive experiments and ablation studies** on large-scale real-world electron density datasets to demonstrate the effectiveness of our model. The QM9 density dataset totals up to more than 1TB of data, and the Cubic dataset contains 84 different types of atoms. In this way, the superior performance can indeed demonstrate the generalizability of InfGCN.\n\nAdditional results and details of **EGNN** as another baseline are also available in the **rebuttal to all reviewers**. In conclusion, EGNN was able to achieve better performance than DimeNet/DimeNet++ on QM9 and Cubic. However, there is still a significant gap between InfGCN. We believe this result further demonstrate the effectiveness of our proposed model.\n\n\n## Q2 Complexity analysis\n\nWe made an explicit complexity analysis in **Appendix D.3** in the original manuscript, and we will briefly summarize the content here. We are aware of the high complexity of using higher-degree spherical harmonics, which has an overall complexity of $O(|\\mathcal{E}|KC(\\ell+1)^4)$ (with optimization), where $\\mathcal{E}$ is the edge set, $K$ is the number of the sampled grid (the same for all models except for CNN), and $C$ is the number of radial channels. We noted that in our experiment setting ($\\ell=7,C=16$), this factor is of approximately the same order as GNN-based models like DeepDFT with a hidden size of 256. In practice, though, we noticed that interpolation nets are slightly quicker than InfGCN and neural operators.\n\nAs InfGCN essentially constructs the graph based on the coordinates of the discrete structure, it inherits the **localization assumption** of normal GNNs. In this way, scalability to larger graphs can be naturally achieved with a reasonable cut-off distance. Furthermore, in Fig.3 and Tab.2, we noted that the performance improvement from using $L=4$ to $L=7$ is not significant as a result of using the residual operator layer. Therefore, we may use a smaller $L$ for larger graphs to accelerate the calculation without significantly damaging the overall performance.\n\n\n\n## Q3 Formatting\n\nThanks for pointing out this problem. We will modify the title of Sec.3.3 to \"Equivariant Message Passing\" to avoid confusion and ambiguity.\n\n\n\nWe value your review to help make our work more comprehensive and solid, and we hope the above clarification will properly address your concerns and problems. Please comment if you have further questions. We are more than happy to address any further concerns."}, {"review_id": "fqtWXxtsdg", "reviewer": "Reviewer_cyPX", "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.\n", "questions": "See the weakness section.", "limitations": "There is no negative societal impact as far as I can tell.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Generalizing neural operator learning to 3D Euclidean space with equivariance guarantee is an interesting and underexplored research question.\n2. The authors provide theoretical explanation to the model by generalizing graph convolution to graphons.\n3. 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:\n    - Why the target feature function $\\rho$ can be expressed as the equation in line 117? What is the rationale behind the equation?\n    - 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?\n    - How is the spherical tensor in section 3.3 connected with the basis function in the construction of $S$?\n    - What is the motivation of using residual operator layer? Why it can mitigate the expressivity issue?\n2. 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.\n3. 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.\n4. 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.\n\n[1] E (n) equivariant graph neural networks", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal to Reviewer cyPX\n\nWe appreciate your recognition of our work's **novel equivariant operator learning** scheme and **rigorous mathematical formulation and proof**. We also thank you for acknowledging our **extensive experiments** on large-scale datasets. We will address your concerns and questions as follows.\n\n## Q1 Expansion of $\\rho$ \n\n We apologize that the equation on line 117 should be instead a **definition**. Here, we are aware of the crucial role of the **underlying discrete structure** that has a significant physical connotation, e.g., atoms in electron density (line 19 & 74). Therefore, we took each discrete node as the center of coefficient expansion and summed over all nodes to obtain the final predicted function $\\hat{\\rho}$. In this way, we may mitigate the error introduced by using a finite basis set.\n\n## Q2 Basis set\n\nWe understand your concerns regarding the choice of basis set in Eq.4, and we will elaborate on the reasons that we chose basis functions in this form.\n\n- They are **widely used in quantum chemical** calculation (known as the **Gaussian-type Orbital**, GTO). The family of solutions of the Schrödinger equation of the hydrogen-like atoms shares a very similar formula as Eq.4. We also mentioned in Sec.4 (line 204) that a claim from Tsubaki et al. can be adapted for our proposed model to be interpreted as the linear combination of atomic orbitals (LCAO).\n- In our settings, the **discrete structure** also plays a vital role in determining the feature function. Therefore, it is reasonable to assume the function is more \"clustered\" around the node and use the node-centered Gaussians as the radial function that decays exponentially.\n- They are relatively **easy to evaluate**. It is possible to use a complete orthonormal basis, as we mentioned in Sec.7. For example, Guseinov (Int. J. Quantum Chem., 2002) proposed such a basis for $L^2(\\mathbb{R}^3)$ whose time complexity of evaluating and backpropagation makes it unsuitable and unnecessary to use them in ML scenarios.\n\nWe further point out that the choice of the radial part in the basis function **does not influence the equivariance** as long as the spherical harmonics part remains. Conceptually, the radial part and the angular part are independent, with the spherical harmonics completely describing the angular part. In Sec.7, we have noted this fact and proposed **potential future work** by using a different basis set that is more task-tailored. \n\n\n\n## Q3 Construction of $S$\n\nWe point out that the overlap integral $S$ is a continuous function of $\\mathbf{r}$. Therefore, we can apply the same **expansion** scheme to represent $S$ as a series of coefficients (as explained in the **rebuttal to all reviewers**). Such an expansion to a spherical tensor was primarily to leverage the fact that the **tensor product of two spherical tensors is equivariant**. Therefore, representing the overlap integral as a spherical tensor, we can design an equivariant message-passing scheme for continuous feature functions.\n\n## Q4 Motivation of the residual operator layer\n\nThe proposed residual operator layer is one of the major contributions of our work, and we will elaborate on the motivation as follow:\n\n- We can only use a **finite set of basis**, which can not be complete for the infinite-dimensional space $L^2(\\mathbb{R}^3)$. This problem limits the expressiveness of the overall model, which we refer to as the \"finite approximation error\" in the manuscript.\n- The proposed residual operator layer resembles the \"interpolation nets\" and takes **query-specific coordinates** as input to produce a **per-coordinate scalar** to \"finetune\" the density. In this sense, this layer can approximate arbitrary residue values between the coefficient-based output and the ground truth given the model is expressive enough.\n- The experiment results have demonstrated the **effectiveness** of such a residual scheme. In our ablation studies, performance improved considerably with the residual operator layer (Tab.2, \"no-res\"). \n\n\n\n## Q5 Preliminary & background\n\nThanks for your valuable advice. We will add an additional introduction to related topics like quantum chemistry and TFN to make our work more self-contained.\n\n\n\n## Q6 Comparison with prior art\n\nWe already summarized our contribution to the ML realm in line 67-72, and we will also briefly discuss the comparison of our ML-based model with traditional quantum chemical calculation. As mentioned in Sec.1 the current art are various *ab initio* methods that scale as $O(N^7)$. They are often regarded as the ground truth density provided by the dataset but take considerable computation time. Other quicker quantum chemical methods like KS-DFT have an $O(N^3)$ scaling factor but less accuracy. In the DeepDFT papers (Jørgensen, 2020; *npj Comput Mater* 2022), the authors made a systematic comparison and claimed that DeepDFT outperformed KS-DFT in terms of accuracy. As our model also outperformed DeepDFT, we may obtain the qualitative conclusion that our model also makes more accurate predictions than KS-DFT.\n\n## Q7 Datasets & baselines\n\nWe further justification for using the density datasets:\n\n- The electron density provides a fundamental description of a molecule in quantum chemistry theory. The predicted density can be further utilized to **improve the performance** on other molecule-related tasks.\n- The electron density is **widely encountered** and **large-scale real-world** datasets are available. For example, the QM9 dataset contains 134k molecules, totaling up more than 1TB of data. Therefore, we believe these large-scale datasets suffice to demonstrate the effectiveness of our proposed model.\n- Other data are less explored. For example, most current work on the neural operator is tested on generated PDE as toy data or 2D data. 3D voxelized data are even scarcer, as they require significantly larger space to store.\n\nResults of **EGNN** are provided in the **rebuttal to all reviewers**."}, {"review_id": "S6L3FYtMLz", "reviewer": "Reviewer_SJUC", "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). ", "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)?\n\nLine 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.\n\nLine 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.\n\nLine 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.\n\nLine 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. \nYou did not discuss the setting properly. What is the data structure/discretization that you work with?\n\nEquation (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.\nMoreover, 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).\n\nLine 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.\n\nLine 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).\n\nFormula (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.\n\nFormula (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. \nWhere is the learnable part of this message passing scheme? This seems like a definition of aggregation or shift operator.\n\nSection 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?\n", "limitations": "Limitations are discussed.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "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.\n\nThe 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.  \nThere 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).\nSee more details in “questions.”\n\nThe 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.\n\nIt seems like this paper has potential, but it is not ready for publication at its current state.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Rebuttal to Reviewer SJUC\n\nWe appreciate your recognition of our work's **underlying physical connection** as well as our **comprehensive experiments**. We are willing to address your concerns and questions. Ambiguity regarding SE(3)/SO(3) is addressed in the **rebuttal to all reviewers**.\n\n\n\n## Q1 Basis function\n\nIn our setting, the basis was **pre-defined for the whole dataset** and was not a learnable part of the model. It can be shown that the basis set defined in Eq.4 indeed spans a subspace in $L^2(\\mathbb{R}^3)$, as the functions are square-integrable over $\\mathbb{R}^3$ and any finite set of them is linearly independent. We are aware that, however, they are neither orthogonal nor complete, as we already mentioned on line 172. We proposed multicentric approximation (Sec.3.2) and the residual operator layer to address this problem (Sec.3.4).\n\nWe explained the **basis expansion** in this sense in the **rebuttal to all reviewers**, and we further explain the \"**equivariant basis**\" in line 132. Your understanding was correct that we want to find a basis set such that the rotation operator commutes with the neural nets on the coefficients. We will rephrase the sentence to avoid confusion.\n\n\n\n## Q2 Construction of graphon\n\nWe appreciate that you mentioned other work on spectral convolution on graphon, e.g., **WNN** (Ruiz, NeurIPS 2020), **cMPNN** (Maskey, NeurIPS 2022), and **MNN** (Wang, ACSSC 2022). These papers provided additional insights into the problem formulation. Though sharing similar ideas of utilizing continuous geometric structures, our method is fundamentally different from these methods in the following aspects:\n\n1. **Construction of graphon**. In the above work, the authors either explicitly construct the graphon kernel $W$(WNN, MNN) or directly estimated the Fredholm integral (cMPNN) in a similar fashion as various neural operators. Our approach, on the other hand, **implicitly constructs the graphon and parameterizes them in the spectral domain**. We mentioned in line 101-105 that $W$ does not have a canonical form in our setting and Monte Carlo estimation is prohibitive for large voxels. Instead, we defined a basis setting and demonstrated that transformation on the coefficients can be also viewed as graphon convolution equivalently (Appendix B.2).\n\n   In this way, we implicitly assume that there exists **a different graphon for each data sample** defined by their discrete structure and their categorical information. In our experiments, they are the atom coordinates and atom types. Nonetheless, the parameterization of graphon was done with the same graphon convolution for the whole dataset, as we expected this information to be generalizable across different samples.\n\n2. **Data structure & task**. In our settings, we further assume there exists an **underlying discrete structure** that has a significant physical connotation, e.g., atoms in electron density (line 19 & 74). The datasets on electron density we experimented with are **real-world** data and are **significantly larger** than those used in previous graphon-related work.\n\n3. **Model design**. Based on the above difference in data structure and tasks, we designed a new network architecture that is different from the work before. We approximated the graphon convolution with NNs on the coefficients instead of the feature function itself, and we also proposed the **residual operator layer** to mitigate the finite approximation error. Also, we **extended the definition of rotation-equivariance** to continuous functions and provided rigid proof that our model achieves such a desirable property by design.\n\nIn conclusion, we did not claim to have invented the graphon convolution, and our reference to graphon might be better viewed as an interpretation of our proposed model from the graph spectral aspect. Nonetheless, we are aware of the values of previous work on graphon in the sense that their theoretical results on convergence and transferability can be adapted for our proposed model to make it more concrete and solid. We will briefly mention these papers in the revised manuscript.\n\n\n\n## Q3 Other questions\n\nQuestions regarding **line 88 & 142 and Eq.4** are addressed in **the rebuttal to all reviewers**. We will also follow your advice on **line 99 & 146 and Eq.5** to avoid ambiguity.\n\n- **line 104 & 105**. In this claim, the rationale can be elaborated as, 1) most current Monte Carlo-based neural nets approximate the integral **using the full (2D) grid** which will be prohibitive for 3D voxels; and 2) Monte Carlo approximation requires considerably **large samples to make an accurate estimation**, as claimed in the GNO paper (one of our baselines). As for the dimensionality problem, we wanted to **expand 2D operator learning to 3D scenarios**. In our experiment settings, for example, we used 1024 basis functions, whereas the number of voxels can easily scale up to 20M, making it impractical to use the full grid. In this sense, we claimed that coefficient-based learning is more efficient and accurate than the Monte Carlo-based method.\n- **line 115**. As we mentioned in line 19-20 and 74, the discrete graph is an **additional input to our model** (and all the baselines). We mentioned in Q2 that it was one of the aspects that distinguish our work from other graphon-based work.\n- **Eq.5**. We will use $\\gets$ (\"gets\") to avoid ambiguity. $\\varphi_{J}^{\\ell k}:\\mathbb{R}^+\\to\\mathbb{R}$ are learnable neural nets that embed the radial length $r=\\|\\mathbf{r}\\|$. Here, the radial index is also omitted for clarity. The full deduction of Eq.5, including the learnable $\\varphi_{J}^{\\ell k}$ can be found in Appendix A.\n\n\n\nWe value your review to help make our work more comprehensive and solid, and we hope the above clarification will properly address your concerns and problems. Please comment if you have further questions. We are more than happy to address any further concerns."}]

Author rebuttal:
# Rebuttal to All Reviewers

Dear Reviewers, we would like to first express our sincere gratitude for your valuable reviews to help make our work more comprehensive and solid. We thank your recognition of the innovation of our work, and we are willing to address your concerns and questions. Here we address some **common questions** and other review-specific questions and concerns are addressed separately.



## Q1 More baselines

We notice the concern of Review sQju & cyPX regarding using more recent baselines. We would like to point out that many of our baselines are newly emerged models, often with SOTA performance in their tasks. Most noticeably, **DeepDFT2** was built upon **PaiNN** (Schütt, ICML 2021), an equivariant GNN with carefully designed tensorial operation. As demonstrated in their paper, PaiNN outperformed EGNN on QM9 and DeepDFT2 indeed provided a very strong baseline. Other models like **FNO** (Li, ICLR 2021), **GNO** (Li, 2020) are also relatively new with the claimed SOTA performance. 

Nonetheless, we also ran the **EGNN** model as another baseline on all datasets, as per request by cyPX. The EGNN code was adapted from the official repo by Satorras et al. In our experiment setting, we used a larger model than the original setting for a fair comparison. We used 4 layers of EGNN layers with an input dimension of 128, hidden and output dimension of 256. This resulted in 2.27M trainable parameters. Other training hyperparameters are identical to DimeNet. The results are shown below:

| Dataset          | NMAE (%) |
| ---------------- | -------- |
| QM9 rotated      | 12.13    |
| QM9 unrotated    | 11.92    |
| Cubic            | 11.74    |
| MD-ethanol       | 13.90    |
| MD-benzene       | 13.49    |
| MD-phenol        | 13.59    |
| MD-resorcinol    | 12.61    |
| MD-ethane        | 15.17    |
| MD-malonaldehyde | 12.37    |

EGNN was able to achieve better performance than DimeNet/DimeNet++ on QM9 and Cubic. However, there is still a significant gap between InfGCN.



## Q2 Equivariance wrt SE(3)/SO(3)

We are sorry for the confusion between SE(3) and SO(3) in the manuscript, and we will clarify it here. We first point out our proposed model is **translation-invariant**, as we only use the relative displacement vector (line 155). Therefore, together with SO(3) equivariance (Appendix A), our model is SE(3) equivariant. We acknowledge that for most of the part in Sec.2-4, we referred to SO(3) (rotation) when we mentioned "equivariance" in our formulation, primarily to be on par with other classical equivariance papers including SE(3)-Transformer (Fuchs, NeurIPS 2020) and EGNN (Satorras, 2021), where they also used SE(3)/E(3) formulation but mainly dealt with SO(3)/O(3) in the context. We will explicitly mention the difference between SE(3) and SO(3) in the introduction section and use SO(3) in the method section to avoid potential ambiguity.



## Q3 Overlap integral $S$ & Expansion

We also notice some ambiguity regarding $S$ in line 142. $S$ should have 6 indices $n_1\ell_1m_1n_2\ell_2m_2$ where $n_1\ell_1m_1$ are basis indices of the first expansion center and $n_2\ell_2m_2$ of the second. We omitted them for clarity as the arguments suit for all overlap integral between any of the two basis functions. As for the **expansion** of a feature function $f$, we slightly abused the notation to indicate a set of coefficients that minimize the norm:
$$
\arg\min_{c_{n\ell m}}\left\\|f(\mathbf{x})-\sum_{n\ell m}c_{n\ell m}\psi_{n\ell m}(\mathbf{x})\right\\|_2
$$
Therefore, line 142 should be viewed as a **finite basis approximation**. We believe such an approximation was reasonable because the overlap integral also decays exponentially with respect to the distance between centers. We often wrote the expansion as $f=\mathsf{f}$, where $\mathsf{f}$ is the coefficient spherical tensor indexed by $n\ell m$. In Sec 3.3 and Appendix A, we also omitted the radius $n$ as it is independent of rotation equivariance, and further used the notation of $f=\mathsf{f}=\\{\mathbf{f}^\ell,\ell\ge 0\\}$. We will use $\approx$ and "approximate" for a more rigorous description. 



We value your reviews and comments and hope the above clarification will properly address your concerns and problems. Please comment if you have further questions. We are more than happy to address any further concerns.


Meta-review: The reviewers strongly recommend acceptance. The paper provides a method to develop SE(3) neural operators to output functions in 3D. The authors show significant improvement over prior art that do not preserve SE(3).
Decision: Accept (spotlight)