On Divergence Measures for Training GFlowNets

Authors: tiago  da
silva  eliezer de souza da silva
diego  mesquita
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 14:22:42

Abstract: Generative Flow Networks (GFlowNets) are amortized samplers of unnormalized distributions over compositional objects with applications to causal discovery, NLP, and drug design. Recently, it was shown that GFlowNets can be framed as a hierarchical variational inference (HVI) method for discrete distributions. Despite this equivalence, attempts to train GFlowNets using traditional divergence measures as learning objectives were unsuccessful. Instead, current approaches for training these models rely on minimizing the log-squared difference between a proposal (forward policy) and a target (backward policy) distributions. In this work, we first formally extend the relationship between GFlowNets and HVI to distributions on arbitrary measurable topological spaces. Then, we empirically show that the ineffectiveness of divergence-based learning of GFlowNets is due to large gradient variance of the corresponding stochastic objectives. To address this issue, we devise a collection of provably variance-reducing control variates for gradient estimation based on the REINFORCE leave-one-out estimator. Our experimental results suggest that the resulting algorithms often accelerate training convergence when compared against previous approaches. All in all, our work contributes by narrowing the gap between GFlowNet training and HVI, paving the way for algorithmic advancements inspired by the divergence minimization viewpoint.

Main paper: 
 introduction: 
 The approximation of intractable distributions is one of the central issues in machine learning and modern statistics [7, 35] . In reinforcement learning (RL), a recurring goal is to find a diverse set of high-valued state-action trajectories according to a reward function. This problem may be cast as sampling trajectories proportionally to the reward, which is generally an intractable distribution over the environment [3, 9, 38, 50] . Similarly, practical Bayesian inference and probabilistic models computations involve assessing intractable posterior distributions [36, 78, 102] . In the variational inference (VI) approach, circumventing this intractability involves searching for a tractable approximation to the target distribution within a family of parametric models. Conventionally, the problem reduces to minimizing a divergence measure, such as Kullback-Leibler (KL) divergence [7, 36, 92] or Renyi-α divergence [51, 70] , between the variational approximation and the target.In particular, Generative Flow Networks (GFlowNets) [3, 4, 48] are a recently proposed family of variational approximations well-suited for distribution over compositional objects (e.g., graphs and texts). GFlowNets have found empirical success within various applications from causal discovery [15, 16] , NLP [30] , and chemical and biological modeling [3, 32] . In a nutshell, a GFlowNet learns an iterative generative process (IGP) [26] over an extension of the target's support, which, for sufficiently expressive parameterizations of transition kernels, yields independent and correctly distributed samples [3, 48] . Remarkably, training GFlowNets typically consists of minimizing the log-squared difference between a proposal and target distributions over the extended space via SGD [4, 55] , contrasting with divergence-minimizing algorithms commonly used in VI [7, 72] .38th Conference on Neural Information Processing Systems (NeurIPS 2024). Indeed, Malkin et al. [56] suggests that trajectory balance (TB) loss training for GFlowNets leads to better approximations of the target distribution than directly minimizing the reverse and forward KL divergence, particularly in setups with sparser rewards. Nevertheless, as we highlight in Section 3, these results are a potential consequence of biases and high variance in gradient estimates for the divergence's estimates, which can be overlooked in the evaluation protocol reliant upon sparse target distributions. Therefore, in Section 5, we present a comprehensive empirical investigation of the minimization of well-known f -divergence measures (including reverse and forward KL), showing it is an effective procedure that often accelerates the training convergence of GFlowNets relative to alternatives. To achieve these results, we develop in Section 4 a collection of control variates (CVs) [63, 71] to reduce the variance without introducing bias on the estimated gradients, improving the efficiency of the optimization algorithms [77, 85] . In summary, our main contributions are:1. We evaluate the performance of forward and reverse KL- [47] , Renyi-α [74] and Tsallis-α [91] divergences as learning objectives for GFlowNets through an extensive empirical campaign and highlight that they frequently outperform traditionally employed loss functions. 2. We design control variates for the gradients of GFlowNets' divergence-based objectives. Therefore, it is possible to perform efficient evaluations of the optimization objectives using automatic differentiation frameworks [69] , and the resulting experiments showcase the significant reduction in the variance of the corresponding estimators. 3. We developed a theoretical connection between GFlowNets and VI beyond the setup of finitely supported measures [56, 112] , establishing results for arbitrary topological spaces.2 Revisiting the relationship between GFlowNets and VI Initially, we review Lahlou et al. [48] 's work on GFlowNets for distributions on topological spaces, a perspective applied consequentially to obtain the equivalence between GFlowNets training and VI divergence minimization in a more generic setting. Finally, we describe standard variance reduction techniques for solving stochastic optimization problems.Notations. Let (S, T ) be a topological space with topology T and Σ be the corresponding Borel σ-algebra. Also, let ν : Σ → R + be a measure over Σ and κ f , κ b : S × Σ → R + be transition kernels over S. For each (B 1 , B 2 ) ∈ Σ × Σ, we denote by ν ⊗ κ(B 1 , B 2 ) := B1 ν(ds)k(s, B 2 ).Likewise, we recursively define the product kernel as κ ⊗0 (s, •) = κ(s, •) and, for n ≥ 1, κ ⊗n (s, •) = κ ⊗n-1 (s, •) ⊗ κ for a transition kernel κ and s ∈ S. Note, in particular, that κ ⊗n is a function from S × Σ ⊗n+1 to R + , with Σ ⊗n+1 representing the product σ-algebra of Σ [1, 96] . Moreover, if µ is an absolutely continuous measure relatively to ν, denoted µ ≪ ν, we write dµ /dν for the corresponding density (Radom-Nikodym derivative) [1] . Furthermore, we denote by P(A) = {S : S ⊆ A} the power-set of a set A ⊂ S and by [d] = {1, . . . , d} the first d positive integers.GFlowNets. A GFlowNet is, in its most general form, built upon the concept of a measurable pointed directed acyclic graph (DAG) [48] , which we define next. Intuitively, it extends the notion of a flow network to arbitrary measurable topological spaces, replacing the directed graph with a transition kernel specifying how the underlying states are connected.Definition 1 (Measurable pointed DAG [48] ). Let ( S, T , Σ) be a measurable topological space endowed with a reference measure ν and forward κ f and backward κ b kernels. Also, let s o ∈ S and s f ∈ S be distinguished elements in S, respectively called initial and final states, and S = S \ {s f }.We assume {s f } is open. A measurable pointed DAG is then a tuple (S, T , Σ, κ f , κ b , ν) satisfying:1. (Terminality) If κ f (s, {s f }) > 0, then κ f (s, {s f }) = 1 ∀s ∈ S. Also, κ f (s f , •) = δ s f .  
divergence measures for learning gflownets: 
 in which P B is typically fixed and P F 's density p θ F is parameterized by θ.Renyi-α [74] and Tsallis-α [91] are families of statistical divergences including, as limiting cases, the widespread KL divergence (Section 3.2) [58] ; see Definition 5. These divergences have been successfully applied to both variational inference [51] and policy search for model-based reinforcement learning [18] . Moreover, as we highlight in Section 5, their performance is competitive with, and sometimes better than, traditional learning objectives for GFlowNets based on minimizing log-squared differences between proposal and target distributions.Definition 5 (Renyi-α and Tsallis-α divergences). Let α ∈ R. Also, let p F θ and p B be GFlowNet's forward and backward policies, respectively. Then, the Renyi-α divergence between P F and P B isR α (P F ||P B ) = 1 α -1 log P S p F θ (τ |s o ) α p B (τ ) 1-α κ f (s o , dτ ).Similarly, the Tsallis-α divergence between P F and P B isT α (P F ||P B ) = 1 α -1 P S p F θ (τ |s o ) α p B (τ ) 1-α κ f (s o , dτ ) -1 .From Definition 5, we see that both Renyi-α and Tsallis-α divergences transition from a masscovering to a mode-seeking behavior as α ranges from -∞ to ∞. Regarding GFlowNet-training, this flexibility suggests that R α and T α are appropriate choices both, e.g., for carrying out Bayesian inference [16] -where interest lies in obtaining an accurate approximation to a posterior distribution -, and for combinatorial optimization [109] -where the goal is to find a few high-valued samples. Additionally, the choice of α provides a mechanism for controlling which trajectories are preferentially sampled during training, with larger values favoring the selection of trajectories leading to highprobability terminal states, resembling the effect of ε-greedy [55] , thompson-sampling [73] , localsearch [40] , and forward-looking [65, 66] techniques for carrying out off-policy training of GFlowNets [56] .To illustrate the effect of α on the learning dynamics of GFlowNets, we show in Figure 1 an early stage of training to sample from a homogeneous mixture of Gaussian distributions by minimizing Renyi-α divergence for different values of α; see Section 5.1 for details on this experiment. At this stage, we note that the GFlowNet covers the target distribution's modes but fails to separate them when α is large and negative. In contrast, a large positive α causes the model to focus on a single high-probability region. Therefore, the use of an intermediate value for α = 0.5 culminates in a model that accurately approximates the target distribution. Also, our early experiments suggested the persistence of such dependence on α for diverse learning tasks, with α = 0.5 leading to the best results. Thus, we fix α = 0.5 throughout our experimental campaign. Importantly, we need only the gradients of R α and T α for solving the optimization problem in Equation 3 and, in particular, learning the target distribution's normalizing constant is unnecessary, as we underline in the lemma below. This property distinguishes such divergence measures from both TB and DB losses in Equation 1 and, in principle, simplifies the training of GFlowNets.Lemma 1 (Gradients for R α and T α ). Let θ be the parameters of p F θ in Definition 5 and, for τ ∈ P S , g(τ, θ) = ( p B (τ |x)r(x) /p F θ (τ |so;θ)) 1-α . The gradient of R α wrt θ is∇ θ R α (p F θ ||p B ) = E[∇ θ g(τ, θ) + g(τ, θ)∇ θ log p F θ (τ |s o ; θ)] (α -1)E[g(τ, θ)] ;the expectations are computed under P F . Analogously, the gradient of T α wrt θ is∇ θ T α (p F θ ||p B ) C = E[∇ θ g(τ, θ) + g(τ, θ)∇ θ log p F θ (τ |s o ; θ)] (α -1) ,in which C = denotes equality up to a multiplicative constant. Lemma 1 uses the REINFORCE method [97] to compute the gradients of both R α and T α , and we implement Monte Carlo estimators to approximate the ensuing expectations based on a batch of trajectories {τ 1 , . . . , τ N } sampled during training [56] . Also, note that the function g is computed outside the log domain and, therefore, particular care is required to avoid issues such as numerical underflow of the unnormalized distribution [3, 87] . In our implementation, we sample an initial batch of trajectories {τ i } N i=1 and compute the maximum of r among the sampled terminal states in log space, i.e., log r = max i log r(x i ). Then, we consider log r(x) = log r(x)log r as the target's unnormalized log density. In Section 4, we will consider the design of variance reduction techniques to decrease the noise level of gradient estimates and possibly speed up the learning process.The KL divergence [47] is a limiting member of the Renyi-α and Tsallis-α families of divergences, derived when α → 1 [70] , and is the most widely deployed divergence measure in statistics and machine learning. To conduct variational inference, one regularly considers both the forward and reverse KL divergences, which we review in the definition below. Definition 6 (Forward and reverse KL). The forward KL divergence between a target P B and a proposalP F is D KL [P B ||P F ] = E τ ∼P B (s f ,•) [log p B (τ ) /p F θ (τ |so)]. Also, the reverse KL divergence is defined by D KL [P F ||P B ] = E τ ∼P F (so,•) [log p F θ (τ |so) /p B (τ )].Remarkably, we cannot use a simple Monte Carlo estimator to approximate the forward KL due to the presumed intractability of P B (which depends directly on R). As a first approximation, we could estimate D KL [P B ||P F ] via importance sampling w/ P F as a proposal distribution as in [56] :EQUATIONand subsequently implement a REINFORCE estimator to compute ∇ θ D KL [P B ||P F ]. Nevertheless, as we only need the divergence's derivatives to perform SGD, we apply the importance weights directly to the gradient estimator. We summarize this approach in the lemma below.Lemma 2 (Gradients for the KL divergence). Let θ be the parameters of P F and s(τ ; θ) = log p F θ (τ |s o ; θ). Then, the gradient of D KL [P F ||P B ] relatively to θ satisfies Correspondingly, the gradient of∇ θ D KL [P F ||P B ] = E τ ∼P F (so,•) ∇ θ s(τ ; θ) + log p F θ (τ |s o ) p B (τ |x)r(x) ∇ θ s(τ ; θ)D KL [P B ||P F ] wrt θ is ∇ θ D KL [P B ||P F ] C = -E τ ∼P F (so,•) p F θ (τ |s o ) p B (τ |x)r(x) ∇ θ s(τ ; θ) .Crucially, choosing an appropriate learning objective is an empirical question that one should consider on a problem-by-problem basis -similar to the problem of selecting among Markov chain simulation techniques [27] . In particular, a one-size-fits-all solution does not exist; see Section 5 for a thorough experimental investigation. Independently of the chosen method, however, the Monte Carlo estimators for the quantities outlined in Lemma 2 are of potentially high variance and may require a relatively large number of trajectories to yield a reliable estimate of the gradients [97] . The following sections demonstrate that variance reduction techniques alleviate this issue.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: The statements of the propositions and lemma include the assumptions, as well as in the notations and definitions presented in Section 4, and proofs provided in Appendix C.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.• Theorems and Lemmas that the proof relies upon should be properly referenced. 
control variates for low-variance gradient estimation: 
 Control variates. We first review the concept of a control variate. Let f : P S → R be a real-valued measurable function and assume that our goal is to estimate E τ ∼π [f (τ )] according to a probability measure π on Σ P (see Section 2 to recall the definitions). The variance of a naive Monte Carlo estimator for this quantity is Varπ(f (τ )) /n. On the other hand, consider a random variable (RV) g : P S → R positively correlated with f and with known expectation E π [g(τ )]. Then, the variance of a naive Monte Carlo forEQUATIONwhich is potentially smaller than 1 n Var π (f (τ )) if the covariance between f and g is sufficiently large. Under these conditions, we choose the value of a that minimizes Equation 5 [94] , namely, a = Covπ(f (τ ),g(τ )) /Varπ(g(τ)). We then call the function g a control variate [63] . Also, although the quantities defining the best baseline a are generally unavailable in closed form, one commonly uses a batchbased estimate of Cov π (f (τ ), g(τ )) and Var π (g(τ )); the incurred bias is generally negligible relatively to the reduced variance [71, 79, 85] . For vector-valued RVs, we let a be a diagonal matrix and exhibit, in the next proposition, the optimal baseline minimizing the covariance matrix's trace.Proposition 2 (Control variate for gradients). Let f, g : P S → R d be vector-valued functions and π be a probability measure on P S . Consider a baseline a ∈ R and assume E π [g(τ )] = 0. Then,arg min a∈R Tr Cov π [f (τ ) -a • g(τ )]= E π [g(τ ) T (f (τ ) -E π [f (τ ′ )])] E π [g(τ ) T g(τ )] .Note that, when implementing the REINFORCE gradient estimator, the expectation we wish to estimate may be generally written asE P F (so,•) [∇ θ f (τ ) + f (τ )∇ θ log p F θ (τ )].For the second term, we use a leave-one-out estimator [85] ; see below. For the first term, we use ∇ θ log p F θ as a control variate, which satisfiesE P F (so,•) [∇ θ log p F θ (τ |s o ; θ)] = 0.Importantly, estimating the optimal baseline a ⋆ in Proposition 2 cannot be done efficiently due to the non-linear dependence of the corresponding Monte Carlo estimator on the sample-level gradients [2] ; i.e., it cannot be represented as a vector-Jacobian product, which is efficient to compute in reverse-mode automatic differentiation (autodiff ) frameworks [8, 69] . Consequently, we consider a linear approximation of both numerator and denominator defining a ⋆ in Proposition 2, which may be interpreted as an instantiation of the delta method [83, Sec. 7.1.3] . Then, given a batch {τ 1 , . . . , τ N } of trajectories, we instead useEQUATIONas the REINFORCE batch-based estimated baseline; ⟨•, •⟩ represents the inner product between vectors. Intuitively, the numerator is roughly a linear approximation to the covariance between ∇ θ log p F θ and ∇ θ f under P F . In contrast, the denominator approximately measures the variance of ∇ θ log p F θ , and ϵ > 0 is included to enhance numerical stability. As a consequence, for the reverse KL divergence, ∇ θ f (τ ) = ∇ θ log p F θ (τ ), â ≈ 1 and the term corresponding to the expectation of ∇ θ f (τ ) vanishes. We empirically find that this approach frequently reduces the variance of the estimated gradients by a large margin (see Figure 2 above and Section 5 below).Leave-one-out estimator.We now focus on obtaining a low-variance estimate ofE τ ∼P F (so,•) [f (τ )∇ θ log p F θ (τ )].As an alternative to the estimator of Proposition 2, Shi et al. [85] and Salimans and Knowles [82] proposed a sample-dependent baseline of the form a(τ i ) =1 N -1 1≤n≤N,n̸ =i f (τ n ) for i ∈ {1, . . . , N }. The resulting estimator, δ = 1 N N n=1   f (τ n ) - 1 N -1 N j=1,j̸ =i f (τ j )   ∇ θ log p F θ (τ n ), is unbiased for E [f (τ )∇ θ log p F θ (τ )] due to the independence between τ i and τ j for i ̸ = j. Strikingly, δ can be swiftly computed with autodiffEQUATIONwith sg as the stop-gradient operation (e.g., lax.stop_gradient in JAX [8] and torch.detach in PyTorch [69] ). Importantly, these techniques incur a minimal computational overhead to the stochastic optimization algorithms relative to the considerable reduction in variance they enact.Relationship with previous works. Importantly, Malkin et al. [56] used â = 1 N n f (τ n ) as baseline and an importance-weighted aggregation to adjust for the off-policy sampling of trajectories, introducing bias in the gradient estimates and relinquishing guarantees of the optimization procedure. A learnable baseline independently trained to match b was also considered. This potentially entailed the inaccurate conclusion that the TB and DB are superior to standard divergence-based objectives. Indeed, the following section underlines that such divergence measures are sound and practical learning objectives for GFlowNets for a range of tasks.Illustration of the control variates' effectiveness. We train the GFlowNets using increasingly larger batches of {2 i : i ∈ [ [5, 10] ]} trajectories with and without CVs. In this setting, Figure 2 showcases the drastic reduction in the variance, represented by the covariance matrix's trace, of the estimated learning objectives' gradients w.r.t. the model's parameters promoted by the CVs. Impressively, as we show in Figure 6 , this approach significantly increases the efficiency of the underlying stochastic optimization algorithm. See Section 5 and Appendix D for further details.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: Details are provided in Section 5, Appendix B and Appendix E.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.• If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.• Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
(finite absorption): 
 There is a N ∈ N such that κ ⊗N f (s, •) = δ s f for every s ∈ S. We designate the corresponding DAG as finitely absorbing.In this setting, the elements in the set X = {s ∈ S \ {s f } : κ f (s, {s f }) > 0)} are called terminal states. Illustratively, when S is finite and ν is the counting measure, the preceding definition corresponds to a connected DAG with an edge from s ∈ S to s ′ ∈ S iff κ f (s, {s ′ }) > 0, with condition 5) ensuring acyclicity and condition 2) implying connectivity. A GFlowNet, then, is characterized by a measurable pointed DAG, a potentially unnormalized distribution over terminal states X and learnable transition kernels on S (Definition 2). Realistically, its goal is to find an IGP over S which, starting at s o , samples from X proportionally to a given positive function.Definition 2 (GFlowNets [48] ). A GFlowNet is a tuple (G,P F ,P B ,µ) composed of a measurable pointed DAG G, a σ-finite measure µ ≪ ν, and σ-finite Markov kernels P F ≪ κ f and P B ≪ κ b , respectively called forward and backward policies.Training GFlowNets. In practice, we denote by p F θ : S × S → R + the density of P F relative to κ f , which we parameterize using a neural network with weigths θ. Similarly, we denote by p B the density of P B wrt k b . Our objective is, for a given target measure R ≪ µ on X with r = dR /dµ, estimate the θ for which the distribution over X induced by P F (s o , •) matches R, i.e., for every B ∈ Σ,n≥0 S n p ⊗n F θ (s o , s 1:n , s f )1 sn∈B κ ⊗n f (s o , ds 1:n ) = R(B) R(X ) .Importantly, the above sum contains only finitely many non-zero terms due to the finite absorption property of κ f . To ensure that p F θ abides by this equation, Lahlou et al. [48] showed it suffices that one of the next balance conditions are concomitantly satisfied by P F and P B .Definition 3 (Trajectory balance condition). For all n ≥ 0 and µ ⊗n -almost surely ∀s ) ∈ S 2 , u(s)p F θ (s, s ′ ) = u(s ′ )p B (s ′ , s) and u(x) = r(x) for x ∈ X .To enforce the trajectory balance (TB) or detailed balance (DB) conditions, we conventionally define a stochastic optimization problem to minimize the expected log-squared difference between the leftand right-hand sides of the corresponding condition under a probability measure ξ supported on an appropriate space [4, 15, 48, 49, 52, 55, 67] . For TB, e.g., we let ξ be defined on Σ ⊗N with support supp(µ ⊗N ). Then, we estimate the GFlowNet's parameters θ by minimizingEQUATIONwith τ = s o:N and h = max{i : s i ̸ = s f } being the last non-final state's index in τ . As denoted by the subscript on Z θ , using the TB loss incurs learning the target's normalizing constant. While tuning p B during training is also possible, the common practice is to keep it fixed. Henceforth, we will consider the measurable space of trajectories (P S , Σ P ), with P S = {(s, s 1 , . . . , s n , s f ) ∈ S n+1 × {s f } : 0 ≤ n ≤ N -1} and Σ P as the σ-algebra generated by N +1 n=1 Σ ⊗n . For notational convenience, we use the same letters for representing the measures and kernels of (S, Σ) and their natural product counterparts in (P S , Σ P ), which exist by Carathéodory extension's theorem [96] ; for example,ν(B) = ν ⊗n (B) for B = (B 1 , . . . , B n ) ∈ Σ ⊗n and p F θ (τ |s o ; θ) is the density of P ⊗n+1 F (s o , •) for τ = (s o , s 1 , . . . , s n , s f ) relatively to µ ⊗n .In this case, we will write τ for a generic element of P S and x for its terminal state (which is unique by Definition 1). For a comprehensive overview of GFlowNets, please refer to [48, 55] .GFlowNets and VI. GFlowNets can be interpreted as hierarchical variational models by framing the forward policy p F θ (τ |s o ; θ) in (P S , Σ P ) as a proposal to r(x) Z p B (τ |x). Malkin et al. [56] demonstrated that, for discrete target distributions, the TB loss in (1) aligns with the KL divergence in terms of expected gradients. Extending this, our Proposition 1 establishes that this relationship also holds for distributions over arbitrary topological spaces.Proposition 1 (TB loss-and KL divergence gradients for topological spaces). Let L T B (τ ; θ) = (log Zp F θ (τ |so;θ) /r(x)p B (τ |x))2 and p B (τEQUATIONwhereD KL [p F θ ||p B ] = E τ ∼P F (so,•) [log p F θ (τ |so;θ) /p B (τ )]is the KL divergence between P F and P B .This proposition shows that minimizing the on-policy TB loss is theoretically comparable to minimizing the KL divergence between P F and P B in terms of convergence speed. Since the TB loss requires estimating the intractable R(X ), the KL divergence, which avoids this estimation, can be a more suitable objective. Our experiments in Section 5 support this, with proofs provided in Appendix C. Extending this result to general topological spaces broadens the scope of divergence minimization strategies, extending guarantees for discrete spaces to continuous and mixed spaces. This generalization aligns with advances in generalized Bayesian inference [45] and generalized VI in function spaces [95] , via optimization of generic divergences. We make the method theoretically firm and potentially widely applicable by proving the equivalence in these broader contexts.Variance reduction. A naive Monte Carlo estimator for the gradient in Equation 2 has high variance [19] , impacting the efficiency of stochastic gradient descent [97] . To mitigate this, we use control variates-random variables with zero expectation added to reduce the estimator's variance without bias [63, 71] . This method, detailed in Section 4, significantly reduces noise in gradient estimates and pragmatically improves training convergence, as shown in the experiments in Section 5. Our experiments serve two purposes. Firstly, we show in Section 5.2 that minimizing divergencebased learning objectives leads to competitive and often better approximations than the alternatives based on log-squared violations of the flow network's balance. This underlines the effectiveness of well-established divergence measures for training GFlowNets [56, 112] . Secondly, we highlight in Section 5.3 that the reduction of variance enacted by our control variates critically accelerates the convergence of GFlowNets. We consider widely adopted benchmark tasks from GFlowNet literature, described in Section 5.1, contemplating both discrete and continuous target distributions. Please refer to Appendix B and Appendix E for additional information on the experimental setup.Below, we provide a high-level characterization of the generative tasks used for synthetic data generation and training. For a more rigorous description in the light of Section 2, see Appendix B. [56] , there is no single best loss function always conducting to the fastest convergence rate, and minimizing well-known divergence measures is often on par with or better than minimizing the TB loss in terms of convergence speed. Results were averaged across three different seeds. Also, we fix α = 0.5 for both Tsallis-α and Renyi-α divergences.Set generation [3, 34, 65, 66] . A state s corresponds to a set of size up to a given S and the terminal states X are sets of size S; a transition corresponds to adding an element from a deposit D to s. The IGP starts at an empty set, and the log-reward of a x ∈ X is d∈x f (d) for a fixed f : D → R.Autoregressive sequence generation [32, 55] . Similarly, a state is a seq. s of max size S and a terminal state is a seq. ended by an end-of-sequence token; a transition appends d ∈ D to s. The IGP starts with an empty sequence and, for x ∈ X , log r(x) = i=1...|x| g(i)f (x i ) for functions f, g.Bayesian phylogenetic inference (BPI) [111] . A state s is a forest composed of binary trees with labeled leaves and unlabelled internal nodes, and a transition amounts to joining the roots of two trees to a newly added node. Then, s is terminal when it is a single connected tree -called a phylogenetic tree. Finally, given a dataset of nucleotide sequences, the reward function is the unnormalized posterior over trees induced by J&C69's mutation model [37] and a uniform prior.Hypergrid Mixture of Gaussians (GMs) [48, 110] . The IGP starts at 0 ∈ R d and proceeds by sequentially substituting each coordinate with a sample from a real-valued distribution. For aK-component GM, the reward of x ∈ R d is defined as k= α k N (x|µ k , Σ k ) with α k ≥ 0 and k α k = 1.Banana-shaped distribution [57, 76] . We use the same IGP implemented for a bi-dimensional GM. For z ∈ R 2 , we set r(x) to a normal likelihood defined on a quadratic function of x, see Equation 8in the supplement. We use HMC samples as ground truth to gauge performance on this task.Next, we provide evidence that minimizing divergence-based objectives frequently leads to faster convergence than minimizing the standard TB loss [55] .Experimental setup. We compare the convergence speed in terms of the rate of decrease of a measure of distributional error when using different learning objectives for a GFlowNet trained to sample from each of the distributions described in Section 5.1. For discrete distributions, we adopt the evaluation protocols of previous works [3, 54, 55, 66] and compute the L 1 distance between the learned p T (x; θ) and target r(x), namely, x∈X |p T (x; θ)r(x) /Z|. To approximate p T , we use a Monte Carlo estimate of p T (x;θ) = E τ ∼P B (x,•) [ p F θ (τ |so;θ) /p B (τ |x)].For continuous distributions, we echo [48, 110] and compute Jensen-Shannon's divergence between P T (x; θ) and R(x): with M (B) = 1 /2 (P T (B) + R(B) /R(X )) being the averaged measure of P T and R and m its corresponding density relatively to the reference measure µ. Remarkably, for the GMs distribution, we can directly sample from the target to estimate D KL [R||M ] , and the autoregressive nature of the generative process ensures that p T (x) = p F θ (τ |s o ) for the unique trajectory τ starting at s o and finishing at x. Hence, we get an unbiased estimate of D KL [P T ||M ]. Finally, let X t be the first t terminal states encountered during training and {x (1) , . . . , x (k) } be an descending ordering of X t according to r. Then, we select a threshold τ ∈ R and an integer K to report NoM(X t ) = |{r(x) : r(x) ≥ τ ∧ x ∈ X t }|, called number of modes, and TopK(X t ) = AVG {r(x (i) ) : 1 ≤ i ≤ K} , referred to as top-K average reward. Both NoM and TopK are metrics of substantial interest in the GFlowNet literature [3, 55, 56, 64, 65] . For. KL 0.22±0.04 0.23±0.12 0.03±0.00 0.09±0.10 Renyi-α 0.22±0.03 0.23±0.10 0.03±0.00 0.19±0. 13 Tsallis-α 0.21±0.04 0.22±0.09 0.03±0.00 0.21±0.11D JS [P T ||R] = 1 /2 (D KL [P T ||M ] + D KL [R||M ]) = E x∼P T [log p T (x) /m(x)] + E x∼R [log r(x) /Zm(x)] ,Results. Figure 3 shows that the procedure minimizing divergence-based measures accelerates the training convergence of GFlowNets, whereas Figure 5 (for the banana-shaped distribution) and Table 1 highlight that we obtain a more accurate model with a fix compute budget. The difference between learning objectives is not statistically significant for the BPI task. Also, we may attribute the superior performance of reverse KL compared to the forward in the sequence generation task to the high variance of the importancesampling-based gradient estimates. Indeed, the observed differences disappear when we increase the batch of trajectories to reduce the estimator's variance (see Figure 8 in Appendix D). In conclusion, our empirical results based on experiments testing diverse generative settings and expanding prior art [48, 56, 112] , shows that training methods based on minimizing f -divergence VI objectives with adequate CVs implemented are practical and effective in many tasks. Correlatively, Figure 4 the fact that minimizing divergence-based objectives frequently implies better coverage of the target's high-probability regions; the only exception is the (extremely sparse) hypergrid task [55] .Figure 2 demonstrates that implementing CVs for the REINFORCE estimator reduces the noise level of gradient estimates significantly. This reduction in variance also accelerates training convergence.To illustrate this, we use the same experimental setup from Section 5.2 and analyze the learning curves for each divergence measure with and without control variates.Results. Figure 6 shows that the implemented gradient reduction techniques significantly enhance the learning stability of GFlowNets and drastically accelerate training convergence when minimizing the reverse KL divergence. Our results indicate that well-designed CVs for gradient estimation can greatly benefit GFlowNets training. Notably, similar improvements have been observed in the context of Langevin dynamics simulations [22, 31, 44] and policy gradient methods for RL [68, 100] .6 Conclusions, limitations and broader impact Discussion. We showed in a comprehensive range of experiments that divergence measures common in VI -forward KL, reverse KL, Renyi-α, and Tsallis-α -are effective learning objectives for training GFlowNets, performing competitively with or better than their balance-based counterparts. To achieve this, the introduction of efficacious control variates for low-variance gradient estimation of the divergence-based objectives was crucial, which is a key distinction between our work an prior art [55, 112] . Additionally, we developed the theoretical connection between GFlowNets and VI beyond the setting of finitely supported measures, establishing results for arbitrary topological spaces.Limitations. Albeit comprehensive and on par with the wider literature, our empirical evaluation was performed on problems of relatively small size due to the intractability of probing a GFlowNet's distributional accuracy on very large state spaces. That said, we acknowledge that an assessment on the domains of natural language processing [30] and drug discovery [3] based on context-specific metrics would strengthen our conclusions; we leave these tasks to future endeavors. Similarly, while we observed promising results for α = 0.5, there might be different choices of α that, depending on the application, might strike a better explore-exploit tradeoff and incur faster convergence. Thus, thoroughly exploring different α might be especially useful to practitioners.Broader impact. Overall, our work highlights the potential of the once-dismissed VI-inspired schemes for training GFNs, paving the way for further research towards improving the GFlowNets by drawing inspiration from the VI literature. For instance, one could develop χ-divergence-based losses for GFNs [19] , combine GFNs with MCMC using Ruiz and Titsias [81] 's divergence, or employ an objective similar to that of importance-weighted autoencoders [10] . Finally, although an ϵ-greedy off-policy sampling scheme can be easily incorporated into a divergence-minimizing algorithm through an importance-sampling correction, it remains elusive whether this would be possible for more sophisticated sampling techniques such as replay buffer [15] and local search [40] .Generative Flow Networks. GFlowNets [3, 4] were initially proposed as a reinforcement learning algorithm for finding diverse high-valued states in a discrete environment by sampling from a distribution induced by a reward function. Shortly after, they were extended to sample from complex distributions in arbitrary topological spaces [48] . Remarkably, this family of models was successfully applied to problems as diverse as structure learning and causal discovery [13, 15, 16] , discrete probabilistic modeling and graphical models [24, 29, 107, 108] , combinatorial optimization and stochastic control [109, 110] , drug discovery [3, 33, 62] , design of biological sequences [32] , natural language processing [30] , and aerial scene classification [23] . Concomitantly to these advances, there is a growing literature concerned with the development of more efficient training algorithms for GFlowNets [4, 40, 55, 84] -primarily drawing inspiration from existing techniques in the reinforcement learning literature [60, 65, 66, 87] . In the same spirit, Tiapkin et al. [87] showed it is possible to frame GFlowNets as an entropy-regularized reinforcement learning. In a study closely related to ours, Malkin et al. [56] proved the equivalence between GFlowNets and hierarchical variational inference (HVI) for discrete distributions; however, when training GFlowNets using divergence-based methods from the VI literature, the authors found no improvement relatively to the traditional flow-matching objectives. Thus, extending beyond discrete distributions, this work provides a definitive analysis of training GFlowNets by directly optimizing a set of divergences typically employed in variational inference training, given a clear context and conditions for effective use of divergence objectives for efficient learning procedures applied on GFlowNets models.Reinforcement Learning as Inference. Reinforcement Learning (RL) has been studied as a form of probabilistic inference extensively, generating relevant insights in the literature, and alternatively referred to as control as inference. Todorov [88] demonstrates a duality between estimation and optimal control, establishing conditions where estimation algorithms could applied for control problems. Kappen et al. [38] demonstrated that optimal control problems could be framed as inference problems in graphical models, providing a unified perspective for solving control tasks. Levine [50] presents a complete and modern RL formulation, linking with VI in particular. Rudner et al. [80] integrates even further RL with VI methods, demonstrating the conceptual and algorithmic gains of leveraging outcome-driven RL with variational inference to optimize policy distributions. Developing further, Toussaint and Storkey [89] applies approximate probabilistic inference methods to solve Markov Decision Processes (MDPs) with discrete and continuous states. The approach also aligns with model-based RL techniques, such as PILCO, which utilizes probabilistic models to enhance data efficiency in policy search [14] . Recent work by Deleu et al. [17] positions discrete probabilistic inference as a control problem in multi-path environments, highlighting the synergy between control theory and probabilistic modeling in the context of GFlowNets. This body of works relates to the approach presented in this paper, comparing optimization of trajectory balance and flow-matching losses related to sequential decisions modeled by the GFlowNet with f -divergence measures minimization procedures -related to approximated variational inference and generalized posterior inference [36, 45, 51, 92, 95] .Divergence measures and gradient reduction for VI. Approximate inference via variational inference (VI) methods [6, 7, 36, 92] initially relied on message passing and coordinate ascent methods to minimize the KL divergence of an unnormalized distribution and a proposal in a parameterized tractable family of distributions. Despite the initial generality of the optimization perspective, the concrete implementation of algorithms often requires specialized update equations and learning objectives for specific classes of models. On the other hand, the development of algorithms and software for automatic differentiation [2] and stochastic gradient estimators [59] unlocked the potential application of generic gradient-based optimization algorithms in inference and learning tasks for a comprehensive class of models. Seminal works such as Black-Box VI (BBVI) [71] , using the REINFORCE/score function estimator, and Automatic Differentiation VI (ADVI) [46] , using reparameterization and change-of-variables, demonstrated practical algorithms for Bayesian inference in generic models, including models combining classical statistical modeling with neural networks. Overall, Mohamed et al. [59] explain the development of the main gradient estimators: the score function [11, 71, 97, 103] , and the pathwise gradient estimator, also known as the parametrization trick [42, 43, 75] . The vanilla REINFORCE/score function estimator has notoriously high variance [11, 71, 77, 97] , which prompted a body of work exploring variance reduction techniques. In the original BBVI proposal, Ranganath et al. [71] explored Rao-Blackwellization, combining iterated conditional expectations and control variates, using the score function estimator (given its zero expectation) as a control variate. Subsequent works have continued to refine these techniques; Liu et al. [53] uses Rao-Blackwellized stochastic gradients for discrete distributions, while Kim et al. [39] and Wang et al. [93] explored joint control variates and provable linear convergence in BBVI. Additionally, Domke [21] and Domke [20] provided smoothness and gradient variance guarantees, further enhancing the robustness of score function estimator for VI methods. Our work demonstrates that effective variance reduction techniques applied to a f -divergence minimization training can significantly enhance the convergence speed and stability of the procedure. In theory and practice, we observed high compatibility between our results of variance-reduced f -divergence GFlowNets training and the body of work of variance-reduced score-function estimators for VI. Furthermore, by showing that these techniques apply to a broad class of models and optimization objectives, including continuous and mixed structured supports, we move GFlowNets' f -divergence minimization training closer to recent notions of generalized Bayesian inference and generalized VI [45] and variational inference in function spaces [95] -with the common thread of casting posterior inference as an optimization problem guided by some divergence measure. This generalization can enable applications of GFlowNets to a diverse range of machine learning tasks, enhancing their versatility and practical relevance.Set generation [3, 34, 65, 66] . S contains the sets of size up to N with elements extracted from a fixed deposit D of size D ≥ N and s o = ∅. For s ∈ S with |s| < N , κ f (s, •) is a counting measure supported at (the σ-algebra induced by) {s ∪ {d} : d ∈ D \ s}; for |s| = N , κ f (s, •) = δ s f . Then, X = {s ∈ S : |s| = N }. Similarly, κ b (s, •)'s support is {s \ {d} : d ∈ D} for s ̸ = s o . We define the unnormalized target's density by log r(x) = d∈x f (d) for a fixed function f : D → R. We parameterize p F (s, •) as a deep set [105] and fix p B (s, •) as a uniform density for s ∈ S.Autoregressive sequence generation [32, 55] Bayesian phylogenetic inference (BPI) [111] . A (rooted) phylogeny is a complete binary tree with labeled leaves and weighted edges; each leaf corresponds to a biological species, and the edges' weights are a measurement of evolutionary time. Formally, we let B be the set of observed biological species and V B be the set of |B| + 1 unobserved species. Next, we represent a phylogeny over B as a weighted directed tree G B = (B ∪ V B , E B , ω B ) with edges E B featured with a weight assignment ω B ; we denote its root by r(G B ). Under these conditions,we define S = K k=1 G F k : k F k = B ∧ G F kis a tree as the set of forests built upon phylogenetic trees over disjoint subsets of B;represents a disjoint union of non-empty sets and s o = b∈B G {b} is the forest composed of singleton trees G {b} . Also, we define the amalgamation of phylogenies G F k and G Fj , A(G F k , G Fj ), as the tree obtained by joining their roots r(G F k ) and r(G Fj ) to a new node r(G F k ∪ G Fj ), with a corresponding extension of the edges' weights. In contrast, the dissolution of a tree G F , R(G F ), returns the union of the two subtrees obtained by removing r(G F ) from G F . Then, κ f (s, •) is the counting measure supportedat K k=1,k̸ =i,j G F k ∪ A(G Fi , G Fj ) : (i, j) ∈ [K] 2 , i ̸ = j with s = K k=1 G F k and K ≥ 2; if s = G B , κ f (s, •) = δ s f . Hence, X is the set of phylogenies over B. Likewise, κ b (s, •)'s support is K k=1,k̸ =j G F k ∪ R(G Fi ) : i ∈ [K] ∧ r(G Fi ) / ∈ B for s = K k=1 G F k and K ≤ |B|.Finally, the unnormalized target is the posterior distribution defined by JC69's mutation model [37] given a data set of genome sequences of the species in B. More specifically, we let the prior be a uniform distribution and compute the model-induced likelihood function by the efficient Felsenstein's algorithm [25] . We assume the edges' weights are constant. See [101] for further details. We parameterize p F (s, •) with GIN [99] and fix p B (s, •) as an uniform distribution. [48, 110] . The training of GFlowNets in continuous spaces is challenging, and the problem of designing highly expressive models in this setting is still unaddressed [16, 48] . However, as we show in Section 5.2, divergence-based measures seem to be very effective learning objectives for autoregressive sampling of a sparse mixture of Gaussians. For a d-dimensional Gaussian distribution, S = {{(0, 0), (x i , i) :1 ≤ i ≤ n}, : n ≤ d, x ∈ R n } ⊂ {(0, 0)} ∪ P(R × [d]) and s o = (0, 0); note S is isomorphic to R d . Also, for s = {(x i , i)} n i=1 , κ f (s, •) is Lebesgue's measure at {s ∪ (x, n + 1) : x ∈ R} if n < d and κ f (s, •) = δ s f otherwise.In particular, X = {s ∈ S : max (x,i)∈s i = d}. Moreover, κ b (s, •) is a measure on {s \ (x, |s|) : x ∈ R}, which is isomorphic to R, which is a singleton. We define the target's density with a homogeneous mixture of Gaussian distributions, 1 K K i=1 N (µ i , σ 2 I) with µ i ∈ R d . We similarly define P F (s, •) as a mixture of one-dimensional Gaussians with mean and variance learned via an MLP [48] .Banana distribution. [57, 76] We consider sampling from the banana distribution, defined byEQUATION)Given its geometry and shape, this distribution is a common baseline in the approximate Bayesian inference literature [90, 104, 106] . This task is identical to sampling from a mixture of Gaussian distributions, except for the different target density specified by the model in Equation (8) . Also, we rely on the implementation Hamiltonian Monte Carlo (HMC) [5, 61] provided by Stan [12] to obtain accurate samples from (8) .A similar description may be utilized for the hypergrid and structure learning domains.We will consider the measurable space of trajectories (P S , Σ P ), with P S = {(s, s 1 , . . . , s n , s f ) ∈ S n+1 × {s f } : 0 ≤ n ≤ N -1} and Σ P as the σ-algebra generated by N +1 n=1 Σ ⊗n . For notational convenience, we use the same letters for representing the measures and kernels of (S, Σ) and their natural product counterparts in (P S , Σ P ), which exist by Carathéodory extension's theorem [96] ; for example, ν(B) = ν ⊗n (B) for B = (B 1 , . . . , B n ) ∈ Σ ⊗n and p F θ (τ |s o ; θ) is the density of P ⊗n+1 F (s o , •) for τ = (s o , s 1 , . . . , s n , s f ) relatively to µ ⊗n . In this case, we will write τ for a generic element of P S and x for its terminal state (which is unique by Definition 1).We will show that the gradient of the expected on-policy TB loss matches the gradient of the KL divergence between the forward and backward policies. Firstly, note that∇ θ D KL [P F ||P B ] = ∇ θ E τ ∼P F (so,•) log p F (τ |s o ; θ) p B (τ ) = ∇ θ τ log p F (τ |s o ; θ) p B (τ ) dP F (s o , dτ ) = ∇ θ τ log p F (τ |s o ; θ) p B (τ ) p F (τ |s o ; θ)dκ f (s o , dτ ) = τ ∇ θ log p F (τ |s o ; θ) p B (τ ) P F (s o , dτ ) + τ log p F (τ |s o ; θ) p B (τ ) ∇ θ p F (τ |s o ; θ)dκ f (s o , dτ )by Leibniz's rule for integrals and the product rule for derivatives. Then, sinceEQUATIONwe omitted the dependency of P F (and of p F thereof) on the parameters θ for conciseness. On the other hand,EQUATIONby the chain rule for derivatives. Thus,EQUATIONensuring that the equivalence between L T B and D KL in terms of expected gradients holds in a context broader than that of finitely supported distributions [56] .Henceforth, we will recurrently refer to the score estimator for gradients of expectations [97] , namely,EQUATIONwhich can be derived using the arguments of the preceding section. In this context, the Renyi-α's divergence satisfies∇ θ R α (P F ||P B ) = ∇ θ E τ ∼P F (so,•) [g(τ, θ)] (α -1)E τ ∼P F (so,•) g(τ, θ) ,with g(τ ; θ) = ( p B (τ |x)r(x) /p F (τ |so;θ)) 1-α and α ̸ = 1; similarly, the Tsallis-α's divergence abides byEQUATIONThe statement then follows by substituting ∇ θ E τ ∼P F (so,•) [g(τ, θ)] with the corresponding score estimator given by Equation (12) .Forward KL divergence. The gradient of D KL [P B ||P F ] is straightforwardly obtained through the application of Leibniz's rule for integrals,∇ θ D KL [P B ||P F ] = -E τ ∼P B (s f ,•) [∇ θ log p F (τ |s o ; θ)] ,since the averaging distribution P B do not depend on the varying parameters θ. However, as we compute Monte Carlo averages over samples of P F , we apply an importance reweighting scheme [63, Chapter 9] to the previous expectation to infer that, up to a positive multiplicative constant,∇ θ D KL [P B ||P F ] C = -E τ ∼P F (so,•) p B (τ |x)r(x) p F (τ |s o ; θ) ∇ θ log p F (τ |s o ; θ) , with C= denoting equality up to a positive multiplicative constant. We emphasize that most modern stochastic gradient methods for optimization, such as Adam [41] and RMSProp [28] , remain unchanged when we multiply the estimated gradients by a fixed quantity; thus, we may harmlessly compute gradients up to multiplicative constants.Reverse KL divergence. We verified in Equation ( 9) that∇ θ D KL [P F ||P B ] = E τ ∼P F (so,•) log p F (τ |s o ) p B (τ ) ∇ θ s θ (τ ) . Since p B (τ ) = p B (τ |x) r(x) /Z and E τ ∼P F (so,•) ∇ θ s θ (τ ) = 0, the quantity ∇ θ D KL [P F ||P B ] may be rewritten as ∇ θ D KL [P F ||P B ] = E τ ∼P F (so,•) log p F (τ |s o ) p B (τ ) ∇ θ s θ (τ ) + E τ ∼P F (so,•) [(log Z)∇ θ s θ (τ )] = E τ ∼P F (so,•) log p F (τ |s o ) p B (τ |x)r(x) ∇ θ s θ (τ ) ,in which x is the terminal state corresponding to the trajectory τ . Thus proving the statement in Lemma 2.We will derive an expression for the optimal baseline of a vector-valued control variate. For this, let f be the averaged function and g : τ → g(τ ) be the control variate. Assume, without loss of generality, that E π [g] = 0 for the averaging distribution π over the space of trajectories. In this case, the optimal baseline for the control variate a ⋆ is found byEQUATIONThus,a ⋆ = arg min a∈R Tr -2a • Cov π [f (τ ), g(τ )] + a 2 Cov π (g(τ )) ,which is a convex optimization problem solved byEQUATIONin which we used the circular property of the trace. This equation exactly matches the result in Proposition 2. In practice, we use g(τ ) = ∇ θ log p F (τ ) for both the reverse KL-and α-divergences, rendering a baseline a ⋆ that depends non-linearly on the sample gradients and is hence difficult to compute in a GPU-powered autodiff framework efficiently. We thus use Equation ( 6) to estimate a ⋆ . Gradient variance for flow-network-based objectives. Figure 7 shows the learning curve for the TB loss in each of the generative tasks. Notoriously, it underlines the low variance of the optimization steps -which, contrarily to their divergence-based counterparts, do not rely on a score function estimator -and suggests that the design of control variates for estimating the gradients of these objectives would not be significantly helpful. Also, the gradient of L T B depends non-linearly on the score function log p F and, consequently, it is unclear how to implement computationally efficient variance reduction techniques in this case. Forward KL for sequence generation. Figure 3 shows that alternative approaches in terms of convergence speed outperformed a GFlowNet trained to minimize the forward KL. One possible cause of this underperformance is the high variance induced by the underlying importance sampling estimator. To verify this, we re-run the corresponding experiments, increasing the size of the batch of trajectories for the forward KL estimator to 1024. Figure 8 presents the experiment's results, with an increased batch size corresponding to an estimator of smaller variance that accelerates the GFlowNet's training convergence. More broadly, this suggests that the design of GFlowNet-specific variance reduction techniques, which we leave to future endeavors, may further improve this family of models.Influence of α on the learning of GFlowNets. As mentioned earlier, divergence-based measures fall short compared to their balance-based counterparts for the hypergrid navigation task. For this extremely sparse problem, the benefits from off-policy exploration during training seem to supersede the statistical efficiency enacted by the minimization of divergences, which fail to properly cover the high-probability regions of the target distribution. In this scenario, Figure 9 suggests that this issue can be mildly circumvented by choosing a sufficiently negative α for the Renyi-α divergences, thereby imposing a mass-covering behavior to the learned model. However, these results should be substantiated by further experimental analysis to be conclusive. Currently, we would suggest a practitioner to stick to the balance-based objectives when dealing with very sparse target distributions.The following paragraph provides further implementation details. Regarding open access to the code, we will make the code public upon acceptance. Shared configurations. For every generative task, we used the Adam optimizer [41] to carry out the stochastic optimization, employing a learning rate of 10 -1 for log Z θ when minimizing L T B and 10 -3 for the remaining parameters, following previous works [48, 55, 56, 66] . We polynomially annealed the learning rate towards 0 along training, similarly to [86] . Also, we use LeakyReLU [98] as the non-linear activation function of all implemented neural networks.Set generation. We implement an MLP of 2 64-dimensional layers to parameterize the policy's logits log p F (s, •). We train the model for 512 epochs with a batch of 128 trajectories for estimating the gradients. Also, we let D = 32 and N = 16 be the source's and set's sizes, respectively.Autoregressive sequence generation. We parameterize the logits of the forward policy with a MLP of 2 64-dimensional layers; we pad the sequences to account for their variable sizes. We respectively consider D = 8 and N = 6 for the source's and sequence's sizes. To approximate the gradients, we rely on a batch of 128 sequences.Bayesian phylogenetic inference. We parameterize the logits of the forward policy with a 2-layer GIN [99] with a 64-dimensional latent embedding, which is linearly projected to log p F . Moreover, we simulated the JC69 model [37] to obtain 25-sized sequences of nucleotides for each of the 7 observed species, setting λ = 0.3 for the instantaneous mutation rate; see [101] for an introduction to computational phylogenetics and molecular evolution. To estimate the gradients, we relied on batches of 64 trajectories.Hypergrid navigation. We consider a H = 12 for Figure 3 and Figure 4 and H = 9 for Figure 9 . In both cases, d = 2 and R o = 10 -3 ; see [55, Section 5.1] . To parameterize the policy, we used a 2-layer 256-dimensional MLP with ReLU activations. We trained the models for 10000 epochs.Bayesian structure learning. We simulated a 100-sized 5-variable data set X from a randomly parameterized linear Gaussian structural model based on a fixed Bayesian network. We implemented a 2-layer 256-dimensional MLP with ReLU activations for the policy network, which received the flattened adjacency matrix of the current state as input. Training lasted for 4000 epochs.Mixture of Gaussian distributions. We consider a mixture of 9 2-dimensional Gaussian distributions centered at µ ij = (i, j) for 0 ≤ i, j ≤ 2, each of which having an isotropic variance of 10 -1 ; see Figure 1 . We use an MLP of 2 64-dimensional layers to parameterize the forward policy.Banana-shaped distribution. The model is specified by Equation (8) . We also consider an MLP of 2 64-dimensional layers to parameterize the forward policy.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: Provided in a zip file during the review period, with the plan for public release once the paper is made public.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips.cc/public/guides/CodeSubmissionPolicy) for more details. • The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Section 5 in the main text presents information about the experiments with full details provided in Appendix B and Appendix E, and supplemented zip-file with the code.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: Plots count on error bars and tables count on standard deviation.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: First paragraph of Appendix E.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: Our submission follow the NeurIPS ethical guidelines.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: We provide a perspective on broader impacts in Section 6, but do not foresee any direct negative societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: No experiments with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.• According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: No experiments with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: This paper f-divergence based objectives for GFlowNets (GFNs). These provide an alternative to the standard trajectory-balance objective for GFN and define generalizations of variational objectives for GFNs based on KL divergences. The authors consider Renyi-$\alpha$, Tsalli-$\alpha$, as well as the forward and reverse KL. The authors derive gradients for these objectives in the case of a fixed backward policy and on-policy samples, and develop control variates that reduce the variance of gradient estimates. Experiments demonstrate that using these objectives can accelerate convergence and enhance stability of training. 

Reviewers did not fully arrive at consensus regarding this submission. Overall reviewers were sympathetic to the exploration of alternative objectives for GFNs presented in the paper, as well the contributions in terms of developing gradient estimators and control variables. However, but reviewers also noted limitations of the approach in terms of the setting with fixed backward policy and on-policy samples, as well the experimental evaluation, which primarily focuses on relatively simple toy problems. Overall two reviewers (HRNm and 3sBX) support acceptance, with 3sBX noting during the AC discussion phase that in spite of shortcomings, they consider the paper to make enough contributions to warrant acceptance. Out of the more critical reviewers, Gn1e raised their score after the author response. Reviewer NozV, who was the most critical of the submission in their review, did not engage with the author response or the AC. As far as the AC can determine, at least some of their concerns were addressed directly by the author response, and as such the AC will somewhat discount the score from this reviewer. 

The AC took a quick look at the submission, which appears to be well-presented piece of work that makes a reasonable set of contributions to this space. This seems above the bar for acceptance. Do please incorporate the revisions that were discussed in the author response.

Reviews: review: paper_summary: paper_summary: This paper investigates divergence measures as learning objectives for Generative Flow Networks, which are amortized inference models designed for sampling from unnormalized distributions over composable objects. The authors review four divergence measures - Renyi-\alpha, Tsallis-\alpha, reverse and forward Kullback-Leibler divergences, and design estimators for their stochastic gradients in the context of training enerative Flow Networks. The authors verify that minimizing these divergences yields correct and empirically effective training schemes on several toy environment, and show that it often lead to faster convergence than previously proposed optimization methods. The authors also design control variates based on REINFORCE and score-matching estimators to reduce gradient variance.
strength_weakness: strength_weakness: 1. While the paper provides empirical evidence for the effectiveness of divergence-based objectives, a more extensive comparison with traditional GFlowNet training methods across a wider range of datasets and applications would strengthen the claims.

2. The choice of the \alpha parameter in Renyi-\alpha and Tsallis-\alpha divergences is not extensively explored. More insights into the impact of \alpha on the learning dynamics and guidance on selecting an appropriate value would be beneficial.

3. The computational overhead introduced by the control variates and their impact on training time is not explicitly discussed.
questions: questions: 1. How do the proposed divergence-based objectives perform on more complex and high-dimensional datasets commonly used in fields such as drug discovery and natural language processing?

2. Can the control variate techniques be extended to other GFlowNet training objectives beyond the divergence measures considered in this paper?

3. How does the choice of the \alpha parameter affect the learned GFlowNet's ability to capture multi-modal target distributions or discover diverse high-quality samples in combinatorial optimization tasks?

4. Are there any theoretical guarantees or bounds on the sample complexity or convergence rates when using the proposed divergence-based objectives for training GFlowNets?
limitations: limitations: Yes.
paper_summary: paper_summary: This paper investigates the potential of using a variety of divergence measures directly as training losses for GFlowNets, relying on many of the connections made between GFlowNets and variational inference. Training GFlowNets essentially consists in enforcing balance/flow-matching conditions between a proposal and a target distribution given some data. Given the links between GFlowNets and variational inference, training the GFlowNet to directly minimize the KL divergence was observed to result in biased and high-variance estimators in general. This paper aims to verify the latter claim, for different divergence measures (KL, reverse KL, Tsallis-\$\alpha\$, Renyi-\$\alpha\$). The paper alleviates the afore-mentioned limits by proposing low-variance gradient estimates to that of the latter quantities through control variates. Finally, the paper verifies the proposed methods experimentally through an extensive set of experiments.
strength_weakness: strength_weakness: See questions.
questions: questions: - How is the diversity of samples impacted? 
- There should be experiments that show that the different behaviors of each method with respected to the number of discovered modes throughout training (that should be compared to TB too!). For instance, similarly to Figure 4 in (Malkin et al., https://arxiv.org/pdf/2201.13259).
limitations: limitations: Limitations adequately addressed throughout the paper.
paper_summary: paper_summary: GFlowNets are a probabilistic framework for training amortized samplers for high-dimensional compositional spaces. The samplers are typically trained using local consistency objectives which are squared log losses. This paper examines alternatives to these obejctives in the form of general statistical f-divergence measures. The authors consider forward and reverse KL, Renyi-$\alpha$ and Tsallis-$\alpha$ divergences. The authors derive gradients for the divergences in the case of a fixed $p_B$ and on-policy samples. These gradient computations rely on REINFORCE-style estimators and consequently can suffer from high-variance gradient estimates which affect the optimization procedure. For variance reduction, the authors derive control variates for their gradient estimators. Through a series of experiments on a variety of tasks, the authors demonstrate the effectiveness of these divergences for training GFlowNets.
strength_weakness: strength_weakness: * The paper considers the on-policy setting for training GFlowNets and the proposed learning objectives based on the divergences assume on-policy smaples. However, existing flow-based objectives are all off-policy, and the advantage of GFlowNets on a lot of tasks (specifically challenging tasks with multi-modal target distributions) comes from the ability to train on off-policy trajectories (e.g. replay buffer[1] to local search [2]). So while the proposed learning objectives empirically perform better than TB on tasks where on-policy sampling is enough, it lacks the flexibility of accomodating off-policy training.
* The authors also assume that the backward policy $P_B$ is fixed (L153). While this is true in some cases, learning $P_B$ results in significant improvements to the learned sampler[3,4]. As far as I can tell, modifying the proposed objectives to accomodate learning the P_B is non-trivial.
* I appreciate the diversity of problems studied by the authors in their experiments, but there are some gaps in the empirical analysis. Specifically, the authors only include a TB baseline and not other objectives such as SubTB, DB which can perform better than TB (and have better training stability) in some cases. Additionally, the paper also does not include the VarGrad-style objective [5] which does away the need for estimating $Z$ in TB.
* Moreover, the tasks consider relatively small tasks so it is not clear how scalable the proposed objectives are. 


[1] Towards Understanding and Improving GFlowNet Training. Max W. Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, Tommaso Biancalani. ICML 2023. 

[2] Local Search GFlowNets. Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, Jinkyoo Park. ICLR 2024. 

[3] Trajectory balance: Improved credit assignment in GFlowNets. Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, Yoshua Bengio. NeurIPS 2022. 

[4] A theory of continuous generative flow networks. Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-García, Léna Néhale Ezzine, Yoshua Bengio, Nikolay Malkin. ICML 2023.

[5] Robust Scheduling with GFlowNets. David W. Zhang, Corrado Rainone, Markus Peschl, Roberto Bondesan. ICLR 2023.
questions: questions: In addition to the points in Weaknesses:

* L135 says TB requires estimating $Z$ but KL doesn't but I am not sure that is correct? Since even in the KL you need the normalizing constant in the $P_B$ term.

* L146: [1] would be a more appropriate reference here I think?

* L157: Missing reference to [2]

* What is the computational performance (in terms of runtime) of the proposed objectives relative to TB? 


[1] Rubinstein, R. Y. (1981). Simulation and the Monte Carlo Method. In Wiley Series in Probability and Statistics. Wiley. https://doi.org/10.1002/9780470316511

[2] f-Divergence Variational Inference. Neng Wan, Dapeng Li, Naira Hovakimyan. https://arxiv.org/abs/2009.13093.
limitations: limitations: The authors do not explicitly address the limitations of their approach (discussed in the weaknesses) section, though the assumptions are mentioned briefly in Section 2 and 3. The authors only mention the choice of $\alpha$ as a limitation.

There is also no discussion of broader impacts (the reference in the checklist is broken too).
paper_summary: paper_summary: This paper investigates alternative training methods for Generative Flow Networks (GFlowNets) by evaluating various divergence measures, including Renyi-α, Tsallis-α, reverse, and forward Kullback-Leibler (KL) divergences. Traditional methods focusing on minimizing log-squared differences are shown to lead to biased and high-variance estimators. The authors propose efficient estimators for the stochastic gradients of these divergences and introduce control variates to reduce gradient variance. Empirical results across diverse tasks demonstrate that minimizing these divergence measures significantly accelerates training convergence and enhances stability compared to traditional methods. The paper establishes theoretical connections between GFlowNets and variational inference, extending these insights to arbitrary topological spaces, and highlights the practical effectiveness of control variates in improving training efficiency.
strength_weakness: strength_weakness: - **Theory** The theoretical contribution builds upon previous works [4,5]. However, I do not think the theory itself has enough contribution. Because the Measurable pointed DAG is from previous work [5] and the main theoretical claim, Proposition 1, has limited novelty compared with Proposition 1 in [4]. Therefore, unlike previous works [4,5], they have their novel theoretical contributions and synthetic experiments would suffice. Thus, for this paper, stronger experiments are required, as will be discussed in the next.  
- **Experiments - Part 1** This paper does not contain enough real-world tasks, such as fragment-based molecule generation [1], graph combinatorial optimization [2], and RNA sequence generation [3] to illustrate its main contribution. These are standard tasks in evaluating GFN performances and are necessary to include. For the only real-world task in the paper, i.e., the BPI task, this paper admits that "not statistically significant" in line 334. Therefore, it is unclear whether the proposed methods to use other divergence measures will be meaningful in real-world scenarios. 
- **Experiments - Part 2** For the synthetic tasks, the plots in Figure 3 are also missing some lines. For example, why are there only two lines in the **Sets** plots? Also, there is no single best loss function that is uniformly better than the KL baseline. Therefore, additional divergence might seem unnecessary. There should be stronger evidence to support the use of other measures. Or, the authors can provide a guideline on how to choose the best measures given prior information about the tasks. 

- I would raise my scores if additional experiments are included and the proposed methods are indeed beneficial in more complex tasks. 


[1] Jain, Moksh, et al. "Biological sequence design with gflownets." International Conference on Machine Learning. PMLR, 2022.

[2] Zhang, Dinghuai, et al. "Let the flows tell: Solving graph combinatorial problems with GFlowNets." Advances in Neural Information Processing Systems 36 (2024).

[3] Kim, Minsu, et al. "Local search gflownets." arXiv preprint arXiv:2310.02710 (2023).

[4] Malkin, Nikolay, et al. "GFlowNets and variational inference." arXiv preprint arXiv:2210.00580 (2022).

[5] Lahlou, Salem, et al. "A theory of continuous generative flow networks." International Conference on Machine Learning. PMLR, 2023.
questions: questions: See **Weaknesses**.
limitations: limitations: See **Weaknesses**.

Comments: title: Response
Response to rebuttal
Author-Reviewer Discussion Phase (ends Aug 13 midnight AOE)
comment: Dear reviewers and conference chairs,

We would like to thank you for your service and dedication to the peer-reviewing process. We understand that some reviewers could not go through our rebuttal during the discussion period. Nonetheless, we made our best effort to incorporate all suggestions and believe the reviewers' feedback was paramount to strengthening our work and clarifying its importance.

Best regards,

The authors
Thank you for your continued engagement.

While we acknowledge many studies consider problems on a larger scale, we note that our empirical analysis primarily focuses on assessing GFlowNet’s distributional accuracy during training.

For problems of much larger scale, such as molecule generation (with approximately $10^{16}$ terminal states), it is not feasible to accurately measure the goodness-of-fit of a trained GFlowNet. This is because it would require exhaustively enumerating the target’s support, which is necessary to assess the total variation distance between the learned and target distributions.
Therefore, we have constrained our analysis to the environments discussed in the rebuttal PDF, which we believe are sufficiently diverse.

We will discuss this matter at the end of the revised manuscript. Thank you very much for your detailed and constructive feedback.
Sorry for the delay in my response, but I appreciate the detailed reply. 

> learned $p_B$

Thanks for sharing these results. I am a bit confused that on the set generation experiments for TB uniform does better than learned $p_B$. But indeed these results are interesting and could be an interesting avenue for future work. 

> For the DAG task, we considered graphs with 6 nodes. 

I appreciate the authors including these two tasks but I should emphasize that as with all the other experiments in the paper, these tasks are significantly smaller (and thus potentially easier) than prior work. For instance, in the TB paper the hypergrids considered were $8\times 8\times 8\times 8$ and $64\times 64$, not to mention much larger sequence design and molecular design tasks. I stand by my original review that the major weakness of the paper is the scale and difficulty of tasks considered in the empirical evaluation.
Thank you very much for engaging in the discussion and acknowledging our rebuttal and the novelty of our work.
Thank you for engaging in the discussion! 

> However recent work has also illustrated the importance of off-policy samples for training GFlowNets in challenging domains [2,3] which might make the objectives like TB better suited.

We agree! We will include our discussion regarding the trade-off between off- and on-policy learning in the revised manuscript. 

> Looking forward to the results.

Below, we report the results comparing the accuracy of the GFlowNet when $\log p_{B}$ is either learned or fixed, respectively, for the tasks of set generation and structure learning (Tables 2 and 3). We observed related results for the task of BPI - with similar performances for learnable and uniform $p_{B}$’s - and we will include these experiments in the revised manuscript. Although they show that jointly minimizing the learning objective wrt both $\log p_{B}$ and $\log p_{F}$ is a sound strategy for divergence-based measures, they do not reveal clear benefits favoring its implementation. 

Table 2: Set generation. Results are averaged across 3 runs. 
| | learn | unif |
|---|---|---|
| TB | 0.29 ± 0.02 | 0.13 ± 0.00 |
| Reverse KL | 0.09 ± 0.01 | 0.03 ± 0.00 |
| Renyi-0.5 | 0.09 ± 0.01 | 0.03 ± 0.00 |

Table 3: DAGs. Results are averaged across 3 runs. 
| | learn | unif |
|---|---|---|
| TB | 0.43 ± 0.11 | 0.47 ± 0.13 |
| Reverse KL | 0.31 ± 0.14 | 0.32 ± 0.14 |
| Renyi-0.5 | 0.14 ± 0.02 | 0.14 ± 0.02 |

To further investigate the shape of the learned backward policy, we computed the expected entropy of $p_{B}(\tau | x)$ under the learned marginal $p_{T}$ over terminal states, i.e., $\mathbb{E}\_{x \sim p\_{T}}[\mathbb{E}\_{\tau \sim p\_{B}(\cdot | x)}[ - \log p\_{B}(\tau | x)]]$. Intuitively, a highly entropic policy is closer to a uniform policy, following the analysis of Shen et al. [1]. In this sense, the results in Tables 4 and 5 below suggest that the learned backward policy closely resembles an uniform distribution. 


Nonetheless, we believe that a deeper investigation of the potential advantages of learning $\log p_{B}$ is a relevant and interesting research direction. 

Table 4: Set generation 
| | learn | unif |
|---|---|---|
| TB | 30.61 ± 0.01 | 30.67 ± 0.00 |
| Reverse KL | 30.61 ± 0.01 | 30.67 ± 0.00 |
| Renyi | 30.61 ± 0.01 | 30.67 ± 0.00 |

Table 5: DAGs
| | learn | unif |
|---|---|---|
| TB | 18.53 ± 3.61 | 18.55 ± 3.61 |
| Reverse KL | 19.26 ± 3.52 | 19.27 ± 3.52 |
| Renyi | 19.29 ± 3.42 | 19.31 ± 3.43 |

[1] Towards Understanding and Improving GFlowNet Training. Shen et al. ICML 2023

> The compute overhead is indeed not high, but my comment was more about the scale of the problems used in the paper.

Apologies for the oversight. For the DAG task, we considered graphs with 6 nodes; the target distribution’s support contains approximately 3.5M graphs. For the hypergrid task, we considered a 12 x 12 grid in Figure 1 and a 9 x 9 grid in Figure 3 of the rebuttal PDF; similarly to [2, 3], we set $R_{o} = 10^{-3}$. We will include these details in the revised manuscript.

[2] GFlowNets and Variational Inference. Malkin et al. ICLR 2023. 

[3] Trajectory balance: Improved Credit Assignment in GFlowNets. Malkin et al. NeurIPS 2022. 

>  This nuance is critical and I hope the authors avoid making sweeping claims.

Thanks for the advice. We only meant to flesh out the diversity of tasks and baselines of our experimental campaign in the rebuttal. 

We are very grateful for your detailed feedback and contribution towards strengthening our paper!
Thanks for your clarifications. I think the design of variance reduction in divergence-based objectives is a quite novel contribution. Admittedly, the design of appropriate divergence is an open question in VI literature. However, since the current experiments still do not include real-world environment like molecules, I can only raise my score to 5.
Thanks for the comments. I have a few follow-ups 

> amenable to a degree of off-policy learning by implementing an importance sampling estimator

Indeed one can always use importance sampling to do off-policy training with an on-policy objective (e.g. as noted in [1]) but IS also introduces other challenges (e.g. high variance). So as you note there is certainly a trade-off. However recent work has also illustrated the importance of off-policy samples for training GFlowNets in challenging domains [2,3] which might make the objectives like TB better suited. 

>  we will run additional experiments for divergence-based training of GFlowNets with a learnable $p_B$

Looking forward to the results. 

> computation overhead incurred by the proposed gradient estimation techniques is negligible

The compute overhead is indeed not high, but my comment was more about the scale of the problems used in the paper. Prior work on learning objectives has considered much larger and challenging problems. I see the authors have added the hypergrid and causal DAG tasks but do not provide details about the size of the problems? What size was the hypergrid and the number of variables considered for the causal DAG task? 

> mathematical definition of TB and KL depends on the constant Z

I agree that the Z does not play a role in the optimization but I would note that in a similar fashion TB can also be optimized without the Z (i.e. VarGrad)

In the common response the authors claim "broadest experimental evaluation of GFlowNet objectives in the literature" but I think this loses a lot of nuance. While the paper does consider 6 tasks, they are all much smaller (and potentially simpler) than tasks considered in prior work (e.g. much larger molecular optimization and sequence design). This nuance is critical and I hope the authors avoid making sweeping claims. 

[1] GFlowNets and variational inference, ICLR 2023. 

[2] Amortizing intractable inference in large language models, ICLR 2024. 

[3] Improved off-policy training of diffusion samplers, arXiv:2402.05098
Dear reviewers,

Thank you for your efforts so far! The authors have performed a number of additional experimental evaluations in response to your reviews. Could you follow up with them to ask any clarifying questions that may further inform your opinion on the paper? 

With thanks,
AC