Generative Flow Networks (GFlowNets) were developed to learn policies for efficiently sampling combinatorial candidates by interpreting their generative processes as trajectories in directed acyclic graphs. In the value-based training workflow, the objective is to enforce the balance over partial episodes between the flows of the learned policy and the estimated flows of the desired policy, implicitly encouraging policy divergence minimization. The policy-based strategy alternates between estimating the policy divergence and updating the policy, but reliable estimation of the divergence under directed acyclic graphs remains a major challenge. This work bridges the two perspectives by showing that flow balance also yields a principled policy evaluator that measures the divergence, and an evaluation balance objective over partial episodes is proposed for learning the evaluator. As demonstrated on both synthetic and real-world tasks, evaluation balance not only strengthens the reliability of policy-based training but also broadens its flexibility by seamlessly supporting parameterized backward policies and enabling the integration of offline data-collection techniques.
Figures & tables
Figure 1: A graphical illustration of a DAG (left) and its graded version (right). Dotted circles represent dummy states, added during the conversion to a graded DAG.
Figure 2
Figure 2: Plots of the means and standard deviations (represented by the shaded area) of DTV for different training methods with parameterized πB and uniform πB on the 256×256 (left) and 128×128 (middle) and 64×64×64 (right) grids, based on five randomly started runs for each method. By default , metric values are recorded every 20 iterations over N=2000 training iterations and smoothed by a sliding window of length 5 for all plotted curves in this paper.
Figure 3: Plots of the mean and standard deviation values (represented by the shaded area) of average reward (left), diversity (right) and FCS (right) of the top 100 unique candidate graphs over 10 nodes, based on five randomly started runs for each method.
Appendix figures & tables18 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Plots of the means and standard deviations (represented by the shaded area) of DJSD for different training methods on the 256×256 (left) 128×128 (middle) and 64×64×64 (right) grids, based on five randomly started runs for each method.
Figure 5: Plots of the means and standard deviations (represented by the shaded area) of DTV (right) for different training methods on the 128×128 (left) 20×20 (right) grids, based on five randomly started runs for Sub-TB-16 Sub-TB-128, Q-much-16 and Q-much 128. Here, “16” and “128” denote the training batch sizes.
Figure 6: Plots of the means and standard deviations (represented by the shaded area) of DTV (left) and DJSD (right) for Sub-EB runs on the 128×128 grid. The hyperparameter λ of Sub-EB ranges from 0.1 to 0.99 . The results are based on five Sub-EB runs for each setup.
Figure 7: Plots of the mean and standard deviation values (represented by the shaded area) of MN for the SIX6 (top left), QM9 (top right), PHO4 (bottom left), and sEH (bottom right) dataset, based on five randomly started runs for each method.
Figure 8: Plots of the mean and standard deviation values (represented by the shaded area) of DTV (left), DJSD (middle), and MA (right) for the SIX6 dataset, based on five randomly started runs for each method.
Figure 9: Plots of the mean and standard deviation values (represented by the shaded area) of DTV (left), DJSD (middle), and MA (right) for the QM9 dataset, based on five randomly started runs for each method.
Figure 10: Plots of the mean and standard deviation values (represented by the shaded area) of DTV (left), DJSD (middle), and MA (right) for the PHO4 dataset, based on five randomly started runs for each method.
Figure 11: Plots of the mean and standard deviation values (represented by the shaded area) of DTV (left), DJSD (middle), and MA (right) for the sEH dataset, based on five randomly started runs for each method.
Figure 12: Plots of the mean and standard deviation values (represented by the shaded area) of average reward (left), diversity (middle) and FCS (right) of the top 100 unique candidate graphs over 15 nodes, based on five randomly started runs for each method.
Figure 13: Plots of the mean and standard deviation values (represented by the shaded area) of DTV (left) and DJSD (right) for the BN learning structure learning task over 5 nodes, based on five randomly started runs for each method.
Figure 14: Four exemplary building blocks. Red and blue dots indicate the available atoms for bonding, where each bonding edge originates from a red dot and terminates at a blue dot.
Figure 15: Plots of the mean and standard deviation values of average reward (left) and diversity (right, Tanimoto diversity of molecular Morgan fingerprints) of top 100 unique molecules for the LogP task, based on five randomly started runs for each method.
Figure 16: Plots of the mean and standard deviation values of average reward (left) and diversity (right, Tanimoto diversity of molecular Morgan fingerprints) of top 100 unique molecules for the JNK3 task, based on five randomly started runs for each method.
Method
Mean
Std
Sub-TB
1h 08m
1.09
Sub-TB-P
1h 14m
10.36
RL
1h 03m
0.92
RL-P
1h 27m
0.57
RL-MLE
1h 09m
1.78
Sub-EB
1h 17m
2.31
Appendix
Table 1: The mean and standard deviation (std) of the total runtimes for each method on the 64×64×64 (left), 128×128×128 (middle), and 256×256×256 (right) grids. Here, ‘h’ denotes hours and ‘m’ denotes minutes, and all std values are reported in minutes.
Method
Mean
Std
Sub-TB
0h 05m
0.21
Sub-TB-B
0h 16m
0.67
RL
0h 04m
0.06
Sub-EB
0h 04m
0.34
Sub-EB-B
0h 17m
1.63
Appendix
Table 2: The mean and standard deviation (std) of the total runtimes for each method on the SIX6 and QM9 datasets. Here, ‘h’ denotes hours and ‘m’ denotes minutes, and all std values are reported in minutes.
Method
Mean Time
Std (min)
Sub-TB
0h 19m
2.05
Sub-TB-B
0h 27m
0.42
RL
0h 18m
0.10
Sub-EB
0h 18m
0.15
Sub-EB-B
0h 30m
1.09
Appendix
Table 3: The mean and standard deviation (std) of the total runtimes for each method on the PHO4 and sEH datasets. Here, ‘h’ denotes hours and ‘m’ denotes minutes, and all std values are reported in minutes.
Method
Mean
Std
Sub-TB
0h 57m
1.11
Sub-TB-B
1h 23m
4.27
RL
0h 56m
0.46
Sub-EB
0h 49m
0.87
Sub-EB-B
1h 07m
1.43
Q-Much
0h 59m
3.57
Appendix
Table 4: The mean and standard deviation (std) of the total runtimes for each method on the 5-node, 10-node and 15-node BN tasks. Here, ‘h’ denotes hours and ‘m’ denotes minutes, and all std values are reported in minutes. The runtimes for the 5-node cases are the largest due to the explicit computation of DTV for performance comparison during training.
Method
Mean Time
Std (min)
RL
0h 9m
0.36
RLEval
0h 13m
0.17
Sub-TB
0h 13m
0.93
Q-Much
0h 19 m
3.50
Appendix
Table 5: The mean and standard deviation (std) of the total runtimes for each method on the LogP and JNK3 tasks. Here, ‘h’ denotes hours and ‘m’ denotes minutes, and all std values are reported in minutes.
Generative Flow Networks (GFlowNets) learn policies that sample discrete objects in proportion to a positive reward. They are trained by minimizing balance objectives, whose exact satisfaction guarantees correct sampling. However, a comprehensive theory linking nonzero balance loss and approximate stationarity to sampling accuracy, especially under evolving sampling distributions and backward policies, is still lacking. To bridge this gap, we develop a unified theory for Flow Matching, Detailed Balance, and Trajectory Balance in this regime. We first bound terminal symmetric KL by the square root of the expected balance loss. Then, we show that approximate stationarity controls this distribution error up to a tangent-space residual measuring parameterization incompleteness. Finally, we establish finite-time convergence under fixed or controlled sampling and fixed, balance-trained, or separately updated backward policies. Under standard smoothness assumptions, with fixed or controlled sampling and a fixed or balance-trained backward policy, the expected loss is O(εˉ2+T−1) and expected terminal symmetric KL is O(εˉ+T−1/2) under uniform coverage. For a separately updated backward policy, we characterize the sufficient update decay to approach the residual neighborhood and preserve these rates. These results connect balance objectives to sampling accuracy and clarify how parameterization, coverage, and moving components govern GFlowNet training.
Rui Hu, Yu Chen, Zhuoran Li +1
Institute for Interdisciplinary Information Sciences, Tsinghua University
This paper explores policy gradient algorithms for training stochastic policies to sample from structured discrete probability distributions under the Generative Flow Network (GFlowNet) framework. Building on extensive theoretical connections between GFlowNets and entropy-regularized reinforcement learning, we derive equivalents of standard policy gradient algorithms for training GFlowNets, as well as experimentally explore their various methodological aspects, including baseline training and advantage estimation. Most importantly, our work is the first to derive and successfully apply proximal policy optimization to GFlowNets, showing its improved convergence speed and data efficiency compared to standard GFlowNet training objectives on benchmarks ranging from synthetic energies to molecular graph generation.
Anna Zykova-Myzina, Timofei Gritsaev, Daniil Tiapkin +1
HSE University · Constructor University · CMAP, CNRS, École polytechnique, IPP
Generative Flow Networks (GFlowNets) learn to sample states proportional to an unnormalized reward. Despite their theoretical promise, practical training is often unstable, exhibiting severe loss spikes and mode collapse. To tackle this, we first assess the sensitivity of GFlowNet objectives, demonstrating that a small Total Variation (TV) distance between the learned and target distributions does not preclude unbounded training loss. Motivated by this mismatch, we establish converse guarantees by deriving loss-to-TV bounds that certify global fidelity from bounded trajectory balance losses. Lastly, we propose Stable GFlowNets, an algorithm that leverages our theoretical results to stabilize training, and empirically demonstrate improved training behavior and superior distributional fidelity.
Zengxiang Lei, Ananth Shreekumar, Jonathan Rosenthal +6
Lyles School of Civil and Construction Engineering, Purdue University, IN, USA · Department of Computer Science, Purdue University, IN, USA · Computer Science and Engineering, University of California, Santa Cruz, USA