gfnx是JAX实现的生成流网络库，训练速度比PyTorch快80倍。
gfnx: Fast and Scalable Library for Generative Flow Networks in JAX
- 基于JAX实现高效可扩展的生成流网络训练与评估
- 在序列生成和贝叶斯网络学习任务中分别提速55倍和80倍
- 提供多种基准环境，适合研究生成模型与强化学习者使用
本文介绍gfnx，一个用JAX编写的生成流网络（GFlowNets）快速且可扩展的训练与评估库。gfnx包含丰富的环境与评估指标，提供核心目标函数的单文件实现。涵盖合成超网格、多种序列生成环境（含不同编辑策略）、分子生成特定奖励设计、系统发育树构建、贝叶斯结构学习及伊辛模型采样。在不同任务中，gfnx相比基于PyTorch的基准（如torchgfn库和作者原始实现）显著提升实际运行速度：例如在基于CPU的序列生成环境中最高提速55倍，在基于GPU的贝叶斯网络结构学习设置中最高提速80倍。该库旨在统一评估标准，加速GFlowNets的研究与应用。代码已开源于GitHub（https://github.com/d-tiapkin/gfnx），可通过PyPI安装（https://pypi.org/project/gfnx/），文档见https://gfnx.readthedocs.io。
原文摘要 · Abstract (English)
In this paper, we present gfnx, a fast and scalable package for training and evaluating Generative Flow Networks (GFlowNets) written in JAX. gfnx provides an extensive set of environments and metrics for benchmarking, accompanied with single-file implementations of core objectives for training GFlowNets. We include synthetic hypergrids, multiple sequence generation environments with various editing regimes and particular reward designs for molecular generation, phylogenetic tree construction, Bayesian structure learning, and sampling from the Ising model energy. Across different tasks, gfnx achieves significant wall-clock speedups compared to Pytorch-based benchmarks (such as torchgfn library) and author implementations. For example, gfnx achieves up to 55 times speedup on CPU-based sequence generation environments, and up to 80 times speedup with the GPU-based Bayesian network structure learning setup. Our package provides a diverse set of benchmarks and aims to standardize empirical evaluation and accelerate research and applications of GFlowNets. The library is available on GitHub (https://github.com/d-tiapkin/gfnx) and on pypi (https://pypi.org/project/gfnx/). Documentation is available on https://gfnx.readthedocs.io.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。