Home

Awesome

:walking: NeuralWalker

arXiv PWC PWC PWC

The repository implements the NeuralWalker in Pytorch Geometric described in the following paper

Dexiong Chen, Till Schulz, and Karsten Borgwardt. Learning Long Range Dependencies on Graphs via Random Walks, Preprint 2024.

TL;DR: A novel random-walk based neural architecture for graph representation learning.

:book: Overview

NeuralWalker

NeuralWalker samples random walks with a predefined sampling rate and length, then uses advanced sequence models to process them. Additionally, local and global message passing can be employed to capture complementary information. The main components of NeuralWalker are a random walk sampler, and a stack of neural walker blocks (a walk encoder block + a message passing block). Each walk encoder block has a walk embedder, a sequence layer, and a walk aggregator.

:hammer_and_wrench: Installation

We recommend the users to manage dependencies using miniconda or micromamba:

# Replace micromamba with conda if you use conda or miniconda
micromamba env create -f environment.yaml 
micromamba activate neuralwalker
pip install -e .

[!NOTE] Our code is also compatible with more recent Pytorch versions, your can use micromamba env create -f environment_latest.yaml for development purposes.

[!TIP] NeuralWalker relies on a sequence model to process random walk sequences, such as CNNs, Transformers, or state-space models. If you encounter any issues when installing the state-space model Mamba, please consult its installation guideline.

:bulb: Reproducing results

Running NeuralWalker on Benchmarking GNNs, LRGB, and OGB

All configurations for the experiments are managed by hydra, stored in ./config.

Below you can find the list of experiments conducted in the paper:

# You can replace zinc with any of the above datasets
python train.py experiment=zinc

# Running NeuralWalker with a different model architecture
python train.py experiment=zinc experiment/model=conv+vn_3L

[!TIP] You can replace conv+vn_3L with any model provided in config/experiment/model, or a customized model by creating a new one in that folder.

Running NeuralWalker on node classification tasks

We integrate NeuralWalker with Polynormer, SOTA model for node classifcation. See node_classifcation for more details.

Debug mode

python train.py mode=debug