Pessimistic Backward Policy for GFlowNets

Authors: hyosoon  jang
yunhui  jang
minsu  kim
jinkyoo  park
sungsoo  ahn
  kaist
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 10:37:07

Abstract: This paper studies Generative Flow Networks (GFlowNets), which learn to sample objects proportionally to a given reward function through the trajectory of state transitions. In this work, we observe that GFlowNets tend to under-exploit the high-reward objects due to training on insufficient number of trajectories, which may lead to a large gap between the estimated flow and the (known) reward value. In response to this challenge, we propose a pessimistic backward policy for GFlowNets (PBP-GFN), which maximizes the observed flow to align closely with the true reward for the object. We extensively evaluate PBP-GFN across eight benchmarks, including hyper-grid environment, bag generation, structured set generation, molecular generation, and four RNA sequence generation tasks. In particular, PBP-GFN enhances the discovery of high-reward objects, maintains the diversity of the objects, and consistently outperforms existing methods.

Main paper: 
 introduction: 
 Generative Flow Networks [1, GFlowNets] are models that sample compositional objects from a Boltzmann distribution defined by some reward function. To this end, GFlowNets construct an object through a trajectory of state transitions, e.g., iteratively adding molecular fragments to construct a molecule. They are attractive for their ability to sample a diverse set of high-reward objects, as demonstrated in molecular discovery [2, 3] , biological sequence design [4] , combinatorial optimization [5] , and large language models [6] .In detail, GFlowNets aim to sample from the Boltzmann distribution using a forward policy to decide the state transitions. However, this is challenging since the forward policy induces the distribution over trajectories, while the Boltzmann distribution is only defined on the terminal state of trajectories, i.e., objects. Hence, directly matching the two distributions with respect to the terminal state requires an intractable marginalization of the forward policy over the exponentially sized trajectory space.To circumvent this issue, GFlowNets employ an auxiliary backward policy that lifts the Boltzmann distribution to the trajectories via reversing the state transitions. In particular, the backward policy decomposes the unnormalized Boltzmann density of a terminal state into the unnormalized densities of trajectories, coined backward flow, associated with the terminal state. Then the forward policy learns the Boltzmann distribution by matching its unnormalized density, i.e., reward, coined forward flow, with the backward flow on the observed trajectories. We call this training scheme flow matching. 1The training objective of the flow matching has been investigated such as detailed balance [7] and trajectory balance [8] , and sub-trajectory balance [9] . To facilitate training, improved credit assignment techniques have been explored [10, 11] . Additionally, exploration methods [12] have been proposed to collect more diverse trajectories. Moreover, exploitation methods such as focusing on the higher-reward trajectories from the backward policy [13] and sampling high-reward trajectories with local search [14] have been presented for the collection of higher-reward trajectories. In this work, we point out a pitfall of the flow matching objectives: when only a small portion of object-sharing trajectories is observed, GFlowNets tend to under-exploit the object. This pitfall stems from the under-determination of the forward flow, due to training only on the observed backward flow that partially represents the true high reward. Consequently, the forward policy tends to assign high probabilities to objects with high observed backward flow, rather than the high reward objects, as illustrated in (a) and (b) of Figure 1 . This is counter-intuitive as the forward policy favors objects with low rewards despite possessing the knowledge of other objects with higher rewards. While one could bypass this issue at the cost of observing more trajectories [13] , we pursue an alternative direction in this work.We propose a simple remedy for the under-exploitation problem: a pessimistic backward policy for GFlowNets (PBP-GFN). Our key idea is the maximization of the observed backward flow to align the observed backward flow to the true reward. Consequently, PBP-GFN resolves the underexploitation problem which favors the object with high observed backward flow while neglecting the true reward, as illustrated in (c) of Figure 1 . We also note that our algorithm preserves the asymptotic optimality to induce the target Boltzmann distribution by simply modifying the backward policy while preserving the true rewards [8] . Additionally, we analyze how our algorithm reduces the error bound in estimating the true Boltzmann distribution.We extensively validate PBP-GFN on various benchmarks: hyper-grid benchmark [1] , bag generation [13] , maximum independent set problem [5] , fragment-based molecule generation [1] , and four RNA sequence generation tasks [4] . In these experiments, we observe that PBP-GFN (1) improves the learning of target Boltzmann distribution and (2) enhances the discovery of high-reward objects, while (3) maintaining the diversity of the sampled high-reward objects.To conclude, our contributions can be summarized as follows:• We characterize the under-exploitation problem stemming from an under-determined flow that only learns the observed flow for partially observed trajectories (Example 1). • To resolve this issue, we propose pessimistic training of backward policy that aims to reduce the amount of unobserved flow for the observed objects. • Through extensive experiments, we show that our algorithm consistently improves the performance of GFlowNets compared to prior works for designing the backward policy, even higher than other training algorithms for discovering high-reward objects. 2 
preliminaries: 
 Generative Flow Networks [1, 7, GFlowNets] generate an object x from the object space X through a trajectory τ = (s 0 , s 1 , . . . , s T ) of state transitions, where the terminal state is the object s T = x ∈ X to be generated. Here, a forward policy P F (s t+1 |s t ) makes the transition from the state s t to the next state s t+1 and assigns a probability of P F (τ ) =T -1 t=0 P F (s t+1 |s t ) to the trajectory τ . Next, GFlowNets train the forward policy to sample objects from a Boltzmann distribution defined by a reward function R(x) that satisfies:EQUATIONHere, P ⊤ F (x) is a distribution of an object x marginalized over exponentially sized non-terminal state spaces. To circumvent this intractability, GFlowNets train on the flow matching objectives.Flow matching for training GFlowNets. To learn the Boltzmann distribution, the forward policy P F aims to align to a backward policy P B . The backward policyP B (τ |x) = T -1 t=0 P B (s t |s t+1 )decomposes the reward into the unnormalized densities of object-sharing trajectories T (x) for an object x, i.e., R(x) = τ ∈T (x) R(x)P B (τ |x).To be specific, the forward policy learns to match the unnormalized densities to the backward policy, coined flow matching, over all trajectories in the trajectory space T :EQUATIONwhere Z θ P F (τ ) is a forward flow defined with a learnable constant Z θ , and R(x)P B (τ |x) is a backward flow. Equation ( 2) induces the forward policy following Boltzmann distribution, i.e.,Z θ P ⊤ F (x) ≈ R(x), by marginalizing trajectory flows over set of trajectories T (x) inducing the objectx, i.e., τ ∈T (x) Z θ P F (τ ) ≈ Z θ P ⊤ F (x) and τ ∈T (x) R(x)P B (τ |x) ≈ R(x).To satisfy Equation (2), GFlowNets minimize various training objectives. One such objective is the trajectory balance [8, TB] , defined as follows:EQUATIONwhich is minimized over trajectories observed during training, e.g., trajectories sampled from the forward policy. The set of observed trajectories stored in the replay buffer inducing the object x is denoted as B(x) ⊂ T (x). Note that training objectives for Equation (2) can also be defined on a transition [7, DB] or sub-trajectories [8, subTB] . 
method: 
 In this section, we introduce our pessimistic backward policy for generative flow networks (PBP-GFN). First, we show that forward policies trained with flow matching tend to under-exploit high-reward object x with partially observed trajectories B(x) when the underdetermined forward flow only learns the small amount of observed backward flow for the high-reward object (Section 3.1). To address this issue, we propose pessimistic training of backward policy that increases the proportion of observed flow for the object, which leads to an accurate estimation of the reward (Section 3.2).First, we explain how conventional flow matching may suffer from the under-exploitation of observed high-reward objects. To this end, we decompose the reward R(x) into two components: (1) observed backward flow R B (x) = τ ∈B(x) R(x)P B (τ |x) assigned to the partially observed trajectories B(x), and (2) unobserved backward flow R(x)-R B (x) assigned to the unobserved trajectories T (x)\B(x).Then (1) and ( 2) are paired with observed forward flow and unobserved forward flow, respectively.In detail, on the one hand, conventional flow matching aligns the observed forward flow to the observed backward flow for the first component, i.e., τ ∈B(x) Z θ P F (τ ) ≈ R B (x). On the other hand, there exists degree of freedom for the unobserved forward flow τ ∈T (x)\B(x) Z θ P F (τ ), as it is challenging to match the flow over unobserved trajectories, i.e. trajectories not in the buffer B. Overall, flow matching induces a forward policy with the marginalized probability P ⊤ F,B (x):EQUATIONHere, our key observation is that, for an observed object x with a high reward R(x) but a small amount of observed backward flow R B (x), the unobserved backward flow R(x) -R B (x) is likely to be much larger than the forward flow of the unobserved trajectories τ ∈T (x)\B(x) Z θ P F (τ ). This leads to the under-exploitation of the high-reward object x, since the forward policy assigns a higher probability to another object x ′ with a lower reward but a larger amount of observed flow R B (x), as illustrated in (a) of Figure 2 . As a result, the marginalized probability P ⊤ F,B (x) may converge to a local optimum that yields a smaller expected reward compared to the target Boltzmann distribution.To further motivate our proposal regarding the under-exploitation problem, we present a failure case of flow matching converged to a local optimum contradicting the observed rewards in Example 1. We construct a particular instance of Equation (4) where the forward policy underestimates the high-reward object compared to the lower one. We depict this example in (a) and (b) of Figure 3 . Example 1. Consider two objects x 1 and x 2 with rewards of 1 and 1 2 , respectively, where x 1 is reached by three trajectories (|T (x 1 )| = 3) and x 2 is reached by one (|T (x 2 )| = 1). Here, one trajectory for each object is observed (|B(x 1 )| = |B(x 2 )| = 1). Then, the probability to induce object x 1 can be assigned as P ⊤ F,B (x 1 ) ∝ 1 3 since the forward flow still matches the backward flows for the observed trajectory τ 1 , i.e., P ⊤ F,B (x 1 ) = P F,B (τ 1 ) ∝ R(x 1 )P B (τ 1 |x 1 ) = 1 3 . This is lower than P ⊤ F,B (x 2 ) ∝ 1 2 assigned with fully observed trajectories.Example 1 is counter-intuitive, as a higher probability is assigned to the lower reward object x 2 despite observing the higher reward object x 1 . The forward policy P F (τ ) also assigns zero probability to unobserved trajectories. Consequently, the probability P ⊤ F,B (x) cannot be corrected even more trajectories are sampled from policy τ ∼ P F,B (τ ). This hints at the necessity of the remedy for the under-exploitation of objects due to the small amount of observed flow, with the fixed observation B.Here, we propose a pessimistic training method for the backward policy in GFlowNets, coined PBP-GFN, which aims to resolve the under-exploitation problem of the flow matching with partially observed trajectories introduced in Section 3.1. To address this challenge, the backward policy is trained to reduce the amount of unobserved backward flow, being pessimistic about unobserved trajectories inducing the observed object. It is notable that the total backward flow for the object, i.e., reward, is preserved by shifting the unobserved backward flow into the observed backward flow. Sample a batch of trajectories {τ (k) } K k=1 from the behaviour policy.Update B ← B ∪ {τ (k) } K k=1 .5:for n = 1, . . . , N do ▷ Learning pessimistic backward policy 6:Update P B to minimize ℓ PBP over B with stochastic gradients.end for 8:Update P F , Z θ to minimize L TB with {τ (b) } K k=1 . 9: until converged To be specific, given a replay buffer B, the pessimistic training of backward policy P B aims to increase the backward flow for the observed trajectories ending with the object x. Specifically, it aims to maximize R B (x) = τ ∈B(x) R(x)P B (τ |x), thereby aligning the observed backward flow R B (x) to the true reward R(x), reaching the upper bound R B (x) ≈ R(x). Consequently, flow matching with such a backward flow for partially observed trajectories induces an observed forward flow that accurately estimates the true reward, thereby preventing the under-exploitation of the rewards due to the small amount of observed flow (Example 1), as illustrated in (b) of Figure 2 and (c) of Figure 3 . Furthermore, it is worth noting how PBP-GFN better estimates the Boltzmann density, i.e., reward. The high-level idea is that, given the fixed total flow, maximizing the observed forward and backward flows with PBP-GFN naturally minimizes the unobserved forward and backward flows, thereby reducing a flow matching error for the unobserved flows effectively. The detailed error bound in estimating the Boltzmann density is described in Appendix A.Pessimistic training of backward policy. We train the parameterized backward policy P B to increase the backward trajectory flows in observed trajectories, R(x) τ ∈Bx P B (τ |x), by assigning higher probabilities to the backward transitions P B (τ |x) = T -1 t=0 P B (s t |s t+1 ) of observed trajectories, i.e., τ ∈B(x) P B (τ |x) ≈ 1. We achieve this by minimizing the negative log-likelihood:EQUATIONwhere x is the object induced by the trajectory τ . It is notable that our approach only modifies the relative backward trajectory flows among trajectories inducing the same object and does not alter the total amount of backward flows, thereby preserving the asymptotic optimality of flow matching for learning the target Boltzmann distribution. Note that the training of the pessimistic backward policy is practical in most cases, as it only requires computing the stochastic gradients to minimize ℓ PBP .Subsequently, we train the GFlowNets with the learned pessimistic backward policy. The pessimistic backward policy is learned online with the forward policy, as new trajectories are observed for the training in each round. The training algorithm is described in Algorithm 1. 3 Note that the pessimistic training is agnostic to the choice of flow matching objectives [7] [8] [9] .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: We provide a detailed proof of how PBP-GFN improves the learning of target Boltzmann distribution in Appendix A 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. 
related work: 
 Generative Flow Networks (GFlowNets). GFlowNets [1, 7] train a forward policy that sequentially constructs objects sampled from a Boltzmann distribution. They are closely related to reinforcement learning in soft Markov Decision Processes (soft MDPs) [15] [16] [17] and variational inference [18] .Recently, there has been a surge in research on improving the training of GFlowNets, such as introducing novel flow matching objective functions [8, 9, 12, 13] , enhancing off-policy exploration [14, [19] [20] [21] , incorporating order information for enabling preference-based optimization [22] , and improving credit assignment [10, 11] . Moreover, GFlowNets are increasingly applied across a wide range of fields such as molecular optimization [2, 3, 23] , biological sequence design [4, 24] , probabilistic modeling and inference [25, 26] , combinatorial optimization [5, 27, 28] , continuous stochastic control [29] [30] [31] , and large language models [6] .Despite the advancements in training GFlowNets, there still exists the challenge of dealing with the vast number of trajectories. The number of trajectories grows exponentially with the increase in the number of state spaces and actions, making it impractical to observe all trajectories during training. This issue can be partially addressed by facilitating the discovery of unobserved trajectories [12] . However, the problem of probability not matching the rewards remains unless sampled trajectories comprehensively cover all possible flows.Training GFlowNets with auxiliary backward policy. GFlowNets train a forward policy to align with the auxiliary backward policy, which inverts the construction process of the object. Therefore, the choice of the backward policy directly impacts the training of GFlowNets and is vital to the improvement of the sampling performance. Despite its crucial role, the choice of backward policy has gained limited attention with only a few works [8, 13, 15] , and none of these works tackle the under-exploitation of high-reward objects caused by unobserved backward flow.For instance, while the uniform [8] and the MaxEnt [15] backward policies assign a fixed probability to the backward transition for enhancing exploration, our pessimistic backward policy learns the backward transition probability for enhancing exploitation. Next, conventional [8] and sub-structure [13] backward policies may enhance the exploitation by learning the backward flow to align with the forward flow or to improve the credit assignments. However, they do not directly reduce the unobserved backward flow and do not resolve the under-exploitation stemming from that.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)? • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [Yes] Justification: We discussed the positive and negative social impacts in last of Appendix. 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). • 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. 
experiment: 
 We evaluate our method on various domains, including a hyper-grid [1] , bags [13] , structured sets [5] , molecules [1] , and RNA sequences [13, 14] . As base metrics, we consider the number of modes, e.g., samples with rewards higher than a specific threshold, and the average top-100 score, which are measured via samples collected during training. We report the performances using three different random seeds. In these experiments, one can observe that:• PBP-GFN improves learning of the target Boltzmann distribution (Figures 4 and 6 ).• PBP-GFN enhances the discovery of high-reward objects (Figures 5, 7 and 8 ).• PBP-GFN maintains the diversity of sampled high-reward objects and promotes the discovery of distinct diverse modes (Figures 7(c ) and 9).In synthetic environments, i.e., hyper-grid environment, bag generation, and maximum independent set, we first show how our method (PBP-GFN) improves the performance compared to the prior methods that proposed various designs of the backward policy, on both the trajectory balance [ (e) PBP-GFN (our)Figure 4 : The target distribution and empirical distributions of each model trained with 10 5 trajectories. The empirical distributions are computed as rescaled products of the distribution over three runs. Our method (PBP-GFN) consistently discovers all modes over three runs and learns the target Boltzmann distribution correctly within the relatively small number of trajectories. backward policy trained with the TB or DB [7] , the uniform backward policy [8] , and the maximum entropy backward policy [15, MaxEnt] .Hyper-grid [1] . We first consider the hyper-grid, where the target Boltzmann distribution is defined over the 16 × 16 × 16 grid illustrated in Figure 4 (a). We also consider the 20 × 20 × 20 × 20 hyper-grid. The actions are incrementing one coordinate by one or terminating. The high-reward regions, i.e., modes, are defined as near the corners of the grid that are separated by regions with very small rewards. In this task, we consider the TB-based implementation following the prior work [8] . The detailed experimental settings are described in Appendix B. In this task, we measure the L1 distance between the target Boltzmann distribution and the empirical distribution of P ⊤ F (x), with the measurable likelihood of the Boltzmann distribution.Bag generation [13] . We next consider a simple bag generation task, where the action is adding an item to a bag. The bag yields a high reward when seven repeated items are included, i.e., modes. We apply our method to the prior TB-based implementation on this task [13] and compare it with TB and MaxEnt. The detailed setting is described in Appendix B.Maximum independent set [5] . We also consider solving maximum independent set problems, where the action is selecting a node and the reward is the size of the independent set. At each epoch, the GFlowNets train with the set of training graphs, and sample 20 solutions for each validation graph and measure the average reward and the maximum reward following Zhang et al. [5] . We apply our method to the prior DB-based implementation of this task [5] . The experimental setting is described in Appendix B. Note that the MaxEnt is equivalent to the uniform backward policy in this task.Results. In Figure 4 and Figure 6 , we depict the empirical sampling distribution and the L1 distance from the target Boltzmann distribution for each method in the hyper-grid environment. Here, one can see that our method (PBP-GFN) captures all modes and converges to the target Boltzmann distribution faster than the baselines. These results can be attributed to the capabilities of PBP-GFN, which enables us to effectively learn from the large amount of correct forward flow. In both bag generation and maximum independent set problems, one can see that our approach also shows (1) superior performance (2) or faster convergence compared to the baselines as illustrated in Figure 5 . One can reason this result stems from the capabilities of PBP-GFN that facilitate the learning of correct Boltzmann distribution. Furthermore, it is worth noting that our method makes improvements over both TB and DB-based implementations.Next, we evaluate our method in the fragment-based molecule generation [1] , where the action is adding a molecular building block. The reward is the binding energy between the molecule and the target protein computed by a pre-trained oracle [1] . Here, the mode is defined as a high-reward molecule with a low Tanimoto similarity [32] measured against previously accepted modes.We consider a TB-based implementation for our method and compare with various baselines including GFlowNets and reinforcement learning algorithms: DB, sub trajectory balance [9, subTB] , TB, TB defined with uniform and MaxEnt backward policies, generative augmented flow networks [10, GAFN] , and Advantage Actor-Critic [33, A2C] . For the evaluation metric, we analyze the trade-off between the average score of the top 100 samples and the diversity of these samples. Additionally, to measure diversity, we compute the average pairwise Tanimoto similarity following prior works. The detailed setting is described in Appendix B.Results. We depict the results in Figure 7 . One can see that our method, i.e., PBP-GFN, outperforms the baselines in enhancing the average score of unique top-100 molecules and the number of modes found during training. These results highlight that PBP-GFN also can make improvements for environments with a huge state space. Furthermore, one can see that our method yields low Tanimoto similarities between top-100 molecules with respect to the average reward. This verifies that our algorithm not only generates high-scoring samples but also diverse molecules.We consider four RNA sequence generation tasks that aim to discover diverse and promising sequences that bind to human transcription factors [4, 34, 35] , where the action is appending or prepending an amino acid. As baselines, we consider the same baselines as in the fragment-based molecule generation. In this task, we conduct experiments on the following four benchmarks. TFBind8. This is the task to generate length-eight RNA sequences. The reward is computed by wet-lab measured DNA binding activity to Sine Oculis Homeobox Homolog 6 [34] . The mode is determined based on whether it is included in a predefined set of promising RNA sequences [13] .RNA-Binding. This task generates length-14 RNA sequences. In this task, we consider three different target transcriptions: RNA-A, RNA-B, and RNA-C [14, 36] . The mode is defined as an RNA sequence with a reward higher than the threshold. In this task, we also consider the 2-hamming ball modes [36] , which is defined as the local maximum among its intermediate neighborhoods defined by modifying n components of the sequence.Results. The results are presented in Figure 8 . One can see that FBP-GFN shows faster convergence or superior performance compared to the considered baselines in enhancing the average score of unique top-100 sequences and the number of modes during training. Furthermore, in Figure 9 , one can see that our method better discovers the diverse distinct modes that are separated far from each other, compared to the baselines.Comparing overall generated sample quality. To further analyze the overall sample quality, we provide the relative mean error [13] which measures the distance between the mean values of the empirical generative distribution and the target Boltzmann distribution. We present the results in Figure 10 . One can see that our method yields the lowest errors. Active round (×10 Comparison with other exploitation methods. We further validate PBP-GFN by comparing or combining the pessimistic backward policy with local search [14] and reward-prioritized buffer [13] that are off-policy sampling methods and orthogonal to our pessimistic backward policy. We present the experimental results in Figure 11 . One can observe that our method (1) shows similar performance compared to them and (2) consistently improves the performance when combined with them. 
conclusion: 
 In this work, we identify the under-exploitation problem in flow matching due to the large amount of unobserved flow lifted by the backward policy. To resolve this, we introduce a pessimistic training of the backward policy for GFlowNets (PBP-GFN) that reduces the probabilities for unobserved backward trajectories leading to observed trajectories. Our PBP-GFN shows a successful alternative to the prior backward policies and has demonstrated improved performance across eight benchmarks.Limitation. As an exploitation method, our PBP-GFN makes a trade-off between obtaining high-reward trajectories and diversified trajectories, i.e., there is no free lunch in the exploitationexploration trade-off. Although our approach maintains the diversity of high-reward sampled objects in the considered benchmarks, this may not hold for some environments where exploration is significant. We discuss such a setting in Appendix C. To relax this issue, one can reduce the learning rate for the pessimistic backward policy or incorporate an exploration-focused off-policy sampling method. One can further control the trade-off by interpolating PBP-GFN with explorative GFNs, e.g., MaxEnt, which can be an interesting future work direction. 
 
 Appendix: We demonstrate that PBP-GFN yields a lower error bound between the true marginalized distribution P ⊤ F (x) and the target Boltzmann distribution P ⊤ B (x) = R(x) Z defined with the Z = x∈X R(x), compared to the conventional GFlowNets. To this end, we derive the following bound:EQUATION)where the ϵ depends on the error in the likelihood for trajectories over the observed trajectories, τ ∈B(x) (P F (τ ) -P B (τ )). In this equation, our pessimistic backward policy maximizes τ ∈B(x) P B (τ ) by maximizing the likelihood of τ ∈B(x) P B (τ |x). This better reduces the error bound compared to the conventional backward policy.The error bound in Equation ( 6) can be derived as follows:x P ⊤ F (x) -P ⊤ B (x) = x τ ∈T (x) P F (τ ) - τ ∈T (x) P B (τ ) ≤ x τ ∈B(x) P F (τ ) - τ ∈B(x) P B (τ ) + x   P ⊤ F (x) - τ ∈B(x) P F (τ )   -   P ⊤ B (x) - τ ∈B(x) P B (τ )   ≤ x τ ∈B(x) P F (τ ) - τ ∈B(x) P B (τ ) + x   P ⊤ F (x) - τ ∈B(x) P F (τ )   + x   P ⊤ B (x) - τ ∈B(x) P B (τ )   = x τ ∈B(x) P F (τ ) - τ ∈B(x) P B (τ ) + x   P ⊤ F (x) - τ ∈B(x) P F (τ )   + x   P ⊤ B (x) - τ ∈B(x) P B (τ )   = x τ ∈B(x) P F (τ ) - τ ∈B(x) P B (τ ) + 2 - x   τ ∈B(x) P F (τ ) + τ ∈B(x) P B (τ )   = ϵ + 2 -2 τ ∈B(x) P B (τ )where the error ϵ is associated with the errors in trajectory flow matching over the observed trajectories τ ∈B(x) (P F (τ ) -P B (τ )).In all experiments, the backward policy is designed to have the same architecture as the forward policy, e.g., a feedforward network with the same hidden dimensions, but does not share the parameters with the forward policy. For all experiments, we set the learning rate for the pessimistic training of the backward policy as 1e-3. Our overall implementations for each benchmark follow the prior studies.Note that we consider both on-policy and off-policy settings. We use a single GPU of NVIDIA GeForce RTX 3090.Hyper-grid environment. The implementations follow the prior study by Malkin et al [8] . We consider 16 × 16 × 16 hyper-grid, where the starting state is (0, 0, 0) and the action is incrementing one coordinate or terminating. The reward is computed as follows:R(s) = R 0 + 0.5 D d=1 I s d H -1 -0.5 ∈ (0.25, 0.5) + 2 D d=1 I s d H -1 -0.5 ∈ (0.3, 0.4) ,where H = 16, D = 3 and R 0 is 1e-3 in our settings.The forward policy is implemented with the feed-forward neural network that consists of two layers with 256 hidden dimensions and is trained with a learning rate of 1e-3. The learning rate for Z θ is 0.1. In this task, the training is on-policy. The GFlowNets are trained with the 64 trajectories sampled from the current policy in each round. We train the pessimistic backward policy to minimize ℓ PBP for these 64 trajectories in each round.Bag generation. The implementations follow the prior study by Shen et al. [13] . The action includes one of seven types of entities in the current bag with a maximum capacity of 15. If it contains seven or more repeats of any items, it has a reward 10 with 75% chance, and 30 otherwise. The threshold for determining the mode is 30.The forward policy is implemented with the feed-forward neural network that consists of two layers with 16 hidden dimensions and is trained with a learning rate of 1e-4. The learning rate for Z θ is 1e-2. In this task, the training is off-policy and consists of online and offline rounds. The online round uses 32 trajectories sampled from the forward policy with an exploration rate 0.1 and the offline round uses 32 trajectories sampled from the backward policy conditioned on the high-reward objects [13] . We store these trajectories (sampled during online and offline rounds) into the buffer B. In each round, we train the pessimistic backward policy to minimize ℓ PBP with trajectories sampled from the buffer B where the N in Algorithm 1 is eight. The buffer stores trajectories sampled during the previous 20 rounds.Maximum independent set. The implementations follow the prior study by Zhang et al. [5] . In this task, the action is selecting a node to construct the maximum independent set. The reward is the set size and the temperature. The number of training graphs and validation graphs are 4000 and 500, respectively. The graph contains around 200 to 300 nodes. Furthermore, the reward is re-scaled with the temperature which is annealed during training, starting from 1 and ending at 500.The forward policy is implemented with the graph isomorphism neural network [37] that consists of five layers with 256 hidden dimensions and is trained with a learning rate of 1e-3. In this task, the training is on-policy transition-based training [5] . In each training step, the GFlowNets are trained with the 64 transitions s → s ′ in trajectories sampled from the current policy. We train the pessimistic backward policy to minimize the negative log-likelihood -log P B (s|s ′ ) for these transitions within each training step.Molecule generation. The overall settings is similar to the prior study [1] , and our implementations are built upon the released codes. 45 This task aims to generate a molecule, where the action is adding a fragment. The number of available fragments is 72. The reward is computed by a pre-trained function [1] . The reward is scaled to have the maximum value nearing 1.0, and the reward exponent is set to 64.0. The mode is defined as a molecule with a reward higher than 0.97 and a Tanimoto similarity lower than 0.65, measured against previously accepted modes.The forward policy is implemented with the graph attention transformer that consists of four layers with 128 hidden dimensions and two attention heads, which is trained with a learning rate of 1e-4.The learning rate for Z θ is 1e-3. In this task, the training is off-policy, where 64 trajectories are sampled from the lagged forward policy P F ′ whose parameters are updated as F ′ = 0.95F ′ + 0.05F in each round. We store these trajectories into the buffer B. In each round, we train the pessimistic backward policy to minimize ℓ PBP with trajectories sampled from the buffer B where the N in Algorithm 1 is eight. The buffer stores trajectories sampled during the previous 20 rounds.TFBind8. The implementations follow the prior study by Shen et al. [13] . The action appending or prepending an amino acid to the sequence with a maximum length of eight. The number of amino acids is four. The reward is pre-computed based on wet-lab measured DNA binding activity to Sine Oculis Homeobox Homolog 6 [34] , which is scaled between 0.001 to 1.0. The reward exponent is set to 3.0. The mode is determined based on whether it is included in a predefined set of promising RNA sequences [13] .The forward policy is implemented with the feed-forward neural network that consists of two layers with 128 hidden dimensions and is trained with a learning rate of 1e-4. The learning rate for Z θ is 1e-2. In this task, the training is off-policy and consists of online and offline rounds. The online round uses 16 trajectories sampled from the forward policy with an exploration rate 0.01 and the offline round uses 16 trajectories sampled from the backward policy conditioned on the high-reward objects [13] . We store these trajectories (sampled during online and offline rounds) into the buffer B.In each round, we train the pessimistic backward policy to minimize ℓ PBP with trajectories sampled from the buffer B where the N in Algorithm 1 is eight. The buffer stores trajectories sampled during the previous 20 rounds.RNA-Binding. The implementations follow the prior study by Kim et al. [14] . The action appending or prepending an amino acid to the sequence with a maximum length of 15. The reward is scaled between 0.001 to 1.0, and the reward exponent is set to 8.0. The mode is determined based on whether it is included in a predefined set of promising RNA sequences [21] .The forward policy is implemented with the feed-forward neural network that consists of two layers with 128 hidden dimensions and is trained with a learning rate of 1e-4. The learning rate for Z θ is 1e-2. In this task, the training is off-policy and consists of online and offline rounds. The online round uses 32 trajectories sampled from the forward policy with an exploration rate 0.01 and the offline round uses 32 trajectories sampled from the backward policy conditioned on the high-reward objects [13] . We store these trajectories (sampled during online and offline rounds) into the buffer B.In each round, we train the pessimistic backward policy to minimize ℓ PBP with trajectories sampled from the buffer B where the N in Algorithm 1 is eight. The buffer stores trajectories sampled during the previous 20 rounds. As the case where our PBP-GFN may reduce exploration, we consider a scenario where an unobserved high-reward trajectory may largely overlap with an observed low-reward trajectory. Then, to explore the high-reward trajectory, one should assign a relatively high probability to the low-reward trajectory, i.e., the opposite of the case requiring exploitation. We exemplify and analyze this scenario in Figure 12 , which illustrates how pessimistic training may reduce exploration by enhancing the exploitation of observed high-reward trajectories.Despite the potential reduction of exploration, we would like to clarify that our method is still effective as exploitation is significant in most environments. There is no free lunch in the exploitationexploration trade-off. One can further control the trade-off by interpolating PBP-GFN with explorative GFNs, e.g., MaxEnt, which can be an interesting future work direction.Improving the performance of GFLowNets can significantly influence various domains, particularly in biology (e.g., molecule and RNA sequence generation as discussed in Section 5). However, these improvements also pose potential risks including the creation of harmful drugs and the misuse of synthesized molecules.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We discuss the limitations of the work in Section 6. In summary, our PBP-GFN makes a trade-off between exploitation and exploration, i.e., obtaining high-reward trajectories and diverse trajectories Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • 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.In this work, we refer to flow matching as the learning scheme that aligns the forward and backward flow.38th Conference on Neural Information Processing Systems (NeurIPS2024).Code: https://github.com/hsjang0/Pessimistic-Backward-Policy-for-GFlowNets.We describe the detailed implementations in Appendix B.https://github.com/recursionpharma/gflownetMIT License, Copyright (c) 2020 Recursion Pharmaceuticals 


Decision: This paper explores the under-exploitation problem in the flow matching of GFlowNets and proposes PBP-GFN.

There were some concerns about the lack of theoretical analysis, which should not be a reason for rejection. The authors addressed most of the concerns raised by the reviewers in their rebuttal. Most importantly, the authors compared their method with a couple of baselines suggested by the reviewers. Even though the proposed method achieves comparable performance to these baselines, they are orthogonal, and the authors show that one can gain more by combining PBP-GFN with these baselines. I recommend an acceptance.

Reviews: review: paper_summary: paper_summary: This paper addresses the under-exploitation of high-reward objects in Generative Flow Networks (GFlowNets) due to the limited observation of trajectories during training. The authors propose PBP-GFN (Pessimistic Backward Policy for GFlowNets), which modifies the backward policy to maximize the observed backward flow, aligning it closer to the true reward. They argue that this pessimistic training scheme encourages the forward policy to assign higher probabilities to high-reward objects, even if they are under-represented in the observed trajectories. Experiments across eight benchmarks, including hypergrid, molecule, and RNA sequence generation, demonstrate that PBP-GFN improves the discovery of high-reward objects while maintaining object diversity.
strength_weakness: strength_weakness: * While the intuition behind PBP-GFN is clear, the paper lacks a strong theoretical analysis to support its claims. The "error bound" analysis in Appendix A is limited in scope and doesn't offer a comprehensive understanding of the algorithm's convergence properties or its impact on the target Boltzmann distribution. 
* The core idea of modifying the backward policy in GFlowNets has been explored in prior works [1, 2, 15]. PBP-GFN, while achieving promising empirical results, doesn't present a significant conceptual departure from these existing approaches.  Its primary contribution appears to be a specific strategy for maximizing the observed backward flow, but the paper lacks a detailed comparison and analysis of how this strategy differs from or improves upon existing techniques.

* The paper primarily focuses on metrics like the number of modes discovered and the average top-100 score. While these metrics are relevant for measuring exploration and exploitation, they provide a limited view of the overall quality and diversity of generated samples.  Evaluating PBP-GFN on more comprehensive downstream task-specific metrics would strengthen the paper's claims.

*  While the paper claims that PBP-GFN maintains object diversity, the pessimistic training scheme inherently introduces a bias towards the observed trajectories. This bias could potentially lead to a reduction in exploration and the generation of degenerate or less diverse samples, especially in domains with significant uncertainty or where observed trajectories are not truly representative of the target distribution.  A deeper analysis and empirical evaluation of this potential bias would be useful.
 
* The performance of PBP-GFN heavily relies on the quality and representativeness of the observed trajectories.  In scenarios where the initial observed trajectories are biased or incomplete, PBP-GFN could amplify these biases, hindering the discovery of truly novel and high-reward objects. The paper doesn't address this sensitivity or propose mitigation strategies. 

* The paper mainly compares PBP-GFN with other backward policy design methods. However, it doesn't provide a thorough comparison with other exploitation-focused techniques in GFlowNets, such as local search GFlowNets [3] or those focusing on higher-reward trajectories [2]. This limited comparison weakens the paper's claim of achieving superior performance in discovering high-reward objects.

[1] Trajectory balance: Improved credit assignment in gflownets
[2] Towards understanding and improving gflownet training.
[3] Local search gflownets
[4] Maximum entropy gflownets with soft q-learning
questions: questions: Please refer to weaknesses section above.
limitations: limitations: N/A.
paper_summary: paper_summary: This paper proposes a pessimistic backward policy for GFlowNets (GFN), which maximizes the expectation of observed backward trajectories. This paper points out the under-exploitation of high-reward objects for previous GFN training methods and provides a deep analysis of this problem. Extensive experiments validate the superior performance of the proposed method.
strength_weakness: strength_weakness: Does the auxiliary objective introduced in Eq. 5 affect the original GFN training objective? Since higher probabilities are assigned to the backward transitions of observed trajectories, could this impact the original assumptions of TB, DB, and other objectives? Do you observe instability during training after some episodes, once the model has seen a sufficient number of samples?
questions: questions: - How do you estimate Eq. 5 in practice? This is an important component of the method, but it is not described.
- What's the meaning of $N$ in algorithm 1? Why does this method need multi-round gradient updates?
limitations: limitations: Limitations have been discussed in this paper.
paper_summary: paper_summary: The authors present a problem with GFlowNets training that originates from not having seen backward trajectories for a particular terminal state. The authors show that because of the lack of observed trajectories, the backward flows underestimate the probabilities for the observed flows, resulting in a (forward) policy that may not match the reward distribution. The authors propose to mitigate the problem by increasing backward flows for the observed trajectories. The authors provide experimental results in eight different environments, showing the superior performances of their method in terms of mode discovery and distribution fitting.
strength_weakness: strength_weakness: 1. The discussion about lacking backward flow given the unseen trajectories is valid. In practice, one can mitigate this by having a uniform backward policy with **a reward-prioritized replay buffer** [1], which also can increase the probability a trajectory yielding high rewards, thereby tackling the mentioned problem. I feel adding this baseline will improve the paper.

2. As an exploitation method, at least one experiment with a larger state-space should be useful to show its performance in larger state-spaces. Hypergrid with a larger dimension and horizon can be one option for this.


[1] Shen, Max W., et al. "Towards understanding and improving gflownet training." International Conference on Machine Learning. PMLR, 2023.
questions: questions: 1. What is the hypermeters for GAFN in Figure 9? Were the hyperparameters tuned for the experiments?
2. The example in Figure 3 raises a question: if $P_B(x|\tau)[i]$ is maximized, $P_B(x|\tau)[j] \approx 0; j \neq i$, thereby matching it will cause $P_F(x|\tau)[j] \approx 0; j \neq i$. In smaller state-space, it may not be a problem, but in a larger state-space, this should be an issue as it may cause $p(x) \approx 0$ where x necessitates action j whose probability we've just made approximately 0. Curious to hear what you think about this/whether you've done experiments in large-scale environments.
limitations: limitations: Yes. Limitations are adequately addressed.
paper_summary: paper_summary: The paper successfully identifies and addresses the under-exploitation problem in the flow matching of GFlowNets. The proposed method, Pessimistic Backward Policy (PBP-GFN), adjusts the probabilities of backward trajectories to improve exploration and exploitation, achieving superior performance in various benchmarks.
strength_weakness: strength_weakness: - The PBP-GFN introduces additional complexity to the training process of GFlowNets. Pessimistic training of backward policy might be computationally expensive or hard to tune in different settings. The author should discuss this in the updated version of the paper. 
- In the experiment sections, this paper does not compare with some new GFlowNets methods, such as [1, 2], that claim improvement over the baselines. I wonder (1) Are PBP-GFN's results better than [1, 2]? (2) Can we combine the pessimistic backward policy with strategies, such as those in [1, 2], to further improve performance? I understand that each paper deals with GFlowNets from a different perspective, therefore, I am ok if not all answers are yes. 
- There are some missing references that are also quite relevant to the topic. [3,4]
- I would raise my scores if my concerns are resolved. 

[1] Kim, M., Yun, T., Bengio, E., Zhang, D., Bengio, Y., Ahn, S., & Park, J. (2023). Local search gflownets. arXiv preprint arXiv:2310.02710.

[2] Jang, H., Kim, M., & Ahn, S. (2023). Learning Energy Decompositions for Partial Inference of GFlowNets. arXiv preprint arXiv:2310.03301.

[3] Chen, Y., & Mauch, L. (2023). Order-Preserving GFlowNets. arXiv preprint arXiv:2310.00386.

[4] Tiapkin, Daniil, et al. "Generative flow networks as entropy-regularized rl." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.
questions: questions: - When minimizing the negative log-likelihood in Eq.5, how to ensure, or do we need to ensure that the sum of the backward probability on a state is 1, and why the total amount of backward flows will not change claimed in line 157? In the provided codebase "PBP.train_proxy", it seems the sum of log_pb_actions will change, and the total number of flows will change. 

- This paper includes MIS from [1]. Can you include more difficult graph combinatorial optimization problems from [1], such as maximum cut? (Maximum clique is not necessary since it is related to MIS)

[1] Zhang, D., Dai, H., Malkin, N., Courville, A., Bengio, Y., & Pan, L. (2023). Let the flows tell: Solving graph combinatorial optimization problems with gflownets. arXiv preprint arXiv:2305.17010.
limitations: limitations: Please see the **Weaknesses** part.

Comments: comment: Dear Reviewer u41F,

We appreciate your constructive feedback for improving our paper in many aspects. We have provided the requested analysis, discussions, and experiments. We are curious whether our rebuttal has resolved your concerns. Thank you again for your time and effort.
Dear reviewer WPBt,

We are happy to hear that our efforts have addressed your concerns! We also appreciate your insightful comments on our works.
Dear reviewer CTvM,

Thanks you for your response. We appreciate your insightful and constructive feedback for improving our paper in many aspects.

Additionally, we would like to emphasize that we have already empirically analyzed the reduction of exploration (as an exploitation method) in large state spaces by measuring the diversity of high-score objects (**Figure 7(c)**). However, we observed no particular reduction in practice. In our future manuscript, we will also incorporate an analysis of the case where our PBP-GFN may reduce exploration (**Figure B** of rebuttal PDF) and discuss mitigation strategies with contents from (**u41F-W4,5**). We hope that this addresses your concerns.
Dear reviewer smvU,

We also appreciate your insightful comments. Thank you again for your time and effort for improving our paper!
Thank you for your rebuttal. I think the method’s promise is weakened because of Q2. Since it is an exploration method, further analysis of this phenomenon in a large state space is required. Hence, I would like to keep my score.
Thank you for the detailed clarifications and additional experiments. I decide to raise my score to 6.
Thanks for your response and hard work. After reading your rebuttal and other reviews, I've decided to maintain my score.