gfnx: Fast and Scalable Library for Generative Flow Networks in JAX

2025-11-20

AI Summary

Existing GFlowNets implementations (e.g., torchgfn) suffer from suboptimal performance, hindering efficient training and evaluation across diverse tasks—including molecular generation, phylogenetic tree inference, Bayesian structure learning, and Ising model sampling. To address this, we introduce the first high-performance, open-source GFlowNets library built on JAX. Leveraging just-in-time compilation, functional parallelization, and hardware-agnostic optimizations, our library significantly improves scalability and computational efficiency. It provides a unified benchmarking framework encompassing multiple standard environments, reward formulations, and evaluation metrics. Empirical results demonstrate speedups of up to 55× on CPU and 80× on GPU over PyTorch-based baselines, substantially reducing training and inference latency. This infrastructure establishes a robust foundation for deploying GFlowNets in complex discrete-structure modeling tasks.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies

Abstract

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.

Problem

• Providing fast scalable library for training Generative Flow Networks
• Establishing standardized benchmarks for GFlowNets across diverse domains
• Accelerating GFlowNet research with significant performance improvements

Innovation

• Fast scalable GFlowNets library built with JAX
• Provides diverse environments and core training objectives
• Achieves significant speedups over PyTorch implementations