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 .
Generative Marginalization Models (MaMs) have been recently introduced as efficient neural sampling models for any-order autoregressive modelling of discrete distributions. By learning both the marginal and conditional probabilities of a persistent-block Gibbs sampler, MaMs enable fast posterior evaluation with a single neural network forward pass. While prior work has considered MaMs to be distinct from Generative Flow Networks (GFlowNets), a well-established paradigm for inference in discrete stochastic models, we show that they are equivalent. Then, we also extend MaMs' sampling strategy to non-autoregressive generative processes. In particular, we describe an automatic criterion for full-state rejuvenation of the Gibbs sampler, derived from the Gelman-Rubin statistic, which plays a key role in speeding up learning convergence. Our experiments show that our method, called Particle GFlowNets, markedly accelerates training in large combinatorial spaces.
Tiago da Silva, Diego Mesquita, Salem Lahlou
MBZUAI · School of Applied Mathematics, Getulio Vargas Foundation
Generative Flow Networks (GFlowNets) are a framework for sampling structured objects via stochastic trajectories in a directed graph. In this work, we establish a theoretical connection between non-acyclic GFlowNets and optimal transport (OT). We show that fixing the initial flow distribution in a minimum-flow GFlowNet reduces its objective to a Kantorovich OT problem with graph-induced shortest path costs. At the optimum, the learned GFlowNet policy therefore encodes an optimal transport plan from the source distribution to the target distribution: we show that sampling trajectories from the minimum-flow GFlowNet recovers the corresponding optimal coupling. Our formulation enables applying the GFlowNet learning framework to OT problems on large graphs via edge flows and neural parameterization. Experiments confirm agreement with exact OT solvers and demonstrate that GFlowNets can learn high-quality transport plans.
Generative Flow Networks (GFlowNets) have shown promising potential to generate high-scoring candidates with probability proportional to their rewards. As existing GFlowNets freely explore in state space, they encounter significant convergence challenges when scaling to large state spaces. Addressing this issue, this paper proposes to restrict the exploration of actor. A planner is introduced to partition the entire state space into overlapping partial state spaces. Given their limited size, these partial state spaces allow the actor to efficiently identify subregions with higher rewards. A heuristic strategy is introduced to switch partial regions thus preventing the actor from wasting time exploring fully explored or low-reward partial regions. By iteratively exploring these partial state spaces, the actor learns to converge towards the high-reward subregions within the entire state space. Experiments on several widely used datasets demonstrate that \modelname converges faster than existing works on large state spaces. Furthermore, \modelname not only generates candidates with higher rewards but also significantly improves their diversity.
Xuan Yu, Xu Wang, Rui Zhu +2
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