SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathcal{G} = (\mathcal{V},\mathcal{E})$ , the graph Fourier transform (GFT) is defined as
$$
\hat {\mathbf {x}} = U ^ {\top} \mathbf {x} \tag {25}
$$
where  $S = U\Lambda U^{\top}$  is the eigenvalue decomposition of the graph shift operator  $S$ . A graph shift operator is a diagonalizable matrix  $S \in \mathbb{R}^{N \times N}$  satisfying  $S_{ij} = 0$  for  $i \neq j, (i,j) \notin \mathcal{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 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:
$$
H \mathbf {x} = \sum_ {k = 0} ^ {K} w _ {k} L ^ {k} \mathbf {x} \tag {26}
$$
where  $w_{k}$  are the learnable parameters. In the GCN formulation, the authors used the diagonalized matrix  $\Lambda$  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  $\mathcal{T}_W$  defined in Eq.(2). As  $W$  is symmetric and square-integrable,  $\mathcal{T}_W$  is a self-adjoint operator that can be decomposed as
$$
\mathcal {U} \mathcal {T} _ {W} = \Lambda \mathcal {U} \tag {27}
$$
where  $\mathcal{U}$  is some unitary operator and  $\Lambda$  is a multiplication operator, i.e., there exists a function  $\xi (\mathbf{x})$  such that for all  $f(\mathbf{x}),\Lambda f(\mathbf{x}) = \xi (\mathbf{x})f(\mathbf{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
$$
\hat {f} = \mathcal {U} f \tag {28}
$$
Following the polynomial approach of approximating the graphon filter, we first define the power series of  $\mathcal{T}_W$  as in Eq.(11) and use the Chebyshev polynomials of  $\mathcal{T}_W$  to approximate the graphon filter  $\mathcal{H}$  as  $\mathcal{H}f\approx \theta_1f + \theta_2\mathcal{T}_Wf$ . Either way, parameterization and evaluation of  $\mathcal{T}_Wf$  is required. As mentioned in Sec.4, our model essentially operates on the eigenvalues of the operator  $\mathcal{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  $\mathcal{T}_W$  to deduce Eq.3. For any complete orthonormal basis  $\{\psi_k\}_{k=1}^{\infty}$  of  $L^2(\mathcal{D})$ , any square-integrable function  $f$  can be expanded as  $f = \sum_{k=1}^{\infty} f_k \psi_k$  where  $f_k = \int_{\mathcal{D}} f(\mathbf{x}) \psi_k(\mathbf{x}) d\mathbf{x}$ . We can then arrange the transform  $g = \mathcal{T}_W f$  as the following matrix-vector form  $\mathbf{g} = \mathbf{W} \mathbf{f}$ , where
$$
W _ {i j} = \int_ {\mathcal {D}} \psi_ {i} (\mathbf {x}) \int_ {\mathcal {D}} W (\mathbf {x}, \mathbf {y}) \psi_ {j} (\mathbf {y}) d \mathbf {x} d \mathbf {y} \tag {29}
$$
If  $\{\psi_k\}_{k=1}^{\infty}$  coincide with the eigenfunctions  $\{\phi_k\}_{k=1}^{\infty}$  which satisfy
$$
\mathcal {T} _ {W} \phi_ {k} = \lambda_ {k} \phi_ {k} \tag {30}
$$
We have  $W_{ij} = \lambda_j\int_{\mathcal{D}}\phi_i(\mathbf{x})\phi_j(\mathbf{x})d\mathbf{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
$$
\hat {\rho} (\mathbf {x}) = \sum_ {u \in \mathcal {V}} \sum_ {i = 1} ^ {\infty} f _ {i, u} \psi_ {i} \left(\mathbf {x} - \mathbf {r} _ {u}\right) \tag {31}
$$
where  $\mathbf{r}_u$  is the coordinate of center  $u$ . Similarly, considering one center at the origin with the other at  $\mathbf{r}$ , we have
$$
W _ {i j} = \int_ {\mathcal {D}} \phi_ {i} (\mathbf {x}) \mathcal {T} _ {W} \phi_ {j} (\mathbf {x} - \mathbf {r}) d \mathbf {x} = \lambda_ {j} \int_ {\mathcal {D}} \phi_ {i} (\mathbf {x}) \phi_ {j} (\mathbf {x} - \mathbf {r}) d \mathbf {x} \tag {32}
$$
The "overlap integral"  $S_{ij}(\mathbf{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:
$$
f _ {i} \leftarrow \sum_ {j = 1} ^ {\infty} w _ {i j} S _ {i j} (\mathbf {r}) f _ {j} \tag {33}
$$
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.
In this section, we provide more details about the datasets we used in the experiments.
QM9. The densities are calculated with VASP 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 VASP 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 \times N_y \times 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 \times N_y \times N_z, 3)$ .
- shape: A 3D vector representing the discretization sizes of each of the three dimensions.
- cell1: 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.
Table 3: Dataset details.
Dataset	QM9	Cubic	MD
train/val/test split	123835/50/10000	14421/1000/1000	1000(2000)/500(400)
max/min/mean #grid	160/40/87.86	448/32/93.97	20/20/20
max/min/mean #node	29/3/17.98	64/1/10.49	14/8/10.83
max/min/mean length (Bohr)	15.83/4.00/8.65	26.20/1.78/5.82	20.00/20.00/20.00
#node type	5	84	3
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  $\ell \leq 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:
$$
f _ {k} (\mathbf {x}) = \sum_ {u \in \mathcal {V}} \exp \left(- a _ {k} \frac {\left| \mathbf {x} - \mathbf {x} _ {u} \right| ^ {2}}{r _ {u}}\right) \tag {34}
$$
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:
$$
f (\mathbf {x}) \leftarrow \sigma \left(W f (\mathbf {x}) + \frac {1}{| \mathcal {N} (u) |} \sum_ {\mathcal {N} (u)} \mathcal {F} (\mathbf {x}, \mathbf {y}, f (\mathbf {x}), f (\mathbf {y})) f (\mathbf {y})\right) \tag {35}
$$
where  $\mathcal{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:
$$
f (\mathbf {x}) \leftarrow \sigma \left(W f (\mathbf {x}) + \mathcal {F} ^ {- 1} R (\mathcal {F} f) (\mathbf {x})\right) \tag {36}
$$
where  $\mathcal{F},\mathcal{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(\mathbf{x},\mathbf{y}) \coloneqq \sum_{j = 1}^{r}\phi_{j}(\mathbf{x})\psi_{j}(\mathbf{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.
Table 4: Training specifications.
Model	Dataset	cutoff	n_iter	lr	patience	batch_size	lr Decay	train_sample	inf_sample
InfGCN	QM9	3.0	40k	1e-3	10	64
Cubic	5.0	10k	5e-3	5	32	0.5	1024	4096
MD	3.0	2k	5e-3	5	64
CNN	QM9	NA	100k	3e-4	10	4	0.5	NA	NA
MD	4k	1e-3	5	32
DeepDFT/DeepDFT2	QM9	3.0	40k	3e-4	10	64
Cubic	5.0	10k	3e-4	10	32	0.5	1024	4096
MD	3.0	2k	1e-3	5	64
EGNN/DimeNet/DimeNet++	QM9	5.0	40k	3e-4	10
Cubic	5.0	10k	3e-4	10	64	0.5	1024	4096
MD	3.0	2k	1e-3	5
GNO/FNO/LNO	QM9	80k	3e-4	10
Cubic	NA	10k	3e-4	10	32	0.5	1024	4096
MD	2k	1e-3	5
D.3 Complexity Analysis
The naive implementation of a message-passing layer with padding and masking scales to  $O(|\mathcal{E}|C(\ell + 1)^6)$  where  $|\mathcal{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  $(\ell + 1)^2$  spherical harmonics with the degree up to  $\ell$ . If coefficients are compressed into one long vector, the complexity can be reduced to  $O(|\mathcal{E}|C(\ell + 1)^4)$ . During the expansion of the basis functions on the voxels, the time complexity is  $O(|\mathcal{E}|KC(\ell + 1)^2)$  where  $K$  is the number of voxels sampled. In practice, we noticed that a small  $\ell$  suffices so that  $(\ell + 1)^4$  can be viewed as a constant  $(\ell = 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(|\mathcal{E}|D^2)$  where  $D$  is the hidden feature size. Therefore, in our standard-setting ( $C = 16$ ,  $\ell = 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^2R)$ , respectively. The additional coefficients are for the Fourier modes or the rank. For 3D-CNN, the time complexity scales as  $O(C_{\mathrm{in}}C_{\mathrm{out}}N_xN_yN_zk^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.
<sup>3</sup>https://e3nn.org/
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.
Figure 6: Visualization of the predicted density, the density error, and the NMAE for  $\mathrm{Rb_8AsO_3}$  (mp-1103058 in Cubic).
For the density visualization, we used linearly spaced values from 0.05 to 3.5 with 5 isosurfaces for the ground truth density and -0.03 (deep pink), -0.01, 0.01, and 0.03 (deep green) for the density errors for QM9 in Figure 2. We used values of 0.3, 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: