Research map / Probabilistic and scientific modeling / Molecules and proteins
gfnx: Fast and Scalable Library for Generative Flow Networks in JAX
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Presents gfnx, a JAX library of JIT-compiled GFlowNet environments, rewards, metrics and single-file baselines, reporting large wall-clock speedups over PyTorch implementations such as torchgfn.
Problem
GFlowNet environments typically run on the host CPU in PyTorch-based libraries, creating CPU-accelerator transfer bottlenecks and slow experimentation for amortized sampling benchmarks.
Approach
Implement GFlowNet environments, reward modules, and success metrics entirely in JAX so that the full training loop can be JIT-compiled end-to-end, decoupling rewards from dynamics and abstracting backward actions to structural choices.
Key results
- gfnx achieves up to 80x speedup on GPU Bayesian structure learning (58.0 vs 0.73 it/s for MDB).
- On CPU QM9 with TB, gfnx reaches 9061.6 it/s versus 162.3 it/s for the baseline (about 55x).
- On hypergrid 4D CPU, gfnx is at least about five times faster than torchgfn while matching total variation convergence.
Limitations
- Only discrete action spaces are supported.
- No non-acyclic environments or multi-objective support yet.
- Trainer vectorization across seeds and hyperparameters is not implemented.
Experimental setup
Datasets: TFBind8 QM9 DBAASP AMP proxy dataset
Benchmarks: hypergrid bit sequence TFBind8 QM9 AMP phylogenetic trees Bayesian structure learning Ising model
Metrics: iterations per second total variation distance Pearson correlation of log-probability and log-reward top-100 reward and diversity
Baselines: torchgfn author implementations recursionpharma/gflownet alexhernandezgarcia/gflownet
Builds on: torchgfn CleanRL purejaxrl Equinox trajectory balance forward-looking detailed balance
Closest papers by meaning
- Pessimistic Backward Policy for GFlowNets (NeurIPS 2024 · similarity 0.89)
- Beyond Squared Error: Exploring Loss Design for Enhanced Training of Generative Flow Networks (ICLR 2025 · similarity 0.88)
- Flow Factorization for Efficient Generative Flow Networks (AAAI 2025 · similarity 0.88)
- Learning Energy Decompositions for Partial Inference in GFlowNets (ICLR 2024 · similarity 0.88)
- Symmetry-Aware GFlowNets (ICML 2025 · similarity 0.87)
- Generative Flows on Synthetic Pathway for Drug Design (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.