Joint Learning of Label and Environment Causal Independence for Graph Out-of-Distribution Generalization
This is the official code for the implementation of "Joint Learning of Label and Environment Causal Independence for Graph Out-of-Distribution Generalization" which is accepted by NeurIPS 2023. 😄
In this work, we propose to simultaneously incorporate label and environment causal independence (LECI) to release the potential of pre-collected environment information in graph tasks, thereby addressing the challenges faced by prior methods on identifying causal/invariant subgraphs. We further develop an adversarial training strategy to jointly optimize these two properties for causal subgraph discovery with theoretical guarantees.
LECI depends on PyTorch (>=1.6.0), PyG (>=2.0), and RDKit (>=2020.09.5). For more details: conda environment
Note that we currently test on PyTorch (==1.10.1), PyG (==2.0.3), RDKit (==2020.09.5); thus we strongly encourage to install these versions.
git clone https://github.com/divelab/LECI.git && cd LECI
pip install -e .
goodtg --config_path final_configs/GOODHIV/scaffold/covaraite/LECI.yaml --exp_round [1/2/3/4/5/6/7/8/9/10] --gpu_idx [0..9]
goodtg --config_path final_configs/GOODHIV/size/covaraite/LECI.yaml --exp_round [1/2/3/4/5/6/7/8/9/10] --gpu_idx [0..9]
goodtg --config_path final_configs/LBAPcore/assay/covaraite/LECI.yaml --exp_round [1/2/3/4/5/6/7/8/9/10] --gpu_idx [0..9]
goodtg --config_path final_configs/GOODMotif/basis/covaraite/LECI.yaml --exp_round [1/2/3/4/5/6/7/8/9/10] --gpu_idx [0..9]
goodtg --config_path final_configs/GOODMotif/size/covaraite/LECI.yaml --exp_round [1/2/3/4/5/6/7/8/9/10] --gpu_idx [0..9]
goodtg --config_path final_configs/GOODCMNIST/color/covaraite/LECI.yaml --exp_round [1/2/3/4/5/6/7/8/9/10] --gpu_idx [0..9]
goodtg --config_path final_configs/GOODSST2/length/covaraite/LECI.yaml --exp_round [1/2/3/4/5/6/7/8/9/10] --gpu_idx [0..9]
goodtg --config_path final_configs/GOODTwitter/length/covaraite/LECI.yaml --exp_round [1/2/3/4/5/6/7/8/9/10] --gpu_idx [0..9]
Explanations of the arguments can be found in this file.
Valid LECI: The training of LECI is valid only when the optimal discriminator Proposition 3.2 is approximately learned, e.g.,
the environment branch loss at least should not indicate a random prediction when the adversarial training is not applied (or is weak). Note that the adversarial intensity
increases from 0 to self.config.train.alpha
in the code.
How to select the right learning rate? Since the environment labels
How to select the valid hyperparameters? If the EA/LA loss never decreases (invalid LECI), please try decreasing
For more details, please refer to the appendix of our paper.
If you find this repository helpful, please cite our paper/preprint.
@inproceedings{gui2023joint,
title={Joint Learning of Label and Environment Causal Independence for Graph Out-of-Distribution Generalization},
author={Gui, Shurui and Liu, Meng and Li, Xiner and Luo, Youzhi and Ji, Shuiwang},
booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
year={2023},
url={https://openreview.net/forum?id=z3HACY5CMa}
}
- The GOOD datasets are under MIT license.
- The DrugOOD dataset is under GPLv3
- The LECI code are under GPLv3 license, since the code architecture is based on GOOD.
Please submit new issues or start a new discussion for any technical or other questions.
Please feel free to contact Shurui Gui or Shuiwang Ji!