[AAAI 2026] Adaptive Riemannian Graph Neural Networks (ARGNN)

This repository provides the implementation of the AAAI 2026 paper "Adaptive Riemannian Graph Neural Networks".

If you find this repository useful or use it in your research, please cite our paper:

@article{wang2025adaptive,
  title={Adaptive riemannian graph neural networks},
  author={Wang, Xudong and Li, Tongxin and Ding, Chris and Fan, Jicong},
  journal={arXiv preprint arXiv:2508.02600},
  year={2025}
}

Overview

Adaptive Riemannian Graph Neural Networks (ARGNN) is a framework that learns a continuous and anisotropic Riemannian metric tensor field over a graph. Each node adapts its local geometry, enabling the model to more effectively represent the graph’s structural landscape.

Project Structure

.
├── figs/                  
├── data/                  # Dataset directory (auto-downloaded if not exist)
├── GraphEnv.yml          
├── README.md              
├── main.py               # Main training and evaluation script
├── models/                # Model definitions
│   └── argnn_model.py     
│   └── ...                # Other baseline methods implementation
├── utils/                 # Utility functions
│   ├── seed_utils.py      
│   ├── logging_utils.py   
│   ├── data_utils.py      # Dataset loading and split management
│   ├── train_utils.py     
│   ├── eval_utils.py      
│   └── loss_utils.py      
├── splits/                # Please put official GeomGCN's splits in this folder. Here we provide Cora/Citeseer and Pubmed splits as examples. 
│   └── ...                # Please download from https://github.com/bingzhewei/geom-gcn   
└── logs/                  # Training logs and saved models
    └── <dataset>_<timestamp>/  # Timestamped experiment directory
        ├── run_configs.json    # Experiment configuration
        ├── <dataset>_records.pkl  # Full training records
        ├── <dataset>_best_<run>.pth  # Best model per run
        └── <dataset>_global_best_model_<run>.pth  # Overall best model

Usage and Quick Start

Node Classification

The example configurations for node classification on Cora, Citeseer and PubMed datasets for a quick start, you can run the command like the following example:

python main.py --dataset Cora --epochs 200 --model argnn --optimizer adam --lr 0.005 \
                --num_runs 10 --argnn_hidden_dim 128 --argnn_num_layers 3 \
                --argnn_metric_hidden_dim 256 --argnn_metric_reg 0.01 \
                --argnn_smoothness_reg 0.001 --dropout 0.3 \
                --task node_classification --use_splits ./split

Link Prediction

Example configuration for link prediction task:

python main.py --dataset Cora --epochs 500 --model argnn --optimizer adam --lr 0.005 \
                --num_runs 1 --argnn_hidden_dim 128 --argnn_num_layers 2 \
                --argnn_metric_hidden_dim 128 --argnn_metric_reg 0.01 \
                --argnn_smoothness_reg 0.001 --dropout 0.2 \
                --task link_prediction

Parameter Descriptions

• --dataset: Name of the dataset (Cora, Citeseer, etc.)
• --epochs: Number of training epochs
• --model: Model architecture (argnn)
• --optimizer: Optimization algorithm (adam, sgd, etc.)
• --lr: Learning rate
• --num_runs: Number of training runs for statistical significance
• --argnn_hidden_dim: Hidden dimension of ARGNN layers
• --argnn_num_layers: Number of ARGNN layers
• --argnn_metric_hidden_dim: Hidden dimension for metric learning
• --argnn_metric_reg: Regularization coefficient for metric learning
• --argnn_smoothness_reg: Regularization coefficient for graph smoothness
• --dropout: Dropout rate
• --task: Task type (node_classification or link_prediction)
• --use_splits: Directory containing predefined data splits

Requirements

Hardware Requirements

• NVIDIA GPU with CUDA support with 8GB+ GPU memory
• Recommended: 16GB+ system memory

Software Requirements

• Python 3.8+
• CUDA 12.1+

Environment Setup

1. Check NVIDIA GPU and CUDA

First, verify that your system has an NVIDIA GPU available:

nvidia-smi
nvcc --version

If NVIDIA drivers are not installed, follow the official NVIDIA guide to install them.

2. Create Conda Environment

Create a conda environment using the provided configuration file:

conda env create -f GraphEnv.yml

This environment file includes all necessary libraries:

• PyTorch 2.1.0 (with CUDA support)
• PyTorch Geometric
• geoopt (for Riemannian optimization)
• networkx
• scikit-learn
• matplotlib
• seaborn

3. Activate Environment

conda activate Graph

Contact

• Email: xudongwang@link.cuhk.edu.cn