Title: Equivariant Neural Operator Learning with Graphon Convolution
Venue: NeurIPS 2023
Accepted: yes
Keywords: Neural Operator Learning

Spectral Graph Theory

Graphon
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

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.

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 an amplitude of applications. For example, the Schrödinger equation for the wave-like behavior of the electron in a molecule, the Helmholtz equation for the time-independent wave functions, and the Navier-Stokes equation for the dynamics of fluids all output a continuous function spanning over R 3 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 [1, 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 on quantum chemical computation as the approximation method to solve ODEs and PDEs. For example, the ground truth electron density is often obtained with ab initio methods [48, 26] with accurate results but an N 7 computational scaling, making it prohibitive or inefficient for large molecules. Other methods like the Kohn-Sham density functional theory (KS-DFT) [24] has an N 3 computational scaling with a 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 equivariance for a function defined on R 3 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 is independent of the poses of the input structure, thus also serving as 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, 15] and

tensor field networks [50, 10]. We notice the close relationship between the tensor field network (TFN) and the equivariance of the continuous functions and also propose our equivariant architecture based on the tensor product.

In this way, we define our task as equivariant neural operator learning. We roughly summarize previous work on operator learning into the following four classes: 1) voxel-based regression (3DCNN) [47, 40, 4]; 2) coefficient learning with a pre-defined set of basis functions [2, 27, 49, 14]; 3) coordinate-based interpolation neural networks [17, 18]; and 4) neural operator learning [28, 25, 29, 32]. The voxel-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 of basis functions. The coordinate-based networks take the raw coordinates as input and use a learnable model to 'interpolate' them to obtain the coordinate-specific output. They leverage the discrete structure and provide a strong baseline, but a 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 1D or 2D data and have difficulty scaling up to 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-specific information. A graphical overview of our model architecture is shown in Fig.1. 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-art 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 a novel architecture that combines coefficient learning with the coordinate-based residual operator layer, with our model guaranteed to preserve SE(3)-equivariance by design; 2) we provided a theoretical interpretation of our model from the graph spectrum point of view as graphon convolution; and 3) we carried out extensive experiments and ablation studies on large-scale electron density datasets to demonstrate the effectiveness of our proposed model. Our code is publicly available at https://github.com/ccr-cheng/InfGCN-pytorch .

(a)

e)

2(x)

米

r

p(x)=

fent a)em(x) + 2(x)

fm

2 Preliminary

We use G = ( V , E ) to denote the (discrete) graph with the corresponding node coordinates { x i } |V| i =1 . A continuous function over the region D is also provided as the target feature: ρ : D → R . We also assume that there is an initial feature function f in : D → R either obtained from less accurate methods, a random guess, or some learnable initialization. Formally, given f in ∈ L 2 ( D ) , which is a square-integrable input feature function over D , and the target feature function ρ ∈ L 2 ( D ) , we want to learn an operator in the Hilbert space T : L 2 ( D ) → L 2 ( D ) to approximate the target function.

Different from common regression tasks over finite-dimensional vector spaces, the Hilbert space L 2 ( D ) is infinite-dimensional.

2.1 Equivariance

Equivariance describes the behavior of the model when the input data are transformed. Formally, for group G acting on X and group H acting on Y , for a function f : X → Y , if there exists a homomorphism S : G → H such that f ( g · x ) = ( S g ) f ( x ) holds for all g ∈ G, x ∈ X , then f is equivariant. Specifically, if S : g ↦→ e maps every group action to the identity action, we have the definition of invariance : f ( g · x ) = f ( x ) , ∀ g ∈ G, x ∈ X .

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 it in our discussion and focus on rotation (i.e., SO(3)) equivariance. We first define the rotation of a continuous function f ∈ L 2 ( R 3 ) as ( R f )( x ) := f ( R -1 x ) , where R is the rotation matrix associated with the rotation operator R . Note that the inverse occurs because we are rotating the coordinate frame instead of the coordinates. In this way, the equivariance condition of an operator T with respect to rotation can be formulated as

For clarity, we will distinguish equivariance and invariance, and use equivariance for functions satisfying Eq.(1).

3 Method

3.1 Intuition

Intuitively, we would like to follow the message passing paradigm [12] to aggregate information from every other point x ∈ D . In our scenario, however, as the nodes indexed by the coordinate x ∈ D are infinite and continuous, the aggregation of the messages must be expressed as an integral:

where W : D×D→ [0 , 1] is a square-integrable kernel function that parameterizes the source node features f ( y ) . There are two major problems regarding the formulation in Eq.(2): 1) unlike the discrete counterpart in which the number of nodes is finite, parameterization of the kernel W in the continuous setting is hard; and 2) even W is well-defined, the integral is generally intractable. Some NOs [29, 32] directly approximate the integral with Monte Carlo estimation over all grid points, which makes it harder to scale to voxels. Instead, we follow a similar idea in the coefficient learning methods [2, 27] to define a set of complete basis functions { ψ k ( x ) } ∞ k =1 over L 2 ( D ) . In this way, the feature function can be expanded onto the basis as f ( x ) = ∑ ∞ k =1 f k ψ k ( x ) where f k are the coefficients. We can then parameterize the message passing in Eq.(2) as the coefficient learning with truncation to the N -th basis. We call such an expansion method unicentric as there is only one basis set for expansion. In theory, as the size of the basis goes to infinite, the above expansion method can approximate any function y ∈ L 2 ( 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 a multicentric expansion scheme. We use the node coordinates r u in the discrete graph as the centers of basis sets: ˆ ρ ( x ) = ∑ u ∈V ∑ ∞ i =1 f i,u ψ i ( x -r u ) . We demonstrated in Appendix B that with some regularity and locality assumptions, the message passing in Eq.(2) can be parameterized as

where S ij ( r ) = ∫ D ψ i ( x ) ψ j ( x -r ) d x models the interaction between the two displaced basis at centers i, j . The outer summation is over ˜ N ( u ) = N ( u ) ∪ { u } , the set of neighboring centers of u including u , and w ij are learnable parameters. Note that, once the basis functions are assigned, S ij 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 S ij ( r ) is often referred to as the overlap integral in quantum chemistry. The basis functions can be viewed as the atomic orbitals and, in this way, the integral can therefore be interpreted as the overlap between two displaced atom-centered electron clouds. The evaluation of the overlap integral is important in the self-consistent field method (Hartree-Fock method) [48].

3.3 Equivariant Message Passing

We will now consider the functions on the 3-dimensional Euclidean space, i.e., D = R 3 , as they are widely encountered in practical applications and non-trivial to achieve equivariance. It is not easy to find a set of equivariant basis that satisfies Eq.(1). Inspired by the atomic orbitals used in quantum chemistry, we construct the basis function with a Gaussian-based radial function R ℓ n ( r ) and a spherical harmonics Y ℓ m (ˆ r ) :

where r = | r | is the vector length and ˆ r = r /r is the direction vector on the unit sphere. c nℓ are normalizing constants such that ∫ R 3 | ψ nℓm ( r ) | 2 dV = 1 . The degree of the spherical harmonics ℓ takes values of non-negative integers, and the order m takes integers values between -ℓ and ℓ (inclusive). Therefore, there are 2 ℓ +1 spherical harmonics of degree ℓ . In this way, the basis index i, j are now triplets of ( n, ℓ, m ) .

To further incorporate the directional information, we follow the Tensor Field Network [50] to achieve equivariance based on the tensor product. Note that for any index pair ( n 1 ℓ 1 m 1 , n 2 ℓ 2 m 2 ) , the overlap integral S ( r ) can also be expanded onto the basis as S ( r ) = ∑ nℓm s nℓm ψ nℓm ( r ) =: ∑ nℓm φ nℓm ( r ) 1 . For a fixed r and radial index n , the coefficient sequence φ = { φ ℓm : ℓ ≥ 0 , -ℓ ≤ m ≤ ℓ } can be viewed as a spherical tensor . Notice that the node feature f = { f ℓ : ℓ ≥ 0 } can also be viewed as a spherical tensor. In the following discussion, we will omit the radial function index n for clarity as 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:

where Q ℓk Jm is the Clebsch-Gordan matrix of shape (2 ℓ + 2) × (2 k + 1) and φ ℓk J : R + → R are learnable radial nets that constitute part of the edge tensor features. A detailed deduction is provided in Appendix A. Intuitively, as TFN can achieve equivariance for spherical tensors, the output spherical tensor interpreted as the coefficients should also give an equivariant continuous function ˆ ρ ( x ) . Indeed we have

Theorem. Given an equivariant continuous input, the message passing defined Eq. (5) gives an equivariant output when interpreted as coefficients of the basis functions.

A rigorous proof of rotation equivariance can be found in Appendix A. The translation equivariance also trivially holds as we only use the relative displacement vector r uv = x v -x u . The equivariance of this scheme relies on the equivariance of the input feature map f in. Note that the 0-degree features that correspond to pure Gaussians are isotropic, so we can 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 v = u , the message passing can be simplified. As the spherical harmonics are orthogonal, the overlap integral is non-zero only if m 1 = m 2 . Therefore,

The first term is referred to as self-interaction in previous papers [50, 44], but can be naturally inferred from our message passing scheme. For the nonlinearity, we follow [50] to use the vector norm of each degree of vector features:

where σ k are the activation functions. The equivariance holds as the vector norm is invariant to rotation. Also, to avoid over-parametrization and save computational resources, we only consider the interactions within the same radial index: ˆ S nℓm,n ′ ℓ ′ m ′ ( r ) := δ nn ′ S mm ′ ,ℓℓ ′ ( 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 (Sec.5.3).

3.4 Residual Operator Layer

The dimension of the function space is infinite, but in practice, we can only use the finite approximation. 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 at a given query point p at coordinate x . More specifically, the residual operator layer aggregates the neighboring node features to produce an invariant scalar 2 to finetune the final estimation:

This scheme resembles the coordinate-based interpolation nets and was proved effective in our ablation study (Sec.5.3). Therefore, the final output function is

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 to the norm in L 2 ( R 3 ) as L = ∥ ˆ ρ -ρ ∥ 2 2 = ∫ R 3 | ˆ ρ ( x ) -ρ ( x ) | 2 d x .

4 Graph Spectral View of InfGCN

Just as 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 a similar concept of graphon convolution . We will first introduce the (slightly generalized) concept of graphon. Defined on region D , a graphon , also known as graph limit or graph function, is a symmetric square-integrable function:

Intuitively, the kernel W ( x , y ) can be viewed as the probability that an edge forms between the continuously indexable nodes x , y ∈ D . Now, consider the operator T W defined in Eq.(2). As the integral kernel W is symmetric and square-integrable, we can apply the spectral theorem to conclude that the operator T W it induces is a self-adjoint operator whose spectrum consists of a countable number of real-valued eigenvalues { λ k } ∞ k =1 with λ k → 0 . Let { ϕ k } ∞ k =1 be the eigenfunctions such that T W ϕ k = λ k ϕ k . Similarly to the graph convolution for discrete graphs, any transformation on the eigenvalues F : { λ k } ∞ k =1 ↦→ { µ k } ∞ k =1 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 H x = ∑ K k =0 w k L k x where w k are the parameters. Define the power series of T W as:

where I is the identity mapping on D . A graphon filter can be then defined as H f = ∑ ∞ k =0 w k T k W f . We can also follow GCN to use the Chebyshev polynomials to approximate the graphon filter H :

Just as the message passing on discrete graphs can be viewed as graph convolution, we point out here that any model that tries to approximate the continuous analog T W f as defined in Eq.(2) can also be viewed as graphon convolution . This includes InfGCN, all NOs, and coefficient learning nets. A 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. [51] 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.(6) 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 a more intuitive interpretation of the radial network in TFN: it captures the magnitude of the radial part of the overlap integral S ( r ) of the basis in Eq.(4). From the point convolution aspect, the TFN structure can be also considered a special case of our proposed InfGCN model. The discrete input features can be regarded as the summation of Dirac measures over the node coordinates as f in ( x ) = ∑ u f u δ ( x -x u ) .

5 Experiments

We carried out extensive experiments on large-scale electron density datasets to illustrate the state-ofthe-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

Weevaluated our model on three electron density datasets. As computers cannot truly store continuous data, all datasets provide the electron density in a 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], which gives 123835 training samples, 50 validation samples, and 10000 testing samples.

Cubic . This large-scale dataset contains electron densities on 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/validation/test was pre-assigned.

MD . The dataset contains 6 small molecules (ethanol, benzene, phenol, resorcinol, ethane, malonaldehyde) with different geometries sampled from molecular dynamics (MD). The former 4 molecules are from [1] with 1000 sampled geometries each. The latter two are from [2] with 2000 sampled geometries each. The models were trained separately for each molecule.

To evaluate the models, we followed [17] to define the normalized mean absolute error (NMAE) as our evaluation metrics:

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, as there is a canonical way of assigning the primitive cell. Similarly, for the MD dataset, the authors described canonical ways to align the molecules [1, 2], so we also did not rotate them. More details regarding the datasets can be found in Appendix C.

We compared our proposed model with a wide range of different baselines including CNN [40, 4], interpolation networks ( DeepDFT [17], DeepDFT2 [18], EGNN [43], DimeNet [23], DimeNet++ [11]), 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 type-specific initial density function is constructed. A sampling scheme is used for all models except for CNN. All models were trained on a 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 are highlighted in bold. Our model is able to achieve state-of-the-art performance for almost all datasets. For QM9 and Cubic, the performance improvement is about 1% compared to the second-best model on both rotated and unrotated data, which is significant considering the small loss. We also noticed that CNN worked well on small molecules in MD and QM9, but quickly ran out of memory for larger Cubic data samples with even a batch size of 1 (marked 'OOM' in the table). This is because CNN ran on the full grid with a maximal number of 384 3 voxels.

Ground

Truth

InfGCN

CNN

DeepDFT

DeepDFT2

EGNN

DimeNet

DimeNet++

GNO

FNO

LNO

Prediction

Density

Error

NMAE/%

1.46

6.75

2.07

1.98

12.46

12.43

11.81

32.30

49.49

26.16

Ground

InfGCN

DeepDFT

DeepDFT2

EGNN

DimeNet

DimeNet++

GNO

FNO

LNO

Truth

Prediction

Density

Error

NMAE/%

10.37

28.80

26.70

26.47

30.38

28.94

78.56

74.66

70.55

The visualizations of the predicted densities in Fig.2 can provide more insights into the models. On QM9, the error of InfGCN had a regular spherical shape, indicating the smoothness property of the coefficient-based 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 a distinct spatial pattern on the partition boundaries of the grid, as it was demonstrated to be sensitive to the partition. More visualizations are provided in Appendix E on Cubic and QM9 with a wide range of representative molecules to demonstrate the generalizability of InfGCN.

The plots of model sizes versus NMAE on the QM9 dataset are shown in Figure 3. It can be clearly 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 as a sampling scheme is required.

5.3 Ablation Study

RotatedQM9

UnrotatedQM9

50.0

·GNO

40.0

GNO

LNO

FNO

FNO

20.0

%

20.0

LNO

%

10.0

s0

DimeNet++

EGNN

DimeNet

NMAE

NMAE

S1

s0

DimeNet++ooDimeNet

EGNN

5.0

s2

10.0

no-res

S1

DeepDFT

2.0

CNN

$2

no-res

6.0

CNNDeepDFT

S3

5.0

DeepDFT2

fc

1.0

S4S5

InfGCN

DeepDFT2

4.0

0.1

0.2

1.0

2.0

10.0

40.0

0.8

0.1

0.2

1.0

2.0

10.0

40.0

Model size / M

Model size / M

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 2 and are also demonstrated in Figure 3 in blue and green.

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, 10] used only the degree-1 spherical tensor which corresponds to vectors. We ran experiments with the maximal degree of the spherical tensor 0 ≤ L ≤ 7 ( s L columns). Note that s 0 corresponds to atom-centered Gaussian mixtures. It can be shown in Figure 3 (in blue) that the error smoothly drops as the maximal degree increases. Nonetheless, the performance gain is not significant with L ≥ 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 to 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-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 in Sec.3.3, we used a channel-wise tensor product instead a fully connected one that allows inter-channel interaction. We also tried the fully-connect tensor product under fc . It 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.

6 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 a symmetry-adapted Gaussian process regression for coefficient estimation. These

traditional methods are able to produce moderate to good results but are also less flexible and difficult to scale. For machine learning-based 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 a 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-layer 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 was 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 [49] constructed the molecule graph to perform variant message passing schemes with the final query-specific prediction. [17] proposed the DeepDFT model which also considered the graph between query coordinates and [18] further extended it to use a locally equivariant GNN [45]. [20] 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, 44, 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, 10], 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. [22] utilized the charge density and spin density as both the supervising signal and the additional input for predicting molecular energies, which achieved a significant performance improvement compared to traditional DFT-based methods. [51, 1] first projected the density onto the pre-defined basis set and then applied different neural nets on the coefficients to make predictions on 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 component 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 a better radial basis, and therefore, it leaves space for improvement. For example, we may use Slater-type orbitals (STO) instead of Gaussian-type orbitals. 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 a neural net is also a feasible option to provide more flexibility.

References

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):5223, Oct 2020. ISSN 2041-1723. doi: 10.1038/s41467-020-19093-1. URL https://doi.org/10.1038/s41467-020-19093-1 .

F. Brockherde, L. Li, K. Burke, and K. Müller. By-passing the kohn-sham equations with machine learning. CoRR , abs/1609.02815, 2016. URL http://arxiv.org/abs/1609.02815 .

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 , 8(12):2004214, 2021. doi: https://doi.org/10.1002/advs.202004214. URL https://onlinelibrary.wiley.com/doi/abs/10.1002/advs.202004214 .

Ö. Çiçek, A. Abdulkadir, S. S. Lienkamp, T. Brox, and O. Ronneberger. 3d u-net: Learning dense volumetric segmentation from sparse annotation. In S. Ourselin, L. Joskowicz, M. R. Sabuncu, G. B. Ünal, and W. M. W. III, editors, Medical Image Computing and ComputerAssisted Intervention - MICCAI 2016 - 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II , volume 9901 of Lecture Notes in Computer Science , pages 424-432, 2016. doi: 10.1007/978-3-319-46723-8\_49. URL https://doi.org/10.1007/ 978-3-319-46723-8_49 .

B. Coors, A. P. 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 Computer Science , pages 525-541. Springer, 2018. doi: 10.1007/978-3-030-01240-3\_32. URL https://doi.org/10.1007/978-3-030-01240-3_32 .

C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasacchi, and L. J. Guibas. Vector neurons: A general framework for so(3)-equivariant networks. In 2021 IEEE/CVF International Conference on Computer Vision, ICCV 2021, Montreal, QC, Canada, October 10-17, 2021 , pages 1218012189. IEEE, 2021. doi: 10.1109/ICCV48922.2021.01198. URL https://doi.org/10. 1109/ICCV48922.2021.01198 .

H. Deng, T. Birdal, and S. Ilic. Ppfnet: Global context aware local features for robust 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 195-205. Computer Vision Foundation / IEEE Computer Society, 2018. doi: 10.1109/CVPR. 2018.00028. URL http://openaccess.thecvf.com/content_cvpr_2018/html/Deng_ PPFNet_Global_Context_CVPR_2018_paper.html .

A. Edmonds. Angular Momentum in Quantum Mechanics . Investigations in physics. Princeton University Press, 1957. URL https://books.google.com/books?id=1swNAQAAIAAJ .

V. Fanaskov and I. V. Oseledets. Spectral neural operators. CoRR , abs/2205.10573, 2022. doi: 10.48550/arXiv.2205.10573. URL https://doi.org/10.48550/arXiv.2205.10573 .

F. Fuchs, D. E. Worrall, V. Fischer, and M. Welling. Se(3)-transformers: 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 612, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 15231a7ce4ba789d13b722cc5c955834-Abstract.html .

J. Gasteiger, S. Giri, J. T. Margraf, and S. Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop, NeurIPS , 2020.

J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neural message passing for quantum chemistry. In D. 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 1263-1272. PMLR, 2017. URL http://proceedings.mlr.press/v70/gilmer17a.html .

A. Grisafi, D. M. Wilkins, G. Csányi, and M. Ceriotti. Symmetry-adapted machine learning for tensorial properties of atomistic systems. Phys. Rev. Lett. , 120:036002, Jan 2018. doi: 10.1103/PhysRevLett.120.036002. URL https://link.aps.org/doi/10.1103/ PhysRevLett.120.036002 .

G. Hegde and R. C. Bowen. Machine-learned approximations to density functional theory hamiltonians, 2016. URL https://arxiv.org/abs/1609.05737 .

B. Jing, S. Eismann, P. N. Soni, and R. O. Dror. Equivariant graph neural networks for 3d macromolecular structure. CoRR , abs/2106.03843, 2021. URL https://arxiv.org/abs/ 2106.03843 .

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 https://openreview.net/forum?id=1YLJDvSx6J4 .

P. B. Jørgensen and A. Bhowmik. Deepdft: Neural message passing network for accurate charge density prediction. CoRR , abs/2011.03346, 2020. URL https://arxiv.org/abs/ 2011.03346 .

P. 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-022-00863-y. URL https://doi.org/10.1038/ s41524-022-00863-y .

J. Jumper, R. Evans, A. Pritzel, T. Green, M. Figurnov, O. Ronneberger, K. Tunyasuvunakool, R. Bates, A. Žídek, A. Potapenko, A. Bridgland, C. Meyer, S. A. A. Kohl, A. 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-589, Aug 2021. ISSN 1476-4687. doi: 10.1038/s41586-021-03819-2. URL https://doi.org/10.1038/s41586-021-03819-2 .

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. URL https://arxiv.org/abs/2201.08348 .

T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings . OpenReview.net, 2017. URL https:// openreview.net/forum?id=SJU4ayYgl .

J. Kirkpatrick, B. McMorrow, D. H. P. 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. J. Cohen. Pushing the frontiers of density functionals by solving the fractional electron problem. Science , 374(6573):1385-1389, 2021. doi: 10.1126/science.abj6511. URL https://www.science.org/doi/abs/10.1126/ science.abj6511 .

J. Klicpera, J. Groß, 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://openreview.net/forum?id= B1eWbxStPH .

W. Kohn and L. J. Sham. Self-consistent equations including exchange and correlation effects. Phys. Rev. , 140:A1133-A1138, Nov 1965. doi: 10.1103/PhysRev.140.A1133. URL https: //link.aps.org/doi/10.1103/PhysRev.140.A1133 .

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://arxiv.org/ abs/2108.08481 .

H. G. KÜMMEL. A biography of the coupled cluster method. International Journal of Modern Physics B , 17(28):5311-5325, 2003. doi: 10.1142/S0217979203020442. URL https: //doi.org/10.1142/S0217979203020442 .

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):3883-3895, Sept. 2022.

Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. M. Stuart, and A. Anandkumar. Neural operator: Graph kernel network for partial differential equations. CoRR , abs/2003.03485, 2020. URL https://arxiv.org/abs/2003.03485 .

Z. Li, N. B. Kovachki, K. Azizzadenesheli, B. Liu, K. Bhattacharya, A. M. Stuart, and A. Anandkumar. 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://openreview.net/forum?id=c8P9NQVtmnO .

R. Livni, D. Carmon, and A. Globerson. Learning infinite layer networks without the kernel trick. In D. 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. URL http: //proceedings.mlr.press/v70/livni17a.html .

L. Lovász. Large Networks and Graph Limits . American Mathematical Society colloquium publications. American Mathematical Society, 2012. ISBN 9780821890851. URL https: //books.google.com/books?id=FsFqHLid8sAC .

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.1038/s42256-021-00302-5. URL https: //doi.org/10.1038%2Fs42256-021-00302-5 .

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. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022 , pages 1891018919. IEEE, 2022. doi: 10.1109/CVPR52688.2022.01836. URL https://doi.org/10. 1109/CVPR52688.2022.01836 .

S. Maskey, R. Levie, Y. Lee, and G. Kutyniok. Generalization analysis of message passing neural networks on large random graphs. In NeurIPS , 2022. URL http://papers.nips.cc/paper_files/paper/2022/hash/ 1eeaae7c89d9484926db6974b6ece564-Abstract-Conference.html .

B. Mohar and W. Woess. A survey on spectra of infinite graphs. Bulletin of the London Mathematical Society , 21(3):209-234, 1989. doi: https://doi.org/10.1112/blms/21.3.209. URL https: //londmathsoc.onlinelibrary.wiley.com/doi/abs/10.1112/blms/21.3.209 .

M. W. Morency and G. Leus. Graphon filters: Graph signal processing in the limit. IEEE Trans. Signal Process. , 69:1740-1754, 2021. doi: 10.1109/TSP.2021.3061575. URL https: //doi.org/10.1109/TSP.2021.3061575 .

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.1021/ja01355a027. URL https://doi.org/10.1021/ja01355a027 .

M. A. Rahman, Z. E. Ross, and K. Azizzadenesheli. U-NO: u-shaped neural operators. CoRR , abs/2204.11127, 2022. doi: 10.48550/arXiv.2204.11127. URL https://doi.org/10.48550/ arXiv.2204.11127 .

R. Ramakrishnan, P. O. Dral, M. Rupp, and O. A. von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data , 1, 2014.

O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In N. Navab, J. Hornegger, W. M. W. III, and A. 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.1007/978-3-319-24574-4\_28. URL https://doi.org/10.1007/978-3-319-24574-4_28 .

L. Ruddigkeit, R. van Deursen, L. C. Blum, and J.-L. Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of Chemical Information and Modeling , 52(11):2864-2875, 2012. doi: 10.1021/ci300415d. URL https: //doi.org/10.1021/ci300415d . PMID: 23088335.

L. Ruiz, L. F. O. Chamon, and A. Ribeiro. Graphon neural networks and the transferability of graph neural 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 612, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 12bcd658ef0a540cabc36cdf2b1046fd-Abstract.html .

V. G. Satorras, E. Hoogeboom, 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-9332. PMLR, 2021. URL http://proceedings. mlr.press/v139/satorras21a.html .

K. Schütt, P. Kindermans, H. E. S. Felix, S. Chmiela, A. Tkatchenko, and K. Müller. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. In I. Guyon, U. von Luxburg, S. Bengio, H. M. Wallach, R. Fergus, S. V. N. Vishwanathan, and R. Garnett, 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://proceedings.neurips.cc/paper/2017/hash/ 303ed4c69846ab36c2904d3ba8573050-Abstract.html .

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. URL http://proceedings.mlr.press/v139/schutt21a.html .

D. Si, S. A. Moritz, J. 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 , 10(1):4282, Mar 2020. ISSN 2045-2322. doi: 10.1038/s41598-020-60598-y. URL https://doi.org/10.1038/s41598-020-60598-y .

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 https://arxiv.org/abs/1809.02723 .

J. C. Slater. The self consistent field and the structure of atoms. Phys. Rev. , 32:339-348, Sep 1928. doi: 10.1103/PhysRev.32.339. URL https://link.aps.org/doi/10.1103/ PhysRev.32.339 .

M. Su, J.-H. Yang, H.-J. Xiang, and X.-G. Gong. Efficient prediction of density functional theory hamiltonian with graph neural network, 2022. URL https://arxiv.org/abs/2205.05475 .

N. Thomas, T. E. Smidt, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3d point clouds. CoRR , abs/1802.08219, 2018. URL http://arxiv.org/abs/1802.08219 .

M. Tsubaki and T. Mizoguchi. On the equivalence of molecular graph convolution and molecular wave function with poor basis set. 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 612, 2020, virtual , 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 1534b76d325a8f591b52d302e7181331-Abstract.html .

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.1103/PhysRevLett.125.206401. URL https: //link.aps.org/doi/10.1103/PhysRevLett.125.206401 .

F. Q. Wang, K. Choudhary, Y. Liu, J. Hu, and M. Hu. Large scale dataset of real space electronic charge density of cubic inorganic materials from density functional theory (dft) calculations. Scientific Data , 9(1):59, Feb 2022. ISSN 2052-4463. doi: 10.1038/s41597-022-01158-z. URL https://doi.org/10.1038/s41597-022-01158-z .

Z. Wang, L. Ruiz, and A. Ribeiro. Convolutional neural networks on manifolds: From graphs and back. In 56th Asilomar Conference on Signals, Systems, and Computers, ACSSC 2022, Pacific Grove, CA, USA, October 31 - Nov. 2, 2022 , pages 356-360. IEEE, 2022. doi: 10.1109/IEEECONF56349.2022.10051964. URL https://doi.org/10.1109/ IEEECONF56349.2022.10051964 .

L. Zepeda-Núñez, Y. Chen, J. Zhang, W. Jia, L. Zhang, and L. Lin. Deep density: Circumventing the kohn-sham equations via symmetry preserving neural networks. J. Comput. Phys. , 443: 110523, 2021. doi: 10.1016/j.jcp.2021.110523. URL https://doi.org/10.1016/j.jcp. 2021.110523 .

Then, we have

Supplementary Material

A Proof of Rotation Equivariance

In this section, we will give a rigid proof of rotation equivariance of our proposed InfGCN model with finite approximation. Just as 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 in Figure 4.

R

(1)(d2L)

Proposition A.1. Rotation of a spherical harmonic of degree ℓ and order m ( R Y ℓ m )( r ) := Y ℓ m ( R -1 r ) transforms into a linear combination of spherical harmonics of the same degree:

where D ℓ mm ′ ( R ) is 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.1.4 in [8]. Therefore, for a square-integrable function defined on the unit sphere in R 3 , we can also describe the rotation of the function with Wigner D-matrics:

Proposition A.2. Assume f ∈ L 2 ( S 2 ) and the rotation of f have the (infinite) expansions onto the spherical harmonic basis as:

where f ℓ is the coefficient vector of degree ℓ with m = 2 ℓ +1 elements, and D ℓ R is the corresponding Wigner D-matrix of degree ℓ .

Proof. Notice that for each degree ℓ , the coefficients are transformed linearly according to Eq.(14), which concludes the proof.

Define spherical tensor f = { f ℓ : ℓ ≥ 0 } , we can further simplify the notation in Proposition A.2 as

A pictorial illustration of the rotation of the spherical harmonics is provided in Figure 5. 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 a continuous function.

One crucial property of the spherical tensors is that the tensor product is equivariant to rotation R∈ SO (3) :

Proposition A.3. The tensor product of two spherical tensors satisfies the following identity:

-0.237

0.350

0.906

p(x)

(Rp)(x)

R:R=

0.874

-0.331

0.356

0.425

-0.876

0.228

(D%

DR=

D

R

R

0.53

+ 0.94+ 0.12

0.54

0.53

0.67

+0.59

0.63

The proof of this property can be found in the TFN paper [50] in Appendix C. Essentially, it is a natural corollary as the property of irreducible representations. Combining Eq.(14) and (3), we can then design an equivariant message as:

where φ m J ( r ) = φ J ( r ) Y J m (ˆ 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. w ℓk are the learnable weights. The tensor product gives a 4-dimensional tensor. The summation indices J, m correspond to the two dimensions other than ℓ, k . One way to define such a tensor product for two spherical tensors comes from the coupling of angular momentum in physics. The tensor product c = a ⊗ b is defined as

where ⟨ ℓm 1 km 2 | Jm ⟩ are the Clebsch-Gordan coefficients, and are nonzero only when | ℓ -k | ≤ J ≤ ℓ + k, -J ≤ m ≤ J . Substituting Eq.(20) into Eq.(19) and ignoring the zero terms, we will get a summation of

where Q ℓk Jm ( m 1 m 2 ) = ⟨ ℓm 1 km 2 | Jm ⟩ . Coalescing the weight into the learnable radial function ˆ φ ℓk J = w ℓk φ J ( r ) , we have our final message passing scheme defined in Eq.(5). 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, the two spherical tensors in Eq.(19) transform as in Eq.(17). Therefore, we have

w ℓk can be moved inside D R because it is essentially a linear combination of equivariant functions, thus it is also equivariant.

Now let's consider finite approximation where 0 ≤ ℓ ≤ L for a fixed L . We have the following proposition:

Proposition A.5. Let P L : a = { a ℓ : ℓ ≥ 0 } ↦→ P L a = { a ℓ : 0 ≤ ℓ ≤ L } be the projection tensor operator that projects the spherical tensor onto the bases with degree less or equal to L , then P L is rotation equivariant:

Proof. It suffices to note that for each degree ℓ , the corresponding component a ℓ rotates according to Eq.(17), which only depends on the components with the same degree. Therefore, P L is equivariant as the components with degree ℓ ≤ L are preserved on both sides.

Proposition A.6. The composition of two equivariant operators T 1 ◦ T 2 is also equivariant.

Proof.

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 ≤ ℓ ≤ L .

Notice that for degree 0 features (scalars), equivariance is equivalent to invariance, and it can be obtained with projection operator P 0 , we have

Corollary A.8. The residual operator layer defined in Eq. (8) with finite approximation is invariant with respect to the grid 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. (9) 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 a discrete graph G = ( V , E ) , the graph Fourier transform (GFT) is defined as

̸

where S = U Λ U ⊤ is the eigenvalue decomposition of the graph shift operator S . A graph shift operator is a diagonalizable matrix S ∈ R N × N satisfying S ij = 0 for i = j, ( i, j ) ̸∈ E . The graph Laplacian L = I -A and the normalized version L = I -D -1 / 2 AD -1 / 2 as was used in GCN [21] where A is the adjacency matrix are such operators. As clear as this definition of a GFT is, it remains computationally prohibitive to implement the filtering operation on a large graph in this way. To filter the graph frequencies, a polynomial approach is thus adopted on the eigenvalues:

where w k are the learnable parameters. In the GCN formulation, the authors used the diagonalized matrix Λ instead of L , which is essentially equivalent.

We now switch to the continuous graphon setting. As defined in Eq.(10), a graphon is a symmetric square-integrable function. The kernel W induces an operator T W defined in Eq.(2). As W is symmetric and square-integrable, T W is a self-adjoint operator that can be decomposed as

where U is some unitary operator and Λ is a mulplication operator , i.e., there exists a function ξ ( x ) such that for all f ( x ) , Λ f ( x ) = ξ ( x ) f ( x ) . This directly follows the result of the spectral theorem for self-adjoint operators. In this way, we may similarly define the graphon Fourier transform as

Following the polynomial approach of approximating the graphon filter, we first define the power series of T W as in Eq.(11) and use the Chebyshev polynomials of T W to approximate the graphon filter H as H f ≈ θ 1 f + θ 2 T W f . Either way, parameterization and evaluation of T W f is required. As mentioned in Sec.4, our model essentially operates on the eigenvalues of the operator T W . In the graph spectral point of view, the eigenvalues are the spectrum of the operator. Therefore, any spectral filtering can 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 T W to deduce Eq.3. For any complete orthonormal basis { ψ k } ∞ k =1 of L 2 ( D ) , any square-integrable function f can be expanded as f = ∑ ∞ k =1 f k ψ k where f k = ∫ D f ( x ) ψ k ( x ) d x . We can then arrange the transform g = T W f as the following matrix-vector form g = Wf , where

If { ψ k } ∞ k =1 coincide with the eigenfunctions { ϕ k } ∞ k =1 which satisfy

We have W ij = λ j ∫ D ϕ i ( x ) ϕ j ( x ) d x . For the unicentric setting, W ij is non-zero only when i = j (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

where r u is the coordinate of center u . Similarly, considering one center at the origin with the other at r , we have

The 'overlap integral' S ij ( r ) arises here and we further parameterize the integral with w ij as the basis transformation also involves index i . Therefore, using the above matrix-vector formulation, we have the following parameterization:

If we also assume the locality of the interaction between two expansion centers, we can sum over all neighboring nodes to give the result in Eq.(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.

B.3 Graphon Convolution and Continuous MPNN

Previous work regarding continuous message-passing 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 message-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 fundamentally 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 a similar fashion as various neural operators. Our approach, on the other hand, implicitly constructed the graphon and parameterized them in the spectral domain. We noted in Sec.3.1 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 a 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 a different graphon.

In terms of the problem formulation, we further assume there exists an underlying discrete structure that has a significant physical connotation, e.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-related 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 Eq.2, the theoretical results on convergence and transferability could be adapted for our model to make it 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 V ASP using the Perdew-Burke-Ernzerhof (PBE) functional [17, 18]. The grid coordinates are guaranteed to be orthogonal but are generally different for different molecules.

Cubic . The densities are calculated with V ASP using the projector augmented wave (PAW) method [53]. As the crystal structure satisfies the periodic boundary condition (pbc), the volumetric data are given for a primitive cell with translation periodicity. We only focus on 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 (fcc) structures are given in its non-orthogonal primitive rhombohedral cell.

MD . Densities from [1] 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 3 . The result volumetric density is represented in Fourier basis, so we first converted it into the Cartesian form for all models.

The raw data of these datasets come in different formats. We defined a unified data interface to facilitate experiments and easy extension to 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 N x × N y × N z . Note that it was flattened to have the order of X-Y-Z.

grid_coord : Coordinates of the grid points where densities are sampled, with size ( N x × N y × N z , 3) .

shape : A 3D vector representing the discretization sizes of each of the three dimensions.

cell : A 3-by-3 matrix representing the cell vectors.

Other common statistics of the datasets are summarized in Table 3. The MD dataset does not have a validation split. The number of grids and grid lengths is for a single dimension so the number of voxels scales cubically with respect to it.

D Model and Training Specification

In this section, we provide our model specifications as 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.

InfGCN . We used spherical degree ℓ ≤ 7 and the number of radial bases n = 16 with the Gaussian parameters a k starting at 0.5 Bohr and ending at 5.0 Bohr. The distance was first embedded in a 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.

CNN [40, 4]. We used a 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:

where r u is the covalent radius of atom u and { a k } are pre-defined Gaussian parameters that contribute to the feature channel. The initial feature map was built with 16 Gaussian parameters a k starting at 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.

DeepDFT [17] and DeepDFT2 [18]. DeepDFT is a GNN-based network that models the interaction between the atom vertices and the query vertices for which the charge density is predicted. As DeepDFT only takes the invariant features of atom types and edge distance as input, it is also globally equivariant. DeepDFT2 uses PaiNN [45] as the GNN architecture. PaiNN designs equivariant interaction between scalar and vectorial features. Therefore, DeepDFT2 is locally equivariant. We simply followed the original model architectures which gave models of 2.04M and 2.93M trainable parameters, respectively.

EGNN [43]. 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 in a 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 [23] and DimeNet++ [11]. DimeNet uses spherical 2D 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 a query-specific residual GNN similar to InfGCN. All the other model hyperparameters are the same as the original models. As a result, the two models have 2.31M and 2.02M parameters, respectively.

The following 3 baselines are neural operators . They directly try to parameterize the Fredholm operator in Eq.(2) using various approaches. Same as CNN, they cannot automatically capture the grid information, so we also use the feature in Eq.(34) as the initial feature. The initial feature map for these models is built with 32 Gaussian parameters a k 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.

GNO [28]. 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:

where F is a 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:

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 low-frequency 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 ( x , y ) := ∑ r j =1 ϕ j ( x ) ψ j ( y ) , similar to the unstacked DeepONet proposed in [32]. We used a feature size of 384 with a 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 a small portion of the grid is randomly sampled

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 a 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 [28], 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 is 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 4. 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 a single NVIDIA A100 GPU. For efficiency, the testing for QM9 was done on the last 1600 samples, so larger molecules were tested. For Cubic and MD, testing was done on all the test samples.

D.3 Complexity Analysis

The naïve implementation of a message-passing layer with padding and masking scales to O ( |E| C ( ℓ + 1) 6 ) 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 Clebsch-Gordan coefficients have their six indices all summed. However, note that there are only ( ℓ +1) 2 spherical harmonics with the degree up to ℓ . If coefficients are compressed into one long vector, the complexity can be reduced to O ( |E| C ( ℓ +1) 4 ) . During the expansion of the basis functions on the voxels, the time complexity is O ( |E| KC ( ℓ +1) 2 ) where K is the number of voxels sampled. In practice, we noticed that a small ℓ suffices so that ( ℓ + 1) 4 can be viewed as a constant ( ℓ = 7 corresponds to 4096). Also, the Clebsch-Gordan coefficients can be pre-computed and stored, and the summation can be efficiently done by index manipulation. Our implementation was based on the e3nn 3 package which implements efficient spherical vector manipulations.

In comparison, most GNN-based layers scale as O ( |E| D 2 ) where D is the hidden feature size. Therefore, in our standard-setting ( C = 16 , ℓ = 7 ), the time complexity is approximately of the same order (with D = 256 ). For GNO, FNO, and LNO, one layer scales as O ( KD 2 ) , O ( KD 3 ) , O ( KD 2 R ) , respectively. The additional coefficients are for the Fourier modes or the rank. For 3D-CNN, the time complexity scales as O ( C in C out N x N y N z k 3 ) 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.

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 in Figure 6.

Ground Truth

InfGCN

DeepDFT

DeepDFT2

EGNN

DimeNet

DimeNet++

GNO

FNO

LNO

Prediction

Density

Error

NMAE (%)

8.15

21.38,

9.35

8.89

15.01

12.87

49.09

46.21

44.01

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, 1, 3, and 8 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 a broad range of different types of molecules, we manually picked some representative molecules and evaluated all the models. The results are provided in Table 5. In order to provide a finer-grained comparison between different models, we used finer-grained 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 a 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 molecules. 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 linear nonane, the branched t -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 a 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 a chemical aspect, a 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 4 . For example, the carbon atom in CF4 has a 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, dit -butyl ether, and isoamyl acetate.

The predictions were even worse for a conjugated and aromatic system where electrons are highly delocalized 5 . Delocalization of electrons allows for long-distance interactions which are harder to capture. The presence of electron-donating groups (EDGs) and electronwithdrawing groups (EWGs) contributes greatly to the conjugated system. For example,

NH,

?

NH,

+

+

+

+

+

+

+

the amino group -NH2 is an EDG when bonding to a conjugated system like benzene, facilitating orthoand paraelectrophilic reactions. In contrast, the nitro group -NO2 is an EWG that facilitates orthoand paranucleophilic reactions (See Fig.7). It can be seen from the visualization that the orthoand parapositions of aniline are 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 at pH 7, 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 is underestimated and that of nitrogen is overestimated.

Name

File

ID

Structure

Ground

Truth

InfGCN

CNN

DeepDFT DeepDFT2

EGNN DimeNet DimeNet++

GNO

FNO

LNO

ammonia

2

NH3

urea

20

acetone

oxime

49

furan

52

tetrafluoro-

methane

184

CF4

glycerol

397

cyclohexane

658

3.16

22.79

5.34

2.02

12.27

11.45

15.22

57.87

157.19

36.00

²HN

0.99

3.56

3.33

1.24

9.63

9.83

9.94

44.72

29.73

22.31

0.92

2.72

2.90

1.02

12.54

13.04

12.60

31.79

30.03

21.52

1.40

3.81

2.86

1.68

12.05

12.51

12.38

30.68

39.98

25.38

1.76

4.83

1.95

1.00

9.79

6.48

7.75

34.19

36.31

27.36

0.79

2.52

2.77

0.88

11.73

11.36

11.03

33.05

29.21

23.78

0.54

1.95

1.75

0.78

13.21

12.65

12.60

30.48

30.80

19.33

Continued on next page

aniline

940

cytosine

4318

cubane

19116

purine

24537

di-

t

-butyl

ether

57520

isoamyl

acetate

60424

asparagine

61439

nonane

114514

nitrobenzene

131915

fluorouracil

132701

1.22

4.67

2.85

1.30

12.62

12.80

12.16

35.29

31.35

24.60

1.25

5.62

2.96

1.37

10.35

10.18

10.26

40.31

24.58

20.80

0.74

8.22

2.76

1.27

11.82

13.03

12.31

28.61

32.88

23.46

1.51

16.34

2.90

1.44

11.41

11.25

10.98

25.98

22.07

20.45

0.78

2.28

2.40

1.04

12.37

12.79

11.88

27.71

29.73

0.67

2.29

2.67

0.87

12.41

12.13

11.65

2.51

27.16

30.92

1.02

3.71

3.12

1.11

10.25

10.05

10.10

44.62

25.34

20.53

0.65

2.54

2.38

0.85

12.86

12.53

12.41

37.23

26.02

20.92

1.44

22.93

3.41

1.45

11.00

10.87

11.16

47.66

31.38

22.41

0.98

4.59

2.61

0.97

10.06

9.97

9.90

51.36

23.77

16.10

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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).\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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\nQUESTIONS:\nParagraph 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\nLIMITATIONS:\nLimitations are discussed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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\n\nQUESTIONS:\nSee the weakness section.\n\nLIMITATIONS:\nThere is no negative societal impact as far as I can tell.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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\nWEAKNESSES:\nSome 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\nQUESTIONS:\nsee weakness\n\nLIMITATIONS:\nWhile 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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?\n\nQUESTIONS:\n* 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\nLIMITATIONS:\nThe authors have communicated limitations of their work", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

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

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

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

QUESTIONS:
- Are there any advantages for explainability from the framework?

LIMITATIONS:
The authors discuss the limitations of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-13