SlowGuess/NeurIPS_2023_OCR huggingface dataset:
GRAPHPATCHER: Mitigating Degree Bias for Graph Neural Networks via Test-time Augmentation
Mingxuan Ju $^{1}$ , Tong Zhao $^{2}$ , Wenhao Yu $^{1}$ , Neil Shah $^{2}$ , Yanfang Ye $^{1}$
<sup>1</sup>University of Notre Dame, <sup>2</sup>Snap Inc.
$^{1}\{mju2,wyu1,yye7\} \text{@nd.edu}; ^{2}\{tzhao,nshah\} \text{@snap.com}$
Recent studies have shown that graph neural networks (GNNs) exhibit strong biases towards the node degree: they usually perform satisfactorily on high-degree nodes with rich neighbor information but struggle with low-degree nodes. Existing works tackle this problem by deriving either designated GNN architectures or training strategies specifically for low-degree nodes. Though effective, these approaches unintentionally create an artificial out-of-distribution scenario, where models mainly or even only observe low-degree nodes during the training, leading to a downgraded performance for high-degree nodes that GNNs originally perform well at. In light of this, we propose a test-time augmentation framework, namely GRAPHPATCHER, to enhance test-time generalization of any GNNs on low-degree nodes. Specifically, GRAPHPATCHER iteratively generates virtual nodes to patch artificially created low-degree nodes via corruptions, aiming at progressively reconstructing target GNN's predictions over a sequence of increasingly corrupted nodes. Through this scheme, GRAPHPATCHER not only learns how to enhance low-degree nodes (when the neighborhoods are heavily corrupted) but also preserves the original superior performance of GNNs on high-degree nodes (when lightly corrupted). Additionally, GRAPHPATCHER is model-agnostic and can also mitigate the degree bias for either self-supervised or supervised GNNs. Comprehensive experiments are conducted over seven benchmark datasets and GRAPHPATCHER consistently enhances common GNNs' overall performance by up to  $3.6\%$  and low-degree performance by up to  $6.5\%$ , significantly outperforming state-of-the-art baselines. The source code is publicly available at https://github.com/jumxglhf/GraphPatcher.
Graph Neural Networks (GNNs) have gained significant popularity as a powerful approach for learning representations of graphs, achieving state-of-the-art performance on various predictive tasks, such as node classification [22, 38, 9], link prediction [50, 53], and graph classification [43, 47, 11]. These tasks further form the archetypes of many real-world applications, such as recommendation systems [45, 3], predicative user behavior models [31, 52], and molecular property prediction [51, 48].
While existing GNNs are highly proficient at capturing information from rich neighborhoods (i.e., high-degree nodes), recent studies [13, 25, 36, 55] have revealed a significant performance degradation of GNNs when dealing with nodes that have sparse neighborhoods (i.e., low-degree nodes). This observation can be attributed to the fact that GNNs make predictions based on the distribution of node neighborhoods [27]. According to this line of theory, GNNs struggle with low-degree nodes due to the limited amount of available neighborhood information, which may not be able to precisely depict the learned distributions. Empirically, as shown in Figure 1, the classification accuracy of GCN [22] proportionally decays as the node degree decreases, resulting in a performance gap of  $\sim 20\%$  accuracy. Furthermore, the sub-optimal performance of GNNs on low-degree nodes can be
Node Degree vs. Accuracy across SoTA
Figure 1: The classification accuracy of GCN and SoTA frameworks that mitigate degree biases.
aggravated by the power-law degree distribution commonly observed in real-world graphs, where the number of low-degree nodes significantly exceeds that of high-degree nodes [36].
To bridge this gap, several frameworks have been proposed to specifically improve GNNs' performance on low-degree nodes [36, 25, 13, 55, 49]. These frameworks either introduce designated architectures or training strategies specifically for low-degree nodes. For examples, Tail-GNN [25] enhances latent representations of low-degree nodes by incorporating high-degree structural information; whereas Coldbrew [55] retrieves a set of existing nodes as virtual neighbors for low-degree nodes. However, these approaches suffer from two significant drawbacks. Firstly, while benefiting low-degree nodes, they inadvertently create an artificial out-of-distribution scenario during training [42], where models primarily observe low-degree nodes, leading to a downgraded performance for high-degree nodes that GNNs originally perform well on. Secondly, deploying these frameworks often requires changing model architectures, which can be impractical in real-world scenarios where the original models are well-trained due to the expensive re-training cost (on large-scale graphs) and the shared usage of it across different functionalities in production.
In light of these drawbacks, we propose a test-time augmentation framework for GNNs, namely GRAPHPATCHER. Given a well-trained GNN, GRAPHPATCHER mitigates the degree bias by patching corrupted ego-graphs with multiple generated virtual neighbors. Notably, GRAPHPATCHER not only enhances the performance of low-degree nodes but also maintains (sometimes improves) GNNs performance on high-degree nodes. This behavior is empirically important because practitioners can universally apply GRAPHPATCHER to all nodes without, like previous works, manually discovering a degree threshold that differentiates the low- and high-degree nodes. To achieve so, we first generate a sequence of ego-graphs corrupted with increasing strengths. Then, GRAPHPATCHER recursively generates multiple virtual nodes to patch the mostly corrupted graph, such that the frozen GNN gives similar predictions for the patched graph and the corresponding corrupted ego-graph in the sequence. Through this scheme, GRAPHPATCHER not only learns how to patch low-degree nodes (i.e., heavily corrupted) but also maintains GNNs original superior performance on high-degree nodes (i.e., lightly corrupted). As a test-time augmentation framework, GRAPHPATCHER is parameterized in parallel with the target GNN. Hence, GRAPHPATCHER is model-agnostic and requires no updates on the target GNN, enabling practitioners to easily utilize it as a plug-and-play module to existing well-established infrastructures. Overall, our contributions are summarized as:
- We study a more practical setting of degree biases on graphs, where both the performances on low- and high-degree nodes are considered. In this case, a good framework is required to not only improve the performance over low-degree nodes but also maintain the original superior performance over high-degree nodes. We evaluate existing frameworks in this setting and observe that many of them trade off performance on high-degree nodes for that on low-degree nodes.
- To mitigate degree biases, we propose GRAPHPATCHER, a novel test-time augmentation framework for graphs. Given a well-trained GNN, GRAPHPATCHER iteratively generates multiple virtual nodes and uses them to patch the original ego-graphs. These patched ego-graphs not only improve GNNs' performance on low-degree nodes but also maintains that over high-degree nodes. Moreover, GRAPHPATCHER is applied at the testing time for GNNs, a plug-and-play module that is easily applicable to existing well-established infrastructures.
- We conduct extensive evaluation of GRAPHPATCHER along with six state-of-the-art frameworks that mitigate degree biases on seven benchmark datasets. GRAPHPATCHER consistently enhances the overall performance by up to  $3.6\%$  and low-degree performance by up to  $6.5\%$  of multiple GNNs, significantly outperforming state-of-the-art baselines.
Graph Neural Networks. Graph Neural Networks (GNNs) have become one of the most popular paradigms for learning representations over graphs [22, 38, 9, 43, 23, 17, 4]. GNNs aim at mapping the input nodes into low-dimensional vectors, which can be further utilized to conduct either graph-level or node-level tasks. Most GNNs explore a layer-wise message passing scheme, where a node iteratively extracts information from its first-order neighbors, and information from multi-hop neighbors can be captured by stacked layers. They achieved state-of-the-art performance on various tasks, such as node classification [22, 44, 12, 35], link prediction [50, 53, 8], node clustering [2, 37], etc. These tasks further form the archetypes of many real-world applications, such as recommendation systems [45, 3], predictive user behavior models [31, 52], question answering [18], and molecular property prediction [51, 48, 7, 24].
Degree Bias underlying GNNs. Recent studies have shown that GNNs exhibit strong biases towards the node degree: they usually perform satisfactorily over high-degree nodes with rich neighbor information but suffer over low-degree nodes [13, 25, 36, 55]. Existing frameworks that mitigate degree biases derive either designated architectures or training strategies specifically for low-degree nodes. For instance, Tail-GNN [25] enhances low-degree nodes' latent representations by injecting high-degree structural information learned from high-degree nodes; Coldbrew [55] retrieves a set of existing nodes as virtual neighbors for low-degree nodes; TuneUp [13] fine-tunes the well-trained GNNs with pseudo labels and heavily corrupted graphs. Though effective for low-degree nodes, they unintentionally create an artificial out-of-distribution scenario [42], where models only observe low-degree nodes during the training, leading to downgraded performance for high-degree nodes that GNNs originally perform well at.
Test-time Augmentation. While data augmentations during the training phase have become one of the essential ingredients for training machine learning models [54], the augmentation applied during the testing time is far less studied, especially for the graph learning community. It has been moderately researched in the computer vision field, aimed at improving performance or mitigating uncertainties [34, 21, 41, 1]. They usually corrupt the same sample by different augmentation approaches and aggregate the model's predictions on all corrupted samples. Whereas in the graph community, GTrans [15] proposes a test-time enhancement framework, where the node feature and graph topology are modified at the test time to mitigate potential out-of-distribution scenarios.
In this work, we specifically focus on the node classification task. Let  $G = (V,E)$  denote a graph, where  $V$  is the set of  $|V| = N$  nodes and  $E \subseteq V \times V$  is the set of  $|E|$  edges between nodes.  $\mathbf{X} \in \mathbb{R}^{N \times d}$  represents the feature matrix, with  $i$ -th row representing node  $v_{i}$ 's  $d$ -dimensional feature vector.  $\mathbf{Y} \subseteq \{0,1\}^{N \times C}$  denotes the label matrix, where  $C$  is the number of total classes. And  $\mathbf{Y}^{(L)}$  denotes the label matrix for training nodes. We denote the ego-graph of node  $v_{i}$  as defined as  $\mathcal{G}(v_i) = (V_i,E_i)$  with  $V_{i} = \mathcal{N}_{k}(v_{i})$ , where  $\mathcal{N}_k(v_i)$  stands for all nodes within the  $k$ -hop neighborhood of  $v_{i}$  including itself and  $E_{i}$  refers to the edges in-between  $\mathcal{N}_k(v_i)$ . A well-trained GNN  $f_{g}(\cdot ;\pmb{\theta}): G \to \mathbb{R}^{N\times C}$  parameterized by  $\pmb{\theta}$  takes  $G$  as input and maps every node in  $G$  to a  $C$ -dimensional class distribution. Formally, we define test-time node patching as the following:
Definition 1 (Test-time Node Patching). Given a GNN  $f_{g}(\cdot ;\pmb{\theta})$  and a graph  $G$ , a test-time node patching framework  $f(\cdot ;\phi):G\to G$  takes  $G$  and outputs the patched graph  $\hat{G}$  with generated nodes and edges, such that the performance of  $f_{g}$  over nodes in  $G$  is enhanced when  $\hat{G}$  is utilized:
$$
\arg \min  _ {\boldsymbol {\phi}} \mathcal {L} \left(f _ {g} (f (G; \boldsymbol {\phi}); \boldsymbol {\theta} ^ {*}), \mathbf {Y}\right), \quad \text {w h e r e} \quad \boldsymbol {\theta} ^ {*} = \arg \min  _ {\boldsymbol {\theta}} \mathcal {L} \left(f _ {g} (G; \boldsymbol {\theta}), \mathbf {Y} ^ {(L)}\right), \tag {1}
$$
where  $\mathcal{L}$  refers to the loss function evaluating the GNN (e.g., cross-entropy or accuracy).
In this work, we aim at mitigating the degree bias via test-time node patching. To achieve so, two challenges need to be addressed: (1) how to optimize and formulate  $f(\cdot; \phi)$ , such that the graphs patched by  $f(\cdot; \phi)$  enhance the performance of  $f_{g}(\cdot; \pmb{\theta}^{*})$  over low-degree nodes; and (2) how to derive a unified learning scheme that allows  $f(\cdot; \phi)$  to not only improve low-degree nodes but also maintain the GNN's original superiority over high-degree nodes.
Figure 2: GRAPHPATCHER is presented ego-graphs corrupted by increasing strengths (i.e., the top half of the figure). From the most corrupted graph, it iteratively generates patching nodes to the anchor node, such that the target GNN behaves similarly given the currently patched graph or the corrupted graph next in the hierarchy (i.e., the bottom half of the figure).
3.2 The Proposed Framework: GRAPHPACKER
Our proposed GRAPHPatcher is a test-time augmentation framework for GNNs to mitigate their degree biases. As shown in Figure 2, GRAPHPatcher is presented a sequence of ego-graphs corrupted by increasing strengths. Starting from the most corrupted graphs, GRAPHPatcher iteratively generates patching nodes to augment the anchor nodes. Compared with the corrupted graphs next in the hierarchy, the patched graphs should allow the target GNN to deliver similar outputs. Through this scheme, GRAPHPatcher not only learns how to patch low-degree nodes while preserving the superior performance over high-degree nodes.
3.2.1 Patching Ego-graphs via Prediction Reconstruction
In order to patch low-degree nodes, a straightforward approach is to corrupt high-degree nodes into low-degree nodes, and allowing the learning model to patch the corrupted nodes to restore their original properties [25, 13]. However, patching low-degree nodes not only affects their own representations but also those of their neighbors due to the message-passing mechanism of GNNs as well as the non-i.i.d. property of nodes in a graph. Besides, modeling over the entire graphs requires the learning model to consider all potential circumstances, whose overheads grow quadratically w.r.t. the number of nodes. Consequently, it becomes challenging to simultaneously determine both features and neighbors of the patching nodes given the entire graph.
To reduce the complexity of the optimization process, instead of working over the entire graph, we conduct node patching over ego-graphs and regard each ego-graph as an i.i.d. sample of the anchor node [56, 19]. For each node  $v_{i}$ , we have  $f_{g}(G;\pmb{\theta})[v_{i}] = f_{g}(\mathcal{G}(v_{i});\pmb{\theta})[v_{i}]$  if  $k$  equals to the number of layers in  $f_{g}(\cdot ;\pmb{\theta})$ . To further simplify the optimization process, we directly wire the generated virtual nodes to the anchor node (i.e., the generated virtual nodes are the first-order neighbors of the anchor node). This implementation is simple yet effective, because we no longer consider the location to place the patching node: any modification that affects the latent representation of the anchor node can be achieved by patching nodes (with different features) directly to the anchor nodes.
We start explaining GRAPHPatcher by the most basic case where we only conduct node patching once. Specifically, given the a trained GNN  $f_{g}(\cdot ;\pmb{\theta}^{*})$ , an anchor node  $v_{i}$ , and a corruption function  $T(\cdot ;t)$  with strength  $t$  (i.e., first-order neighbor dropping with probability  $t$  to simulate a low-degree
scenario), that is,  $\mathcal{G}'(v_i) = (V'(v_i), E'(v_i)) = \mathcal{T}(\mathcal{G}(v_i), t)$ . GRAPHPATCHER  $f(\cdot; \phi)$  takes the corrupted ego-graph  $\mathcal{G}'(v_i)$  as input and outputs the augmented ego-graph  $\hat{\mathcal{G}}(v_i)$  with a patching node  $v_p$  and its feature  $\mathbf{x}_p$ , which is directly connected to  $v_i$ . That is,
$$
\hat {\mathcal {G}} (v _ {i}) = f \left(\mathcal {G} ^ {\prime} (v _ {i}); \phi\right), \text {w h e r e} \hat {V} = V ^ {\prime} (v _ {i}) \cup \left\{v _ {p} \right\}, \hat {E} = E ^ {\prime} (v _ {i}) \cup \left\{e _ {(i, p)} \right\}, \tag {2}
$$
where  $e_{(i,p)}$  refers to the edge connecting  $v_{i}$  and  $v_{p}$  and  $V^{\prime}(v_{i})$  and  $E^{\prime}(v_{i})$  refer to the nodes and edges in  $\mathcal{G}'(v_i)$ , respectively. To optimize  $f(\cdot :\phi)$  such that  $f_{g}(\cdot ;\boldsymbol{\theta}^{*})$  gives similar predictions to  $\hat{\mathcal{G}}'(v_i)$  and  $\mathcal{G}(v_i)$ , we minimize the Kullback-Leibler divergence between the frozen GNN's predictions on these two ego-graphs, which is defined as:
$$
\arg \min  _ {\phi} \sum_ {v _ {i} \in V _ {\mathrm {t r}}} \operatorname {K L - D i v} \left(f _ {g} \left(\mathcal {G} \left(v _ {i}\right); \boldsymbol {\theta} ^ {*}\right) [ v _ {i} ], f _ {g} \left(f \left(\mathcal {G} ^ {\prime} \left(v _ {i}\right); \phi\right); \boldsymbol {\theta} ^ {*}\right) [ v _ {i} ]\right), \tag {3}
$$
where  $\mathrm{KL - Div}(\mathbf{y}_1,\mathbf{y}_2) = (\mathbf{y}_1 + \epsilon)\cdot \bigl (\log (\mathbf{y}_2 + \epsilon) - \log (\mathbf{y}_1 + \epsilon)\bigr)^1$  with  $\epsilon >0$  and  $V_{\mathrm{tr}}$  refers to the set of anchor nodes for training. Intuitively, the reconstruction process above enforces GRAPHPATCHER to remedy the corrupted neighborhood caused by  $\mathcal{T}(\cdot ;t)$  via adding a patching node directly to the anchor node. It is philosophically similar to the existing works (e.g., TuneUp [13] and Tail-GNN [25]), where models gain better generalization over low-degree nodes via the corrupted high-degree nodes. Empirically, we observe that this branch of approaches can effectively enhance performance over lowdegree nodes. Though promising, according to our empirical studies, it falls short on the high-degree node that original GNNs perform well at. This phenomenon may be attributed to the unintentionally created out-of-distribution scenario [42], wherein models primarily encounter nodes with low degrees during the training. Consequently, the performance of GNNs, which is typically proficient with high-degree nodes, is adversely affected and downgraded.
3.2.2 Iterative Patching to Mitigate Degree Bias
In this work, we emphasize that: mitigating degree bias should not focus specifically on the low-degree nodes: trading off performance on high-degree nodes for that on low-degree nodes simply creates a new bias towards high-degree nodes. Therefore, besides enhancing the performance on low-degree nodes, maintaining GNN's original superiority on high-degree nodes is equally critical. This behavior is empirically desirable because practitioners can universally apply GRAPHPATCHER to all nodes without, like previous works do, manually discovering the degree threshold that differentiates the low- and high-degree nodes. Furthermore, the fact that these frameworks are applicable only to low-degree nodes indicates a lack of robustness: further remedying a neighborhood that is informative enough to deliver a good classification result should not jeopardize the performance.
To mitigate the degree bias, we propose a novel training scheme for GRAPHPATCHER such that it observes both low- and high-degree nodes simultaneously during the optimization. Specifically, given a node  $v_{i}$ , we firstly create a sequence of  $M$  corrupted ego-graphs of  $v_{i}$ , denoted as  $S(v_{i}) = [\mathcal{G}'(v_{i})_{m} = \mathcal{T}(\mathcal{G}(v_{i}),t_{m})]_{m = 1}^{M}$ , with decreasing corruption strength (i.e.,  $\forall m,n\in \{1,\dots ,M\}$ ,  $t_m > t_n$  if  $m < n$ ). Instead of the one-step patching to match the prediction on the original ego-graph as described in Section 3.2.1, GRAPHPATCHER traverses  $S(v_{i})$  and recursively patches the corrupted ego-graph to match the target GNN's prediction on the ego-graph next in the sequence. As also illustrated in Figure 2, this optimization process is formulated as:
$$
\arg \min  _ {\phi} \sum_ {v _ {i} \in V _ {\mathrm {t r}}} \sum_ {m = 1} ^ {M - 1} \operatorname {K L - D i v} \left(f _ {g} \left(\mathcal {G} ^ {\prime} \left(v _ {i}\right) _ {m + 1}; \boldsymbol {\theta} ^ {*}\right) [ v _ {i} ], f _ {g} \left(\hat {\mathcal {G}} \left(v _ {i}\right) _ {m}; \boldsymbol {\theta} ^ {*}\right) [ v _ {i} ]\right), \tag {4}
$$
$$
\begin{array}{l} \text {s . t .} \hat {\mathcal {G}} (v _ {i}) _ {m} = f (\hat {\mathcal {G}} (v _ {i}) _ {m - 1}; \phi), \end{array}
$$
where  $\hat{\mathcal{G}}(v_i)_m = (\hat{V}_m, \hat{E}_m)$  with  $\hat{V}_m = V_1'(v_i) \cup \{v_p\}_{p=1}^m$ ,  $\hat{E}_m = E_1'(v_i) \cup \{e_{(i,p)}\}_{p=1}^m$ , and  $\hat{\mathcal{G}}(v_i)_0 = \mathcal{G}'(v_i)_1$ .
The one-step patching described in Section 3.2.1 remedies low-degree anchor nodes directly to the distributions of high-degree nodes. During this process, the model does not observe distributions of high-degree nodes and hence delivers sub-optimal performance. Therefore, we design GRAPH-PATCHER to be an iterative multi-step framework. At each step, it takes the previously patched
$^{1}$ KL divergence used here is equal to the regularized cross-entropy. It is strongly convex and Lipschitz continuous due to the incorporation of  $\epsilon$ . These two properties are required for the derivation of Theorem 1.
ego-graph as input and further remedies the partially patched ego-graph to match the GNN's prediction on the ego-graph next in the sequence. This scheme enables GRAPHPACHER to learn to patch low-degree nodes in early steps when the ego-graphs are heavily corrupted (e.g., low-degree case in Figure 2) and maintain the original performance in later steps when ego-graphs are lightly corrupted (e.g., high-degree case in Figure 2). Specifically, at the  $m$ -th patching step, the currently patched ego-graph  $\hat{\mathcal{G}}(v_i)_m$  reflects the neighbor distribution of ego-graphs corrupted by a specific strength of  $t_{m+1}$ . GRAPHPACHER takes  $\hat{\mathcal{G}}(v_i)_m$  as input and further generates another patching node  $v_{m+1}$  to approach the neighbor distribution of ego-graphs corrupted by a slightly weaker strength of  $t_{m+2}$ . This process iterates until GRAPHPACHER traverses  $S(v_i)$ . Intuitively, the incorporation of  $v_{m+1}$  enriches the neighbor distribution by an amount of  $t_{m+2} - t_{m+1}$  corruption strength. This optimization scheme allows GRAPHPACHER to observe neighbor distributions with varying corruption strengths and makes our proposal applicable to both low- and high-degree nodes.
However, the target distribution at each step (i.e.,  $f_{g}\big(\mathcal{G}^{\prime}(v_{i})_{m + 1};\pmb {\theta}\big)[v_{i}]$  in Equation (4)) is not deterministic due to the stochastic nature of the corruption function  $\mathcal{T}$ . Given an ego-graph  $\mathcal{G}(v_i)$  and a corruption strength  $t$ , one can at most generate  $\binom{|V_i|}{(1-t)|V_i|}$  different corrupted ego-graphs. With a large corruption strength (e.g., ego-graphs early in the sequence  $S(v_{i})$ ), two corrupted ego-graphs generated by the same exact priors might exhibit completely different topologies. Such differences could bring high variance to the supervision signal and instability to the optimization process. To alleviate the issue above, at each step we sample  $L$  ego-graphs with the same corruption strength and let GRAPHPATCHER approximate multiple predictions over them, formulated as:
$$
\mathcal {L} _ {\text {p a t c h}} = \sum_ {v _ {i} \in V _ {\mathrm {t r}}} \sum_ {m = 1} ^ {M - 1} \sum_ {l = 1} ^ {L} \operatorname {K L - D i v} \left(f _ {g} \left(\mathcal {G} ^ {\prime} \left(v _ {i}\right) _ {m + 1} ^ {l}; \boldsymbol {\theta} ^ {*}\right) [ v _ {i} ], f _ {g} \left(\hat {\mathcal {G}} \left(v _ {i}\right) _ {m}; \boldsymbol {\theta} ^ {*}\right) [ v _ {i} ]\right), \tag {5}
$$