Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation
Year: 2021

Abstract:
This paper is about the problem of learning a stochastic policy for generating an object (like a molecular graph) from a sequence of actions, such that the probability of generating an object is proportional to a given positive reward for that object. Whereas standard return maximization tends to converge to a single return-maximizing sequence, there are cases where we would like to sample a diverse set of high-return solutions. These arise, for example, in black-box function optimization when few rounds are possible, each with large batches of queries, where the batches should be diverse, e.g., in the design of new molecules. One can also see this as a problem of approximately converting an energy function to a generative distribution. While MCMC methods can achieve that, they are expensive and generally only perform local exploration. Instead, training a generative policy amortizes the cost of search during training and yields to fast generation. Using insights from Temporal Difference learning, we propose GFlowNet, based on a view of the generative process as a flow network, making it possible to handle the tricky case where different trajectories can yield the same final state, e.g., there are many ways to sequentially add atoms to generate some molecular graph. We cast the set of trajectories as a flow and convert the flow consistency equations into a learning objective, akin to the casting of the Bellman equations into Temporal Difference methods. We prove that any global minimum of the proposed objectives yields a policy which samples from the desired distribution, and demonstrate the improved performance and diversity of GFlowNet on a simple domain where there are many modes to the reward function, and on a molecule synthesis task.

heading: 1 Introduction
text: The maximization of expected return R in reinforcement learning (RL) is generally achieved by putting all the probability mass of the policy π on the highest-return sequence of actions. In this paper, we study the scenario where our objective is not to generate the single highest-reward sequence of actions but rather to sample a distribution of trajectories whose probability is proportional to a given positive return or reward function. This can be useful in tasks where exploration is important, i.e., we want to sample from the leading modes of the return function. This is equivalent to the problem of turning an energy function into a corresponding generative model, where the object to be generated is obtained via a sequence of actions. By changing the temperature of the energy function (i.e., scaling it multiplicatively) or by taking the power of the return, one can control how selective the generator should be, i.e., only generate from around the highest modes at low temperature or explore more with a higher temperature.
A motivating application for this setup is iterative black-box optimization where the learner has access to an oracle which can compute a reward for a large batch of candidates at each round, e.g., in drug-discovery applications. Diversity of the generated candidates is particularly important when the oracle is itself uncertain, e.g., it may consist of cellular assays which is a cheap proxy for clinical
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
trials, or it may consist of the result of a docking simulation (estimating how well a candidate small molecule binds to a target protein) which is a proxy for more accurate but more expensive downstream evaluations (like cellular assays or in-vivo assays in mice).
When calling the oracle is expensive (e.g. it involves a biological experiment), a standard way (Angermueller et al., 2020) to apply machine learning in such exploration settings is to take the data already collected from the oracle (say a set of (x, y) pairs where x is a candidate solution an y is a scalar evaluation of x from the oracle) and train a supervised proxy f (viewed as a simulator) which predicts y from x. The function f or a variant of f which incorporates uncertainty about its value, like in Bayesian optimization (Srinivas et al., 2010; Negoescu et al., 2011), can then be used as a reward function R to train a generative model or a policy that will produce a batch of candidates for the next experimental assays. Searching for x which maximizes R(x) is not sufficient because we would like to sample for the batch of queries a representative set of x’s with high values of R, i.e., around modes of R(x). Note that alternative ways to obtain diversity exist, e.g., with batch Bayesian optimization (Kirsch et al., 2019). An advantage of the proposed approach is that the computational cost is linear in the size of the batch (by opposition with methods which compare pairs of candidates, which is at least quadratic). With the possibility of assays of a hundred thousand candidates using synthetic biology, linear scaling would be a great advantage.
In this paper, we thus focus on the specific machine learning problem of turning a given positive reward or return function into a generative policy which samples with a probability proportional to the return. In applications like the one mentioned above, we only apply the reward function after having generated a candidate, i.e., the reward is zero except in a terminal state, and the return is the terminal reward. We are in the so-called episodic setting of RL.
The proposed approach views the probability assigned to an action given a state as the flow associated with a network whose nodes are states, and outgoing edges from that node are deterministic transitions driven by an action (not to be confused with normalizing flows; Rezende and Mohamed (2016)). The total flow into the network is the sum of the rewards in the terminal states (i.e., a partition function) and can be shown to be the flow at the root node (or start state). The proposed algorithm is inspired by Bellman updates and converges when the incoming and outgoing flow into and out of each state match. A policy which chooses an action with probability proportional to the outgoing flow corresponding to that action is proven to achieve the desired result, i.e., the probability of sampling a terminal state is proportional to its reward. In addition, we show that the resulting setup is off-policy; it converges to the above solution even if the training trajectories come from a different policy, so long as it has large enough support on the state space.
The main contributions of this paper are as follows:
• We propose GFlowNet, a novel generative method for unnormalized probability distributions based on flow networks and local flow-matching conditions: the flow incoming to a state must match the outgoing flow.
• We prove crucial properties of GFlowNet, including the link between the flow-matching conditions (which many training objectives can provide) and the resulting match of the generated policy with the target reward function. We also prove its offline properties and asymptotic convergence (if the training objective can be minimized). We also demonstrate that previous related work (Buesing et al., 2019) which sees the generative process like a tree would fail when there are many action sequences which can lead to the same state.
• We demonstrate on synthetic data the usefulness of departing from seeking one mode of the return, and instead seeking to model the entire distribution and all its modes.
• We successfully apply GFlowNet to a large scale molecule synthesis domain, with comparative experiments against PPO and MCMC methods.
All implementations are available at https://github.com/bengioe/gflownet.

heading: 2 Approximating Flow Network generative models with a TD-like objective
text: Consider a discrete set X and policy π(a|s) to sequentially build x ∈ X with probability π(x) with
π(x) ≈ R(x) Z
= R(x)∑
x′∈X R(x ′)
(1)
where R(x) > 0 is a reward for a terminal state x. This would be useful to sample novel drug-like molecules when given a reward function R that scores molecules based on their chemical properties. Being able to sample from the high modes ofR(x) would provide diversity in the batches of generated molecules sent to assays. This is in contrast with the typical RL objective of maximizing return which we have found to often end up focusing around one or very few good molecules. In our context, R(x) is a proxy for the actual values obtained from assays, which means it can be called often and cheaply. R(x) is retrained or fine-tuned each time we acquire new data from the assays.
What method should one use to generate batches sampled from π(x) ∝ R(x)? Let’s first think of the state space under which we would operate.
Let S denote the set of states and X ⊂ S denote the set of terminal states. Let A be a finite set, the alphabet, A(s) ⊆ A be the set of allowed actions at state s, and let A∗(s) be the set of all sequences of actions allowed after state s. To every action sequence ~a = (a1, a2, a3, ..., ah) of ai ∈ A, h ≤ H corresponds a single x, i.e. the environment is deterministic so we can define a function C mapping a sequence of actions ~a to an x. If such a sequence is ‘incomplete’ we define its reward to be 0. When the correspondence between action sequences and states is bijective, a state s is uniquely described by some sequence ~a, and we can visualize the generative process as the traversal of a tree from a single root node to a leaf corresponding to the sequence of actions along the way.
However, when this correspondence is non-injective, i.e. when multiple action sequences describe the same x, things get trickier. Instead of a tree, we get a directed acyclic graph or DAG (assuming that the sequences must be of finite length, i.e., there are no deterministic cycles), as illustrated in Figure 1. For example, and of interest here, molecules can be seen as graphs, which can be described in multiple orders (canonical representations such as SMILES strings also have this problem: there may be multiple descriptions for the same actual molecule). The standard approach to such a sampling problem is to use iterative MCMC methods (Xie et al., 2021; Grathwohl et al., 2021). Another option is to relax the desire to have p(x) ∝ R(x) and to use non-interative (sequential) RL methods (Gottipati et al., 2020), but these are at high risk of getting stuck in local maxima and of missing modes. Indeed, in our setting, the policy which maximizes the expected return (which is the expected final reward) generates the sequence with the highest return (i.e., a single molecule).

heading: 2.1 Flow Networks
text: In this section we propose the Generative Flow Network framework, or GFlowNet, which enables us to learn policies such that p(x) ∝ R(x) when sampled. We first discuss why existing methods are inadequate, and then show how we can use the metaphor of flows, sinks and sources, to construct adequate policies. We then show that such policies can be learned via a flow-matching objective.
With existing methods in the bijective case, one can think of the sequential generation of one x as an episode in a tree-structured deterministic MDP, where all leaves x are terminal states (with reward R(x)) and the root is initial state s0. Interestingly, in such a case one can express the pseudo-value of a state Ṽ (s) as the sum of all the rewards of the descendants of s (Buesing et al., 2019).
In the non-injective case, these methods are inadequate. Constructing π(τ) ≈ R(τ)/Z, e.g. as per Buesing et al. (2019), MaxEnt RL (Haarnoja et al., 2017), or via an autoregressive method (Nash and Durkan, 2019; Shi et al., 2021) has a particular problem as shown below: if multiple action sequences ~a (i.e. multiple trajectories τ ) lead to a final state x, then a serious bias can be introduced in the generative probabilities. Let us denote ~a +~b as the concatenation of the two sequences of actions ~a and~b, and by extension s+~b the state reached by applying the actions in~b from state s. Proposition 1. Let C : A∗ 7→ S associate each allowed action sequence ~a ∈ A∗ to a state s = C(~a) ∈ S. Let Ṽ : S 7→ R+ associate each state s ∈ S to Ṽ (s) = ∑~b∈A∗(s)R(s +~b) > 0, whereA∗(s) is the set of allowed continuations from s and s+~b denotes the resulting state, i.e., Ṽ (s) is the sum of the rewards of all the states reachable from s. Consider a policy π which starts from the state corresponding to the empty string s0 = C(∅) and chooses from state s ∈ S an allowable action a ∈ A(s) with probability π(a|s) = Ṽ (s+a)∑
b∈A(s) Ṽ (s+b) . Denote π(~a = (a1, . . . , aN )) =∏N
i=1 π(ai|C(a1, . . . , ai−1)) and π(s) with s ∈ S the probability of visiting a state s with this policy. The following then obtains: (a) π(s) = ∑ ~ai:C(~ai)=s π(~ai).
(b) IfC is bijective, then π(s) = Ṽ (s) Ṽ (s0) and as a special case for terminal states x, π(x) = R(x)∑ x∈X R(x)
. (c) If C is non-injective and there are n(x) distinct action sequences ~ai s.t. C(~ai) = x, then π(x) = n(x)R(x)∑
x′∈X n(x ′)R(x′) .
See Appendix A.1 for the proof. In combinatorial spaces, such as for molecules, where C is noninjective (there are many ways to construct a molecule graph), this can become exponentially bad as trajectory lengths increase. It means that larger molecules would be exponentially more likely to be sampled than smaller ones, just because of the many more paths leading to them. In this scenario, the pseudo-value Ṽ is “misinterpreting” the MDP’s structure as a tree, leading to the wrong π(x).
An alternative is to see the MDP as a flow network, that is, leverage the DAG structure of the MDP, and learn a flow F , rather than estimating the pseudo-value Ṽ as a sum of descendant rewards, as elaborated below. We define the flow network as a having a single source, the root node (or initial state) s0 with in-flow Z, and one sink for each leaf (or terminal state) x with out-flow R(x) > 0. We write T (s, a) = s′ to denote that the state-action pair (s, a) leads to state s′. Note that because C is not a bijection, i.e., there are many paths (action sequences) leading to some node, a node can have multiple parents, i.e. |{(s, a) | T (s, a) = s′}| ≥ 1, except for the root, which has no parent. We write F (s, a) for the flow between node s and node s′ = T (s, a), F (s) for the total flow going through s1. This construction is illustrated in Fig. 1.
∑ ∑
To satisfy flow conditions, we require that for any node, the incoming flow equals the outgoing flow, which is the total flow F (s) of node s. Boundary conditions are given by the flow into the terminal nodes x, R(x). Formally, for any node s′, we must have that the in-flow
F (s′) = ∑
s,a:T (s,a)=s′
F (s, a) (2)
equals the out-flow F (s′) = ∑ a′∈A(s′) F (s′, a′). (3)
More concisely, with R(s) = 0 for interior nodes, and A(s) = ∅ for leaf (sink/terminal) nodes, we write the following flow consistency equations:∑
s,a:T (s,a)=s′
F (s, a) = R(s′) + ∑
a′∈A(s′)
F (s′, a′). (4)
with F being a flow, F (s, a) > 0 ∀s, a (for this we needed to constrain R(x) to be positive too). One could include in principle nodes and edges with zero flow but it would make it difficult to talk about the logarithm of the flow, as we do below, and such states can always be excluded by the allowed set of actions for their parent states. Let us now show that such a flow correctly produces π(x) = R(x)/Z when the above flow equations are satisfied.
1In some sense, F (s) and F (s, a) are close to V and Q, RL’s value and action-value functions. These effectively inform an agent taking decisions at each step of an MDP to act in a desired way. With some work, we can also show an equivalence between F (s, a) and the “real” Qπ̂ of some policy π̂ in a modified MDP (see A.2).
Proposition 2. Let us define a policy π that generates trajectories starting in state s0 by sampling actions a ∈ A(s) according to
π(a|s) = F (s, a) F (s)
(5)
where F (s, a) > 0 is the flow through allowed edge (s, a), F (s) = R(s) + ∑ a∈A(s) F (s, a) where R(s) = 0 for non-terminal nodes s and F (x) = R(x) > 0 for terminal nodes x, and the flow consistency equation ∑ s,a:T (s,a)=s′ F (s, a) = R(s ′) + ∑ a′∈A(s′) F (s
′, a′) is satisfied. Let π(s) denote the probability of visiting state s when starting at s0 and following π(·|·). Then (a) π(s) = F (s)F (s0) (b) F (s0) = ∑ x∈X R(x) (c) π(x) = R(x)∑ x′∈X R(x ′) .
Proof. We have π(s0) = 1 since we always start in root node s0. Note that ∑ x∈X π(x) = 1
because terminal states are mutually exclusive, but in the case of non-bijective C, we cannot say that∑ s∈S π(s) equals 1 because the different states are not mutually exclusive in general. This notation is different from the one typically used in RL where π(s) refers to the asymptotic distribution of the Markov chain. Then
π(s′) = ∑
(a,s):T (s,a)=s′
π(a|s)π(s) (6)
i.e., using Eq. 5,
π(s′) = ∑
(a,s):T (s,a)=s′
F (s, a)
F (s) π(s). (7)
We can now conjecture that the statement
π(s) = F (s)
F (s0) (8)
is true and prove it by induction. This is trivially true for the root, which is our base statement, since π(s0) = 1. By induction, we then have that if the statement is true for parents s of s′, then
π(s′) = ∑
s,a:T (s,a)=s′
F (s, a)
F (s)
F (s)
F (s0) =
∑ s,a:T (s,a)=s′ F (s, a)
F (s0) = F (s′) F (s0) (9)
which proves the statement, i.e., the first conclusion (a) of the theorem. We can then apply it to the case of terminal states x, whose flow is fixed to F (x) = R(x) and obtain
π(x) = R(x)
F (s0) . (10) Noting that ∑ x∈X π(x) = 1 and summing both sides of Eq. 10 over x we thus obtain (b), i.e.,
F (s0) = ∑ x∈X R(x). Plugging this back into Eq. 10, we obtain (c), i.e., π(x) = R(x)∑ x′∈X R(x ′) .
Thus our choice of π(a|s) satisfies our desiderata: it maps a reward function R to a generative model which generates x with probability π(x) ∝ R(x), whether C is bijective or non-injective (the former being a special case of the latter, and we just provided a proof for the general non-injective case).

heading: 2.2 Objective Functions for GFlowNet
text: We can now leverage our RL intuitions to create a learning algorithm out of the above theoretical results. In particular, we propose to approximate the flows F such that the flow consistency equations are respected at convergence with enough capacity in our estimator of F , just like the Bellman equations for temporal-difference (TD) algorithms (Sutton and Barto, 2018). This could yield the following objective for a trajectory τ :
L̃θ(τ) = ∑
s′∈τ 6=s0  ∑ s,a:T (s,a)=s′ Fθ(s, a)−R(s′)− ∑ a′∈A(s′) Fθ(s ′, a′) 2 . (11)
One issue from a learning point of view is that the flow will be very large for nodes near the root (early in the trajectory) and tiny for nodes near the leaves (late in the trajectory). In high-dimensional spaces where the cardinality of X is exponential (e.g., in the typical number of actions to form an x), the F (s, a) and F (s) for early states will be exponentially larger than for later states. Since we want F (s, a) to be the output of a neural network, this would lead to serious numerical issues.
To avoid this problem, we define the flow matching objective on a log-scale, where we match not the incoming and outgoing flows but their logarithms, and we train our predictor to estimate F logθ (s, a) = logF (s, a), and exponentiate-sum-log theF log θ predictions to compute the loss, yielding the square of a difference of logs:
Lθ, (τ) = ∑
s′∈τ 6=s0 log  +∑ s,a:T (s,a)=s′ expF logθ (s, a) − log  +R(s′) +∑ a′∈A(s′) expF logθ (s ′, a′) 2 (12) which gives equal gradient weighing to large and small magnitude predictions. Note that matching the logs of the flows is equivalent to making the ratio of the incoming and outgoing flow closer to 1. To give more weight to errors on large flows and avoid taking the logarithm of a tiny number, we compare log( +incoming flow) with log( +outgoing flow). It does not change the global minimum, which is still when the flow equations are satisfied, but it avoids numerical issues with taking the log of a tiny flow. The hyper-parameter also trades-off how much pressure we put on matching large versus small flows, and in our experiments is set to be close to the smallest value R can take. Since we want to discover the top modes of R, it makes sense to care more for the larger flows. Many other objectives are possible for which flow matching is also a global minimum.
An interesting advantage of such objective functions is that they yield off-policy offline methods. The predicted flows F do not depend on the policy used to sample trajectories (apart from the fact that the samples should sufficiently cover the space of trajectories in order to obtain generalization). This is formalized below, which shows that we can use any broad-support policy to sample training trajectories and still obtain the correct flows and generative model, i.e., training can be off-policy. Proposition 3. Let trajectories τ used to train Fθ be sampled from an exploratory policy P with the same support as the optimal π defined in Eq. 5 for a consistent flow F ∗ ∈ F∗. A flow is consistent if Eq. 4 is respected. Also assume that ∃θ : Fθ = F ∗, i.e., we choose a sufficiently rich family of predictors. Let θ∗ ∈ argminθEP (τ)[Lθ(τ)] a minimizer of the expected training loss. Let Lθ(τ) have the property that when flows are matched it achieves its lowest possible value. First, it can be shown that this property is satisfied for the loss in Eq. 12. Then
Fθ∗ = F ∗, and Lθ∗(τ) = 0 ∀τ ∼ P (θ), (13)
i.e., a global optimum of the expected loss provides the correct flows. If πθ∗(a|s) = Fθ∗ (s,a)∑ a′∈A(s) Fθ∗ (s,a ′) then we also have
πθ∗(x) = R(x)
Z . (14)
The proof is in Appendix A.1. Note that, in RL terms, this method is akin to asynchronous dynamic programming (Sutton and Barto, 2018, §4.5), which is an off-policy off-line method which converges provided every state is visited infinitely many times asymptotically.

heading: 3 Related Work
text: The objective of training a policy generating states with a probability proportional to rewards was presented by Buesing et al. (2019) but the proposed method only makes sense when there is a bijection between action sequences and states. In contrast, GFlowNet is applicable in the more general setting where many paths can lead to the same state. The objective to sample with probability proportional to a given unnormalized positive function is achieved by many MCMC methods (Grathwohl et al., 2021; Dai et al., 2020). However, when mixing between modes is challenging (e.g., in high-dimensional spaces with well-separated modes occupying a fraction of the total volume) convergence to the target distribution can be extremely slow. In contrast, GFlowNet is not iterative and amortizes the challenge of sampling from such modes through a training procedure which must be sufficiently exploratory.
This sampling problem comes up in molecule generation and has been studied in this context with numerous generative models (Shi et al., 2020; Jin et al., 2020; Luo et al., 2021), MCMC methods (Seff
et al., 2019; Xie et al., 2021), RL (Segler et al., 2017; Cao and Kipf, 2018; Popova et al., 2019; Gottipati et al., 2020; Angermueller et al., 2020) and evolutionary methods (Brown et al., 2004; Jensen, 2019; Swersky et al., 2020). Some of these methods rely on a given set of "positive examples" (high-reward) to train a generative model, thus not taking advantage of the "negative examples" and the continuous nature of the measurements (some examples should be generated more often than others). Others rely on the traditional return maximization objectives of RL, which tends to focus on one or a few dominant modes, as we find in our experiments. Beyond molecules, there are previous works generating data non-greedily through RL (Bachman and Precup, 2015) or energy-based GANs (Dai et al., 2017).
The objective that we formulate in (12) may remind the reader of the objective of control-asinference’s Soft Q-Learning (Haarnoja et al., 2017), with the difference that we include all the parents of a state in the in-flow, whereas Soft Q-Learning only uses the parent contained in the trajectory. Soft Q-Learning induces a different policy, as shown by Proposition 1, one where P (τ) ∝ R(τ) rather than P (x) ∝ R(x). More generally, we only consider deterministic generative settings whereas RL is a more general framework for stochastic environments.
Literature at the intersection of network flow and deep learning is sparse, and is mostly concerned with solving maximum flow problems (Nazemi and Omidi, 2012; Chen and Zhang, 2020) or classification within existing flow networks (Rahul et al., 2017; Pektaş and Acarman, 2019). Finally, the idea of accounting for the search space being a DAG rather than a tree in MCTS, known as transpositions (Childs et al., 2008), also has some links with the proposed method.

heading: 4 Empirical Results
text: We first verify that GFlowNet works as advertised on an artificial domain small enough to compute the partition function exactly, and compare its abilities to recover modes compared to standard MCMC and RL methods, with its sampling distribution better matching the normalized reward. We find that GFlowNet (A) converges to π(x) ∝ R(x), (B) requires less samples to achieve some level of performance than MCMC and PPO methods and (C) recovers all the modes and does so faster than MCMC and PPO, both in terms of wall-time and number of states visited and queried. We then test GFlowNet on a large scale domain, which consists in generating small drug molecule graphs, with a reward that estimates their binding affinity to a target protein (see Appendix A.3). We find that GFlowNet finds higher reward and more diverse molecules faster than baselines.

heading: 4.1 A (hyper-)grid domain
text: Consider an MDP where states are the cells of a n-dimensional hypercubic grid of side length H . The agent starts at coordinate x = (0, 0, ...) and is only allowed to increase coordinate i with action ai (up to H , upon which the episode terminates). A stop action indicates to terminate the trajectory. There are many action sequences that lead to the same coordinate, making this MDP a DAG.The reward for ending the trajectory in x is some R(x) > 0. For MCMC methods, in order to have an ergodic chain, we allow the iteration to decrease coordinates as well, and there is no stop action.
We ran experiments with this reward function: R(x) = R0+R1 ∏ i I(0.25 < |xi/H−0.5|)+R2 ∏ i I(0.3 < |xi/H−0.5| < 0.4)
0 10 20 30 40 50 60
0 10 20 30 40 50 60
with 0 < R0 R1 < R2, pictured when n = 2 on the right. For this choice of R, there are only interesting rewards near the corners of the grid, and there are exactly 2n modes. We set R1 = 1/2, R2 = 2. By varyingR0 and setting it closer to 0, we make this problem artificially harder, creating a region of the state space which it is undesirable to explore. To measure the performance of a method, we measure the empirical L1 error E[|p(x)− π(x)|]. p(x) = R(x)/Z is known in this domain, and π is estimated by repeated sampling and counting frequencies for each possible x. We also measure the number of modes with at least 1 visit as a function of the number of states visited.
We run the above experiment for R0 ∈ {10−1, 10−2, 10−3} with n = 4, H = 8. In Fig. 2 we see that GFlowNet is robust to R0 and obtains a low L1 error, while a Metropolis-Hastings-MCMC based method requires exponentially more samples than GFlowNet to achieve some level of L1 error. This is apparent in Fig. 2 (with a log-scale horizontal axis) by comparing the slope of progress of GFlowNet (beyond the initial stage) and that of the MCMC sampler. We also see that MCMC takes much longer to visit each mode once as R0 decreases, while GFlowNet is only slightly affected, with GFlowNet converging to some level of L1 error faster, as per hypothesis (B). This suggests that
GFlowNet is robust to the separation between modes (represented by R0 being smaller) and thus recovers all the modes much faster than MCMC (again, noting the log-scale of the horizontal axis).
To compare to RL, we run PPO (Schulman et al., 2017). To discover all the modes in a reasonable time, we need to set the entropy maximization term much higher (0.5) than usual ( 1). We verify that PPO is not overly regularized by comparing it to a random agent. PPO finds all the modes faster than uniform sampling, but much more slowly than GFlowNet, and is also robust to the choice of R0. This and the previous result validates hypothesis (C). We also run SAC (Haarnoja et al., 2018), finding similar or worse results. We provide additional results and discussion in Appendix A.6.

heading: 4.2 Generating small molecules
text: Here our goal is to generate a diverse set of small molecules that have a high reward. We define a large-scale environment which allows an agent to sequentially generate molecules. This environment is challenging, with up to 1016 states and between 100 and 2000 actions depending on the state.
We follow the framework of Jin et al. (2020) and generate molecules by parts using a predefined vocabulary of building blocks that can be joined together forming a junction tree (detailed in A.3). This is also known as fragment-based drug design (Kumar et al., 2012; Xie et al., 2021). Generating such a graph can be described as a sequence of additive edits: given a molecule and constraints of chemical validity, we choose an atom to attach a block to. The action space is thus the product of choosing where to attach a block and choosing which block to attach. There is an extra action to stop the editing sequence. This sequence of edits yields a DAG MDP, as there are multiple action sequences that lead to the same molecule graph, and no edge removal actions, which prevents cycles.
The reward is computed with a pretrained proxy model that predicts the binding energy of a molecule to a particular protein target (soluble epoxide hydrolase, sEH, see A.3). Although computing binding energy is computationally expensive, we can call this proxy cheaply. Note that for realistic drug design, we would need to consider many more quantities such as drug-likeness (Bickerton et al., 2012), toxicity, or synthesizability. Our goal here is not solve this problem, and our work situates itself within such a larger project. Instead, we want to show that given a proxy R in the space of molecules, we can quickly match its induced distribution π(x) ∝ R(x) and find many of its modes. We parameterize the proxy with an MPNN (Gilmer et al., 2017) over the atom graph. Our flow predictor Fθ is parameterized similarly to MARS (Xie et al., 2021), with an MPNN, but over the junction tree graph (the graph of blocks), which had better performance. For fairness, this architecture is used for both GFlowNet and the baselines. Complete details can be found in Appendix A.4.
We pretrain the proxy with a semi-curated semi-random dataset of 300k molecules (see A.4) down to a test MSE of 0.6; molecules are scored according to the docking score (Trott and Olson, 2010), renormalized so that most scores fall between 0 and 10 (to have R(x) > 0). We plot the dataset’s reward distribution in Fig. 3. We train all generative models with up to 106 molecules. During training, sampling follows exploratory policy P (a|s) which is a mixture between π(a|s) (Eq. 5), used with probability 0.95, and a uniform distribution over allowed actions with probability 0.05.
Experimental results In Fig. 3 we show the empirical distribution of rewards in two settings; first when we train our model with R(x), then with R(x)β . If GFlowNet learns a reasonable policy π,
this should shift the distribution to the right. This is indeed what we observe. We compare GFlowNet to MARS (Xie et al., 2021), known to work well in the molecule domain, and observe the same shift. Note that GFlowNet finds more high reward molecules than MARS with these β values; this is consistent with the hypothesis that it finds high-reward modes faster (since MARS is an MCMC method, it would eventually converge to the same distribution, but takes more time).
In Fig. 4, we show the average reward of the top-k molecules found so far, without allowing for duplicates (based on SMILES). We compare GFlowNet with MARS, PPO, and JT-VAE (Jin et al., 2020) with Bayesian Optimization. As expected, PPO plateaus after a while; RL tends to be satisfied with good enough trajectories unless it is strongly regularized with exploration mechanisms. For GFlowNet and for MARS, the more molecules are visited, the better they become, with a slow convergence towards the proxy’s max reward. Given the same compute time, JT-VAE+BO generates only about 103 molecules (due to its expensive Gaussian Process) and so does not perform well.
The maximum reward in the proxy’s dataset is 10, with only 233 examples above 8. In our best run, we find 2339 unique molecules during training with a score above 8, only 39 of which are in the dataset. We compute the average pairwise Tanimoto similarity for the top 1000 samples: GFlowNet has a mean of 0.44± 0.01, PPO, 0.62± 0.03, and MARS, 0.59± 0.02 (mean and std over 3 runs). As expected, our MCMC baseline (MARS) and RL baseline (PPO) find less diverse candidates. We also find that GFlowNet discovers many more modes (>1500 with R>8 vs <100 for MARS). This is shown in Fig. 5 where we consider a mode to be a Bemis-Murcko scaffold (Bemis and Murcko, 1996), counted for molecules above a certain reward threshold. We provide additional insights into how GFlowNet matches the rewards in Appendix A.7.

heading: 4.3 Multi-Round Experiments
text: To demonstrate the importance of diverse candidate generation in an active learning setting, we consider a sequential acquisition task. We simulate the setting where there is a limited budget for calls to the true oracle O. We use a proxy M initialized by training on a limited dataset of (x,R(x))
pairs D0, where R(x) is the true reward from the oracle. The generative model (πθ) is trained to fit to the unnormalized probability function learned by the proxy M . We then sample a batch B = {x1, x2, . . . xk} where xi ∼ πθ, which is evaluated with the oracle O. The proxy M is updated with this newly acquired and labeled batch, and the process is repeated for N iterations. We discuss the experimental setting in more detail in Appendix A.5.
Hyper-grid domain We present results for the multi-round task in the 4-D hyper-grid domain in Figure 6. We use a Gaussian Process (Williams and Rasmussen, 1995) as the proxy. We compare the Top-k Return for all the methods, which is defined as mean(top -k(Di))−mean(top -k(Di−1)), where Di is the dataset of points acquired until step i, and k = 10 for this experiment. The initial dataset D0 (|D0| = 512) is the same for all the methods compared. We observe that GFlowNet consistently outperforms the baselines in terms of return over the initial set. We also observe that the mean pairwise L2-distance between the top -k points at the end of the final round is 0.83± 0.03, 0.61± 0.01 and 0.51± 0.02 for GFlowNet, MCMC and PPO respectively. This demonstrates the ability of GFlowNet to capture the modes, even in the absence of the true oracle, as well as the importance of capturing this diversity in multi-round settings.
Small Molecules For the molecule discovery task, we initialize an MPNN proxy to predict docking scores from AutoDock (Trott and Olson, 2010), with |D0| = 2000 molecules. At the end of each round we generate 200 molecules which are evaluated with AutoDock and used to update the proxy. Figure 7 shows GFlowNet discovers molecules with significantly higher energies than the initial set D0. It also consistently outperforms MARS as well as Random Acquisition. PPO training was unstable and diverged consistently so the numbers are not reported. The mean pairwise Tanimoto similarity in the initial set is 0.60. At the end of the final round, it is 0.54 ± 0.04 for GFlowNet and 0.64± 0.03 for MARS. This further demonstrates the ability of GFlowNet to generate diverse candidates, which ultimately helps improve the final performance on the task. Similar to the single step setting, we observe that JT-VAE+BO is only able to generate 103 molecules with similar compute time, and thus performs poorly.

heading: 5 Discussion & Limitations
text: In this paper we have introduced a novel TD-like objective for learning a flow for each state and (state, action) pair such that policies sampling actions proportional to these flows draw terminal states in proportion to their reward. This can be seen as an alternative approach to turn an energy function into a fast generative model, without the need for an iterative method like that needed with MCMC methods, and with the advantage that when training succeeds, the policy generates a great diversity of samples near the main modes of the target distribution without being slowed by issues of mixing between modes.
Limitations. One downside of the proposed method is that, as for TD-based methods, the use of bootstrapping may cause optimization challenges (Kumar et al., 2020; Bengio et al., 2020) and limit its performance. In applications like drug discovery, sampling from the regions surrounding each mode is already an important advantage, but future work should investigate how to combine such a generative approach to local optimization in order to refine the generated samples and approach the local maxima of reward while keeping the batches of candidates diverse.
Negative Social Impact. The authors do not foresee negative social impacts of this work specifically.
Acknowledgments and Disclosure of Funding
This research was enabled in part by computational resources provided by Calcul Québec (www. calculquebec.ca) and Compute Canada (www.computecanada.ca). All authors are funded by their primary academic institution. We also acknowledge funding from Samsung Electronics Co., Ldt., CIFAR and IBM.
The authors are grateful to Andrei Nica for generating the molecule dataset, to Maria Kadukova for advice on molecular docking, to Harsh Satija for feedback on the paper, as well as to all the members of the Mila Molecule Discovery team for the many research discussions on the challenges we faced.
Author Contributions
EB and YB contributed to the original idea, and wrote most sections of the paper. YB wrote the proofs of Propositions 1-3, EB the proof of Proposition 4. EB wrote the code and ran experiments for sections 4.1 (hypergrid) and 4.2 (small molecules). MJ wrote the code and ran experiments for section 4.3 (multi-round) and wrote the corresponding results section of the paper. MK wrote the biochemical framework upon which the molecule experiments are built, assisted in debugging and running experiments for section 4.3, implemented mode-counting routines used in 4.2, and wrote the biochemical details of the paper.
MK, DP and YB provided supervision for the project. All authors contributed to proofreading and editing the paper.

Review summaries:
The paper considers generative processes for which state transitions form a DAG. The authors view such a generative process as a flow in a graph—with a single source and sinks connected to final objects. The authors set objects' probabilities being proportional to the reward function and train a generative model by enforcing flow-matching conditions. The main contribution of the paper is a new generative approach for unnormalized distributions.

The paper tries to tackle a scenario where the goal is not to learn to generate a sequence of actions that has the highest return but rather to generate a diverse set of trajectories with high returns. They do so by using the given positive return function (which is given by the oracle function) to formulate a generative policy that samples with a probability proportional to the return.  They used a flow network to express the MDP and its value function where the value of a state and action pairs is defined by the flow between the two nodes in the flow network, and finally, the flow consistency equations are used as a learning objective.

This paper proposes an approach (called GFlowNet) that uses ideas from reinforcement learning to learn a "policy" for sampling from a discrete distribution over a set of objects ("terminal states"), where these objects can be traversed between using a sequence of "actions". Specifically, the authors consider distributions proportional to some (non-negative) reward function, such that one is more likely to sample better performing points (although I believe the approach is general enough to sample from any unnormalized distribution).  In the experimental section the authors first consider a toy problem to highlight the advantages of their approach before moving on to consider the problem of designing new molecules, where the "state" is the current molecule, and the "actions" are defined using a vocabulary of fragments which one can add to the current state (as well as the stop action).  Here, the authors show GFlowNet is able to find  a set of better and more diverse molecules more quickly than previous MCMC and PPO approaches.   **More specific details about GFlowNet:**    In GFlowNet, states and actions can be described using a flow network (i.e. a transportation network). Here the states are represented using nodes and the actions as edges (the edge weight, ie flow, describes the probability of taking an action). The authors always consider starting from the same state and actions cannot return to previously seen states (ie fragments can be added to molecules but not taken away). Moreover, one can arrive at the same state with different actions (i.e. different fragments could be used to create the same molecule). Overall, this means that the flow network can be represented using a DAG (directed acyclic graph, see Figure 1).   The core idea of GFlowNet is then about how to learn the policy over the actions (ie the flows) such that when sampling actions from this policy, you sample from a predefined distribution over terminal states. To this end, the authors derive a flow matching objective (equation 9 -- the flow incoming to a state must match the outgoing flow) which can be differentiated through to learn the parameters parameterizing the policy. The authors prove that this loss can be used off-policy.