Particle GFlowNets: Rethinking Generative Marginalization Models
Tiago da Silva1
Diego Mesquita2
Salem Lahlou1
1
arXiv:2609.11538v1 [cs.LG] 10 Sep 2026
2
MBZUAI School of Applied Mathematics, Getulio Vargas Foundation
Abstract 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.
1
INTRODUCTION
Generative Flow Networks [GFlowNets; Bengio et al., 2021, 2023] have been used in recent years as flexible samplers for distributions over compositional, discrete objects such as sequences and sets. Notably, they were successfully applied to problems in computational biology [Jain et al., 2022, 2023a, Laajil et al., 2025], combinatorial optimization [Zhang et al., 2023a,b, Zhang and Cao, 2025], and Bayesian inference [Deleu et al., 2022, 2023], several of which seem not to be as easily solvable with other approaches. In fact, given a positive function R(x) defined over a finite and combinatorial space, a GFlowNet can often generate independent samples from the corresponding probability distribution π(x) ∝ R(x) with remarkable accuracy [Shen et al., 2023].
To achieve this, GFlowNets learn a policy function for a Markov Decision Process (MDP) whose marginal distribution over terminal states matches π(x) [Tiapkin et al., 2024]. However, training is frequently bottlenecked by the simulation of long, overlapping trajectories, requiring countless neural network forward passes for each gradient step. Generative Marginalization Models [MaMs; Liu et al., 2024] mitigate this issue by learning both the marginal and conditional distributions of a Gibbs sampler [Geman and Geman, 1984]. In doing so, each learning step requires a constant number of model evaluations—corresponding to a single Gibbs transition—regardless of the state space’s size. Contrarily to conventional understanding, we demonstrate that MaMs are equivalent to permutation-conditioned (PC) GFlowNets, which we introduce as a particular case of the well-established conditional GFlowNet framework [Bengio et al., 2023, Jain et al., 2023b]. Specifically, we show that MaMs’ learning objective corresponds to the traditional detailed balance loss function [Bengio et al., 2023, Example 5] for PC-GFlowNets. With this in mind, the key contribution of MaMs is a computationally efficient Gibbs sampling algorithm that minimizes the number of neural network evaluations during training. Their core assumption, however, is that R(x) is supported on a factorized domain (i.e., {1, . . . , K}d ), which limits MaMs’ applicability in permutation-invariant (e.g., graph-structured) spaces [Liu et al., 2024, Section 3]. The central question we ask in this work is whether such an approach can be adapted to non-factorized domains. We answer it affirmatively. For this, we propose Particle (P) GFlowNets. As in MaMs, a P-GFlowNet maintains a persistent Gibbs sampler over terminal states. Similarly to PC-GFlowNets, a P-GFlowNet optimizes a learning objective that enforces the detailed balance condition of the underlying Markov chain. In contrast to both methods, which were designed for autoregressive modelling, the reverse (backward) move for P-GFlowNets is non-deterministic, as the forward mapping is non-injective. To circumvent this, we use a stochastic reverse transition kernel, which can also be learned. Importantly, once
trained, a P-GFlowNet can generate both independent or correlated samples via MDP simulation or Gibbs sampling, respectively. The crucial difference lies in the per-generation computational cost and the effective sample size of the generated samples, both of which are larger for independent sampling. Given a limited time budget, the optimal approach will likely need to be determined in a case-by-case basis. In practice, we find that samples from P-GFlowNet’s persistent Gibbs chain can become highly correlated during training, reducing the accuracy of gradient estimates and slowing down convergence. To address this limitation, we occasionally refresh the stochastic process based on the Gelman-Rubin statistics Gelman and Rubin [1992], Carpenter et al. [2017], also known as R-hat or R̂, being above a certain threshold. We empirically show that this process, called rejuvenation in the probabilistic programming literature Lew et al. [2022a,b], significantly accelerates learning convergence. Additionally, we evaluate P-GFlowNets on both established benchmarks and novel tasks for which relevant functionals can be efficiently computed and used for reliable performance assessment. Our experiments show that P-GFlowNets significantly accelerate training convergence when the cost of MDP simulation exceeds that of evaluating the target distribution, which is typical for large, long-horizon state spaces. In summary, our contributions are as follows. 1. We show MaMs can be represented as PC GFlowNets, and that their learning objective is equivalent to Bengio et al. [2023]’s expected detailed balance loss function. 2. We introduce P-GFlowNets. When compared to prior methods for GFlowNet training, P-GFlowNets asymptotically reduce the average number of forward passes per gradient step as the MDP’s horizon increases. 3. We empirically demonstrate that P-GFlowNets significantly reduce wall-clock time for learning convergence. The reader only interested in learning about P-GFlowNets may skip Section 3 and focus on Sections A, 4 and 5 instead.
2
PRELIMINARIES
Notations. Let X be a discrete space. We assume X is compositional, i.e., each x ∈ X can be described as a collection of components C, x = {c1 , . . . , cd } ⊆ C, with d ≥ 1 possibly depending on x. Often, we will consider X := [K]d := {1, . . . , K}d for positive integers K and d, in which case we will refer to d as X ’s dimension. In this setting, C = [K] × [d] and each x = {(i1 , 1), . . . , (id , d)} corresponds to the sequence (i1 , . . . , id ) ∈ [K]d . We use these notations interchangeably, i.e., x = {(1, 1), (0, 2), (1, 3)} represents the sequence (1, 0, 1), with x1 = 1, x2 = 0, and x3 = 1. In particular, we write x ≡ {(ij , j)}dj=1 . In general, X ⊆ 2C is a subset of C’s power set, 2C .
Our objectives are to generate x ∈ X in proportion to a given positive measure in X , the probability mass of which we will denote by R : X → R+ , and to evaluate the marginal probability of subsets of x under R. We will call π(x) ∝ R(x) the normalized distribution. As in Bengio et al. [2021], Malkin et al. [2023], we refer to R as a reward function. Additionally, we will define the space S as S = {s ∈ 2C : s ⊂ x for some x ∈ X }.
(1)
Clearly, ∅ ∈ S. As in Bengio et al. [2021, 2023], we call S ∪ X the state space, and define so := ∅ as the initial state. Similarly, we let G := (S ∪ X , E), E = {(s, s′ ) : ∃c ∈ C \ s′ such that s′ = s ∪ {c}}, be the state graph. Intuitively, s → s′ in G if s′ differs from s by a single additional component. We say that the generative process is autoregressive when G is a tree rooted at so ; see Figure 1.
(a) Autoregressive.
(b) Non-autoregressive.
Figure 1: The state graph (G) is a tree in an autoregressive setting (a). When many trajectories lead to the same state, such as so → (s1 , s2 ) → x2 in (b), G is a non-tree DAG. GFlowNets. A GFlowNet learns a forward policy function pF : S × (S ∪ X ) → [0, 1] such that pF (s, ·) is a probability distribution supported on the children of s in the state graph G. By defining so ⇝ x as the set of trajectories in G from so to x ∈ X , we seek a function pF satisfying X p⊤ (x) := pF (so , τ ) ∝ R(x), (2) τ ∈so ⇝x
with τ = (so , . . . , sd−1 , x), sd := x, and (with a slight Qd abuse of notation) pF (so , τ ) = i=1 pF (si−1 , si ); we refer to d as the trajectory’s length. Equation (2) represents the marginal distribution over X induced by pF (so , ·). When the generative process is autoregressive, so ⇝ x contains a single trajectory τx , and p⊤ (x) = pF (so , τx ). Otherwise, the set so ⇝ x might be intractably large; see Figure 1. Under these conditions, we introduce a backward policy pB , which is a forward policy on the transposed state graph, and approximate Equation (2) via importance sampling, − pF (so , → τ) − p⊤ (x) = E← ∝ R(x), τ ∼pB (x,·) − pB (x, ← τ) − − in which we use → τ and ← τ to distinguish the forward and backward trajectories. (We will drop thisP notation in the rest of the text for clarity). By letting Z := x∈X R(x) be the partition function of our target distribution, the above equation may be re-arranged as (dividing both sides by R(x)) Z · pF (so , τ ) Eτ ∼pB (x,·) = 1. (3) R(x) · pB (x, τ )
The condition Z · pF (so , τ ) = R(x) · pB (x, τ ) is known as trajectory balance [TB; Malkin et al., 2022]. In practice, pF (s, ·) is parameterized as a softmax neural network receiving s as input, and we optimize both pF and Z via stochastic gradient descent on the objective function " 2 # Z · pF (so , τ ) log , LTB (pF , pB , Z) = Eτ ∼pE R(x) · pB (x, τ ) with pE as an exploratory (full-support) policy [see, e.g., Madan et al., 2025, Kim et al., 2025]. As the reader may have noticed, pB was chosen arbitrarily; it can be either learned [e.g., Shen et al., 2023, Gritsaev et al., 2025] or fixed [e.g., Deleu et al., 2022, Zhou et al., 2024]. Alternatively, we commonly parameterize the flow functions X F (s) = Z pF (so , τ ), τ ∈so ⇝s
for s ̸= so and F (so ) = Z, and define the detailed balance (DB) learning objective LDB (pF , pB , F ) as 2 X F (s )p (s , s ) i−1 F i−1 i , (4) Eτ ∼pE wi log F (si )pB (si , si−1 ) 1≤i≤d
with τ = (so , . . . , sd ) as in Equation (2), FP (sd ) = R(sd ), n and wi > 0 as positive weights such that i=1 wi = 1. Common choices for wi are either uniform wi = 1/n [Bengio et al., 2023] or wi ∝ exp{λi} for λ > 0, the intuition being that states closer to X should be assigned with larger weights during training [Silva et al., 2025a]. Other learning objectives have also been studied [e.g., Madan et al., 2022, Zhang et al., 2023a]; see Section C for related works. A conditional GFlowNet, on the other hand, models a family of distributions Rω (x) indexed by a parameter ω ∈ Ω that is also used as input for both the policy (pF and pB ) and flow (F ) functions. Prominent applications include multiobjective combinatorial optimization, in which Ω is a simplex and Rω represents a weighted mixture of each objective according to ω [Jain et al., 2023b, Roy et al., 2023, Zhu et al., 2023, Laajil et al., 2025], and distributed learning, with each ω ∈ Ω representing a distinct subgraph of the state graph and Rω (x) simply referring to the restriction of a given R(x) to the corresponding subset of X [Silva et al., 2025b]. MaMs. In the context of sampling, MaMs were designed by Liu et al. [2024] for fast autoregressive modelling of discrete models. Simply put, they concomitantly learn a marginal and conditional distributions, pθ and pϕ , on [K]d satisfying both a correctness and consistency conditions. Given indices I, J ⊆ {1, . . . , d} such that I ⊂ J and |J \ I| = M , we respectively define these conditions as pθ (x) =
R(x) and pθ (xI )pϕ (xJ |xI ) = pθ (xJ ), Z
(5)
for x ∈ X . As in Liu et al. [2024], we assume M = 1, although our analysis can be easily extended to M > 1 by interpreting a group of M adjacent variables as a single variable. As with GFlowNets, pθ and pϕ are parameterized as neural networks. Also, a special symbol △ is used to represent a marginalized variable—e.g., if x = (x1 , . . . , xd ), then x{1,2} = (x1 , x2 , △, . . . , △)—and Z is a learned parameter. Once trained, the probability pθ (xI ) of any x ∈ X and any I ⊆ {1, . . . , d} can be evaluated in a single neural network forward pass. To enforce Equation (5), we minimize KL[pθ ||π] + λConsistencyError(pθ , pϕ , Z), (6) h i θ (x) in which λ > 0, KL[pθ ||π] := Ex∼pθ log pπ(x) is the Kullback-Leibler (KL) divergence between pθ and π, and ConsistencyError is defined as 2 pθ (xσ(<m) ) · pϕ (xm |xσ([m−1]) ) Ex∼q Em Eσ log , pθ (xσ([m]) ) with m ∼ U[d], σ ∼ U(Sd ), q as any distribution over X , U[d] as an uniform distribution over [d] = {1, . . . , d}, (hardcoded) restriction pθ ((△, . . . , △)) = Z, and σ uniformly picked from the space of permutations, which we denote by Sd = {σ : [d] → [d] | σ is bijective}.
(7)
To generate samples from pθ , which are needed to estimate the KL divergence in Equation (6), Liu et al. [2024] implement a persistent-block Gibbs sampling scheme having pϕ as the transition kernel. As we explain next, under the consistency condition in Equation (5), the resulting Markov chain has the marginal pθ as the stationary distribution. Gibbs sampling. Originally designed by [Geman and Geman, 1984] to sample from the Gibbs distribution [Gibbs, 1902], and also known as Successive Substitution Sampling [SSS; Schervish, 2012], Gibbs sampling creates a Markov chain by iteratively modifying each coordinate of a state x according to the conditionals of the target distribution π(x). That is, we iteratively select i ∈ [d] and then x′i ∼ π(xi |x−i ), with x−i = x{1,...,i−1,i+1,...,d} as in Equation (5). This process may be pictured as (x1 , x2 ) → (x′1 , x2 ) → (x′1 , x′2 ) for two-dimensional distributions. When i is picked at random, the algorithm is known as random scan Gibbs sampler; otherwise, as systematic scan Gibbs sampler [Owen, 2013]. We will focus on the former. The reader is invited to notice that the resulting Markov chain is stationary with respect to the target distribution π. In the context of learning and energy-based models [Hinton, 2002, Tieleman, 2008, Liu et al., 2024], a persistent-block Gibbs sampler extends the canonical algorithm by (i) grouping variables into blocks, each of which is updated in a single step, and (ii) maintaining the Markov chain’s state across gradient updates of the underlying neural sampling model. We notice that, in this case, the Markov chain ceases to be stationary until training converges.
3
RETHINKING MAMS AS GFLOWNETS
We now show that MaMs may be interpreted as a conditional GFlowNet (Section 3.1). In a nutshell, the conditioning space Ω := Sd will be the space of permutations defined in Equation (7), and we will demonstrate the consistency loss function is equivalent to the weighted DB loss in Equation (4) in expectation (Section 3.2). 3.1
unchanged under σ (except for re-labeling). Interestingly, however, the permutation-conditioned flow function F σ satisfying the DB condition in Equation (4) can be exactly interpreted as the marginal probability distribution of a sequence under R. This result, outlined below, is crucial in bridging the gap between GFlowNets and MaMs. Proposition 3.1. Let F σ be a flow function abiding by the DB condition of a permutation-conditioned GFlowNet, i.e., F σ (s)pσF (s, s′ ) = F σ (s′ )pσB (s′ , s)
PERMUTATION-CONDITIONED GFLOWNETS
(9)
σ
To understand the connection between MaMs and GFlowNets, we first introduce a family of permutationconditioned (PC) GFlowNets for autoregressive modelling. Again, let X = [K]d . We define a learnable state-embedding function ψ : S → Rh and weights {Wk }K k=1 such that Wk ∈ Rh×d for a given dimension h. Then, recall from Section 2 that the state space S ∪ X is simply a subset of the power set of C := [K] × [d]. As such, the policy function evaluated at s(i) := {(k1 , σ(1)), . . . , (ki , σ(i))} ∈ S and conditioned on a permutation σ ∈ Sd is defined as T pσF (s(i) , s(i+1) ) ∝ exp{wk,σ(i+1) ψ(s(i) )}
(8)
for s = s ∪ {(k, σ(i + 1))} and k ∈ [K]. Intuitively, pσF fills up a sequence up to a prescribed size (d, in this case) according to the ordering imposed by σ. We illustrate this in Figure 2. Given a σ and a s(i+1) ∈ S, there is only one state s(i) such that the transition s(i) → s(i+1) has positive probability under pσF (si , ·). Hence, the only choice for the backward policy abiding by the TB condition in Equation (3) is pσB (s(i+1) , ·) = δs(i) for s(i) = {(kj , σ(j))}ij=1 and i+1 s(i+1) = s(i) ∪ {(ki+1 , σ(i + 1))} with (kj )i+1 ; j=1 ∈ [K] (i) δs(i) is the Dirac delta at s , i.e., δs(i) (s) = 1 if s = s(i) and 0 otherwise. Thus, as pB plays no significant role in learning, we set it aside from most of our analysis until Section 4. Our main claim in this section is that there a bijective correspondence between MaMs and PC GFlowNets. (i+1)
(i)
Claim 3.1. For each MaM, there is a unique and functionally equivalent PC GFlowNet; see Proposition 3.2.
(i)
and F (x) = R(x) for x ∈ X . Then, denoting by s = {(k1 , σ(1)), . . . , (ki , σ(i))} for i ≤ d, X X F σ (s(i) ) = R(x) and F σ (so ) = R(x), (i)
x∈X : xσ([i]) =sσ([i])
x∈X
(10) (i) in which xσ([i]) = sσ([i]) means that x ∈ [K]d agrees with
s(i) on the first i coordinates according to the permutation (i) (i) σ, i.e., xσ(1) = sσ(1) , . . . , xσ(i) = sσ(i) . In other words, F σ (s(i) ) (i) under R. F σ (so ) is the marginal distribution of s
As in Section 2, the state s(i) in the Proposition 3.1 can be (i) (i) written as sj = kj for j ∈ σ([i]) and sj = △ otherwise.
When learning a conditional GFlowNet, we often augment the input of the neural networks parameterizing both F σ and pσF with σ [Bengio et al., 2023, Zhang et al., 2023a]. As MaMs avoid such augmentation, we ask: is it needed for PC GFlowNets? The following corollary shows that, when states s1 and s2 differ solely by the permutation σ, the optimal flow function in Equation (9) does not depend on σ. Corollary 3.1. Let σ1 , σ2 ∈ Sd be permutations, and define s1,(i) = {(k1 , σ1 (1)), . . . , (ki , σ1 (i))} and s2,(i) = {(k1 , σ2 (1)), . . . , (ki , σ2 (i))}. Then, if i satisfies σ1 ([i]) = σ2 ([i]), i.e., σ1 and σ2 coincide on [i] := {1, . . . , i}, F σ1 (s1,(i) ) = F σ2 (s2,(i) ). To evaluate pσF (s(i) , s(i+1) ), however, information about σ(i+1) is needed. As we show below, σ1 and σ2 coincide on [i + 1], then pσF1 (s(i) , s) = pσF2 (s(i) , s) for each s ∈ S ∪ X .
Corollary 3.2. Let σ1 , σ2 ∈ Sd . Define s1,(i) and s2,(i) as in Corollary 3.1. Assume σ1 (≤ i + 1) = σ2 (≤ i + 1). Then, pσF1 (s1,(i) , ·) = pσF2 (s2,(i) , ·),
(a) σ(1) = 1, σ(2) = 2.
(b) σ(1) = 2, σ(2) = 1.
Figure 2: A PC GFlowNet fills up a sequence according to a given permutation (σ). After trained, it supports fast evaluation of any-subset marginals. (△ means no token.) Clearly, the conditioning permutation σ only affects the navigation of S, as both X and R : X → R+ remain
i.e., if s1,(i+1) ≡ s2,(i+1) correspond to the same object in ([K] ∪ {△})d , pσF1 (s1,(i) , s1,(i+1) ) = pσF2 (s2,(i) , s2,(i+1) ). In practice, we parameterize pσF (s(i) , ·) as a neural network returning a matrix P = RK×d with sum(P) = 1 represent+ ing the probability of each variable k being in the ith position, Pk,i , and mask out all columns except the σ(i + 1)-th one, which results in Equation (8). Drawing on the derivations above, Claim 3.1 is only a matter of bookkeeping.
Notation mapping. We recall a MaM is characterized by a marginal pθ and conditional pϕ distributions. Also, for each s = (k1 , . . . , kd ) with kj ∈ [K] or kj = △ (i.e., kj is marginalized), as described in Section 2, there is a permutation σ and an index i for which s ≡ {(j, σ(j))}ij=1 in the PC GFlowNet’s state graph. By Corollary 3.1, we can unambiguously define F σ (s) := pθ (s). Similarly, let J ⊃ I be subsets of {1, . . . , d}, and define I = |I|. (Recall |J \ I| = 1). As before, given x ∈ X , there is a permutation σ such that σ([I]) = I, σ([I + 1]) = J , and xI ≡ {(xσ(j) , σ(j))}Ij=1 and xJ ≡ {(xσ(j) , σ(j))}I+1 j=1 . (11) Under Corollary 3.2, pσF (xI , ·) does not depend on σ(j) for j > I + 1. Again, we can thus define pϕ (xJ |xI ) = pσF (xI , xJ ) without ambiguity. Our central result, stated below, establishes that there is a unique PC GFlowNet for each MaM satisfying both the distributional (i.e., pθ (x) = π(x) for x ∈ X ) and consistency conditions, and vice-versa. Proposition 3.2 (MaMs are PC GFlowNets). Let (pθ , pϕ ) be a consistent MaM satisfying pθ (x) = π(x). Then, there is a unique PC GFlowNet such that F σ (xJ ) = pθ (xJ ) and pσF (xI , xJ ) = pϕ (xJ |xI ) for each triplet I, J , σ as in Equation (11). This PC GFlowNet satisfies the DB condition in Equation (9). Conversely, a PC GFlowNet satisfying DB induces a unique consistent MaM such that pθ (x) = π(x). From an operational viewpoint, both MaMs and PC GFlowNets allow for the evaluation of any-subset marginals with a single neural network forward pass in discrete models. In fact, the only lingering difference between MaMs and PC GFlowNets is their differing learning objectives. As we discuss next, however, the estimator used for MaM’s loss function is biased unless the consistency condition holds. Hence, it might be unsuited for training. We instead derive a novel unbiased estimator that preserves MaMs’ constant number of forward passes per gradient step. 3.2
REVISITING MAMS’ OBJECTIVE
At a basic level, MaMs aim to minimize the KL divergence between pθ and π under the consistency constraint, i.e., min KL[pθ ||π] such that pθ (xI )pϕ (xJ |xI ) = pθ (xJ ) pθ
for each set of indices I, J as in Proposition 3.2. As directly enforcing pθ to be consistent with pϕ is computationally unfeasible, a penalty function—the ConsistencyError—is introduced instead; recall Equation (6). A central question, however, remains unanswered: how to generate samples from pθ to estimate the objective function KL[pθ ||π]?
When the consistency constraint is satisfied, a persistent Gibbs sampling scheme using pϕ as the transition kernel would realize a Markov chain ergodic with respect to pθ , the
samples of which could be used for estimating KL[pθ ||π]. However, pθ is only approximately consistent with pϕ . Thus, the generated samples would produce biased estimates of the KL objective. Based on the connection between MaMs and PC GFlowNets in Proposition 3.2, we propose instead minimizing the consistency-only objective below. I) Definition 3.1. We let pθ (xI ) = Fθ (x for given parametZ ric function Fθ and x ∈ X and indices I ⊆ [d] with Z := F (x∅ ) learnable. The consistency-only objective LCO is 2 Fθ (xσ([m−1]) )pϕ (xσ(m) |xσ([m−1]) ) Ex∼q,m,σ log , Fθ (xσ([m]) )
in which σ ∼ U(Sd ), m ∼ Cat(w) for w ≥ 0 and Pd i=1 wi = 1, and q is a full-support distribution over X . Additionally, we restrict Fθ (x) = R(x) for x ∈ X . The only difference between LCO and MaMs’ consistency error is that we enforce Fθ (x) = R(x) for x ∈ X instead of minimizing KL[pθ ||π].P Clearly, when LCO is globally minimized, Fθ (x∅ ) = x∈X R(x), which allows for direct evaluation of pθ . Moreover, the LCO corresponds to the DB loss LDB in Equation (4) with a specific weight scheme (wi )di=1 . In contrast to LDB , however, a Monte Carlo approximation of LCO requires a constant number of forward passes with respect to the state space’s dimension. Proposition 3.3. Let (F, pσF ) be a PC GFlowNet. (By Corollary 3.1, F does not depend on σ). Let pσE be an exploratory policy such that the marginal of pσE (so , ·) over X matches the distribution q defined in Definition 3.1. Then, 2 σ X F (s )p (s |s ) i−1 i i−1 F wi = E log F (si ) σ∼U (Sd ), τ ∼pσ E (so ,·)
1≤i≤d
" E
σ∼U (Sd ), x∼q, i∼Cat(w)
F (xσ([i−1]) )pσF (xσ([i−1]) , xσ([i]) ) log F (xσ([i]) )
2 #
with τ = (si )di=0 , sd = x, and xσ([i]) = {(xj , σ(j))}ij=1 . Proposition 3.3 establishes that MaMs’ ConsistencyError corresponds to a transition-wise estimator of the conventional DB loss function for conditional GFlowNets. After thoroughly outlining the connection between PC GFlowNets and MaMs, we ask: how can we leverage our results to improve GFlowNet training? We investigate this next.
4
PARTICLE GFLOWNETS
We extend the persistent-block Gibbs sampling scheme to non-autoregressive generative processes, such as set generation, in which multiple trajectories may lead to the same object. As with MaMs, this reduces the number of forward passes from O(d) to O(1) per gradient step. We also develop an automatic criterion for refreshing the Gibbs chain, which improves exploration and accelerates convergence.
,
Particle GFlowNets. We recall from Section 2 that a compositional object x is represented as a collection of components from a set C, i.e., x = {c1 , . . . , cd }. In this scenario, a policy function induces a probability distribution over a subset of C conditioned on a state s ∈ 2C . For PC GFlowNets, C = [K] × [d] and pσF (s(i) , ·) induces a distribution over [K] × {σ(i + 1)} ⊂ C, as in Proposition 3.1. We henceforth assume that each x ∈ X is naturally represented by exactly d components, i.e., X ⊆ Cd . This is often the case for usual applications and benchmarks in the GFlowNet literature, such as set generation Jang et al. [2024], phylogenetic inference [Zhou et al., 2024], design of mRNA sequences [Laajil et al., 2025], causal discovery [Silva et al., 2026], and certain combinatorial optimization tasks [Zhang et al., 2023b].
In practice, we pick k ∼ Cat(w) with w ∈ Rd as the probabilities of a truncated Poisson distribution with average equal to log d. As suggested by Proposition 3.3, this choice provides larger weight to near-terminal states, which has been shown to be beneficial [Silva et al., 2025a], while ensuring coverage of the entire state graph. We then define L̂P = L̂C + L̂I
(12)
as our loss function. We refer to a model trained by minimizing L̂P as a Particle (P) GFlowNet. Notably, we show below that the Markov chain {xt }t≥1 described above is ergodic with respect to the target π when the P-GFlowNet abides by the DB condition. This ensures P-GFlowNets support both correlated and independent sampling, the choice of which to use being a trade off between computational cost with statistical efficiency. Proposition 4.1. Let (pF , pB , F ) be a P-GFlowNet abiding by the detailed balance condition. Define {xt }t≥1 as the Markov chain with the transition kernel depicted in Figure 3. Then, {xt }t≥1 is ergodic with respect to π.
Figure 3: A transition x → y for P-GFlowNets’ persistent Gibbs sampler. At each iteration, we replace a component of the current state x and evaluate the loss function in Equation (12) by averaging over the B stochastic processes. This requires O(1) forward passess, regardless of d. To build intuition, consider Figure 3. We start by generating samples {xbt }B b=1 ⊆ X from an untrained GFlowNet. Each xb can be thought of as a d-sized subset of C, for instance, t,b xbt = {ct,b 1 , . . . , cd }. At each iteration, we remove a comb b b ponent ct,b i from xt according to pB (xt , ·), resulting in st . We then choose a c ∈ C according to the forward policy pF (sbt , ·) and attach it to sbt , generating xbt+1 := sbt ∪ {c}. In conclusion, we may update each model via a gradient step on 2 F (sbt )pF (sbt , xbt+1 ) 1 X L̂C := log . B R(xbt+1 )pB (xbt+1 , sbt ) 1≤b≤B Upon repetition, this produces a coupled Markov chain {{xbt }B b=1 }t≥1 satisfying the DB condition for terminal (xbt ) and near-terminal (sbt ) states. However, independent sampling via MDP simulation can only be achieved if the DB condition is satisfied for every intermediate state. In view of this, we also sample sb,k ∼ pkB (xbt , ·) by removing t b k k components from xt ; pB denotes pB ’s k-fold composition. b,k Then, we select sb,k t+1 ∼ pF (st , ·) and compute !2 b,k b,k F (sb,k 1 X t )pF (st , st+1 ) L̂I := log . b,k b,k B F (sb,k t+1 )pB (st+1 , st ) 1≤b≤B
Importantly, our persistent Gibbs chain might suffer from inadequate state space exploration due to the structural similarity of adjacent states. To mitigate this issue, we occasionally refresh the process with fresh independent samples from the current policy pF (so , ·). This approach is discussed next. Chain rejuvenation. We use the R̂ metric to decide when to restart our stochastic process. Following standard statistical practice, we define R̂ > 1.1 as our condition for rejuvenation [Carpenter et al., 2017]. To understand this, recall that the R̂ measures the discrepancy between the withinand inter-chain variances. When samples are independently generated, R̂ ≈ 1. A large R̂ indicates that the diversity of our batched stochastic process is significantly larger than that of each individual sequence, which suggests inefficient state space exploration. To account for the non-stationarity of our Gibbs sampler, we use Gelman’s split R̂ metric. As we are dealing with discrete state spaces, however, we cannot directly compute meaningful variances. Instead, we use the last layer embeddings of the forward policy’s neural network to compute the split R̂ metric. As we show in the following section, the proposed criterion notoriously improves learning convergence and exploration.
5
EXPERIMENTS
Our experimental campaign addresses the following research questions (RQs) regarding P-GFlowNets. RQ1 Under which conditions do P-GFlowNets improve convergence relative to a standard GFlowNet? RQ2 How effective is our approach for chain rejuvenation?
d = 16
d = 64
0.20
d = 128
L1 error
0.15 0.10 0.05
0.15
0.15
0.10
0.10
0.05 0
10 20 Elapsed Time (s)
0.05
30
0
50 100 Elapsed Time (s)
0
100 200 Elapsed Time (s)
(a) Set generation with log-additive rewards. n = 100
n = 400
n = 700
100 100
100 LPPD
0 0
0 −100 0
50 100 Elapsed Time (s)
−100
−100 0
50 100 Elapsed Time (s)
150
0
50 100 150 Elapsed Time (s)
(b) Bayesian variable selection. P-GFlowNet
c Zk,i =
GFlowNet
Figure 4: P-GFlowNets often converge faster to the target distribution. As the MDP horizon d grows (a), the relative speed up of our method increases. In contrast (b), our method grows more effective as the reward query cost decreases relative to the sampling cost—i.e., as the number n of samples for likelihood evaluation becomes smaller. R̂ threshold 1.1 2 ∞
R̂
L1 error
0.10 10
1
0.05 100 0
10 20 Elapsed time (s)
30
10
20 Elapsed time (s)
30
Figure 5: Chain rejuvenation critically accelerates learning. Left: Equation (13). Right: R̂ throughout training. As P-GFlowNets reduce the number of neural network forward passes for sample generation, we expect it to reduce training time when learning is bottlenecked by policy evaluation (RQ1). This is the case when trajectories are long, and sampling is expensive, and reward queries are cheaper than multiple model inferences. Our experiments confirm this intuition (see Figures 4 and 6). We also show that chain rejuvenation does not only improve state space exploration, but significantly accelerates learning convergence (RQ2). All in all, our observations position P-GFlowNets as a compute-efficient and principled algorithm for GFlowNet training. We provide further details and computer code for our experiments in the supplement. Set generation with log-additive rewards. Our first task consists of generating fixed-size subsets of a given set C = {1, . . . , K}. Let X := {s ⊆ C : |s| = d} be the space of d-sized subsets of C. Similarly, S := {s ⊆ C : |s| < d}, and so = ∅ be the initial state. Given a utility function u : C → R+ , we define the target distribution R : X → R+ as X log R(x) = log u(c). c∈x
This function has several important features. First, despite being factorizable into x’s components, it cannot be arbitrarily well-approximated by a mean-field variational approximation. Second, the function R(x) can be naturally extended to arbitrary d (state graph’s diameter) and K (state graph branching’s factor) with d ≤ K. Third, if 1s denotes s’sPindicator function, both the partition P function Z := x∈X R(x) and the marginals pc := Z1 s∈X R(x) · 1s (c) for c ∈ C can be efficiently computed in O(K · d2 ) through the following dynamic programming algorithm, which may be of independent interest for GFlowNet evaluation. Let firstk (S) denote the first k elements of any S ⊆ C in a fixed order, and define X
Y
X
u(e) and Zk,i =
s⊆firstk (C\{c}) e∈s |s|=i
Y
u(e)
s⊆firstk (C) e∈s |s|=i
for k ∈ {1, . . . , K} and i ∈ {1, . . . d}. It should be Zc c c clear that if Pk,i = u(c) · Zk−1,i then pc = PK,d . Also, k,i Zk+1,i = u(k + 1) · Zk,i−1 + Zk,i , with a similar recurc rence equation for Zk,i obtained by replacing u(k+1) by the appropriate (k +1)-th element of C \{c}. As such, we gauge the accuracy of a trained GFlowNet on this task by drawing N independent states {{x1 , . . . , xN }} ⊆ X and computing 1 X 1 |pc − p̂c | with p̂c = |C| N c∈C
X
1xn (c).
(13)
1≤n≤N
We consider (d, K) ∈ {(32, 64), (64, 128), (128, 256)} and log u(c) ∼ N (0, 1) drawn from a standard Gaussian for c ∈ C. Then, we track the reduction in the above metric in terms of wall-clock time in Figure 4. As d increases, the runtime gap between P-GFlowNet and a standard GFlowNet widens. Additionally, Figure 5 shows our rejuvenation approach is paramount for speeding up training. There, we consider (d, K) = (64, 128) and the conditions R̂ > α for α = 1.1 (default), α = 2, and α = ∞ (i.e., no rejuvenation). Bayesian variable selection. When reward query costs offset the sampling overhead, we expect that allocating more compute per sample should speed up learning convergence. Indeed, this has been empirically observed by Madan et al. [2025], Kim et al. [2025], Dall’Antonia et al. [2026]. To understand how this behavior affects the training of P-GFlowNets, we consider the problem of Bayesian variable selection [George and McCulloch, 1993] with progressively larger sample sizescorresponding to increasingly expensive-to-evaluate posterior distributions. Let X ∈ Rn×K and y ∈ Rn be a dataset with n samples and K-dimensional features, and denote by XF ∈ Rn×|F | the column-filtered data with F ⊆ {1, . . . , K}. Also, let In (resp. I|F | ) is n-dimensional (resp. |F |-dimensional) identity matrix, F ∼ Multinomial({1, . . . , K}, π) indicates each F is sampled with probability π |F | (1 − π)K−|F | for
Average Top 10 log-rewards
d = 100
d = 600
30
75
20
50 0
25 50 75 Elapsed Time (s)
d = 900
125
100
40
100 75 50 0
100 200 Elapsed Time (s)
0
100 200 300 Elapsed Time (s)
Average Top 10 log-rewards
(a) Ising model. d = 32
d = 128
d = 256
20
15
We next consider the tasks of Ising model simulation, also present in [Liu et al., 2024], and of bit generation, a common testbed for GFlowNets [e.g., Viviano et al., 2023]. For each example, X is the space of d-sized sequences with elements from {−1, 1} and {1, 0}, respectively. As is standard practice Malkin et al. [2022], Pan et al. [2023], Kim et al. [2025], we evaluate a model by measuring the average log reward for the top 10 most valuable samples found throughout training.
16 14
18
15 13
14 16 0
20 40 Elapsed Time (s)
60
0
50 100 Elapsed Time (s)
0
100 200 Elapsed Time (s)
(b) Bit sequences. P-GFlowNet
GFlowNet
Figure 6: P-GFlowNets improve exploration. We show the average log-reward of the 10 most rewarding states found in training. As in Figure 4, the runtime gains from P-GFlowNets as the MDP horizon (d) grows. π ∈ [0, 1]. We then consider the linear model y|β, F ∼ N (XF βF , η 2 In ) with
βF |F ∼ N (0, ν 2 I|F | ) and F ∼ Multinomial([K], π). We assume π, η > 0, and ν > 0 are known. Under these conditions, the marginal posterior distribution over the set F is log R(F ) = log N (y|0, ν 2 XF XTF +η 2 In )+|F | log
π , 1−π
with N (y|µ, Σ) denoting the density of a Gaussian distribution with mean µ and covariance Σ. Clearly, evaluating R(F ) costs O(K · n2 ). To assess a GFlowNet, we evaluate the log-predictive posterior density (LPPD) of a held-out dataset (X⋆ , y⋆ ) throughout training. Given samples {{F1 , . . . , FN }}, we approximate the LPPD as LPPD = log
1 N
X 1≤n≤N
p(y⋆ |X⋆ , y, F ),
with p(y |y, F ) = N (y |XF µF , X⋆F ΣF XT + η 2 In ) and µF = η −2 ΣF XTF y and ΣF = (η −2 XTF XF + ν −2 I|F | )−1 ⋆
⋆
as βF posterior’s mean and covariance given F , respectively. In practice, each row of XF is independently sampled from N (0, Γ), with 0 ∈ RK and Γ ∈ RK×K and Γij = γ |i−j| for γ = 0.8. Under this model, X’s columns are highly correlated, and the posterior distribution over F is multimodal. Notably, Figure 4 shows that P-GFlowNets consistently outperform a standard GFlowNet when n small. However, as n grows, the posterior evaluation cost outpaces that of trajectory sampling, and the computational benefits from our model dwindle. As explained above, this confirms our initial assumptions regarding P-GFlowNets.
Ising model. Simply put, let J ∈ Rd×d and h ∈ Rd . We define an energy function as E(x) = − 12 xT Jx − hT x, and p(x) ∝ exp {−E(x)/β } as the probability of a configuration x ∈ {−1, 1}d under a temperature β > 0. In Figure 6, we consider d ∈ {100, 300, 900}. We observe that P-GFlowNets drastically improve exploration when trajectories are long and sampling is consequently expensive. Bit sequences. As in Malkin et al. [2022], Tiapkin et al. [2024], we let M ⊆ {1, 0}d be a set of modes, and define β log R(x) = 1 − minm∈M ρ(m,x)/d, with ρ(x, m) as edit distance between x and m and β > 0 is a temperature parameter. We consider d ∈ {64, 128, 256} and β = 1/20. As we consistently observed in prior experiments, Figure 6 shows P-GFlowNets significantly speed up the discovery of high-reward states as the MDP horizon grows.
6
DISCUSSION
We showed that MaMs [Liu et al., 2024], which were previously thought to be distinct from GFlowNets, can be seen as an instantiation of a conditional GFlowNet using a persistent Gibbs sampler for exploration during training. Based on this, we also demonstrated this strategy generalizes beyond autoregressive modelling, for which MaMs were originally designed, while maintaining its computational benefits. Our experiments highlighted that the resulting method, called Particle GFlowNets, significantly accelerated learning convergence in terms of wall-clock time when compared against conventional GFlowNet training algorithms. From a broader perspective, our work (esp. Propositions 3.3 and 4.1) strengthens the connection between GFlowNets and Markov chain methods, which was also formally studied by Deleu and Bengio [2023]. This raises several questions. How to optimally decide when to rejuvenate the persistent Gibbs sampler? As noted in Silva et al. [2025a], diagnosing GFlowNets is strikingly difficult; can we draw inspirations from the MCMC literature to properly assess the distributional accuracy of GFlowNets? Successful Markov samplers, such as Langevin dynamics-based methods [Welling and Teh, 2011, Girolami and Calderhead, 2011], rely on simulating a latent dynamics for each transition of the underlying stochastic process; is such a technique extensible to GFlowNets, and when does it accelerate training? We believe these to be interesting directions for future research.
Acknowledgements DM acknowledges the support of the Fundação Carlos Chagas Filho de Amparo à Pesquisa do Estado do Rio de Janeiro (FAPERJ) (SEI-260003/020348/2025, SEI260003/020694/2025) and the Conselho Nacional de Desenvolvimento Científico e Tecnológico (CNPq) (404336/20230, 305692/2025-9, 445170/2024-7). References Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. In NeurIPS (NeurIPS), 2021. Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. Journal of Machine Learning Research (JMLR), 2023. James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of statistical software, 2017. Sanghyeok Choi, Sarthak Mittal, Víctor Elvira, Jinkyoo Park, and Nikolay Malkin. Reinforced sequential monte carlo for amortised sampling, 2025. URL https:// arxiv.org/abs/2510.11711. Pedro Dall’Antonia, Tiago da Silva, Daniel Csillag, Salem Lahlou, and Diego Mesquita. Avoid what you know: Divergent trajectory balance for gflownets, 2026. URL https://arxiv.org/abs/2602.17827. Tristan Deleu and Yoshua Bengio. Generative flow networks: a markov chain perspective, 2023. Tristan Deleu, António Góis, Chris Chinenye Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. In UAI, 2022. Tristan Deleu, Mizu Nishikawa-Toomey, Jithendaraa Subramanian, Nikolay Malkin, Laurent Charlin, and Yoshua Bengio. Joint Bayesian inference of graphical structure and parameters with a single generative flow network. In Advances in Neural Processing Systems (NeurIPS), 2023. Walter M. Fitch. Toward defining the course of evolution: minimum change for a specific tree topology. Systematic Zoology, 20(4):406–416, 1971.
Andrew Gelman and Donald B. Rubin. Inference from iterative simulation using multiple sequences. Statistical Science, 1992. Stuart Geman and Donald Geman. Stochastic relaxation, gibbs distributions, and the bayesian restoration of images. IEEE Transactions on Pattern Analysis and Machine Intelligence, PAMI-6(6), November 1984. doi: 10.1109/tpami.1984.4767596. Edward I George and Robert E McCulloch. Variable selection via gibbs sampling. Journal of the American Statistical Association, 88(423):881–889, 1993. Josiah Willard Gibbs. Elementary Principles in Statistical Mechanics: Developed with Especial Reference to the Rational Foundation of Thermodynamics. Charles Scribner’s Sons, New York, 1902. Reprinted by Dover Publications (1960) and others. Mark Girolami and Ben Calderhead. Riemann manifold langevin and hamiltonian monte carlo methods. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 2011. Timofei Gritsaev, Nikita Morozov, Sergey Samsonov, and Daniil Tiapkin. Optimizing backward policies in gflownets via trajectory likelihood maximization, 2025. URL https://arxiv.org/abs/2410.15474. Geoffrey E. Hinton. Training products of experts by minimizing contrastive divergence. Neural Comput., 14 (8):1771–1800, August 2002. ISSN 0899-7667. doi: 10.1162/089976602760128018. URL https://doi. org/10.1162/089976602760128018. Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, and et al. Amortizing intractable inference in large language models, 2023a. Edward J. Hu, Nikolay Malkin, Moksh Jain, Katie Everett, Alexandros Graikos, and Yoshua Bengio. Gflownet-em for learning compositional latent variable models, 2023b. URL https://arxiv.org/abs/2302.06576. Rui Hu, Yifan Zhang, Zhuoran Li, and Longbo Huang. Beyond squared error: Exploring loss design for enhanced training of generative flow networks, 2024. URL https://arxiv.org/abs/2410.02596. Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcia, Jarrid Rector-Brooks, Bonaventure F. P. Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological sequence design with GFlowNets. In International Conference on Machine Learning (ICML), 2022.
Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernandez-Garcia, and Yoshua Bengio. Gflownets for ai-driven scientific discovery. Digital Discovery, 2023a. Moksh Jain, Sharath Chandra Raparthy, Alex HernandezGarcia, Jarrid Rector-Brooks, Yoshua Bengio, Santiago Miret, and Emmanuel Bengio. Multi-objective GFlowNets. In International Conference on Machine Learning (ICML), 2023b. Hyosoon Jang, Minsu Kim, and Sungsoo Ahn. Learning energy decompositions for partial inference in GFlownets. In The Twelfth International Conference on Learning Representations, 2024. Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https://kellerjordan.github. io/posts/muon/. Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local search gflownets. arXiv preprint arXiv:2310.02710, 2023. Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local search gflownets, 2024. URL https://arxiv. org/abs/2310.02710. Minsu Kim, Sanghyeok Choi, Taeyoung Yun, Emmanuel Bengio, Leo Feng, Jarrid Rector-Brooks, Sungsoo Ahn, Jinkyoo Park, Nikolay Malkin, and Yoshua Bengio. Adaptive teachers for amortized samplers. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum? id=BdmVgLMvaf. Aya Laajil, Abduragim Shtanchaev, Sajan Muhammad, Eric Moulines, and Salem Lahlou. Curriculum-augmented gflownets for mrna sequence generation, 2025. URL https://arxiv.org/abs/2510.03811. Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-García, Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of continuous generative flow networks. In ICML, volume 202 of Proceedings of Machine Learning Research, pages 18269–18300. PMLR, 2023. Elaine Lau, Nikhil Vemgal, Doina Precup, and Emmanuel Bengio. Dgfn: Double generative flow networks, 2023. URL https://arxiv.org/abs/2310.19685. Alexander K. Lew, Monica Agrawal, David Sontag, and Vikash K. Mansinghka. Pclean: Bayesian data cleaning at scale with domain-specific probabilistic program-
ming, 2022a. URL https://arxiv.org/abs/ 2007.11838. Alexander K. Lew, Marco Cusumano-Towner, and Vikash K. Mansinghka. Recursive monte carlo and variational inference with auxiliary variables, 2022b. URL https: //arxiv.org/abs/2203.02836. Sulin Liu, Peter J Ramadge, and Ryan P Adams. Generative marginalization models. In International Conference on Machine Learning (ICML) , 2024. Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Cristian Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning gflownets from partial episodes for improved convergence and stability. In International Conference on Machine Learning, 2022. Kanika Madan, Alex Lamb, Emmanuel Bengio, Glen Berseth, and Yoshua Bengio. Towards improving exploration through sibling augmented GFlownets. In The Thirteenth International Conference on Learning Representations, 2025. Idriss Malek, Aya Laajil, Abhijith Sharma, Eric Moulines, and Salem Lahlou. Loss-guided auxiliary agents for overcoming mode collapse in gflownets, 2025. URL https://arxiv.org/abs/2505.15251. Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlownets. In NeurIPS (NeurIPS), 2022. Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. International Conference on Learning Representations (ICLR), 2023. Sean P Meyn and Richard L Tweedie. Markov Chains and Stochastic Stability. Cambridge University Press, Cambridge, 2nd edition, 2009. Cambridge Mathematical Library. Art B. Owen. Monte Carlo theory, methods and examples. 2013. Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GFlowNets with local credit and incomplete trajectories. In International Conference on Machine Learning (ICML), 2023. Ling Pan, Moksh Jain, Kanika Madan, and Yoshua Bengio. Pre-training and fine-tuning generative flow networks. In The Twelfth International Conference on Learning Representations, 2024.
Julien Roy, Pierre-Luc Bacon, Christopher Pal, and Emmanuel Bengio. Goal-conditioned gflownets for controllable multi-objective molecular design. arXiv preprint arXiv:2306.04620, 2023.
Max Welling and Yee Whye Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th International Conference on Machine Learning (ICML-11), 2011.
Mark J Schervish. Theory of statistics. Springer Science & Business Media, 2012.
David W Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan. Robust scheduling with gflownets. In International Conference on Learning Representations (ICLR), 2023a.
Max W. Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso Biancalani. Towards understanding and improving gflownet training. In International Conference on Machine Learning, 2023. Tiago Silva, Rodrigo Barreto Alves, Eliezer de Souza da Silva, Amauri H Souza, Vikas Garg, Samuel Kaski, and Diego Mesquita. When do GFlownets learn the right distribution? In The Thirteenth International Conference on Learning Representations, 2025a.
Dinghuai Zhang, Hanjun Dai, Nikolay Malkin, Aaron Courville, Yoshua Bengio, and Ling Pan. Let the flows tell: Solving graph combinatorial optimization problems with gflownets. In NeurIPS (NeurIPS), 2023b. Ni Zhang and Zhiguang Cao. Hybrid-balance GFlownet for solving vehicle routing problems. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
Tiago Silva, Amauri H Souza, Omar Rivasplata, Vikas Garg, Samuel Kaski, and Diego Mesquita. Generalization and distributed learning of GFlownets. In The Thirteenth International Conference on Learning Representations, 2025b.
Ming Yang Zhou, Zichao Yan, Elliot Layne, Nikolay Malkin, Dinghuai Zhang, Moksh Jain, Mathieu Blanchette, and Yoshua Bengio. PhyloGFN: Phylogenetic inference with generative flow networks. In The Twelfth International Conference on Learning Representations, 2024.
Tiago Silva, Bruna Bazaluk, Eliezer da Silva, António Góis, Salem Lahlou, Dominik Heider, Samuel Kaski, Diego Mesquita, and Adele Ribeiro. Expert-aided causal discovery of ancestral graphs. SSRN, 01 2026. doi: 10.2139/ssrn.6074306.
Mingyang Zhou, Zichao Yan, Elliot Layne, Nikolay Malkin, Dinghuai Zhang, Moksh Jain, Mathieu Blanchette, and Yoshua Bengio. Phylogfn: Phylogenetic inference with generative flow networks, 2023.
Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry Vetrov. Generative flow networks as entropyregularized rl, 2024. Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In Proceedings of the 25th international conference on Machine learning, pages 1064–1071, 2008. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. URL https://arxiv.org/abs/1706.03762. Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, Alexandre Adam, Jarrid Rector-Brooks, Yoshua Bengio, Glen Berseth, and Nikolay Malkin. Amortizing intractable inference in diffusion models for vision, language, and control, 2024. URL https://arxiv.org/abs/2405. 20971. Joseph D Viviano, Omar G Younis, Sanghyeok Choi, Victor Schmidt, Yoshua Bengio, and Salem Lahlou. torchgfn: A pytorch gflownet library. arXiv e-prints, pages arXiv– 2305, 2023.
Yiheng Zhu, Jialu Wu, Chaowen Hu, Jiahuan Yan, ChangYu Hsieh, Tingjun Hou, and Jian Wu. Sample-efficient multi-objective molecular optimization with gflownets, 2023.
Particle GFlowNets: Rethinking Generative Marginalization Models (Supplementary Material)
Tiago da Silva1
Diego Mesquita2
Salem Lahlou1
1
2
A
MBZUAI School of Applied Mathematics, Getulio Vargas Foundation
SUMMARY & PSEUDOCODE Table 1: Summary of our main theoretical results. Result
Name
Statement
Prop. 3.1
Flow = Marginal
The flow at any partial state equals the sum of rewards over all completions consistent with that state.
Cor. 3.1
σ-Independence of Flow
The optimal flow does not depend on the variable ordering σ used to construct the DAG.
Cor. 3.2
σ-Independence of Policy
The forward policy is likewise permutation-independent, so the network need not receive σ as input.
Prop. 3.2
MaMs ≡ PC-GFlowNets
Consistent MaMs and PC-GFlowNets satisfying detailed balance are equivalent.
Prop. 3.3
Consistency Error = DB Loss
The MaM consistency error is an unbiased, O(1)-cost estimator of the weighted detailed balance training loss.
Prop. 4.1
P-GFlowNet Ergodicity
The persistent Markov chain induced by a P-GFlowNet satisfying detailed balance converges to π ∝ R.
The main purpose of our work is to show that MaMs and GFlowNets, previously thought to be distinct, are the same—with MaM differing from conventional GFlowNet implementations solely in how states are sampled during training. Building on this equivalence, we introduce P-GFlowNets, demonstrating its correctness and computational efficiency as a sampling model for discrete, compositional spaces. In this context, Table 1 summarizes our central claims regarding the equivalence between P-GFlowNets and MaMs (Propositions 3.1 and 3.2, Corollaries 3.1 and 3.2) and the correctness of P-GFlowNets (Propositions 3.3 and 4.1). We also provide a pseudocode description of P-GFlowNets in Algorithms 1 and 2. Algorithm 1: Transition kernel for Algorithm 2. Procedure TransitionKernel(x) /* One step of the persistent Gibbs chain (used to produce y (b) in L̂C ) */ 2 Sample s ∼ pB (· | x); // remove one component via backward policy 3 Sample x′ ∼ pF (· | s); // add one component via forward policy 4 return x′ ;
1
Algorithm 2: P-GFlowNet Training Loop Input: Reward function R, batch size B, GR threshold R̂thr , weight vector w (truncated Poisson, mean log d) Output: Trained parameters θ = (pF , pB , F ) /* Initialisation */ (b) B 1 Sample particles {x }b=1 ∼ pF (· | so ) via MDP simulation; 2 while not converged do /* Complete loss L̂C (terminal / near-terminal transitions) */ 3 for b = 1 to B do 4 s(b) ∼ pB (· | x(b) ); // remove one component (b) (b) 5 y ∼ pF (· | s ); // add one component 2 (b) (b) (b) F (s ) pF (s , y ) 1 PB log 6 L̂C ← ; B b=1 R(y (b) ) pB (y (b) , s(b) ) /* Intermediate loss L̂I (interior transitions, depth sampled from w) */ 7 for b = 1 to B do 8 Sample depth k (b) ∼ Cat(w); // truncated Poisson, mean log d (b,k) k(b) (b) 9 s ∼ pB (· | x ); // remove k components 10 s′(b,k) ∼ pF (· | s(b,k) ); // add one component 2 1 PB F (s(b,k) ) pF (s(b,k) , s′(b,k) ) 11 L̂I ← ; b=1 log B F (s′(b,k) ) pB (s′(b,k) , s(b,k) ) /* Gradient step */ 12 L̂P ← L̂C + L̂I ; 13 θ ← θ − η∇θ L̂P (θ); /* Particle update via Gibbs step */ 14 for b = 1 to B do 15 x(b) ← y (b) ; // advance persistent chain /* Rejuvenation (Gelman-Rubin criterion) Compute split R̂ from last-layer embeddings of pF over the 512 most recent states; if R̂ > R̂thr then Resample all particles: x(b) ∼ pF (· | so ) for b = 1, . . . , B;
16 17 18 19
*/
return θ;
B
PROOFS
B.1
PROOF OF PROPOSITION 3.1
We first show that F σ (so ) equals the partition function of Z and, in particular, does not depend on σ. To see this, notice that Equation (9) implies F σ (so )pσF (so , τσ,x ) = F σ (x) := R(x) (14) for the only trajectory τσ,x starting at so and finishing at x ∈ X with positive probability under pσF (as explained earlier, pτB (s, ·) is either 1 or 0). Also, it should be clear that X x∈X
pσF (so , τσ,x ) =
X τ ∈so ⇝X
pσF (so , τ ) = 1,
as there is only one P trajectory from so to each x and, separating the sum according to the terminal state x ∈ X of each τ , this is exactly x∈X p⊤ (x) for the p⊤ introduced in Equation (2). Then, when we sum Equation (14) over x ∈ X , F σ (so ) =
X
R(x) := Z.
x∈X
This shows F σ (so ) does not depend on σ. Similarly, to show that Equation (10) is satisfied, we proceed by induction in i. (i) For this, let T(σ, s(i) ) = {x ∈ X : xσ([i]) = sσ([i]) }. Then, for i = d, s(d) ∈ X , T(σ, s(d) ) = {s(d) } and F σ (s(d) ) = R(x) P (i+1) by definition. Assume, for i < d, that F (s(i+1) ) = . Then, x∈T(σ,s(i+1) ) R(x) for each (i + 1)-sized object s σ summing over the support of pF (si , ·) in the DB condition in Equation (9), we observe that X
F σ (s(i) ) =
k∈[K]
F σ (s(i) ∪ {(k, σ(i + 1))}).
Clearly, s(i),k = s(i) ∪ {(k, σ(i + 1))} corresponds to an (i + 1)-sized sequence. Also, the sets T(σ, s(i),k ) are ′ disjoint for k ∈ [K] as, if x ∈ T(σ, s(i),k ) ∩ T(σ, s(i),k ), then xσ(i+1) = k and xσ(i+1) = k ′ . Additionally, S T(σ, s(i) ) = k∈[K] T(σ, s(i),k ) (see Figure 2). By induction, F σ (s(i) ) =
X
F σ (s(i),k )
k∈[K]
=
X
X
R(x) =
k∈[K] x∈T(σ,s(i),k )
X
R(x).
x∈T(σ,s(i) )
In particular, F σ (s(i) ) = F σ (so )
X x∈T(σ,s
R(x) , Z (i) )
which is exactly the marginal distribution of s(i) under R. B.2
PROOF OF COROLLARY 3.1
This follows from Proposition 3.1. In fact, notice that T(σ1 , s1,(i) ) = T(σ2 , s2,(i) ), as if x ∈ T(σ1 , s1,(i) ), then 1,(i) 2,(i) xσ([i]) = sσ([i]) = sσ([i]) , and so x ∈ T(σ2 , s2,(i) ), and vice-versa. As a consequence, F σ1 (s1,(i) ) =
X x∈T(σ1
=
R(x)
,s1,(i) )
X
R(x) = F σ2 (s2,(i) ).
x∈T(σ2 ,s2,(i) )
B.3
PROOF OF COROLLARY 3.2
This follows from Corollary 3.1. As both pσF1 and pσF2 satisfy the DB condition, for j ∈ {1, 2}, σ
pFj (sj,(i) , sj,(i+1) ) =
F σj (sj,(i+1) ) , F σj (sj,(i) )
with sj,(i+1) = sj,(i) ∪ {(k, σj (i + 1))} for some k ∈ [K], and zero otherwise. As σ1 (≤ i + 1) = σ2 (≤ i + 1), Corollary 3.1 ensures that F σ1 (s1,(i) ) = F σ2 (s2,(i) ) and F σ1 (s1,(i+1) ) = F σ2 (s2,(i+1) ), and then pσF1 (s1,(i) , ·) = pσF2 (s2,(i) , ·).
B.4
PROOF OF PROPOSITION 3.2
This proposition follows directly from Corollaries 3.1 and 3.2, and the notational mapping discussed in Section 3.1. That said, we provide a comprehensive demonstration below. ( =⇒ ) Let (pθ , pϕ ) be a consistent MaM satisfying pθ (x) = π(x) for all x. Then, define a PC-GFlowNet (pσF , F σ ) as follows. For each J ⊆ {1, . . . , d}, let σJ be a permutation for which the first |J | elements are J . For i ∈ {1, . . . , d} \ J , i let σJ be a permutation for which the first |J | elements are J , and the (|J | + 1)-th element is i. σi
In this context, let F σJ (xJ ) = pθ (xJ ). Also, let I = J ∪ {i} and pFJ (xJ , xI ) = pϕ (xI |xJ ), with xI and xJ as in Equation (11). By Corollaries 3.1 and 3.2, the above construction does not depend on the choice of permutation σ. Consequently, (pσF , F σ ) is equivalent to the MaM (pθ , pϕ ). i i ( ⇐= ) Conversely, let (pσF , F σ ) be a PC-GFlowNet abiding by the DB condition. Let σJ , σJ , and I be defined as above, and let pθ (xJ ) = F σJ (xJ ) and pϕ (xI |J ) = pF (xJ , xI ). i By Corollaries 3.1 and 3.2, again, the definition above does not depend on permutation σJ and σJ , as long as they satisfy the property that the first |J | elements are J , and the (|J | + 1)-th is i.
Consequently, (pθ , pϕ ) implements the PC-GFlowNet (pσF , F σ ). This shows the equivalence. B.5
PROOF OF PROPOSITION 3.3
To see this, let x be τ ’s terminal state. First, notice that pσE simply denotes a(exploratory) policy such that pE (s, ·) and pσF (s, ·) have the same support for each s ∈ S. By the definition of PC GFlowNets, each τ has size d and si = xσ([i]) for i ∈ {0, . . . , d}. For conciseness, define ∆(x, σ, i) :=
2 F (xσ([i−1]) )pσF (xσ([i]) |xσ([i−1]) ) log . F (xσ([i]) )
Then, the LHS of Proposition 3.3 can be written as X E wi ∆(x, σ, i) = σ τ ∼pE (so ,·)
1≤i≤d
E
[∆(x, σ, i)] ,
x∼q,i∼Cat(w)
as the marginal of pσE matches q. This is exactly the inner expectation of Proposition 3.3’s RHS. B.6
PROOF OF PROPOSITION 4.1
To ensure that {xt } is ergodic, we show that (i) it is irreducible (i.e., every state is reachable from every other state), (ii) aperiodic (i.e., the chain does not return to xt at regular intervals), and (iii) stationary with respect to π [Meyn and Tweedie, 2009]. By definition of both X and S, and since neither pF or pB are degenerate (i.e., they do not assign zero probability to valid transition), the chain {xt } is irreducible. Additionally, for any t, we can return to xt with positive probability after any number of steps. Hence, the chain is aperiodic. To see that the chain is stationary with respect to π, we first recall that the detailed balance implies F (s)pF (s, x) = R(x)pB (x, s) for any x ∈ X and s ∈ S. Denote by κ : X × X → R+ the transition kernel of {xt }. Our objective is to show that R(x)κ(x, x′ ) = R(x′ )κ(x′ , x),
which implies {xt } is stationary with respect to π(x) ∝ R(x). For this, we notice that X κ(x, x′ ) = pB (x, s)pF (s, x′ ). s∈S
Hence, R(x)κ(x, x′ ) =
X
R(x)pB (x, s)pF (s, x′ )
s∈S
=
X
F (s)pF (s, x)pF (s, x′ )
s∈S
=
X
F (s)pF (s, x)pF (s, x′ )
s∈S
=
X
R(x′ )pB (x′ , s)pF (s, x)
s∈S
= R(x′ )κ(x′ , x); we highlight in teal and in blue the terms for which we apply the DB condition. This shows {xt } is stationary with respect to π. Taken together, our results ensure {xt } is ergodic with respect to π.
C
RELATED WORKS
GFlowNets [Bengio et al., 2021, Lahlou et al., 2023, Bengio et al., 2023] are a topic of major interest in the probabilistic modelling literature, providing a clear framework for reasoning about complex hierarchical variational approximations of discrete stochastic models [Malkin et al., 2023, Choi et al., 2025]. The central challenge hampering GFlowNets’ broader applicability, in our opinion, is that they are often notoriously difficult to train. A prominent research direction for mitigating this problem, in fact, has been the development of effective learning objectives [Madan et al., 2022, Malkin et al., 2022, Hu et al., 2024, Pan et al., 2024, 2023] that accelerate training convergence. As learning efficiency is characterized by not only the objective function, but also by which samples are observed throughout training, another significant recent line of research has explored meta-heuristic approaches for enhanced state space exploration [Lau et al., 2023, Kim et al., 2025, Madan et al., 2025, Malek et al., 2025, Dall’Antonia et al., 2026]. In particular, our method operationally resembles Kim et al. [2024], Hu et al. [2023b], both of which introduce a sampling technique based on repeated applications of forward and backward policies; however, only P-GFlowNets asymptotically reduces the number of forward passes per gradient step as a function of the underlying MDP’s horizon. Nonetheless, despite significantly enhancing GFlowNet’s sample efficiency, these approaches often incur in a significant computational cost due to expensive exploration strategies requiring numerous policy evaluations per sample. In fact, our evaluation of SA-GFlowNets, Adaptive Teachers GFlowNets, and ACE Madan et al. [2025], Kim et al. [2025], Dall’Antonia et al. [2026], alongside Local Search GFlowNets Kim et al. [2023], suggested an increase of up to 4× in the per-sample processing time when compared against Madan et al. [2022], Malkin et al. [2022]’s traditional algorithms, which remain standard in the GFlowNet literature, e.g., [Hu et al., 2023a, Venkatraman et al., 2024, Zhou et al., 2023]. As in MaMs [Liu et al., 2024], our work addresses a fundamentally different problem: contrarily to prior approaches, we assume reward querying is cheap, and exploration cost is dominated by policy evaluation. This is the dominant setting for Bayesian inference over complex models in large state spaces, e.g., Deleu et al. [2022], wherein the policy pF is frequently parameterized with inference-intensive models such as a transformer Vaswani et al. [2023]. With this in mind, as discussed in Section 6, we believe the ideal algorithm would adaptively provision the appropriate amount of computation for each batch of samples. How such an approach would have to be implemented, however, remains open.
D
EXPERIMENTAL DETAILS
Our models were implemented in JAX [Bradbury et al., 2018]. All our experiments were run in a Apple MacBook Pro, Apple M4 (10-core: 4P + 6E), 16 GB unified memory, macOS 26.2 (Tahoe). In Section 5, the GFlowNet was trained by minimizing the TB loss Malkin et al. [2022]. We used Muon optimizer Jordan et al. [2024] for minimizing the learning objectives for both GFlowNets and P-GFlowNets. As in Madan et al. [2022], we used a learning rate of 10−3 for pF and of 10−2 for F and Z; pB was fixed as an uniform policy. For each experiment, we implemented a 2-layer MLP with 256 hidden units for parameterizing the policy network, and fixed B = 64 for the batch size. We trained each model with a fixed time budget shown in Figures 4 to 6. In particular, we also set ν = η = 10−1 for the task of Bayesian variable selection and,
Runtime for phylogenetic inference tasks
15
Runtime per 100 training steps (s)
Runtime per 100 training steps (s)
Runtime for set generation tasks TB-GFlowNet ACE AT-GFlowNet P-GFlowNet
10
5
0
25
50
75 d
100
100
60 40 20 0
125
TB-GFlowNet ACE AT-GFlowNet P-GFlowNet
80
12
14
16
18
20
d
Figure 8: Per-training step runtime for ACE, AT, TB, and P-GFlowNets (ours). By circumventing complete trajectory sampling during training, P-GFlowNets reduce per-step computation cost by several orders of magnitude. Convergence for set generation tasks (d = set size) d = 16
n = 100
d = 64
n = 400
0.2 0
0 −200 LPPD
L1 error
0.2 ACE AT-GFlowNets P-GFlowNets
0.1 0.1
−500
−400 −600
−1000
−800
−1500
−1000 0.0
0
5
10
15
20
25
0
0.0
30
0
20
40
60
80
100
120
20
40 60 80 Elapsed Time (s)
100
120
0
25
50 75 100 Elapsed Time (s)
125
150
Runtime (s)
(a) L1 error for set generation.
(b) LPPD for variable selection.
Figure 9: P-GFlowNets converge faster than AT, ACE, and TB GFlowNets in the set generation and variable selection tasks, achieving more accurate marginals (a) and larger log-predictive density (b), respectively, within a shorter time span. for the bit sequence generation, followed Malkin et al. [2022]’s approach for constructing the set of modes M. All plots show the average across 3 independent runs; error bars represent one standard deviation from the average. We periodically evaluated the R̂ asynchronously based on the 512 most recently observed states, following the implementation in Stan [Carpenter et al., 2017], rejuvenating the chain when the computed R̂ exceeded 1.1.
E
ADDITIONAL EXPERIMENTS Phylogenetic trees: Parsimony (d = number of nodes)
Parsimony score
d = 12
d = 15
−1210 −1540
−1220 −1230 0
20
40 60 Elapsed Time (s)
80
−1560
GFlowNets ACE AT-GFlowNets P-GFlowNets
0
25
50 75 Elapsed Time (s)
100
To further evaluate P-GFlowNets, we compare it against the recently proposed Adaptive Teachers (AT) Kim et al. [2025] and the Adaptive Complementary Exploration (ACE) Dall’Antonia et al. [2026] training algorithms. We also confirm P-GFlowNets’ effectiveness in the phylogenetic inference task Zhou et al. [2024].
125
Comparison against AT and ACE. In contrast to P-GFlowNets, whose focus lies on reducing the per-step computational cost for both Figure 7: P-GFlowNets finds more parsitraining and inference, AT and ACE aim at improving a GFlowNet’s monious phylogenetic trees than ACE, AT sample efficiency by training an exploratory model to search for highly GFlowNets and TB GFlowNets. informative (e.g., unvisited, high-probability under the target) regions during learning. In doing so, however, training cost often increases multifold due to the evaluation of both the exploratory and target GFlowNets on both forward and backward trajectories; see Figure 8. From this perspective, we found that P-GFlowNets converge significantly faster than AT and ACE with respect to wall-clock time. We show illustrate this in Figure 9 for both the set generation and variable selection tasks. In both cases, we followed the implementations of Dall’Antonia et al. [2026] for both AT and ACE GFlowNets. The main setting for which these artificial curiosity-inspired strategies would be appropriate, in our opinion, is when the reward function is extremely expensive to evaluate relatively to the policy network and this expensiveness cannot be reduced by exploiting the shared structure of adjacent states through caching, as explained next. Phylogenetic inference. To further evaluate P-GFlowNets, we also consider the problem of phylogenetic inference using the parsimony score as the log-reward function Fitch [1971]. We adopt the algorithm suggested by Zhou et al. [2024]. Given
a phylogenetic tree T with leaves L annotated with {A, T, C, G}, let en ∈ {1, 0}4 be the one-hot encoding of node n. For n (1) (2) (j) in L, en = 1 if n is annotated with A; en = 1, if T , and etc. Otherwise, en = 0. Then, by letting LC(n) and RC(n) denote the left and right child of n, the parsimony score is recursively defined as ( P4 eLC(n) ∧ eRC(n) if i=1 (eLC(n) ∧ eRC(n) )(i) ≥ 1 en = eLC(n) ∨ eRC(n) , otherwise. ParScore(n, T) = ParScore(LC(n), T) + ParScore(RC(n), T) + [en ̸= eLC(n) ∧ eRC(n) ],
with the boundary condition ParScore(n, T) = 0 for n ∈ L and u, v 7→ [u ̸= v] being 1 if u and v match element-wise and 0 otherwise. The parsimony score of a tree whose leaves are annotated with a sequence of {A, T, C, G} is the sum of the parsimony score of a tree whose leaves are solely annotated with each element in this sequence, corresponding to the usual bag-of-words assumption in phylogenetic statistical models. The intuition is that a tree is as parsimonious (having low parsimony score) as the number of disagreements between a parent and its children. We define R(T) = exp{−ParScore(r(T), T)} as our reward function, with r(T) as the root of T. Importantly, the modular nature of ParScore ensures that, in transitioning from T to T′ using P-GFlowNets’ backwardforward kernels, most of the computation required for ParScore(T′ ) can be reused from ParScore(T), reducing the cost of reward querying. Based on this, we compare P-GFlowNets against AT and ACE GFlowNets on their state space exploration capabilities during training in Figure 7. The initial exploration phase, highlighted as a horizontal line for P-GFlowNets, corresponds to the period before which samples are collected for computing R̂, which is used to decide whether the chain should be rejuvenated. As in the set generation and variable selection tasks, P-GFlowNets improve upon both AT and ACE given a similar wall-clock time budget. That said, extending P-GFlowNets to mixed, discrete and continuous, spaces—as often required in phylogenetic inference based on stochastic evolution models—remains an interesting venue for future research.