SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We carried out extensive experiments on large-scale electron density datasets to illustrate the state-of-the-art performance of our proposed InfGCN model over the current baselines. Multiple ablation studies were also carried out to demonstrate the effectiveness of the proposed architecture.
5.1 Datasets and Baselines
We evaluated our model on three electron density datasets. As computers cannot truly store continuous data, all datasets provide the electron density in 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:
$$
\mathrm {N M A E} = \frac {\int_ {\mathbb {R} ^ {3}} | \hat {\rho} (\mathbf {x}) - \rho (\mathbf {x}) | d \mathbf {x}}{\int_ {\mathbb {R} ^ {3}} | \rho (\mathbf {x}) | d \mathbf {x}} \tag {13}
$$
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. Equivalence 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.
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.
Table 1: NMAE (%) on QM9, Cubic, and MD datasets.
Dataset/Model	InfGCN	CNN	Interpolation Net	Neural Operator
DeepDFT	DeepDFT2	EGNN	DimeNet	DimeNet++	GNO	FNO	LNO
QM9	rotated	4.73	5.89	5.87	4.98	12.13	12.98	12.75	46.90	33.25	24.13
unrotated	0.93	2.01	2.95	1.03	11.92	11.97	11.69	40.86	28.83	26.14
Cubic	8.98	OOM	14.08	10.37	11.74	12.51	12.18	53.55	48.08	46.33
MD	ethanol	8.43	13.97	7.34	8.83	13.90	13.99	14.24	82.35	31.98	43.17
benzene	5.11	11.98	6.61	5.49	13.49	14.48	14.34	82.46	20.05	38.82
phenol	5.51	11.52	9.09	7.00	13.59	12.93	12.99	66.69	42.98	60.70
resorcinol	5.95	11.07	8.18	6.95	12.61	12.04	12.01	58.75	26.06	35.07
ethane	7.01	14.72	8.31	6.36	15.17	13.11	12.95	71.12	26.31	77.14
malonaldehyde	10.34	18.52	9.31	10.68	12.37	18.71	16.79	84.52	34.58	47.22
Figure 2: Visualization of the predicted density, the density error, and the NMAE. Up: Indole (file ID 24492 from QM9). Down:  $\mathrm{Cr_4CuNiSe_8}$  (mp-1226405 from Cubic). The colors of the points indicate different atom types, and the isosurfaces indicate different density values. The pink and green isosurfaces in the error plots represent the negative and positive errors, respectively.
The visualizations of the predicted densities in Fig.2 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.
Figure 3: Plot of model sizes vs NMAE for different models and ablation studies on the QM9 dataset.
To further demonstrate the effectiveness of InfGCN, we also carried out extensive ablation studies on various aspects of the proposed architecture on the QM9 dataset. The results are summarized in Table 2 and are also demonstrated in Figure 3 in blue and green.
Table 2: NMAE (%) and the parameter count of different model settings on the QM9 dataset.
Model	InfGCN(s7)	s6	s5	s4	s3	s2	s1	s0	no-res	fc
QM-rot (%)	4.73	4.77	4.76	4.77	4.86	6.95	9.56	12.62	6.14	4.95
QM-unrot (%)	0.93	1.01	1.11	1.08	1.46	4.65	8.07	12.05	3.72	1.36
Parameters (M)	1.20	0.85	0.58	0.39	0.26	0.17	0.13	0.11	1.16	17.42
Number of spherical basis. For coefficient learning models, using more basis functions will naturally lead to a more expressive power of the model. For discrete tasks, [50, 10] used only the degree-1 spherical tensor which corresponds to vectors. We ran experiments with the maximal degree of the spherical tensor  $0 \leq L \leq 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 \geq 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 of a fully connected one that allows inter-channel interaction. We also tried the fully-connect tensor product under f.c. 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.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.
[51] 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 6-12, 2020, virtual, 2020. URL https://proceedings.neurips.cc/paper/2020/bitnet/1534b76d325a8f591b52d302e7181331-AbsAbstract.html.
[52] M. Tsubaki and T. Mizoguchi. Quantum deep field: Data-driven wave function, electron density generation, and atomization energy prediction and extrapolation with machine learning. Phys. Rev. Lett., 125:206401, Nov 2020. doi: 10.1103/PhysRevLett.125.206401. URL https://link.aps.org/doi/10.1103/PhysRevLett.125.206401.
[53] 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.
[54] 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.
[55] 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.
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.
Figure 4: An illustration of rotation equivariance of the electron density function of benzene on  $\mathbb{R}^3$ .
Proposition A.1. Rotation of a spherical harmonic of degree  $\ell$  and order  $m\left(\mathcal{R}Y_{\ell}^{m}\right)(\mathbf{r}) := Y_{\ell}^{m}(R^{-1}\mathbf{r})$  transforms into a linear combination of spherical harmonics of the same degree:
$$
\left(\mathcal {R} Y _ {\ell} ^ {m}\right) (\mathbf {r}) = \sum_ {m ^ {\prime} = - \ell} ^ {\ell} D _ {m m ^ {\prime}} ^ {\ell} (\mathcal {R}) Y _ {\ell} ^ {m ^ {\prime}} (\mathbf {r}) \tag {14}
$$
where  $D_{mm'}^{\ell}(\mathcal{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  $\mathbb{R}^3$ , we can also describe the rotation of the function with Wigner D-matrics:
Proposition A.2. Assume  $f \in L^{2}(S^{2})$  and the rotation of  $f$  have the (infinite) expansions onto the spherical harmonic basis as:
$$
f (\mathbf {r}) = \sum_ {\ell m} f _ {m} ^ {\ell} Y _ {\ell} ^ {m} (\mathbf {r})
$$
$$
\mathcal {R} f (\mathbf {r}) = \sum_ {\ell m} ^ {\ell m} g _ {m} ^ {\ell} Y _ {\ell} ^ {m} (\mathbf {r}) \tag {15}
$$
$$
\mathbf {g} ^ {\ell} = D _ {\mathcal {R}} ^ {\ell} \mathbf {f} ^ {\ell} \tag {16}
$$
where  $\mathbf{f}^{\ell}$  is the coefficient vector of degree  $\ell$  with  $m = 2\ell +1$  elements, and  $D_{\mathcal{R}}^{\ell}$  is the corresponding Wigner  $D$ -matrix of degree  $\ell$ .
Proof. Notice that for each degree  $\ell$ , the coefficients are transformed linearly according to Eq.(14), which concludes the proof.
Define spherical tensor  $\mathsf{f} = \{\mathbf{f}^{\ell}:\ell \geq 0\}$ , we can further simplify the notation in Proposition A.2 as
$$
\mathcal {R} f = D _ {\mathcal {R}} (\mathrm {f}) \tag {17}
$$
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  $\mathcal{R} \in SO(3)$ :
Proposition A.3. The tensor product of two spherical tensors satisfies the following identity:
$$
D _ {\mathcal {R}} (\mathsf {a} \otimes \mathsf {b}) = D _ {\mathcal {R}} (\mathsf {a}) \otimes D _ {\mathcal {R}} (\mathsf {b}) \tag {18}
$$
Figure 5: An illustration of rotation equivariance of linear combination of spherical harmonics as a continuous function on the unit sphere. For better visualization, the radial value used to plot the spherical harmonics is the squared density  $|\rho|^2$ . The calculation, however, is still done on the original (real-valued) spherical harmonics.
The proof of this property can be found in the TFN paper [50] in Appendix C. Essentially, it is a natural corollary as the property of irreducible representations. Combining Eq.(14) and (3), we can then design an equivariant message as:
$$
\mathrm {f} _ {u} = \sum_ {v \in \tilde {\mathcal {N}} (u)} \sum_ {\ell k} w _ {\ell k} \sum_ {J m} \varphi (\mathrm {r} _ {u v}) \otimes \mathrm {f} _ {v} \tag {19}
$$
where  $\varphi_J^m (\mathbf{r}) = \varphi_J(r)Y_J^m (\hat{\mathbf{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_{\ell 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  $\ell ,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  $\mathsf{c} = \mathsf{a}\otimes \mathsf{b}$  is defined as
$$
C _ {J m} = \sum_ {m _ {1} = - \ell} ^ {\ell} \sum_ {m _ {2} = - k} ^ {k} a _ {\ell m _ {1}} b _ {k m _ {2}} \left\langle \ell m _ {1} k m _ {2} \right| J m \rangle \tag {20}
$$
where  $\langle \ell m_1km_2|Jm\rangle$  are the Clebsch-Gordan coefficients, and are nonzero only when  $|\ell -k|\leq J\leq \ell +k, - J\leq m\leq J.$  Substituting Eq.(20) into Eq.(19) and ignoring the zero terms, we will get a summation of
$$
\mathbf {f} _ {u} ^ {\ell} = \sum_ {v \in \tilde {\mathcal {N}} (u)} \sum_ {k \geq 0} \sum_ {J = | k - \ell |} ^ {k + \ell} w _ {\ell k} \varphi_ {J} (r) \sum_ {m = - J} ^ {J} Y _ {J} ^ {m} (\hat {\mathbf {r}}) Q _ {J m} ^ {\ell k} \mathbf {f} _ {v} ^ {k} \tag {21}
$$
where  $Q_{Jm}^{\ell k}(m_1m_2) = \langle \ell m_1km_2|Jm\rangle$ . Coalescing the weight into the learnable radial function  $\hat{\varphi}_J^{\ell k} = w_{\ell k}\varphi_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
$$
\begin{array}{l} \mathcal {T} \left(\mathcal {R} \mathrm {f} _ {u}\right) = \sum_ {v \in \tilde {\mathcal {N}} (u)} \sum_ {\ell k} w _ {\ell k} \sum_ {J m} \mathcal {R} \varphi \left(\mathbf {r} _ {u v}\right) \otimes \mathcal {R} \mathrm {f} _ {v} \\ = \sum_ {v \in \tilde {\mathcal {N}} (u)} \sum_ {\ell k} w _ {\ell k} \sum_ {J m} D _ {\mathcal {R}} \varphi (\mathbf {r} _ {u v}) \otimes D _ {\mathcal {R}} \mathbf {f} _ {v} \\ = \sum_ {v \in \tilde {\mathcal {N}} (u)} \sum_ {\ell k} w _ {\ell k} \sum_ {J m} D _ {\mathcal {R}} \left(\varphi \left(\mathbf {r} _ {u v}\right) \otimes \mathbf {f} _ {v}\right) \tag {22} \\ = D _ {\mathcal {R}} \left[ \sum_ {v \in \tilde {\mathcal {N}} (u)} \sum_ {\ell k} w _ {\ell k} \sum_ {J m} \varphi (\mathbf {r} _ {u v}) \otimes \mathbf {f} _ {v} \right] \\ = \mathcal {R} (\mathcal {T f} _ {u}) \\ \end{array}
$$
$w_{\ell k}$  can be moved inside  $D_{\mathcal{R}}$  because it is essentially a linear combination of equivariant functions, thus it is also equivariant.
Now let's consider finite approximation where  $0 \leq \ell \leq L$  for a fixed  $L$ . We have the following proposition:
Proposition A.5. Let  $\mathcal{P}_L: \mathfrak{a} = \{\mathbf{a}_{\ell} : \ell \geq 0\} \mapsto \mathcal{P}_L \mathfrak{a} = \{\mathbf{a}_{\ell} : 0 \leq \ell \leq L\}$  be the projection tensor operator that projects the spherical tensor onto the bases with degree less or equal to  $L$ , then  $\mathcal{P}_L$  is rotation equivariant:
$$
\mathcal {P} _ {L} (\mathcal {R} \mathbf {a}) = \mathcal {R} \left(\mathcal {P} _ {L} \mathbf {a}\right) \tag {23}
$$
Proof. It suffices to note that for each degree  $\ell$ , the corresponding component  $\mathbf{a}^{\ell}$  rotates according to Eq.(17), which only depends on the components with the same degree. Therefore,  $\mathcal{P}_L$  is equivariant as the components with degree  $\ell \leq L$  are preserved on both sides.
Proposition A.6. The composition of two equivariant operators  $\mathcal{T}_1\circ \mathcal{T}_2$  is also equivariant.
$$
\left(\mathcal {T} _ {1} \circ \mathcal {T} _ {2}\right) (\mathcal {R} \mathfrak {a}) = \mathcal {T} _ {1} \left(\mathcal {T} _ {2} (\mathcal {R} \mathfrak {a})\right) = \mathcal {T} _ {1} \left(\mathcal {R} \left(\mathcal {T} _ {2} \mathfrak {a}\right)\right) = \mathcal {R} \left(\mathcal {T} _ {1} \mathcal {T} _ {2} \mathfrak {a}\right) \tag {24}
$$
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 \leq \ell \leq L$ .
Notice that for degree 0 features (calvars), equivariance is equivalent to invariance, and it can be obtained with projection operator  $\mathcal{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