<h1 align='center'>Enhanced Expert Merging for Mixture-of-Experts in Graph Foundation Models</h1>

## Environment Setup
Download the data files at [this link](https://huggingface.co/datasets/hkuds/AnyGraph_datasets), which are provided by the AnyGraph authors. And fill in your own directories for data storage at function `get_data_files(self)` of class `DataHandler` in the file `data_handler.py`.


**Packages**: Our experiments were conducted with the following package versions:
* python==3.10.15
* torch==2.3.1
* numpy==1.26.4
* scipy==1.14.1
* cuda-version==12.6

**Device Requirements**: The training and testing of KDEM/PPEM requires only one GPU with 24G memory (e.g. RTX 3090, 4090). Using larger input graphs may require devices with larger memory.

## Datasets

<img src='imgs/data_status.png' />

The statistics for the experimental datasets are presented in the table above. We categorize them into distinct groups as below. Note that Link1 and Link2 include datasets from different sources, and the datasets do not share the same feature spaces. This separation ensures a robust evaluation of true zero-shot performance in graph prediction tasks.

| Group | Included Datasets |
| ----- | ----- |
| Link1 | Products-tech, Yelp2018, Yelp-textfeat, Products-home, Steam-text, Amazon-text, Amazon-book, Citation-2019, Citation-20Century, Pubmed-link, Citeseer, OGB-PPA, P2P-Gnutella06, Soc-Epinions1, Email-Enron |
| Link2 | Photo, Goodreads, Fitness, Movielens-1M, Movielens10M, Gowalla, Arxiv, Arxiv-t, Cora, CS, OGB-Collab, Proteins-0, Proteins- 1, Proteins-2, Proteins-3, OGB-DDI, Web-Stanford, RoadNet-PA |
| Ecommerce | Products-tech, Yelp2018, Yelp-textfeat, Products-home, Steam-text, Amazon-text, Amazon-book, Photo, Goodreads, Fitness, Movielens-1M, Movielens10M, Gowalla |
| Academic | Citation-2019, Citation-20Century, Pubmed-link, Citeseer, OGB-PPA, Arxiv, Arxiv-t, Cora, CS, OGB-Collab |
| Others | P2P-Gnutella06, Soc-Epinions1, Email-Enron, Proteins-0, Proteins- 1, Proteins-2, Proteins-3, OGB-DDI, Web-Stanford, RoadNet-PA |
| Node | Cora, Arxiv, Pubmed, Home, Tech |

## Usage
### Train
To train KDEM, run:
```
python main.py --dataset link1+link2 --topk_expert 3 --lr 1e-4 --ensemble_steps 100 --ensemble_loss_coef 1e-2 --save pretrain_link1_top3_lr1e-4_esb100_coef1e-2
```
To train PPEM, run:
```
python main.py --dataset link1+link2 --topk_expert 5 --lr 1e-4 --ema_every_steps 100 --alpha 0.999 --save pretrain_link1_top5_lr1e-4_ema100_alpha_0.999
```
### Test
To test the pre-trained models, run the following command lines:
```
# Test on the entire link2 dataset group.
# KDEM
python main.py --load pretrain_link1_top3_lr1e-4_esb100_coef1e-2 --topk_expert 3 --epoch 0 --dataset link2
# PPEM
python main.py --load pretrain_link1_top5_lr1e-4_ema100_alpha_0.999 --topk_expert 5 --epoch 0 --dataset link2

# Test on the Ecommerce datasets in the Link2 group.
# Testing on Academic and Others datasets are conducted similarily.
python main.py --load pretrain_link1_top3_lr1e-4_esb100_coef1e-2 --topk_expert 3 --epoch 0 --dataset ecommerce_in_link2
python main.py --load pretrain_link1_top5_lr1e-4_ema100_alpha_0.999 --topk_expert 5 --epoch 0 --dataset ecommerce_in_link2
```

