Higher-order Transformer (PyTorch)

Transformers Generalize DeepSets and Can be Extended to Graphs and Hypergraphs
Jinwoo Kim, Saeyoon Oh, Seunghoon Hong @ KAIST
NeurIPS 2021

We present a generalization of Transformers to sets, graphs, and hypergraphs, and reduce its computational cost to linear to input size.

Powerful operations, involving both local and global interactions over input elements
Translation between different-order graphs (e.g., set-to-graph, graph-to-set, graph-to-vector)
Theoretically and empirically stronger than MPNNs, even with the same linear complexity
Works well on large datasets!

In this repository, we provide the PyTorch implementation of:

Up-to-second-order Transformers: {set, graph} input × {set, graph, vector} output
A space of their variants: {dense, sparse} data × {softmax, kernel} attention
A special extension for higher-order hyperedge prediction: set input, k-hyperedges output
Corresponding equivariant linear layers as baselines

In practice, we recommend to use the sparse kernel variant as it scales linearly to data size while still powerful.

How do I use Higher-order Transformers for my project?

Go to your repository and execute the following.

# this is enough
pip install torch>=1.8.1
# add this repository as a submodule
git submodule add https://github.com/jw9730/hot hot
git submodule init
# run basic tests
cd hot
python3 run_tests.py
cd ..

For basic use cases, please see run_perf_tests.py and run_tests.py.

To update the code, execute the following:

git submodule update --remote --merge

Setting up experiments

Using Docker

git clone https://github.com/jw9730/hot.git hot
cd hot
docker build --no-cache --tag hot:latest .
docker run -it --gpus all --ipc=host --name=hot -v /home:/home hot:latest bash
# upon completion, you should be at /hot inside the container

Using pip

git clone https://github.com/jw9730/hot.git hot
cd hot
pip3 install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html
pip install -r requirements.txt
pip install torch-sparse==0.6.9 torch-scatter==2.0.6 -f https://data.pyg.org/whl/torch-1.8.0+cu111.html
pip install torch-geometric==1.6.3

Running experiments

Runtime and memory analysis

python3 run_perf_tests.py

PCQM4M-LSC graph regression

cd regression/examples/pcqm4m-lsc
# Second-order Transformers (sparse kernel)
bash enc.sh
bash enc-small.sh
# Comparison with more baselines
bash enc-short.sh
bash enc-small-short.sh
# Second-order MLP
bash mlp-short.sh
# Vanilla Transformer + Laplacian PE
bash laplacian-short.sh

Set-to-graph prediction

cd set-to-graph/main
python download_jets_data.py
cd ../examples
# Delaunay (50)
cd delaunay-a
bash enc-kernel.sh
bash enc-softmax.sh
bash s2g.sh
# visualize
bash visualize-kernel.sh
bash visualize-softmax.sh
# Delaunay (20-80)
cd ../delaunay-b
bash enc-kernel.sh
bash enc-softmax.sh
bash s2g.sh
# visualize
bash visualize-kernel.sh
bash visualize-softmax.sh
# Jets
cd ../jets
bash enc-kernel.sh
bash enc-softmax.sh
# test
bash test-enc-kernel.sh
bash test-enc-softmax.sh

k-uniform hyperedge prediction

cd k-uniform-hyperedge/examples
# GPS
cd gps
bash enc.sh
bash s2g+.sh
bash hyper-sagnn-e.sh
bash hyper-sagnn-w.sh
# MovieLens
cd ../movielens
bash enc.sh
bash s2g+.sh
bash hyper-sagnn-e.sh
bash hyper-sagnn-w.sh
# Drug
cd ../drug
bash enc.sh
bash s2g+.sh
bash hyper-sagnn-e.sh
bash hyper-sagnn-w.sh

References

This implementation uses code from the following repositories:

Performer for FAVOR+ attention kernel
Graph Transformer for Laplacian PE of baseline model
Graphormer for graph regression experiment
Set2Graph for set-to-graph prediction experiment
Hyper-SAGNN for k-uniform hyperedge prediction experiment

Citation

If you find our work useful, please consider citing it:

@article{kim2021transformers,
  author    = {Jinwoo Kim and Saeyoon Oh and Seunghoon Hong},
  title     = {Transformers Generalize DeepSets and Can be Extended to Graphs and Hypergraphs},
  booktitle = {2021 Conference on Neural Information Processing Systems (NeurIPS)},
  year      = {2021}
}