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

• URL: http://arxiv.org/abs/2511.16592v1
• Date: Thu, 20 Nov 2025 17:44:45 GMT
• Title: gfnx: Fast and Scalable Library for Generative Flow Networks in JAX
• Authors: Daniil Tiapkin, Artem Agarkov, Nikita Morozov, Ian Maksimov, Askar Tsyganov, Timofei Gritsaev, Sergey Samsonov,
• Abstract summary: We present gfnx, a package for training and evaluating Generative Flow Networks (GFlowNets) written in JAX.<n>gfnx provides an extensive set of environments and metrics for benchmarking, accompanied with single-file implementations of core objectives for training GFlowNets.<n>Across different tasks, gfnx achieves significant wall-clock speedups compared to Pytorch-based benchmarks.
• Score: 12.748861864193623
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Comgra: A Tool for Analyzing and Debugging Neural Networks [35.89730807984949]
  We introduce comgra, an open source python library for use with PyTorch. Comgra extracts data about the internal activations of a model and organizes it in a GUI. It can show both summary statistics and individual data points, compare early and late stages of training, focus on individual samples of interest, and visualize the flow of the gradient through the network.
  arXiv  Detail & Related papers  (2024-07-31T14:57:23Z)
• Benchmarking Predictive Coding Networks -- Made Simple [48.652114040426625]
  We tackle the problems of efficiency and scalability for predictive coding networks (PCNs) in machine learning.<n>We propose a library, called PCX, that focuses on performance and simplicity, and use it to implement a large set of standard benchmarks.<n>We perform extensive tests on such benchmarks using both existing algorithms for PCNs, as well as adaptations of other methods popular in the bio-plausible deep learning community.
  arXiv  Detail & Related papers  (2024-07-01T10:33:44Z)
• Green AI: A Preliminary Empirical Study on Energy Consumption in DL Models Across Different Runtime Infrastructures [56.200335252600354]
  It is common practice to deploy pre-trained models on environments distinct from their native development settings. This led to the introduction of interchange formats such as ONNX, which includes its infrastructure, and ONNX, which work as standard formats.
  arXiv  Detail & Related papers  (2024-02-21T09:18:44Z)
• torchgfn: A PyTorch GFlowNet library [44.94532429787822]
  We present torchgfn, a PyTorch library that aims to address this need.<n>Its core contribution is a modular and decoupled architecture which treats environments, neural network modules, and training objectives as interchangeable components.<n>This provides users with a simple yet powerful API to facilitate rapid prototyping and novel research.
  arXiv  Detail & Related papers  (2023-05-24T00:20:59Z)
• PARTIME: Scalable and Parallel Processing Over Time with Deep Neural Networks [68.96484488899901]
  We present PARTIME, a library designed to speed up neural networks whenever data is continuously streamed over time. PARTIME starts processing each data sample at the time in which it becomes available from the stream. Experiments are performed in order to empirically compare PARTIME with classic non-parallel neural computations in online learning.
  arXiv  Detail & Related papers  (2022-10-17T14:49:14Z)
• Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining [58.10436813430554]
  Mini-batch training of graph neural networks (GNNs) requires a lot of computation and data movement. We argue in favor of performing mini-batch training with neighborhood sampling in a distributed multi-GPU environment. We present a sequence of improvements to mitigate these bottlenecks, including a performance-engineered neighborhood sampler. We also conduct an empirical analysis that supports the use of sampling for inference, showing that test accuracies are not materially compromised.
  arXiv  Detail & Related papers  (2021-10-16T02:41:35Z)
• Transparent FPGA Acceleration with TensorFlow [1.0828616610785522]
  We propose a toolflow for developers who want to make use of a new deep learning accelerator. On the backend we use an FPGA, which is addressable via a runtime environment. This can be achieved by our HSA toolflow, since the hardware is not statically configured with the structure of the network.
  arXiv  Detail & Related papers  (2021-02-02T06:49:33Z)
• Efficient Graph Deep Learning in TensorFlow with tf_geometric [53.237754811019464]
  We introduce tf_geometric, an efficient and friendly library for graph deep learning. tf_geometric provides kernel libraries for building Graph Neural Networks (GNNs) as well as implementations of popular GNNs. The kernel libraries consist of infrastructures for building efficient GNNs, including graph data structures, graph map-reduce framework, graph mini-batch strategy, etc.
  arXiv  Detail & Related papers  (2021-01-27T17:16:36Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.