/FedDAG

[TMLR23] FedDAG: Federated DAG Structure Learning

Primary LanguagePython

FedDAG: Federated DAG Structure Learning

This repository contains an implementation of the federated DAG structure learning methods described in "FedDAG: Federated DAG Structure Learning".

If you find it useful, please consider citing:

@article{
gao2023feddag,
title={Fed{DAG}: Federated {DAG} Structure Learning},
author={Erdun Gao and Junjia Chen and Li Shen and Tongliang Liu and Mingming Gong and Howard Bondell},
journal={Transactions of Machine Learning Research},
year={2023},
url={https://openreview.net/forum?id=MzWgBjZ6Le},
note={}
}

Requirements

  • CUDA 10.0
  • cuDNN 7.4
  • Anaconda

Use conda env create -f environment.yml to create a feddag_test_env conda environment.

Use 10.0-cudnn7-devel-ubuntu18.04 to create a container if running from docker.

Examples

After creating a new conda environment, you can run bash test_run.sh to test the codes. Our logging results of the test_run.sh is recorded in log_example.txt. If you want to test our method with other scales of graphs, change the hyper-parameters according to settings provided in the Appendix.

Acknowledgments

  • Our implementation is highly based on the existing Tool-chain named gcastle pip link and paper link, which includes many gradient-based DAG structure learning methods.
  • Our implementation is also highly based on NOTEARS-tensorflow and MCSL.

Recommendations

  • Notice that NOTEARS-ADMM is a concurrent and interesting work that also considers the same problem with our FedDAG. In NOTEARS-ADMM, ADMM is leveraged to jointly learn the graph.
  • The baseline method of our FedDAG is MCSL. Please read this paper if you have concerns about the basic modules of FedDAG.

You are highly recommended to read these two papers and to cite them.

@inproceedings{Ng2022federated,
  author = {Ng, Ignavier and Zhang, Kun},
  title = {Towards Federated Bayesian Network Structure Learning with Continuous Optimization},
  booktitle = {International Conference on Artificial Intelligence and Statistics},
  year = {2022},
}

@inproceedings{ng2022masked,
  title={Masked gradient-based causal structure learning},
  author={Ng, Ignavier and Zhu, Shengyu and Fang, Zhuangyan and Li, Haoyang and Chen, Zhitang and Wang, Jun},
  booktitle={Proceedings of the 2022 SIAM International Conference on Data Mining (SDM)},
  pages={424--432},
  year={2022},
  organization={SIAM}
}