This is the official code for the paper: On the Expressivity of Stable Positional Encodings for Graphs.

Feel free to contact [email protected] if there is any question.

In this work, we present SPE, a Laplacian-based graph positional encodings that are provably stable and expressive. The key insight is to perform a soft and learnable ``partition" of eigensubspaces in an eigenvalue dependent way, hereby achieving both stability (from the soft partition) and expressivity (from dependency on both eigenvalues and eigenvectors).

Our SPE method processes eigenvectors $V\in\mathbb{R}^{n\times d}$ and eigenvalues $\lambda\in\mathbb{R}^d$ into node positional encodings as follows: $$\text{SPE}(V, \lambda)=\rho(V\text{diag}{\phi_1(\lambda)}V^{T}, V\text{diag}{\phi_2(\lambda)}V^{T}, ..., V\text{diag}{\phi_m(\lambda)}V^{T}),$$ where $\rho:\mathbb{R}^{n\times n\times m}\to\mathbb{R}^{n\times p}$ and $\phi_i:\mathbb{R}^{d}\to\mathbb{R}^d$ are permutational equivariant functions w.r.t. $n\times n$ and $d$ axes respectively.

See requirements.txt for necessary python environment.

Download all required datasets from here. The downloaded 'data' directory should be placed in the root direcotry. For example, './data/drugood', etc.

To reproduce experiments on ZINC, cd to ./zinc and run

python runner.py --config_dirpath ../configs/zinc --config_name SPE_gine_gin_mlp_pe37.yaml --seed 0

To reproduce experiments on Alchemy, cd to ./alchemy and run

python --config_dirpath ../configs/alchemy --config_name SPE_gine_gin_mlp_pe12.yaml --seed 0

To reproduce experiments on DrugOOD, cd to ./drugood and run

python --config_dirpath ../configs/assay --config_name SPE_gine_gin_mlp_pe32_zeropsi.yaml --dataset assay --seed 0
python --config_dirpath ../configs/scaffold --config_name SPE_gine_gin_mlp_pe32_standard_dropout.yaml --dataset scaffold --seed 0
python --config_dirpath ../configs/scaffold --config_name SPE_gine_gin_mlp_pe32_standard_dropout.yaml --dataset size --seed 0

To reproduce substructures counting, cd to ./count and run

bash run.sh