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

In the graph domain, deep graph networks based on Message Passing Neural Networks (MPNNs) or Graph Transformers often cause over-smoothing of node features, limiting their expressive capacity. Many upsampling techniques involving node and edge manipulation have been proposed to mitigate this issue. However, these methods are often heuristic, resulting in extensive manual labor and suboptimal performance and lacking a universal integration strategy. In this study, we introduce UniGAP, a universal and adaptive graph upsampling framework to mitigate over-smoothing in node classification tasks. Specifically, we design an adaptive graph upsampler based on condensed trajectory features, serving as a plug-in component for existing GNNs to mitigate the over-smoothing problem and enhance performance. Moreover, UniGAP serves as a representation-based and fully differentiable framework to inspire further exploration of graph upsampling methods. Through extensive experiments, UniGAP demonstrates significant improvements over heuristic data augmentation methods in various datasets and metrics. We analyze how graph structure evolves with UniGAP, identifying key bottlenecks where over-smoothing occurs, and providing insights into how UniGAP addresses this issue. Lastly, we show the potential of combining UniGAP with large language models (LLMs) to further improve downstream performance. Our code is available at: https://github.com/wangxiaotang0906/UniGAP

💡 Research Summary

UniGAP (Universal and Adaptive Graph Upsampling) tackles the pervasive over‑smoothing problem in deep graph neural networks (GNNs) by learning to insert intermediate nodes on edges in a task‑driven, differentiable manner. The authors first observe that node representations across layers form “trajectories” that encode how each node’s features evolve and gradually converge to a common vector as depth increases. These trajectories are pre‑computed using one of three strategies—zero initialization, non‑parametric message passing (pure powers of the adjacency matrix), or a pretrained GNN that provides richer feature embeddings.

Once the layer‑wise trajectories are collected, a Multi‑View Condensation (MVC) encoder compresses the L‑step trajectory tensor (shape L × N × d) into a compact per‑node vector (N × d). Two concrete MVC designs are explored: (1) Trajectory‑MLP‑Mixer, which learns hop‑wise attention weights and mixes them via MLP layers, and (2) Trajectory‑Transformer, which treats each hop as a token and applies self‑attention to capture long‑range dependencies. The resulting condensed features capture a node’s propensity to over‑smooth.

The Adaptive Upsampler then predicts, for every original edge, a probability of inserting an intermediate node. This probability is produced by a small MLP that consumes the condensed trajectory features of the two incident nodes. A differentiable sampling step (e.g., Gumbel‑Softmax) decides whether to actually insert a node, ensuring gradients can flow back to both the MVC encoder and the upsampler. Inserted nodes are linked to the two original endpoints, and their initial attributes are derived from a weighted combination of neighboring features.

The upsampled graph—now larger in both node count and edge set—is fed into any downstream GNN (GCN, GraphSAGE, GAT, etc.) for the target task, typically node classification. The task loss is back‑propagated through the entire pipeline, jointly updating the MVC encoder, the upsampler, and the downstream GNN. After an initial “warm‑up” epoch that uses the original graph (to avoid cold‑start issues), trajectories are recomputed from the refined GNN, and the process iterates until convergence.

Key contributions include:

1. A universal, plug‑and‑play upsampling module that replaces heuristic rules with learnable, data‑driven probabilities.
2. Full differentiability, enabling end‑to‑end training with any GNN backbone.
3. Interpretability: the locations of inserted intermediate nodes directly reveal structural bottlenecks where over‑smoothing is most severe, offering visual diagnostics.
4. Demonstrated synergy with large language models (LLMs); the authors show that LLM‑generated textual cues or prompt‑guided adjustments to insertion probabilities can further boost performance.

Extensive experiments on citation networks (Cora, Citeseer, PubMed) and large‑scale Open Graph Benchmark datasets (ogbn‑arxiv, ogbn‑products) validate UniGAP. When combined with various GNNs, UniGAP consistently improves classification accuracy by 3–7 percentage points over strong baselines, especially as the number of layers exceeds 8, where traditional GNNs suffer severe performance drops due to over‑smoothing. Computational overhead remains modest because the upsampling step is lightweight and the MVC encoder reduces the trajectory dimensionality before the upsampler.

The paper positions UniGAP as a new paradigm: instead of merely tweaking message‑passing mechanisms or applying static graph rewiring, it treats graph topology itself as a learnable representation that can be adaptively refined to counteract the diffusion of information. Future directions suggested include extending the framework to more complex graph transformations (subgraph insertion, multi‑node expansions), exploring unsupervised trajectory extraction, and deepening the integration with LLMs for domain‑specific knowledge injection. Overall, UniGAP offers a principled, versatile, and empirically validated solution to one of the most stubborn challenges in deep graph learning.