This repository contains the official implementation of our paper "A unified gradient regularization method for heterogeneous graph neural networks", published in Neural Networks (2025).

Short Abstract:
Heterogeneous Graph Neural Networks (HGNNs) often suffer from over-smoothing and unstable training.
We introduce Grug, a unified gradient regularization framework that jointly regularizes node-type and message-level gradients during message passing.
Grug provides theoretical guarantees on stability and diversity, unifies existing dropping and adversarial methods, and achieves state-of-the-art performance on six public datasets with minimal overhead.

📄 Read the paper: https://www.sciencedirect.com/science/article/abs/pii/S0893608025009840

• torch >= 1.12.1
• torch-geometric >= 2.2.0
• dgl >= 0.9.1

• ACM
• DBLP
• IMDB
• ogbn-mag
• Amazon
• LastFM

• RGCN
• RGAT
• HGT
• HAN
• SimpleHGN

• DropEdge
• DropNode
• DropMessage
• FLAG

You can directly run the source code of Grug:

python main.py -m <model> -t <task> -d <dataset> -a <alpha> -b <beta> -M <iterations>

Arguments:

• -m model: Model name (e.g., RGCN, RGAT, HGT).
• -t task: Task name (node_classification or link_prediction).
• -d dataset: Dataset name (ACM, DBLP, IMDB, ogbn-mag, etc.).
• -a alpha: Coefficient for message matrix regularization.
• -b beta: Coefficient for node matrix regularization.
• -M iterations: Number of regularization iterations.

Example:

python main.py -m RGCN -t node_classification -d ACM -a 0.01 -b 0.01 -M 3

If you find this work helpful, please consider citing our paper:

@article{yang2025unified,
  title = {A Unified Gradient Regularization Method for Heterogeneous Graph Neural Networks},
  author = {Yang, Xiao and Zhao, Xuejiao and Shen, Zhiqi},
  journal = {Neural Networks},
  pages = {108104},
  year = {2025},
  publisher = {Elsevier}
}