Embarrassingly Parallel GFlowNets
<|reference_start|>Embarrassingly Parallel GFlowNets: GFlowNets are a promising alternative to MCMC sampling for discrete compositional random variables. Training GFlowNets requires repeated evaluations of the unnormalized target distribution or reward function. However, for large-scale posterior sampling, this may be prohibitive since it incurs traversing the data several times. Moreover, if the data are distributed across clients, employing standard GFlowNets leads to intensive client-server communication. To alleviate both these issues, we propose embarrassingly parallel GFlowNet (EP-GFlowNet). EP-GFlowNet is a provably correct divide-and-conquer method to sample from product distributions of the form $R(\cdot) \propto R_1(\cdot) ... R_N(\cdot)$ -- e.g., in parallel or federated Bayes, where each $R_n$ is a local posterior defined on a data partition. First, in parallel, we train a local GFlowNet targeting each $R_n$ and send the resulting models to the server. Then, the server learns a global GFlowNet by enforcing our newly proposed \emph{aggregating balance} condition, requiring a single communication step. Importantly, EP-GFlowNets can also be applied to multi-objective optimization and model reuse. Our experiments illustrate the EP-GFlowNets's effectiveness on many tasks, including parallel Bayesian phylogenetics, multi-objective multiset, sequence generation, and federated Bayesian structure learning.<|reference_end|>
@article{da silva2024embarrassingly,
    title={Embarrassingly Parallel GFlowNets},
    author={Tiago da Silva, Luiz Max Carvalho, Amauri Souza, Samuel Kaski, Diego
  Mesquita},
    journal={arXiv preprint arXiv:2406.03288},
    year={2024},
    archivePrefix={arXiv},
    eprint={2406.03288},
    primaryClass={cs.LG stat.ML}
}