JAX中的GFlowNet生成流网络快速可扩展库

Daniil Tiapkin; Artem Agarkov; Nikita Morozov; Ian Maksimov; Askar Tsyganov; Timofei Gritsaev; Sergey Samsonov

JAX中的GFlowNet生成流网络快速可扩展库

机器学习 2025-11-21 v1

作者: Daniil Tiapkin , Artem Agarkov , Nikita Morozov , Ian Maksimov , Askar Tsyganov , Timofei Gritsaev , Sergey Samsonov

在 arXiv 上查看 ↗ PDF ↗

摘要

本文我们介绍了gfnx，一个用于训练和评估生成流网络 (GFlowNets) 的快速可扩展软件包，基于JAX实现。gfnx提供了广泛的环境和指标，用于基准测试，附带核心训练目标的单文件实现。我们包括合成高维格、多序列生成环境及各种编辑方案和特定奖励设计，用于分子生成、 phylogenetic 树构建、贝叶斯结构学习以及从Ising模型能量中进行采样。在不同任务中，gfnx相较于基于PyTorch的基准（如torchgfn库）和作者实现，实现了显著的 wall-clock速度提升。例如，在基于CPU的序列生成环境中，gfnx实现了最高可达55倍的速度提升；在基于GPU的贝叶斯网络结构学习 setup中，速度提升可达80倍。该软件包提供了多样化的基准，旨在标准化经验评估并加速GFlowNets的研究和应用。该库可在GitHub (https://github.com/d-tiapkin/gfnx) 和pypi (https://pypi.org/project/gfnx/) 上获取。在 https://gfnx.readthedocs.io 处提供文档。

关键词

生成流网络 JAX 机器学习分子生成贝叶斯结构学习

引用

@article{arxiv.2511.16592,
  title  = {gfnx: Fast and Scalable Library for Generative Flow Networks in JAX},
  author = {Daniil Tiapkin and Artem Agarkov and Nikita Morozov and Ian Maksimov and Askar Tsyganov and Timofei Gritsaev and Sergey Samsonov},
  journal= {arXiv preprint arXiv:2511.16592},
  year   = {2025}
}

备注

GitHub: https://github.com/d-tiapkin/gfnx | Documentation: https://gfnx.readthedocs.io

JAX中的GFlowNet生成流网络快速可扩展库

摘要

关键词

引用

备注

相关论文