SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$\mathbf{LD}^2$ : Scalable Heterophilous Graph Neural Network with Decoupled Embeddings
Ningyi Liao
liao0090@e.ntu.edu.sg
Siqiang Luo*
siqiang.luo@ntu.edu.sg
xiangli@dase.ecnu.edu.cn
Jieming Shi
jieming.shi@polyu.edu.hk
Heterophilous Graph Neural Network (GNN) is a family of GNNs that specializes in learning graphs under heterophily, where connected nodes tend to have different labels. Most existing heterophilous models incorporate iterative non-local computations to capture node relationships. However, these approaches have limited application to large-scale graphs due to their high computational costs and challenges in adopting minibatch schemes. In this work, we study the scalability issues of heterophilous GNN and propose a scalable model,  $\mathrm{LD}^2$ , which simplifies the learning process by decoupling graph propagation and generating expressive embeddings prior to training. Theoretical analysis demonstrates that  $\mathrm{LD}^2$  achieves optimal time complexity in training, as well as a memory footprint that remains independent of the graph scale. We conduct extensive experiments to showcase that our model is capable of lightweight minibatch training on large-scale heterophilous graphs, with up to  $15\times$  speed improvement and efficient memory utilization, while maintaining comparable or better performance than the baselines. Our code is available at: https://github.com/gdmnl/LD2
Graph Neural Networks (GNNs) combine graph management techniques and neural networks to learn from graph-structured data, and have shown remarkable performance in diverse graph processing tasks, including node classification [1, 2], link prediction [3, 4], and community detection [5, 6]. Common GNN models rely on the principle of homophily, which assumes that connected nodes tend to be similar to each other in terms of classes [7]. This inductive bias introduces additional information from the graph structure and improves model performance in appropriate tasks [8].
However, this assumption does not always hold in practice. A broad range of real-world graphs are heterophilous, where class labels of neighboring nodes usually differ from the ego node [9]. In such cases, the aggregation mechanism employed by conventional GNNs, which only passes messages from a node to its neighbors, may mix the information from non-homophilous nodes and cause them to be less discriminative. Consequently, the locality-based design is considered less advantageous or even potentially harmful in these applications [10, 11]. Various solutions have been proposed to address the heterophily problem, giving rise to a class of specialized GNNs known as heterophilous GNNs. Common strategies to address heterophily include discovering multi-hop or global graph relations [12, 13, 14, 15, 16, 17], and retrieving expressive node information through enhanced network architectures [18, 19, 20, 21, 22, 23].
Table 1: Time and memory complexity of homophilous and non-homophilous models with respect to precomputation, training, and inference stages. "Batch" refers to the minibatch availability of the original model design. Training and inference memory indicate the GPU usage for storing and updating representation and weight matrices, while precomputation is mostly on RAM. Training and inference time complexity represent the forward-passing computational operations on respective node sets.
Model	Batch	Pre. Mem.	Training Mem.	Inference Mem.	Pre. Time	Training Time	Inference Time
MLP	Y	-	O(LnbF + LF2)	O(LnbF + LF2)	-	O(ILnF2)	O(LnF2)
GCN	N	-	O(LnbF + LF2)	O(LnbF + LF2)	-	O(ILmF + ILnF2)	O(LmF + LnF2)
GSAINT	Y	-	O(LPbNbF + LF2)	O(LnbF + LF2)	-	O(ILPmLnF + ILnF2)	O(LmF + LnF2)
APPNP	Y	O(m)	O(LnbF + LF2 + nb)	O(LnbF + LF2 + nb)	O(m)	O(ILPmF + ILnF2)	O(LPmF + LnF2)
PPRGo	Y	O(n/δ)	O(LnbF + LF2)	O(LnbF + LF2)	O(m/δ)	O(Inf2 + ILnF2)	O(nF2 + LnF2)
SGC	Y	O(nF)	O(LnbF + LF2)	O(LnbF + LF2)	O(LPmF)	O(ILnF2)	O(LnF2)
GPRGN	N	O(m)	O(LnbF + LF2 + m)	O(LnbF + LF2 + m)	O(m)	O(ILPmF + ILnF2)	O(LPmF + LnF2)
GCNJK	N	-	O(LcNbF + LC2F2)	O(LcNbF + LC2F2)	-	O(ILmF + ILnF2)	O(LmF + LnF2)
MixHop	N	-	O(CLNbF + CLF2)	O(CLNbF + CLF2)	-	O(ILPmF + ILnF2)	O(LPmF + LnF2)
LINKX	Y	-	O(LcNbF + LC2F2 + nF)	O(LcNbF + LC2F2 + nF)	-	O(ImF + ILnF2)	O(mF + LnF2)
LD2(ours)	Y	O(CnF)	O(LcNbF + LC2F2)	O(LcNbF + LC2F2)	O(LPmF)	O(ILnF2)	O(LnF2)
Existing heterophilous GNNs are not scalable enough. Scalability has become a prominent concern in GNN studies. The ever-increasing sizes of graph data nowadays can easily exceed the memory limit of devices such as GPUs, rendering these models impractical for large-scale tasks [24, 25]. We observe that this issue is particularly critical in the context of heterophilous GNNs, due to an inherent conflict that most current models have not taken into account: heterophily-oriented designs usually rely on non-local information calculated by certain types of whole-graph operations. As the graph structure is involved, the time and memory overhead escalates substantially with the graph size. A recent investigation [26] reveals that all the evaluated full-graph GNNs run out of 24GB GPU memory when applied to the million-scale graph wiki (1.77M nodes, 243M edges), while minibatch strategies often result in performance loss. It is thus crucial to develop GNNs scalable to large graphs while retaining the capability for heterophily.
$\mathbf{LD}^2$ : a scalable solution via decoupling. In this work, we examine the scalability problem and propose  $\mathbf{LD}^2$ , a scalable GNN model for heterophilous graphs with Low-Dimensional embeddings and Long-Distance aggregation. The model highlights simplicity by decoupling graph dependency from iterative computations and solely learning from a set of precomputed embeddings. Derived from node attributes and graph topology, these novel embeddings are capable of aggregating node relations of varying objectives and distances in the graph into low-dimensional features. To facilitate the decoupled scheme, we specifically propose an algorithm to efficiently estimate all embeddings before training, which enjoys time complexity only linear to the graph scale and a guaranteed precision bound. After the precomputation, a simple but powerful multi-channel neural network is subsequently employed to learn from the extracted node features. Theoretical and empirical results showcase that the combination of embeddings effectively retrieves representations among heterophilous nodes. On the efficiency aspect,  $\mathbf{LD}^2$  is advantageous in its scalable design, including the straightforward minibatch scheme, optimal training and inference times, and the superior memory utilization.
Our contribution. (1) We propose  $\mathrm{LD}^2$  as a scalable GNN under heterophily, which removes the reliance on iterative train-time full-graph computations. The model realizes theoretically optimized training, highlighting time complexity that is only linear to the number of nodes  $O(n)$  and memory overhead independent of the graph scale. To the best of our knowledge,  $\mathrm{LD}^2$  is the first model achieving such optimization in the context of heterophilous GNNs. (2) We design an array of feature and topology embeddings by applying multi-hop discriminative propagation, encoding expressive node representation within a compact size. An end-to-end precomputation algorithm is proposed for efficient embedding calculation with a linear complexity. (3) We conduct comprehensive experiments to evaluate the efficacy and efficiency of  $\mathrm{LD}^2$ . On large-scale datasets, our model demonstrates  $3 - 15\times$  faster minibatch training and inference, and up to  $5\times$  smaller memory footprint than state of the art, with comparable or better accuracy. Particularly, it completes training on the wiki dataset within 1 minute and inference in 0.1 seconds, demanding only 5GB of GPU memory.
2 Preliminary and Related Work
Graph Notation and Heterophily. In an undirected graph  $G = (V, E)$  with node set  $V$  and edge set  $E$ , the number of nodes, number of edges, and average degree are denoted by  $n = |V|$ ,  $m = |E|$ , and  $d = m / n$ , respectively. The neighborhood of an ego node  $u \in V$  is the set  $\mathcal{N}(u) = \{v | (u, v) \in E\}$ , and its degree  $d(u) = |\mathcal{N}(u)|$ . The diagonal degree matrix is  $D = \mathrm{diag}(d(u_1), d(u_2), \dots, d(u_n))$ .
The graph connectivity is represented by the adjacency matrix  $\mathbf{A} \in \mathbb{R}^{n \times n}$ . We adopt the general graph normalization scheme [31, 32] with coefficients  $a, b \in [0,1]$  and  $\tilde{\mathbf{A}} = D^{-a}\mathbf{A}\mathbf{D}^{-b}$ . The normalized adjacency matrix with self-loop edges is also frequently used, which is denoted as  $\tilde{\mathbf{A}} = (\mathbf{I} + \mathbf{D})^{-a}(\mathbf{I} + \mathbf{A})(\mathbf{I} + \mathbf{D})^{-b}$ , and the corresponding graph Laplacian matrix is  $\tilde{\mathbf{L}} = \mathbf{I} - \tilde{\mathbf{A}}$ . Each node  $u \in V$  is represented by an  $F$ -dimensional attribute vector  $\mathbf{X}(u)$ , which composes the attribute matrix  $\mathbf{X} \in \mathbb{R}^{n \times F}$ . The graph heterophily is measured by the node homophily score [14], which is the average proportion of the neighbors with the same class of each node.
Vanilla and Sampling-based GNNs. We summarize our analysis of the time and memory complexity of related GNN models in Table [1]. In general, a GNN recurrently computes the node representation  $\pmb{H}^{(l)}$  in its  $l$ -th layer. For the vanilla GCN [2], the model input is the attribute matrix  $\pmb{H}^{(0)} = \pmb{X}$ , and the layer representation is updated by  $\pmb{H}^{(l+1)} = \sigma(\tilde{\pmb{A}}\pmb{H}^{(l)}\pmb{W}^{(l)})$ ,  $l = 0,1,\dots,L-1$ , where  $\pmb{W}^{(l)}$  is the trainable weight matrix each layer,  $\sigma(\cdot)$  is the activation function, and  $L$  is the number of layers. For simplicity we assume the feature size  $F$  to be constant in all layers. Previous research [33, 31] points out that the operation of dominating expense in both GCN training and inference is the graph propagation  $\tilde{\pmb{A}}\cdot\pmb{H}^{(l)}$ , which can be regarded as repetitive sparse-dense matrix multiplications, resulting in a total complexity  $O(LmF)$ . The overhead for feature transformation by applying  $\pmb{W}^{(l)}$  is  $O(LnF^2)$ . These two procedures are iteratively performed for  $I$  epochs throughout training. In terms of memory usage, GCN typically requires  $O(LnF + LF^2)$  space to store layer-wise node representations and weight matrices, respectively. For large-scale cases where  $n\gg F$ , the overhead of dense node representations  $O(LnF)$  becomes the primary term [24].
The above analysis indicates that the scalability bottleneck of GCN lies in the time complexity of graph propagation as well as the memory overhead of full-graph representation. There is a large scope of GNNs attempting to address the issue by sampling techniques, which simplify the propagation by replacing the entire graph with subgraphs in minibatches [2, 27, 33]. For example, the widely-used GSAINT [27] incorporates  $L_{P}$ -hop random walk sampling, reducing the in-memory representation to  $O(L_{P}Ln_{b}F)$ , where  $n_b$  is the batch size. However, it is not applicable to the inference stage.
Full-graph Heterophilous GNNs. In the context of GNNs under heterophily, a large number of models augment the spatial graph convolution to a full-graph scheme, relying upon the complete graph topology to compute inter-node relationships. For instance,  $\mathrm{H}_2\mathrm{GCN}$  [13] and MixHop [12] incorporate 2-hop propagation  $\tilde{A}^2 H^{(l)}$ , while  $\mathrm{GeomGCN}$  [14] and  $\mathrm{GloGNN}$  [22] exploit hierarchical computation on non-local connections. These high-order calculations are shown to be effective in retrieving information beyond immediate neighbors, but come at the price of more complex propagation operations. Another common practice for non-homophilous design is altering transformation to learn from multiple features, i.e. channels. MixHop [12] mixes its multi-hop representation as  $H^{(l + 1)} = \sigma (H^{(l)}W_0^{(l)}\| \tilde{A} H^{(l)}W_1^{(l)}\| \tilde{A}^2 H^{(l)}W_2^{(l)})$ , where  $(\cdot \| \cdot)$  denotes matrix concatenation. GGCN [19] and ACM [20] select frequency-based filters expressed by  $\tilde{A} HW_{l}$ ,  $(I - \tilde{A})HW_{h}$ ,  $IHW_{i}$ , while GCNJK [21] records individual layer representations as channels. Denote the number of channels as  $C$  and  $L_{C} = L + C$ , employing multi-channel learning increases the memory budget for node representations and weight matrices by a factor related to  $C$  or  $L_{C}$ .
Alternatively, a line of approaches choose to apply learnable spectral filters for flexibility on nonhomophilous data, deriving specific propagation in the spectral domain for various graphs. GPRGNN [17] generalizes the multi-hop propagation as  $T = \sum_{l=0}^{L_P} \theta_l \tilde{A}^l$ , which inevitably involves the entire adjacency matrix for acquiring weight factors  $\theta_l$ . BernNet [34] and ChebNetII [35] respectively utilize order- $L_P$  Bernstein and Chebyshev polynomials with respect to  $\tilde{L}$  to approximate their spectral filters. These models focus relatively more on the expressiveness of graph propagation rather than the overall GNN architectural design, usually demanding additional cost for determining spectral filters and updating the learnable basis.
In spite of their advantageous capabilities, above full-graph designs hardly address the scalability bottleneck in GNN propagation, which can be observed from the  $O(n)$  or  $O(m)$  terms in their complexity. Recent studies also discover that heterophilous GNNs are naturally unsuitable for enforcing scalable training via sampling-based minibatching, since their distant or full-graph information is heavily overlooked in batches built on locality [11]. Evaluations show that simply fitting these models to learn from induced subgraph samples causes performance degradation [26]. Therefore, we believe these full-graph models targeting effectiveness are orthogonal to our study on scalability.
Decoupled GNNs. As the scalability of GNN is closely tied to the graph propagation, a promising approach to simplify the process is decoupling it from iterative training and efficiently computing it in advance. The representative two-stage model SGC [30] encodes graph information with  $\mathbf{X}$  into an embedding matrix  $\mathbf{P} = \tilde{\mathbf{A}}^{L_P} \cdot \mathbf{X}$ , which is then input to a Multi-Layer Perceptron  $\mathbf{H}^{(L)} = \mathrm{MLP}(\mathbf{P})$ . [36, 32, 25] further generalize the embedding aggregation. Such decoupled GNN is considered optimal for training efficiency, as its time complexity  $O(ILnF)$  is identical to the simple MLP [31].
However, applying the decoupling technique to heterophilous GNNs is non-trivial due to the full-graph relationships. To our knowledge, LINKX [26] is the only model conceptually similar to this scheme, removing graph-related propagation during training iterations and enabling solely node-wise minibatching. It exploits a simple architecture  $\pmb{H}^{(L)} = \mathrm{MLP}(\pmb{X}\pmb{W}_X\| \pmb{A}\pmb{W}_A)$ , where the matrix  $\pmb{A}$  is used as an input feature in learning. The major drawback of this design is that it suffers from the  $O(nF)$  term in model size and  $O(mF)$  term in forward time, hindering its scalability to large graphs.
In this section, we first present an overview of our  $\mathrm{LD}^2$  model, then respectively elaborate on the selection of the adjacency and feature embeddings. Lastly, an end-to-end scalable algorithm, namely  $\mathrm{A}^2\mathrm{Prop}$ , is proposed to efficiently and concurrently compute all the embeddings.
3.1  $\mathbf{LD}^2$ : A Decoupled Heterophilous GNN
In order to achieve superior time and memory scalability for heterophilous GNNs, we employ the concept of decoupling, which removes the dependency of graph adjacency propagation in training iterations. The main idea of our model is first generating embeddings from raw features including node attributes and adjacency in a precomputation stage. Then, these embeddings are taken as inputs to learn representations by a simple neural network model. We embrace the multi-channel architecture [37, 20] to enhance flexibility, where the input data is a list consisting of embedding matrices  $[P_1, P_2, \dots, P_C]$ . Each embedding is separately processed and then merged in the network.
$\mathrm{LD}^2$  utilizes diverse embeddings based on pure graph adjacency and node attributes, denoted as  $P_A(A)$  and  $P_X(X,A)$ , respectively. Both types of embeddings can be produced by our precomputation  $\mathrm{A}^2\mathrm{Prop}$  following Algorithm 1. The initial layer of the  $\mathrm{LD}^2$  network applies a separate linear transformation to each embedding input, and the results are concatenated to form the representation matrix. Lastly, an  $L$ -layer MLP is leveraged for the classification task. The high-level framework of  $\mathrm{LD}^2$  is depicted in Figure 1 and can be expressed as follow:
Precompute:  $P_A, P_X = \mathrm{A}^2\mathrm{Prop}(A, X)$ ; Transform:  $H^{(L)} = \mathrm{MLP}(P_A W_A \| P_X W_X)$ . (1)
Training/Inference Complexity. Our decoupled model design enables a simple on-demand mini-batch scheme in training and inference, that only  $n_b$  rows corresponding to the batch nodes in the embedding matrices are loaded into GPU and processed by the network transformation. For  $\mathrm{LD}^2$  with  $C$  channels, the GPU memory footprint is therefore bounded by  $O(L_{C}n_{b}F + L_{C}F^{2})$ . It is worth noting that such complexity does not depend on the graph scale  $n$  or  $m$ . Consequently, the training is freely configurable with an arbitrary GPU memory budget. Regarding computation operations, the time complexity of forward inference through the graph is  $O(LnF^2)$ , being just linear to  $n$ . As
Figure 1:  $\mathrm{LD}^2$  framework: decoupled precomputation and training. Figure 2: Propagations under heterophily.
Constant 2-hop Adjacency Propagation
the model complexity only contain essential operations of MLP-like transformation on nodes in the graph with no additional expense, this is the optimal scale with respect to the iterative training of GNN architectures.
3.2 Low-Dimensional Adjacency Embedding
Several studies reveal that, despite the feature information of nodes, the pure graph structure is equally or even more important in the context of heterophilous GNNs [9, 14, 26]. Particularly, the most informative aspects are often associated with 2-hop neighbors, i.e., "neighbors of neighbors" of ego nodes. [13] proves that even under heterophily, the 2-hop neighborhood is expected to be homophily-dominant. [38] also verifies that the 2-hop similarity is strongly relevant to GNN performance. We thence intend to explicitly model such topological information.
The 2-hop relation can be described by the 2-hop adjacency matrix  $A^2$ . Note that as the sparse matrix  $A$  has  $m$  entries, the number of entries in  $A^2$  is at the scale of  $O(md)$ , which indicates that directly applying 2-hop graph propagation in the training stage will demand even more expensive time and memory overhead to be scaled up. We instead propose an approximate scheme that seeks to prevent the 2-hop adjacency from explicit processing, and retrieves a low-dimensional but expressive embedding prior to training in the precomputation stage. In other words, we utilize the embedding to resemble 2-hop information which can be directly learned by the neural network transformation. Denote the  $F$ -dimensional embedding matrix as  $P_A \in \mathbb{R}^{n \times F}$ . We aim to minimize its approximation error in Frobenius norm  $(\| \cdot \|_F)$ :
$$
\boldsymbol {P} _ {A} = \underset {\boldsymbol {P} \in \mathbb {R} ^ {n \times F}} {\arg \min } \| \boldsymbol {A} ^ {2} - \boldsymbol {P} \boldsymbol {P} ^ {T} \| _ {F} ^ {2}. \tag {2}
$$
The solution to Eq. (2) can be derived from the eigendecomposition of the symmetric matrix  $A^2$ , that  $P_A^* = U|\Lambda|^{1/2}$ , where  $\Lambda = \mathrm{diag}(\lambda_1, \dots, \lambda_F)$  is the diagonal matrix with top- $F$  eigenvalues  $\lambda_1 \geq \lambda_2 \geq \dots \geq \lambda_F$ , and  $U \in \mathbb{R}^{n \times F}$  is the matrix consisting of corresponding orthogonal eigenvectors. The eigenvalues are also called frequencies of the graph, and large eigenvalues of the adjacency matrix refer to low-frequency signals in the graph spectrum.
Spectral Analysis. Let  $A_{2}(u,v)$  be the entry  $(u,v)$  of matrix  $A^{2}$ . Its diagonal degree matrix is  $D_{2} = \mathrm{diag}(d_{2}(u_{1}), d_{2}(u_{2}), \dots, d_{2}(u_{n}))$ , where  $d_{2}(u) = \sum_{v \in V} A_{2}(u,v)$ . Denote  $P_{A}(u)$  as the  $F$ -dimensional adjacency embedding vector of node  $u$ . We show that the embedding matrix  $P_{A}^{*}$  defined by Eq. (2) is also the solution to the following optimization problem:
$$
\boldsymbol {P} _ {A} = \underset {\boldsymbol {P} \in \mathbb {R} ^ {n \times F}, \boldsymbol {P} ^ {\top} \boldsymbol {D} _ {2} \boldsymbol {P} = \boldsymbol {\Lambda}} {\arg \min } \sum_ {u, v \in V} A _ {2} (u, v) \| \boldsymbol {P} (u) - \boldsymbol {P} (v) \| ^ {2}. \tag {3}
$$
This is because  $\sum_{u,v} A_2(u,v) \| P(u) - P(v) \|^2 = 2\sum_u d_2(u) \| P(u) \|^2 - 2\sum_{u,v} A_2(u,v) P(u)P(v) = 2\operatorname{tr}(P^\top D_2P - P^\top A^2P)$ . As  $P^\top D_2P$  is fixed, finding the minimum of Eq. (3) is equivalent to optimizing  $\max_P P^\top A^2 P$ , of which the solution is exactly  $P_A^*$  according to the property of eigenvectors. Equation (3) implies that, 2-hop neighbors  $(u,v), t \in \mathcal{N}(u), v \in \mathcal{N}(t)$  in the graph will share similar embeddings  $P_A(u)$  and  $P_A(v)$ .
In fact, the low-dimensional embedding  $P_A^*$  can be interpreted as the adjacency spectral embedding of the 2-hop graph  $A^2$ . Graph spectral embedding is a technique concerning the low-frequency spectrum of a graph, and is employed in tasks such as graph clustering [39]. As  $P_A$  corresponds to the dominant eigenvalues of  $A^2$ , the embedding provides an approximate representation of the 2-hop neighborhoods based on the overall graph topology. Alternatively, if we regard the adjacency information solely as features input into the network like LINKX,  $P_A$  correlates to the uncentered principal components of matrix  $A$ . Therefore, learning a linear transformation  $P_A W_A$  with weight matrix  $W_A \in \mathbb{R}^{F \times F}$  in  $\mathrm{LD}^2$  is the same expressive as the rank-  $F$  approximation of  $A W_{A0}$  in LINKX, where  $W_{A0} \in \mathbb{R}^{n \times F}$ , but with a less computational cost independent to the graph scale. Compared to other works attempting to generate graph embeddings based on graph geometric or similarity measures [40, 41, 14, 16, 19, 42], our approach offers the advantages of lower dimensionality and efficient calculation as demonstrated in Section 3.4
3.3 Long-Distance Feature Embedding
Decoupling the node features through approximate propagation has been extensively studied in regular GNNs with various schemes [28, 30, 36, 29, 43, 32, 44]. Nonetheless, these approaches are
based on the homophily assumption and focus on local neighborhoods. In order to apply decoupled propagation to heterophilous graphs and exploit the multi-channel ability of our model, we formulate the general form of approximate propagation as the weighted sum of powers of a propagation matrix applied to the input feature, i.e.,  $P_{X} = \sum_{l=1}^{L_{P}} \theta_{l} T^{l} X$ . Examples of propagation matrix  $T$  include  $\tilde{A}$  and  $\tilde{L}$ , which respectively correspond to aggregative and discriminative operations.
$\mathrm{LD}^2$  utilizes the following channels jointly as input embeddings: (1) inverse summation of 1-hop improved Laplacian propagations  $P_{X,H} = \frac{1}{L_{P,H}}\sum_{l=1}^{L_{P,H}}(I + \tilde{L})^l X$ , ( $\theta_l = 1$ ,  $T = I + \tilde{L}$ ); (2) constant summation of 2-hop adjacency propagations  $P_{X,L2} = \frac{1}{L_{P,L2}}\sum_{l=1}^{L_{P,L2}}\bar{A}^{2l}X$ , ( $\theta_l = 1$ ,  $T = \bar{A}^2$ ); (3) raw node attributes  $P_{X,0} = X$ .