Learning an ensemble of GFlowNets to sample from a discrete target distribution has become a common approach for achieving better state space exploration and convergence than that of a monolithic sampler. However, these methods often add a substantial runtime overhead to the base model, and their conceptual connection remains elusive. To address this, we first propose a general-purpose theoretical framework for describing a mixture of GFlowNets, which we specialize into continuously (CI) and discretely indexed (DI) collections. On the one hand, we show CI GFlowNets can be interpreted through the lens of a random features expansion, provably boosting the sampler's expressivity in graph-structured tasks and reducing learning instability via spectral shifting. On the other hand, we demonstrate DI GFlowNets encompass prior approaches for GFlowNet training and provide the foundation for the newly proposed Stratum-Conditioned (SC) GFlowNets. This method, which is inspired by the Doob's h-transform of Markov chains, decomposes the state space according to a prescribed modular function and restricts each component to sample from a distinct subset of it. Importantly, SC GFlowNets support centralized and component-wise embarrassingly parallel training, and we show both of them significantly speed up learning convergence and mode coverage without introducing any non-negligible extra computation.
Figures & tables
Figure 1 : An ensemble of GFlowNets that, when averaged, matches the target.
Figure 2 : SC GFlowNets for the Lines environment ( N=256 ). (Left) We show the TV distance between the learned and target distributions, emphasizing the faster convergence promoted by stratum-conditioning. (Middle) The marginal distribution of both models after 1000 training steps, and the target. (Right) The marginal distribution of SC GFlowNets for each stratum k∈{1,2,3} .
Figure 3
Figure 5 : SC GFlowNets converge substantially faster than GFlowNets in the Hypergrid domain, with distinct partitioning strategies ( Length & Octant ) yielding similar benefits.
Figure 6 : SC GFlowNets enhance exploration (left) and speed up learning (right) for the Ancestral Graphs domain. We denote by H the number of variables (nodes).
Appendix figures & tables21 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7 : Any α∈[1,2] solves the assignment problem in this state graph. The solution for uniform pB is α=1.5 .
Figure 8 : KL between pB and pU for the Hypergrid domain.
Figure 9 : State graph for Example C.3 .
Figure 10 : Random features act as a regularizer for GFlowNet training.
Figure 11 : SC GFlowNets run both faster (left) and more compute-efficiently (right) than Boosted GFlowNets for the Set Generation domain. Throughout this section, runtimes represent per-step averages over 3 independent runs, each with 600 iterations.
Figure 12 : SC GFlowNets are more compute-efficient than Boosted GFlowNets in the Hypergrid domain, requiring substantially fewer FLOPs per gradient step during training.
Figure 13 : SC GFlowNets is far more efficient than Boosted GFlowNets in the Lazy Random Walk domain ( Rings corresponds to the Rings target; GM , to Gaussian Mixture ).
Figure 14 : SC GFlowNets converge substantially faster than Boosted GFlowNets in the Set Generation domain, particularly for larger state spaces. See also Figure 4 for Erritem ’s definition.
Figure 15 : SC GFlowNets exhibit faster learning convergence when compared to Boosted GFlowNets, while also being far more compute-efficient ( Figure 12 ).
Figure 16 : On top of being more computationally efficient ( Figure 13 ), SC GFlowNets also learn a more accurate distributional approximation to the target relatively to Boosted GFlowNets in both the Rings and Gaussian Mixture variants of the Lazy Random Walk domain.
Figure 17 : SC GFlowNets significantly accelerate learning convergence when approximating the generative process in Equation 15 , with greater improvements being observed for larger state spaces. TVitem and TVlength are defined in Equation 17 .
Figure 18 : Embarrassingly parallel counterpart of Figure 17 , comparing ASC GFlowNet against the best of K randomly initialized and independently trained GFlowNets with matched computational cost. ASC GFlowNet accelerates learning convergence.
Figure 19 : Accuracy of SC GFlowNets as a function of the number of components in the mixture.
Figure 20 : SC GFlowNets improve goodness-of-fit for the Lazy Random Walk task, exhaustively covering the target distribution’s high probability regions even for an imperfect state space partitioning (as for the Gaussian Mixture ; see Figure 21 ).
Figure 21 : SC GFlowNet’s marginal mF(k)(so,x) (recall Equation 1 ) in each x∈X for both the (a) Rings and (b) Gaussian Mixture target distributions of the Lazy Random Walk domain.
Figure 22 : Learned log-partition functions by ASC GFlowNets in the Lazy Random Walk domain. We recall that, for Rings , k=3 is the starved component, containing negligible probability mass; for Gaussian Mixture , it is k=1 . As expected, ASC GFlowNets learn the correct values (a, c); however, the approximation of logZk for regions with near-zero probability is imperfect (b, d), as sizeable deviations in logZk in this case have little effect on distributional accuracy.
Figure 23 : KL between prior mixing distributions.
Figure 24 : ASC GFlowNets converge faster than their synchronous counterparts in the Lazy Random Walk task, achieving a more accurate approximation in fewer training iterations.
Figure 25 : SC GFlowNets improve learning convergence (left) and accelerate the discovery of high-probabiltiy states (right) for the Bayesian variable selection task; H represents the number of variables (columns) in the dataset.
Figure 26 : SC GFlowNets’ distributional accuracy for distinct numbers of partitions.
Figure 27 : Component-wise learned distributions for both considered SC GFlowNets in Figure 5 .
University of Science and Technology of China (USTC), Hefei, China · Suzhou Institute for Advanced Research, USTC, Suzhou, China · State Key Laboratory of Precision and Intelligent Chemistry, USTC, Hefei, China