amazon-science/cusp

amazon-science/cusp

Description: Official implementation of Spectro-Riemannian Graph Neural Networks (ICLR 2025)

Language: Python

Stars: 16

Forks: 1

Open issues: 4

Created: 2024-10-18T23:20:41Z

Pushed: 2026-06-10T23:03:25Z

Default branch: main

Fork: no

Archived: no

README:

Spectro-Riemannian Graph Neural Networks

Official implementation of Spectro-Riemannian Graph Neural Networks (ICLR 2025).

🛠 Dependencies and Installation

• geoopt 0.5.0 (For Riemannian optimization and implementation of $\kappa$-stereographic model)
• GraphRicciCurvature 0.5.3.2 (For computing Ollivier-Ricci Curvature)
• torch 2.4.0 (Implementation in PyTorch)
• torch_geometric 2.6.1
• Other required packages in requirements.txt

# git clone this repository
git clone https://github.com/amazon-science/cusp.git

# install python dependencies
pip3 install -r requirements.txt

💽 Datasets

We have performed extensive experimentation and ablation studies across 8 datasets $-$ Homophilic (Cora, Citeseer, PubMed) and Heterophilic (Chameleon, Actor, Squirrel, Texas, Cornell). The datasets are downloaded automatically from torch_geometric and saved in data/ directory. Consider the following script from train.py.

datasets = {
"Cora": Planetoid(root="data/Cora", name="Cora", transform=T.ToUndirected()),
"Citeseer": Planetoid(root="data/Citeseer", name="Citeseer", transform=T.ToUndirected()),
"PubMed": Planetoid(root="data/PubMed", name="PubMed", transform=T.ToUndirected()),
"Chameleon": WikipediaNetwork(root="data/WikipediaNetwork", name="chameleon", transform=T.ToUndirected()),
"Actor": Actor(root="data/Actor", transform=T.ToUndirected()),
"Squirrel": WikipediaNetwork(root="data/WikipediaNetwork", name="squirrel", transform=T.ToUndirected()),
"Texas": WebKB(root="data/WebKB", name="Texas", transform=T.ToUndirected()),
"Cornell": WebKB(root="data/WebKB", name="Cornell", transform=T.ToUndirected())}

Specify the dataset to be used using the --dataset argument while running train.py.

🔂 Training

Key Training Arguments

The following table describes some important optional arguments that control different ablations of the CUSP model. The entire list of training arguments are present in train.py.

| Argument | Description | |---------------------------|-------------| | use_curvature_encoding | Enables functional curvature encoding for improved representation learning. | | use_cusp_pooling | Uses curvature-based positional encoding in Cusp Pooling, allowing hierarchical attention over embeddings. | | euclidean_variant | Forces the Euclidean variant of the model, where all manifolds are Euclidean. | | use_cusp_laplacian | Uses the Cusp Laplacian (default). If not set, the model falls back to using the standard graph Laplacian. | |K| Number of filters in the filterbank.| |manifold_config| Product manifold signature to use, for learning representations.| |d_f|Dimension of the functional curvature encoding.|

Example Usage

To run the full CUSP model for the task of node classification using Riemannian Adam optimizer (radam), with K=10 filters, on the Cora dataset, for the product manifold signature $\mathbb{H}^{16} \times \mathbb{H}^{16}\times \mathbb{S}^{16} \times\mathbb{E}^{16}$, use the below script:

python train.py --dataset Cora \
--epochs 100 \
--model cusp \
--optimizer radam \
--lr 0.01 \
--num_runs 1 \
--use_curvature_encoding \
--use_cusp_laplacian \
--use_cusp_pooling \
--K 10 \
--d_f 64 \
--ricci_alpha 0.5 \
--manifold_config H16H16S16E16 \
--task node_classification

To train the CUSP model with different ablations use the following commands:

• Euclidean variant without Cusp pooling and curvature encoding for node classification.

python train.py --dataset Cora --model cusp --epochs 30 --lr 4e-3 --num_runs 2 --euclidean_variant --use_cusp_laplacian --manifold_config H16H16S16E16 --K 10

• Full CUSP without Cusp Laplacian for link prediction.

python train.py --dataset Cora --epochs 30 --model cusp --optimizer radam --lr 4e-3 --num_runs 2 --use_curvature_encoding --use_cusp_pooling --K 10 --task node_classification --manifold_config H16H16S16E16

📞 Contact

If you have any questions or issues, please feel free to reach out to Karish Grover at karishg@cs.cmu.edu.

Security

See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more information.

License

This library is licensed under the CC-BY-NC-4.0 License.

✏️ Citation

If you think that this work is helpful, please leave a star ⭐️ and cite our paper:

@misc{grover2025spectroriemanniangraphneuralnetworks,
title={Spectro-Riemannian Graph Neural Networks},
author={Karish Grover and Haiyang Yu and Xiang Song and Qi Zhu and Han Xie and Vassilis N. Ioannidis and Christos Faloutsos},
year={2025},
eprint={2502.00401},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2502.00401},
}