SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chaoran Cheng  
University of Illinois Urbana-Champaign  
chaoran7@illinois.edu
Jian Peng  
University of Illinois Urbana-Champaign  
jianpeng@illinois.edu
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.
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  $\mathbb{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  $\mathbb{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 (3D-CNN) [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.
Figure 1: Overview of the model architecture. (a) The input molecule with node-wise spherical tensor features. (b) The message passing scheme in InfGCN.  $\otimes$  denotes the tensor product of two spherical tensors  $f, \varphi(r)$  (Sec.3.3). (c) Coordinate-specific residual operator layer (Sec.3.4). (d) Spherical harmonics. (e) The final prediction combining the expanded basis functions and the residue.
We use  $\mathcal{G} = (\mathcal{V},\mathcal{E})$  to denote the (discrete) graph with the corresponding node coordinates  $\{\mathbf{x}_i\}_{i = 1}^{|\mathcal{V}|}$ . A continuous function over the region  $\mathbb{D}$  is also provided as the target feature:  $\rho :\mathcal{D}\to \mathbb{R}$ . We also assume that there is an initial feature function  $f_{\mathrm{in}}:\mathcal{D}\rightarrow \mathbb{R}$  either obtained from less accurate methods, a random guess, or some learnable initialization. Formally, given  $f_{\mathrm{in}}\in L^2 (\mathcal{D})$ , which is a square-integrable input feature function over  $\mathcal{D}$ , and the target feature function  $\rho \in L^2 (\mathcal{D})$ , we want to learn an operator in the Hilbert space  $\mathcal{T}:L^2 (\mathcal{D})\to L^2 (\mathcal{D})$  to approximate the target function.
Different from common regression tasks over finite-dimensional vector spaces, the Hilbert space  $L^2(\mathcal{D})$  is infinite-dimensional.
2.1 Equivariance
Equivalence describes the behavior of the model when the input data are transformed. Formally, for group  $G$  acting on  $\mathcal{X}$  and group  $H$  acting on  $\mathcal{Y}$ , for a function  $f: \mathcal{X} \to \mathcal{Y}$ , if there exists a homomorphism  $S: G \to H$  such that  $f(g \cdot \mathbf{x}) = (Sg)f(\mathbf{x})$  holds for all  $g \in G, \mathbf{x} \in X$ , then  $f$  is equivariant. Specifically, if  $S: g \mapsto e$  maps every group action to the identity action, we have the definition of invariance:  $f(g \cdot \mathbf{x}) = f(\mathbf{x}), \forall g \in G, \mathbf{x} \in 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 \in L^{2}(\mathbb{R}^{3})$  as  $(\mathcal{R}f)(\mathbf{x}) \coloneqq f(R^{-1}\mathbf{x})$ , where  $R$  is the rotation matrix associated with the rotation operator  $\mathcal{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  $\mathcal{T}$  with respect to rotation can be formulated as
$$
\mathcal {T} (\mathcal {R} f) = \mathcal {R} (\mathcal {T} f), \forall \mathcal {R} \tag {1}
$$
For clarity, we will distinguish equivariance and invariance, and use equivariance for functions satisfying Eq.(1).
Intuitively, we would like to follow the message passing paradigm [12] to aggregate information from every other point  $\mathbf{x} \in \mathcal{D}$ . In our scenario, however, as the nodes indexed by the coordinate  $\mathbf{x} \in \mathcal{D}$  are infinite and continuous, the aggregation of the messages must be expressed as an integral:
$$
\mathcal {T} _ {W} f (\mathbf {x}) := \int_ {\mathcal {D}} W (\mathbf {x}, \mathbf {y}) f (\mathbf {y}) d \mathbf {y} \tag {2}
$$
where  $W: \mathcal{D} \times \mathcal{D} \to [0,1]$  is a square-integrable kernel function that parameterizes the source node features  $f(\mathbf{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  $\{\psi_k(\mathbf{x})\}_{k=1}^{\infty}$  over  $L^2(\mathcal{D})$ . In this way, the feature function can be expanded onto the basis as  $f(\mathbf{x}) = \sum_{k=1}^{\infty} f_k \psi_k(\mathbf{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 \in L^2(\mathcal{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  $\mathbf{r}_u$  in the discrete graph as the centers of basis sets:  $\hat{\rho} (\mathbf{x}) = \sum_{u\in \mathcal{V}}\sum_{i = 1}^{\infty}f_{i,u}\psi_i(\mathbf{x} - \mathbf{r}_u)$ . We demonstrated in Appendix B that with some regularity and locality assumptions, the message passing in Eq.(2) can be parameterized as
$$
f _ {i, u} = \sum_ {v \in \tilde {\mathcal {N}} (u)} \sum_ {j = 1} ^ {\infty} w _ {i j} S _ {i j} \left(\mathbf {r} _ {u v}\right) f _ {j, v} \tag {3}
$$
where  $S_{ij}(\mathbf{r}) = \int_{\mathcal{D}} \psi_i(\mathbf{x}) \psi_j(\mathbf{x} - \mathbf{r}) d\mathbf{x}$  models the interaction between the two displaced basis at centers  $i, j$ . The outer summation is over  $\tilde{\mathcal{N}}(u) = \mathcal{N}(u) \cup \{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  $\mathbf{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}(\mathbf{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.,  $\mathcal{D} = \mathbb{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}^{\ell}(r)$  and a spherical harmonics  $Y_{\ell}^{m}(\hat{\mathbf{r}})$ :
$$
\psi_ {n \ell m} (\mathbf {r}) = R _ {n} ^ {\ell} (r) Y _ {\ell} ^ {m} (\hat {\mathbf {r}}) = c _ {n \ell} \exp \left(- a _ {n} r ^ {2}\right) r ^ {\ell} Y _ {\ell} ^ {m} (\hat {\mathbf {r}}) \tag {4}
$$
where  $r = |\mathbf{r}|$  is the vector length and  $\hat{\mathbf{r}} = \mathbf{r} / r$  is the direction vector on the unit sphere.  $c_{n\ell}$  are normalizing constants such that  $\int_{\mathbb{R}^3}|\psi_{n\ell m}(\mathbf{r})|^2 dV = 1$ . The degree of the spherical harmonics  $\ell$  takes values of non-negative integers, and the order  $m$  takes integers values between  $-\ell$  and  $\ell$  (inclusive). Therefore, there are  $2\ell + 1$  spherical harmonics of degree  $\ell$ . In this way, the basis index  $i,j$  are now triplets of  $(n,\ell,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\ell_1m_1,n_2\ell_2m_2)$ , the overlap integral  $S(\mathbf{r})$  can also be expanded onto the basis as  $S(\mathbf{r}) = \sum_{n\ell m}s_{n\ell m}\psi_{n\ell m}(\mathbf{r})=: \sum_{n\ell m}\varphi_{n\ell m}(\mathbf{r})^1$ . For a fixed  $\mathbf{r}$  and radial index  $n$ , the coefficient sequence  $\varphi = \{\varphi_{\ell m}:\ell \geq 0, - \ell \leq m\leq \ell \}$  can be viewed as a spherical tensor. Notice that the node feature  $\mathbf{f} = \{\mathbf{f}^{\ell}:\ell \geq 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:
$$
\mathbf {f} _ {u} ^ {\ell} \leftarrow \sum_ {v \in \tilde {\mathcal {N}} (u)} \sum_ {k \geq 0} W ^ {\ell k} \left(\mathbf {x} _ {v} - \mathbf {x} _ {u}\right) \mathbf {f} _ {v} ^ {k}, \quad W ^ {\ell k} (\mathbf {r}) = \sum_ {J = | k - \ell |} ^ {k + \ell} \varphi_ {J} ^ {\ell k} (r) \sum_ {m = - J} ^ {J} Y _ {J} ^ {m} (\hat {\mathbf {r}}) Q _ {J m} ^ {\ell k} \tag {5}
$$
where  $Q_{Jm}^{\ell k}$  is the Clebsch-Gordan matrix of shape  $(2\ell + 2) \times (2k + 1)$  and  $\varphi_J^{\ell k}: \mathbb{R}^+ \to \mathbb{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  $\hat{\rho}(\mathbf{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  $\mathbf{r}_{uv} = \mathbf{x}_v - \mathbf{x}_u$ . The equivariance of this scheme relies on the equivariance of the input feature map  $\mathsf{f}_{\mathrm{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,
$$
\mathbf {f} _ {u} ^ {\ell} = w ^ {\ell} \mathbf {f} _ {u} ^ {\ell} + \sum_ {v \in \mathcal {N} (u)} \sum_ {k \geq 0} W ^ {\ell k} \left(\mathbf {x} _ {v} - \mathbf {x} _ {u}\right) \mathbf {f} _ {v} ^ {k} \tag {6}
$$
<sup>1</sup>If we use an infinite number of complete orthonormal basis, the expansion can be achieved without error. However, if we use some none-complete or finite basis, this should be viewed as an approximation in the subspace spanned by the basis.
The first term is referred to as self-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:
$$
f ^ {0} = \sigma_ {0} \left(f ^ {0}\right), \quad \mathbf {f} ^ {\ell} = \sigma_ {\ell} \left(\| \mathbf {f} ^ {\ell} \| _ {2}\right) \mathbf {f} ^ {\ell} \tag {7}
$$
where  $\sigma_{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:  $\hat{S}_{n\ell m,n^{\prime}\ell^{\prime}m^{\prime}}(\mathbf{r})\coloneqq \delta_{nn^{\prime}}S_{mm^{\prime},\ell \ell^{\prime}}(\mathbf{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  $\mathbf{x}$ . More specifically, the residual operator layer aggregates the neighboring node features to produce an invariant scalar<sup>2</sup> to finetune the final estimation:
$$
z (\mathbf {x}) = \sum_ {v \in \mathcal {N} (p)} \sum_ {k \geq 0} W _ {\text {r e s}} ^ {k} \left(\mathbf {x} _ {v} - \mathbf {x}\right) \mathbf {f} _ {v} ^ {k} \tag {8}
$$
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
$$
\hat {\rho} (\mathbf {x}) = \sum_ {n \ell m} f _ {n \ell m} \psi_ {n \ell m} (\mathbf {x}) + z (\mathbf {x}) \tag {9}
$$
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(\mathbb{R}^3)$  as  $\mathcal{L} = \|\hat{\rho} - \rho\|_2^2 = \int_{\mathbb{R}^3} |\hat{\rho}(\mathbf{x}) - \rho(\mathbf{x})|^2 d\mathbf{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  $\mathcal{D}$ , a graphon, also known as graph limit or graph function, is a symmetric square-integrable function:
$$
W: \mathcal {D} \times \mathcal {D} \rightarrow [ 0, 1 ], \int_ {\mathcal {D} ^ {2}} | W (\mathbf {x}, \mathbf {y}) | ^ {2} d \mathbf {x} d \mathbf {y} <   \infty \tag {10}
$$
Intuitively, the kernel  $W(\mathbf{x}, \mathbf{y})$  can be viewed as the probability that an edge forms between the continuously indexable nodes  $\mathbf{x}, \mathbf{y} \in \mathcal{D}$ . Now, consider the operator  $\mathcal{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  $\mathcal{T}_W$  it induces is a self-adjoint operator whose spectrum consists of a countable number of real-valued eigenvalues  $\{\lambda_k\}_{k=1}^{\infty}$  with  $\lambda_k \to 0$ . Let  $\{\phi_k\}_{k=1}^{\infty}$  be the eigenfunctions such that  $\mathcal{T}_W \phi_k = \lambda_k \phi_k$ . Similarly to the graph convolution for discrete graphs, any transformation on the eigenvalues  $\mathcal{F}: \{\lambda_k\}_{k=1}^{\infty} \mapsto \{\mu_k\}_{k=1}^{\infty}$  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\mathbf{x} = \sum_{k=0}^{K} w_k L^k\mathbf{x}$  where  $w_k$  are the parameters. Define the power series of  $\mathcal{T}_W$  as:
$$
\mathcal {T} _ {W} ^ {n} f (\mathbf {x}) = \mathcal {T} _ {W} \mathcal {T} _ {W} ^ {n - 1} f (\mathbf {x}) = \int_ {\mathcal {D}} W (\mathbf {x}, \mathbf {y}) \mathcal {T} _ {W} ^ {n - 1} f (\mathbf {y}) d \mathbf {y}, \mathcal {T} _ {W} ^ {0} = \mathcal {I} \tag {11}
$$
where  $\mathcal{I}$  is the identity mapping on  $\mathcal{D}$ . A graphon filter can be then defined as  $\mathcal{H}f = \sum_{k=0}^{\infty} w_k \mathcal{T}_W^k f$ . We can also follow GCN to use the Chebyshev polynomials to approximate the graphon filter  $\mathcal{H}$ :
$$
\mathcal {H} f \approx \theta_ {1} f + \theta_ {2} \mathcal {T} _ {W} f \tag {12}
$$
2Note that for scalars, SO(3) equivariance is equivalent to invariance.
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  $\mathcal{T}_Wf$  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(\mathbf{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_{\mathrm{in}}(\mathbf{x}) = \sum_{u}f_{u}\delta (\mathbf{x} - \mathbf{x}_{u})$