This is the implementation accompanying the NIPS 2018 paper, Watch Your Step: Learning Node Embeddings via Graph Attention.

First, create a fresh virtual environment and install the requirements.

# From google-research/
virtualenv -p python3 .
source ./bin/activate

pip install -r graph_embedding/watch_your_step/requirements.txt

Next, download the datasets used in Abu-El-Haija et al, CIKM'17:

# From google-research/
curl http://sami.haija.org/graph/datasets.tgz > datasets.tgz
tar zxvf datasets.tgz
export DATA_DIR=datasets

After downloading the datasets, you can run the code like the following.

# From google-research/
python -m graph_embedding.watch_your_step.graph_attention_learning --dataset_dir ${DATA_DIR}/wiki-vote

The above run should reproduce the results in a few minutes.

To save the output, please use --output_dir. Consider other flags for options. Output file will contain train/test metrics, embeddings, as well as learned context distributions.

If you find Watch Your Step useful in your research, we ask that you cite the following paper:

Abu-El-Haija, S., Perozzi, B., Al-Rfou, R., and Alemi, A. (2018). Watch Your Step: Learning Node Embeddings via Graph Attention. In Neural Information Processing Systems.

@inproceedings{abu2018watch,
 author={Abu-El-Haija, Sami and Perozzi, Bryan and Al-Rfou, Rami and Alemi, Alex}
 title={Watch Your Step: Learning Node Embeddings via Graph Attention},
 booktitle = {Neural Information Processing Systems},
 year = {2018},
}

For questions or comments about the implementation, please contact sami@haija.org and bperozzi@acm.org.