UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node Classification Tasks

Trajectories encompass multi-view information of each node, including one-hop, two-hop, and higher-order information, which are essential for subsequent modules. As the output of node features at each layer of the model, they implicitly reflect the gradual process of over-smoothing as the number of layers increases. Our aim is to derive a high-dimensional representation of the over-smoothing process from the trajectories, thereby enabling us to devise strategies to mitigate over-smoothing. In this section, we will elucidate the process of trajectory computation.

Let the trajectories be denoted as $T\in\mathbb{R}^{L\times N\times d}$, where $L$ represents the number of layers, $N$ denotes the number of nodes, and $d$ signifies the dimensionality of features. Trajectories include the concatenation of multi-view features (hidden states from each layer):

where $\operatorname{concat}$ denotes channel-wise concatenation. $L$ is the number of model layers. $T^{(1)},\cdots,T^{(L)}$ represent multi-view features (1-hop, $\cdots$, L-hop information), we collect them through gathering the hidden representations in each layer:

where the representation of each hop is denoted as $H^{(l)}=f^{(l)}(H^{(l-1)})\in\mathbb{R}^{N\times d}$, $f^{(l)}(\cdot)$ means $l$-th layer in trajectory computation model, and $\operatorname{Norm}({x}_{i})=\frac{{x}_{i}}{\left|\left|{x}_{i}\right|\right|_{2}}$ signifies a channel-wise $L2$ normalization, which aids in regulating the magnitude of features that may escalate exponentially with the number of iterations. $w$ is a hyper-parameter that determines the frequency of normalizing features.

In the initial stage, a specific model is not available for calculating the trajectories of the original graph. Therefore, a trajectory precomputation module is needed to obtain initialized trajectories. Various strategies can be employed to achieve this goal. In this work, we explore three off-the-shelf strategies, considering their efficiency and simplicity: Zero Initialization, Non-parametric Message Passing, and Pretrained GNN. Next, we will discuss how to integrate these strategies into our methods:

1) Zero initialization. We initialize the trajectories with all-zero vectors, resulting in zero multi-view condensed features during the first epoch of training. Consequently, the upsampling phase in this initial epoch does not introduce new edges or nodes, thus utilizing the original graph for downstream tasks. After finishing the first epoch, we store the hidden representations of each layer in downstream model as the trajectories for subsequent epochs. This zero initialization ensures that the UniGAP module learns the graph structure tailored to each downstream task with minimal noise and optimal performance. The zero initialization can be formulated as

2) Non-parametric Message Passing. After removing the non-linear operator in graph neural networks, we can decouple the message passing ($AX$) and feature transformation ($XW$). However, in the initial stage, we have no access to the transformation weights $W$ unless pre-training methods are employed. Therefore, we adopt non-parameteric message passing to collect multi-view information, i.e., $AX,A^{2}X,\cdots A^{L}X$. This approach avoids nonlinear transformations, preserving maximum original feature, graph topology, and trajectories. The formulation is as follows:

where $\mathcal{P}^{(l)}$ denotes the $l$-th layer in $\mathcal{P}$ and $T^{(0)}=G$.

3) Pretrained GNN. The first two methods do not introduce additional parameters, which may limit the expressivity of the information embedded in the raw features. Here, we introduce a more powerful method, though it may incur additional computational costs. We leverage the training data to train a GNN model by designing a pretext task such as unsupervised contrastive loss (You et al., 2020; Zhu et al., 2022, 2024b) or masked prediction (Hou et al., 2022; Zhu et al., 2023) to obtain a generalized feature extractor. An additional advantage of this method is the ability to use pre-trained parameters to initialize the encoding parameters in the downstream model, facilitating faster convergence and reducing training time. In this way, we can kill two birds with one stone. The pretraining procedure can be formulated as follows:

After the pre-training process, the pre-trained GNN model will be employed to compute trajectories:

After obtaining trajectories (multi-view information), we aim to adaptively condense this multi-view information into compact features that reflect the node’s propensity to over-smoothing, while considering efficiency and adaptivity. To achieve this goal, we design a multi-view condensation (MVC) encoder that can adaptively condense trajectories $T\in\mathbb{R}^{L\times N\times d}$ into condensed trajectory features $\hat{T}\in\mathbb{R}^{N\times d}$. This reduces computational cost and adaptively adjusts the weights of information at each hop. We instantiate the MVC encoder using two strategies in this work: Trajectory-MLP-Mixer (TMM) and Trajectory-Transformer.

1) Trajectory-MLP-Mixer (TMM). Considering simplicity and efficiency, we introduce a lightweight model, i.e., Trajectory-MLP-Mixer, to instantiate the MVC encoder as depicted in Figure 2a. Initially, trajectory mixing is applied to the trajectory dimension to adaptively integrate information from each hop. Subsequently, channel mixing is applied to the mixed trajectories to obtain condensed multi-view features:

where $w_{\text{traj}}\in\mathbb{R}^{1\times d}$ is employed for trajectory mixing, $\odot$ denotes the broadcasted Hadamard product with a summation across the feature dimension, and $\sigma$ represents the normalization function. Thus, $\sigma(w_{\text{traj}}\odot T)\in\mathbb{R}^{L\times N}$ signifies the importance weights for each hop information of each node, which will be utilized for the weighted summation $\oplus$ of multi-view information $T$. Additionally, $W_{\text{channel}}\in\mathbb{R}^{d\times d}$ is used for channel mixing.

2) Trajectory-Transformer (TT). We treat the trajectories as a sequence, with each hop information considered a token input, as illustrated in Figure 2b. These tokens are fed into a Transformer encoder to learn the mixed-hop representations. Subsequently, an attention-based readout function is developed to adaptively aggregate mixed-hop features to obtain condensed trajectory features. The steps can be formulated as follows:

where $\operatorname{Tr}(\cdot)$ denotes a shallow Transformer encoder, and $\mathcal{R}$ represents an attention-based readout function.

After obtaining condensed trajectory features $\hat{T}$, we use them to compute the upsampling probability for each existing edge to decelerate message passing, thereby mitigating the over-smoothing problem (theoretical justification can be found in Section 3.5.1). Subsequently, we employ a differentiable sampling method to generate an augmented graph. Specifically, we compute the upsampling logits as follows:

where $P_{ij}$ represents the upsampling logits for the edge between nodes $i$ and $j$, ${W}_{\text{up}}$ is a learnable weight matrix, and $\left[\hat{{T}}_{i}\|\hat{{T}}_{j}\right]$ denotes the concatenation of the condensed trajectory features of nodes $i$ and $j$.

Then we employ Gumbel-Softmax (Jang et al., 2017) with additional tricks to derive the upsampling mask and generate the augmented graph $\hat{G}$:

where $\epsilon\in\mathbb{R}^{|\mathcal{E}|\times 2}$ represents Gumbel noise used to introduce diversity into the sampling process (exploration), preventing the selection of edges solely based on highest probabilities and ensuring differentiability. $\hat{P}$ denotes the upsampling probability, and $\mathds{1}_{\hat{P}[:,0]\leq\hat{P}[:,1]}$ is an indicator function that returns 1 if $\hat{P}[:,0]\leq\hat{P}[:,1]$. $\tilde{P}$ results from truncating the gradient of $\hat{P}$. This approach achieves two objectives: (1) ensures the output is exactly one-hot (by adding and then subtracting a soft value), and (2) maintains the gradient equal to the soft gradient (by keeping the original gradients). Lastly, $M$ represents an active mask used to determine which edges will have a new node inserted. Specifically, if the mask for edge $i\to j$ is 1, a new node $k$ is inserted between nodes $i$ and $j$. The original edge $i\to j$ is then masked and replaced by new edges $i\to k$ and $k\to j$. The initial feature vector of node $k$ can be set to a zero vector, the average of nodes $i$ and $j$, or an adaptive sum:

Then, we obtain the augmented graph $\hat{{G}}=(\hat{\mathcal{V}},\hat{\mathcal{E}},\hat{X})$, which will be used for downstream tasks to mitigate over-smoothing and improve overall performance.

To optimize the augmented graph, we will utilize it in downstream tasks such as node classification. Let the downstream model be denoted as $\mathcal{D}_{\psi}(\cdot)$. The augmented graph will be fed into this model to compute the downstream loss $\mathcal{L}_{\text{down}}$. Additionally, we incorporate over-smoothing metrics, such as Mean Average Distance (MAD) (Chen et al., 2020b) and Dirichlet Energy (Rusch et al., 2023), as regularization terms $\mathcal{L}_{\text{smooth}}$ to ensure that UniGAP explicitly focuses on mitigating the over-smoothing issue. The total loss is defined as:

where $\tilde{y}$ represents the prediction of the downstream model, $\mathcal{N}_{v}$ denotes the set of neighboring nodes associated with node $v$ used in the downstream model, and $\beta$ is the regularization coefficient that controls the impact of the smoothing metric. We empirically set $\beta$ to 1 because a larger value would negatively impact downstream performance. The formula of MAD is represnted as:

Through optimizing the downstream loss, all parameters in our framework will be updated in an end-to-end manner through backpropagation. After warming up the downstream model, it will be used to compute trajectories instead of relying solely on the trajectory precomputation module mentioned in Section 3.3.1. The model training process continues until convergence, aiming to obtain the optimal upsampled graph and model parameters.

Equipped with UniGAP, the optimized upsampled graph can not only enhance downstream performance but also mitigate the over-smoothing problem. In this subsection, we will provide theoretical explanations by answering the following questions:

RQ1: How does UniGAP learn in concert with downstream tasks to get the optimal upsampling graph structure?

RQ2: How does UniGAP mitigate the over-smoothing problem?

As shown in Equation 14,15, UniGAP can learn the message passing dynamics and adaptivly change the nodes’ receptive filed, finally reach a optimal structural representation for downstream task. The derivation and more illustrations can be found in Appendix B.

As shown in Lemma 3.2, the rate of smoothing of original MP is $2k$ due to the $A^{2k}\Sigma$ operator (Keriven, 2022). Then we consider the condition with UniGAP.

Here $\Sigma^{(k)}_{\text{UniGAP}}$ is the approximated covariance of the node features, $p=\mathbb{E}(\hat{P}[:,0])$ and $\hat{P}$ denotes the upsampling probability at the optimal graph structure. The rate of smoothing of MP with UniGAP is $k-1$ due to the $A^{k-1}\Sigma$ operator. The complete proof can be found in Appendix B.

With Lemma 3.2 and 3.3, we can naturally get the Theorem:

By slowing down the smoothing rate of the covariance of node features during message passing, UniGAP reduces the smoothing rate from $2k$ to $k-1$.

In this subsection, we will present a complexity analysis of UniGAP from a theoretical perspective.

Time Complexity. Our framework comprises four primary components: trajectory computation, multi-view condensation, adaptive graph upsampler, and downstream model. The trajectory precomputation module involves complexities of $\mathcal{O}(1)$, $\mathcal{O}(LEd)$, or $\mathcal{O}(LEd+LNd^{2})$ depending on the initialization method: zero initialization, non-parametric message passing, or pre-trained GNN. The MVC Condensation module (Trajectory-MLP-Mixer) operates with a time complexity of $\mathcal{O}(LN+Nd^{2})$, while the adaptive graph upsampler computes with a complexity of $O(Ed)$ to determine upsampling probabilities for adjacent node pairs. The time complexity of downstream model, e.g., GNN, is $\mathcal{O}(LEd+LNd^{2})$. Consequently, the total time complexity of our method is $O(LN+LEd+LNd^{2})$, comparable to traditional L-layer GNNs, i.e., $\mathcal{O}(LEd+LNd^{2})$.

Space Complexity. The space complexity of our method also closely aligns with that of traditional GNNs ($\mathcal{O}(E+Ld^{2}+LNd)$). Specifically, the trajectory precomputation module exhibits space complexities of $\mathcal{O}(1)$, $\mathcal{O}(E)$, or $\mathcal{O}(E+Ld^{2}+LNd)$. The multi-view condensation module requires $\mathcal{O}(d+d^{2}+LN+Nd)$, whereas the adaptive graph upsampler operates with a complexity of $\mathcal{O}(d+E)$. The space complexity of the downstream model, such as GNN, is $\mathcal{O}(E+Ld^{2}+LNd)$. Consequently, the total space complexity of our method is $\mathcal{O}(d+d^{2}+LN+Nd+E+Ld^{2}+LNd)$ which is similar to that of GNN.

In summary, the complexity of UniGAP is on par with traditional GNN methods, maintaining efficiency.

We evaluated our method on 8 node classification benchmark datasets, as shown in Table 1, comprising 4 homophilic and 4 heterophilic datasets. More detailed information is provided in Appendix A.

To compare UniGAP with baselines, we use public split setting for all datasets, and report the mean accuracy along with the standard deviation across twenty different random seeds. UniGAP is integrated as a plug-in method with various SOTA GNN models to evaluate its effectiveness. Besides, we report the MAD metric to observe the phenomenon of over-smoothing as the number of layers increases.

The baselines primarily consist of four categories: (1) traditional GNN methods such as GCN (Kipf and Welling, 2017), GAT (Velickovic et al., 2018), and GraphSAGE (Hamilton et al., 2017), (2) advanced GNN methods including GCNII (Chen et al., 2020c), GGCN (Yan et al., 2021), H2GCN (Zhu et al., 2020b), GPRGNN (Chien et al., 2021), and AERO-GNN (Lee et al., 2023), (3) Graph Transformer methods like GT (Dwivedi and Bresson, 2020), GraphGPS (Rampášek et al., 2023), and NAGphormer (Chen et al., 2023), and (4) other graph upsampling methods, i.e., HalfHop. We integrate UniGAP and HalfHop with various GNN models for evaluation.

For heterophilic datasets, UniGAP significantly enhances traditional GNN models, which heavily rely on the “homophily” assumption. For example, UniGAP achieves an average absolute improvement on four heterophilic datasets of 10.0% on GCN, 13.3% on GAT, and 5.0% on GraphSAGE, respectively. It also enhances other advanced GNN models, with absolute improvements such as 2.3% on GCNII, 1.1% on H2GCN, and 2.9% on GPRGNN. Compared to HalfHop, UniGAP surpasses it in all heterophilic datasets with over 3.9% absolute improvement with GCN, demonstrating the effectiveness of adaptive graph upsampling in identifying the optimal graph structure to enhance the downstream model’s performance.

For homophilic datasets, traditional GNN models benefit from their inductive bias, specifically the homophily assumption, leading to moderate performance. However, UniGAP can still identify graph structures that enhance downstream tasks through adaptive learning, thereby improving the performance of all baselines across different datasets. Average absolute improvements range from 1.1% to 4.8%.

In summary, compared to all the baselines and datasets, UniGAP demonstrates a significant improvement, with an average absolute boost ranging from 1.4% to 6.7%, surpassing the SOTA baselines, showcasing the effectiveness of UniGAP.

To analyze the extent to which UniGAP effectively mitigates the phenomenon of over-smoothing, we conducted experiments on GCN across several datasets, evaluating its accuracy and MAD as the number of layers increases. As depicted in Figure 3, the experimental results demonstrate that UniGAP significantly enhances model performance and alleviates over-smoothing at each layer, surpassing AdaEdge and HalfHop.

Specifically, while HalfHop and AdaEdge can mitigate over-smoothing to some extent, they do not achieve optimal performance as UniGAP does. HalfHop fails to identify crucial elements necessary to mitigate over-smoothing and enhance performance because it inserts nodes randomly rather than adaptively finding optimal positions for the inserted nodes. AdaEdge increases connectivity between intra-class nodes, but effective message propagation between inter-class nodes requires an optimal graph structure. Therefore, adaptive upsampling represents a more sophisticated approach to mitigating over-smoothing and enhancing model performance.

To better analyze the reasons for mitigating over-smoothing, we examine the ratio of intermediate node insertions. This analysis provides valuable insights into the root causes of over-smoothing. As shown in Figure 4, we identify the locations where intermediate nodes are added. We observe that new nodes tend to be inserted between inter-class node pairs, slowing down the message passing between different classes and preventing the features of different classes from collapsing into identical vectors, thus mitigating the over-smoothing problem. For the portion of nodes inserted into intra-class edges, we speculate that some intra-class edges contain noise. Equipped with UniGAP, these harmful noises can be eliminated, thereby improving performance.