Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation
Authors: Emmanuel Bengio
Moksh Jain
Maksym Korablyov
Doina Precup
Yoshua Bengio
Venue: NeurIPS 2021 Poster (2021)
Keywords: sequential generative model
temporal difference
diversity
molecules

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.

Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation
Emmanuel B2 Moksh Jain 15, Maksym Korablyov' Doina Precup-.24 12,4 Yoshua Bengio
Mila, MeGiil University, Université de Montréal, ' *DeepMind, Miccosoot
Abstract
This paper about the problem of learning stochastic policy for generating an object (like a molecular graph) from a sequence of actions, such that the probability of generating an object proportional that object. Whereas standard return maximization tends to converge to single return- muxummiaiiizing 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 problem of approximately converting an energy function a generative distribution. While MCMC methods can achieve that, they are expensive and generally only perform local exploration. Instead, training 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 flow network, making it possible handle the tricky case where different trajectories can yield the same final state, c.g., there are many ways sequentially add atoms to generate some molecular graph. We cast the set of trajectories as a flow and convert the flow consistency equations into learning objective, akin the casting of the Bellman equations into Temporal Difference methods. We prove that any global minimum of the proposed objectives yields policy which samples from the desired distribution, and demonstrate the improved performance and diversity of GFlowNer on a simple domain where there are many modes to the reward function, and on molecule synthesis task.
given positíve reward for
Introduction
The maximization of expected return R in reinforcement learning (RL) is generally achieved by putting all the probability mass 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 sample a distribution of trajectories whose probability 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 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) should be, i.c., only generate from around the highest modes at low temperature or explore more with higher temperature.
by taking the power of the return, one control how selective the generator
motivating application for this setup is iterative black-box optimization where the learner has access to an oracle which can compute a reward for large batch of candidates at each round, e.g., in drug-discover applications. Diversity of the generated candidates particularly important when the oracle is itself uncertain, e.g., it may consist of cellular assays which is cheap proxy for clinical
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

trials, or it may consist of the result of a docking simulation (estimating how well a candidate small molecule binds 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), standard way (Anger- mueller et al., 2020) to apply machine learning in such exploration settings is to take the data already collected from the oracle (say set of (a, y) pairs where a candidate solution an y a scalar evaluation of x from the oracle) and train supervised proxy f (viewed as a simulator) which predicts y from 2. The function f or variant of f which incorporates uncertainty about its value, like in Bayesian optimization (Srinivas et al., 2010; Negoescu et al., 201 1), can then be used as a reward function R. to train a generative model or policy that will produce 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 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 great advantage.
In this paper, we thus focus on the specific machine learning problem of turning given positive reward or return function into a generative policy which samples with a probability propor- tional the return. In applications like the one mentioned above. we only apply the reward function after having generated candidate, i.e., the reward is zero except in terminal state, and the return is the terminal reward. We are in the so-called episodic setting of RI..
The proposed approach views the probability assigned to an action given a state as the flow associated with 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 the sum of the rewards in the terminal states (i.e., 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- wow-matching conditions: the flow incoming 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 GelowNet a large scale molecule synthesis domain, with comparative experiments against PPO and MCMC methods.
All implementations are available at https: //g:::/gtittt wtt.cccc//////eeee hwtm//engio//gffoooee
2 Approximating Flow Network generative models with a TD-like objective
Consider a discrete set X and policy (als) to sequentially build x E X' with probability x(x) with
R(r)
R(x)
(1)
Z Evex R(x*)
2

---

where R(x) > O is reward for a terminal state z. This would be useful to sample novel drug-like molecules when given reward function R that scores molecules based on their chemical properties. Being able to sample from the high modes of R(x) would provide diversity in the batches of generated molecules sent to assays. This is in contrast with the typical RI. objective of maximizing return which R(x) is retrained or fine-tuned each time we acquire new data from the assays.
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.
What method should one use to generate batches sampled from (2x) 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 C S denote the set of terminal states. Let A be finite set, the alphahet, A(s) S A be the set of allowed actions at state 8, and let A*) be the set of all sequences of actions allowed after state To every action sequence ā [01, (2, (3;--., corresponds a single x. i.e. the environment is deterministic so we can define a function C' mapping a sequence of actions ă to an æ. If such a sequence is 'incomplete" we define its reward be 0. When the correspondence between action sequences and states is bijective, state is uniquely described by some sequence ã. and can visualize the generative process the traversal of tree from single root node to a leaf corresponding to the sequence of actions along the way.
-.. .......))))))))..))))..) of 0g A, h S H
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 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 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(z) x R(xx) and to use non-interative (sequential) RL methods (Gottipati et al., 2020), but these are 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.c., a single molecule).
2.1
Flow Networks
In this section we propose the Generative Flow Network framework, or GFlowNet, which enables us to learn policies such that p(z) 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 tree-structured deterministic MDP, where all leaves are terminal states (with reward R(x)) and the root is initial state so. Interestingly, in such a case one can express the pseudo-value of a state (s) as the sum of all the rewards of the descendants of s (Buesing et al., 2019).
In the non non-injective case, these methods are inadequate. Constructing T(T) = R(T)/Z, c.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 shown below: if multiple action sequences ă (i.e. multiple trajectories 7) lead to a final state x, then a serious bias can be introduced in the generative probabilities. Let us denote ă + b as the concatenation of the two sequences of actions ā and b, and by extension + 5 the state reached by applying the actions in 5 from state s.
Proposition 1. Let C. : A + S associate each allowed action sequence ă E A to a state = =C(() E S. Let V S R associate each state S to V(s) = ETEA' R(s + b) > O,
where (s) is the set of allowed continuations from and + denotes the resulting state, i.e (ss is the sum of the rewards of all the states reachable from 8. Consider a policy n which starts from C(() and chooses from state E S an allowable V(ata)
the state corresponding the empty string so
action A(s) with probability π (als)) = EsEAm V(o+b)- Denote π( = (Q1....,aN)) =
TH (aaz/C(ann a;)1)) and (8) with E S the probability of visiting state with this policy.
The following then obtains:
(a) m(s)= = Ea, (d.))
3

---

(b) If is bijective, then (s)
V(s)
and as a special case for terminal states æ, TTx)) "Eeex R(x)
R(x)
(ea)
(c) If C' is non-injective and there are n(z) distinct action sequences s.t. C(ā:) I, then
n(z)((t)) (a)) Exex R(a")
See Appendix A. for the proof. In combinatorial spaces, such as for molecules, where C' is non- injective (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 them. In this scenario, the pseudo-value V is "misinterpreting" the MDP's structure as a tree, leading the wrong (xx)
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 as a sum of descendant rewards, as having single source, the root node (or initial
elaborated below. We define the flow network state) 8o with in-flow Z, and one sink for each leaf (or terminal state) x with out-flow R(x) > O. 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'}l 2 1, except for the root, which has parent. We write F(s, a) for the flow between node s and node s' = T(8, a), F(8) for the total flow going through s! This construction is illustrated in Fig.
nodes have multiple paths from the root
E ass
Σ F(s3, d') a' d'EA(as)
=.
Fso)
Z
S3
Q1
as
044
22
0.7
root
81
54
03
terminal state
Figure I: A flow network MDP. Episodes start at source sq with flow Z. Like with SMILES strings, there are cycles. Terminal states are sinks with out-flow R(s). Exemplar state $3 has parents {(s, (s) a) = 53} = ((1 02), (32, }s)} and allowed actions A(ss) = (a4, 47 84 is a terminal sink state with R(s4) O and only one parent. The goal is to estimate Fss a) such that the flow equations are satisfied for all states: for each node, incoming flow equals outgoing flow.
To satisfy flow conditions, require that for any node, the incoming flow equals the outgoing flow, which is the total flow Fss) 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(6') = Σ
F(s.a)
s,aTT(saa=s*
equals the out-flow
F(s') = Σ F(s',a').
(3)
d'EA(e))
More concisely. with R(s) = 0 for interior nodes, and A(s) Ø for leaf (sink/terminal) nodes, write the following flow consistency equations:
Σ
= R(s')
Σ F(d'.d).
(4)
a.a=T(saa)=6'
d'EA(6))
with F' being a flow, F(s, a) > 0 Vs, 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 now show that such flow correctly produces T (2) = R(x)/2 / when the above flow equations are satisfied.
In some sense, F(s) and F(s,, a) are close to V and Q. RL's value and action-value functions. These etfectively inform an agent taking decisions at each step of an MDP to act in desired way. With some work, we also show equivalence between F'ss, a) and the "real" Q* of some policy à in modified MDP (see A.2).

---

Proposition 2. Let us define a policy that generates trajectories starting in state by sampling actions E A(s) according
F(s,a)
x (als)
(5)
F(s)
where F(s,a) > O the flow through allowed edge (s, a), F(s)= Ens F(s, where
R(s) 0 for non-terminal nodes and F(x) = R(x) > for terminal nodes I, and the flow
consistency equation _=.:T(e.))=) F(s, a) R(s')(ss)))))))))) E Ea'EA(o)) F(s', .') is satisfied. Let x (s)
denote the probability of visiting state s when starting at sy and following T--..). Then
(a) (s)) = F(s)
F(ao)
(b) (sa)) = -mex R(x)
(c) T()) = Ev'xx R(x*)*
R(w)
Proof. We have T(0)) since we always start in root node 6g. Note that E (④x) = 1 because terminal states are mutually exclusive, but in the case of non-bijective C, we cannot say that Exes (s) equals because the different states are not mutually exclusive in general. This notation different from the one typically used in RL where (s)) refers to the asymptotic distribution of the
Markov chain. Then
T(8')
Σ
T(als)t(s)
(6)
(1,a):((s.))m8
i.e. using Eq. 5,
F(s,a) F(s) (8).
(s')
Σ
((.0):
We can now conjecture that the statement
F(s)
T = F(so)
(8)
true and prove by induction. This trivially true for the root, which our base statement, since o (so) = 1. By induction, we then have that if the statement is true for parents s of s', then
F(s.a) F(s) [.s)) F(s,a) (F, F(s) F(so)
F(s')
T(s') = Σ
(9)
F(so)
F(so)
which proves the statement, i.e., the first conclusion (a) of the theorem. We can then apply it to the case of terminal states v, whose flow is fixed to F(x) = R(x)) and obtain
R(x)
(Φx) = F(su)
(10)
Noting that ΣxEx (2) = and summing both sides of Eq. 10 over we thus obtain (b), i.e.,
F(so) = Evex R(x). Plugging this back into Eq. 10, we obtain (c), i.c. T2)) = Erex H(ee))
R(())
Thus choice of क (a]s) satisfies desiderata: it maps reward function R to generative model which generates with probability (x) xx R(x), whether C is bijective or non-injective (the former being a special case of the latter, and just provided proof for the general non-injective case).
2.2 Objective Functions for GFlowNet
We can now leverage our RL intuitions to create 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--iiffernnce (TD) algorithms (Sutton and Barto, 2018). This could yield the following objective for trajectory
[a(T)= = Σ
F -R(s)) -
(II)
S'ETTSO
a'EA(e')
5

---

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 z), the F(s, a) and F(s) for early states will be exponentially larger than for later states. Since we want F(s,a) 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 (s, a) = log F(x, a), and exponentiate-sun-log the F/08 predictions to compute the loss, yielding the square of difference of logs:
log + R(s') + Σ exp f (s', (s, )) (12)
(T. = Σ log + Σexp exp Ao (s,a) a)
s.a:T(s.a)=d'
a'EA(S')
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 I. To give more weight to errors on large flows and avoid taking the logarithm of tiny number, we compare logoeminn flow) with (og(c+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 e 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 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.
Let trajectories T used to train Fo be sampled from an exploratory policy P with the same support as the optimal π defined in Eq. 5 for consisteni flow F F Affow is consistent if Eq. is respected. Also assume that 30 : FFg- = F*, i.e., we choose a sufficiently rich family of predictors. Let 6* E argming Eriinngrrre Le()) a minimizer of the expected training loss. Let Lo(T) 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
Proposition 3.
Fg- = F=, and Lg- (T) = 0 VT ~ P(O).
(13)
i.e., a global optimum of the expected loss provides the correct flows. If To. (als) EneEAw (s,a)) then also have
Fg- (a.0)
R(x) Z
m-- (x)
(14)
The proof is in Appendix A. Note that, in RL terms, this method akin to asynchronous dynamic programming (Sutton and Barto, 2018, $4.5), which is an off-policy off- method which converges provided every state is visited infinitely many times asymptotically.
3 Related Work
The objective of training policy generating states with a probability proportional to rewards was presented by Bucsing 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 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 fraction of the total volume) convergence 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 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 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 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-as- inference's ' Soft Q-Learning (Haarnoja et al., 2017), with the difference that we include all the parents of state in the in-flow, whereas Soft Q-Learning only uses the parent contained in the trajec- tory. Soft Q-Learning induces a different policy, as shown by Proposition 1, one where P(r) x R(T) rather than
P(x) 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 clas- sification within existing flow networks (Rahul et al.. 2017; Pektaş 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.
Empirical Results
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 standard MCMC
and RL methods, with its sampling distribution better matching the normalized reward. We find that GelowNet (A) converges (x) x R(xx), (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.
4.1 A (hyper-)grid domain
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 with action a, (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 æ is some R(x)= > O. 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) Ro + Rı II, I(0.25 Ja,/H 0.5l) + R2 II, I(0.3 < Ja:/H < 0.4)
with 0- < Ry < R1 < R2, pictured when n = 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 R = 1122, R2 = 2. By varying Rg and setting it closer to 0, we make this problem artificially harder, creating region of the state space which it is undesirable to explore. To measure the performance of a method, we measure the empirical LI error [[[(()) - . (x)]]. p(x) = R(#)/2 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 visit as a function of the number of states visited.
We run the above experiment for Ro E (10-1.10 10-2, 10-2,1003 10-3) with n 4, H = 8. In Fig. 2 we see that GFlowNet is robust to Ro and obtains a low LI error, while a Metropolis- Cestropois-HHattngs--A MaastigssMCMMe based method requires exponentially more samples than GFlowNet to achieve some level of LI error. This is apparent Fig. 2 (with a log-scale horizontal axis) by comparing the slope of progress of GelowNet (beyond the initial stage) and that of the MCMC sampler. We also see that MCMC takes much longer to visit each mode once as Ro decreases, while GFlowNet is only slightly affected, with GFlowNet converging to some level of L1 error faster, as per hypothesis (B). This suggests that
7

---

GFlowNet is robust to the separation between modes (represented by Ro being smaller) and thus recovers all the modes much faster than MCMC (again, noting the log-scale of the horizontal axis).
To compare to RI., we run PPO (Schulman et al., 2017). To discover all the modes in a reasonable time, we need 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 Ro. This and the previous result validates hypothesis (C). We also run SAC (Haarnoja et al., 2018) 1)), finding similar or worse results. We provide additional results and discussion in Appendix A.6.
the entropy maximization term much higher (0.5) than usual (<< 1). We verify
modes found (max=16)
15
RS 10
empirical L1 emoo
Ro 10-2 2
10
Ro 10
10
ours
5
MCMC
PPO
random
104
10S
106
10a
10a
10c
states visited
states visited
Figure 2: Hypergrid domain. Changing the task difficulty Ro to illustrate the advantage of GFlowNet over others. We see that as Ro gets smaller, MCMC struggles struggles to visit all the modes. PPO also struggles to find all the modes, and requires very large entropy regularization, but is robust to the choice of Rg. We plot means over 10 runs for each setting.
fit the distribution because
4.2 Generating small molecules
Here our goal is to generate diverse set of small molecules that have a high reward. We define large-scale environment which allows an agent to sequentially generate molecules. This environment is challenging, with up to 1026 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 fragment- based drug design (Kumar et al., 2012; Xie al., 2021). Generating such a graph can be described as sequence of additive edits: given a molecule and constraints of chemical validity, we choose an atom to attach block to. The action space 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 molecule to 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 proxy R in the space of molecules, we can quickly match its induced distribution (2) o 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 Fo is parameterized similarly to MARS (Xie al., 2021), with an MPNN, but over the junction tree graph (the graph of blocks), which had better performance. For fairness, this architecture used for both GelowNet and the baselines. Complete details can be found Appendix A.4.
We pretrain the proxy with a semi-curated semi-random dataset of 300k molecules (see A.4) down to 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 106 molecules. During e training, sampling follows exploratory policy P(als) which is a mixture between x (a]s) (Eq. 5), used with probability 0.95, and a uniform distribution over allowed actions with probability 0.05.
Experimental results Fig. 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 reasonable policy
8

---

0.08
outs,
8
ours, =44
avg R of unique top.
MARS
0.06
MARS.
PPO
((R)
MARS. B
JT-VAE BO
0.04
proxy dataset
top
top 100
0.02
top 1000
0.00
2
8
10a
10
10*
10a
10
R(x)
molecules visited
Figure 3: Empirical density rewards. We ver- ify that GelowNet is consistent by training it with Rº 3 = which has the hypothesized effect of shifting the density to the right.
Figure The average reward of the top-k as function of learning (averaged over 3 runs). Only unique hits are counted. Note the log scale. Our method finds more unique good molecules faster.
this should shift the distribution to the right. This indeed what we observe. We compare GelowNet MARS (Xie 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 3 values; this consistent with the hypothesis that it finds high-reward modes faster (since MARS 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 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 while; RL tends to be satisfied with good enough trajectories unless it is strongly regularized with exploration mechanisms. For GelowNet and for MARS, the more molecules convergence towards the proxy's 's max reward. Given the same compute time, JT- TAE+BO generates only about 10a molecules (due to its expensive Gaussian Process) and does not perform well.
visited, the better they become, with slow
The maximum reward in the proxy's dataset is 10, with only 233 examples above 8. our best run, find 2339 unique molecules during training with dataset. We compute the average pairwise Tanimoto similarity for the top 1000 samples: GFlowNet has 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 certain reward threshold. We provide additional insights into how GFlowNet matches the rewards in Appendix A.7.
above 8, only 39 of which are in the
10000
GFlowNet
GFlowNet
# modes with R 7.5
8000
MARS
# of modes with R >>3
MARS
PPO
1000
PPO
6000
4000
500
2000
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
states visited
x10%
states visited
x10"
Figure 5: Number of diverse of the number of molecules seen. Left, T 7.5. Right, 8.
scaffolds found above reward threshold T' as function
4.3 Routi-Round Experiments
To demonstrate the importance of diverse candidate generation in active learning setting, consider a sequential acquisition task. We simulate the setting where there a limited budget for calls to the true oracle O. We use proxy M initialized by training limited dataset of (æ, R(x))

---

pairs Do. where R(z) is the true reward from the oracle. The generative model (TT) is trained to fit to the unnormalized probability function learned by the proxy M. We then sample batch B = (21, 22.... ..k where ~ To, which is evaluated with the oracle O. The proxy M is updated with this newly acquired and labeled batch, and the process repeated for N iterations. We discuss the experimental setting in more detail in Appendix A.5.
1.5
Rg = 10
Gel
8.5
Top-k Retur
Rg - 10-2
Top-k Rewand
MARS
1.0
VAE+BO
Rg 10-3
Random
0.5
GelowNet
7.5
Top 10
MCMC
Top 100
PPO
0.0
10
15
2000
2500
3000
3500
Rounds of acquisition
No. of molccules docked
Figure 6: The top-k return (mean over 3 runs) in the 4-D Hyper-grid task with active learning. GFlowNet gets the highest return faster.
Figure 7: The top-k docking reward (mean over runs) in the molccule task with active learning. GFlowNet consistently generates better samples.
Hyper-grid domain We present results for the multi-round task in the 4-D hyper-grid domain in Figure 6. We use Gaussian Process (Williams and Rasmussen, 1995) the proxy. We compare the Top-k Return for all the methods, which is defined as mean(top-ke ((.))
mean(top-alt -(Dp) DD,1) )),
where D, is the dataset of points acquired until step i, and 10 for this experiment. The initial dataset Do (Dool = 512) is the same for all the methods compared. We observe that GelowNet 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 GelowNet 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 (Trout and Olson, 2010), with Dol 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 Dy- 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.
5 Discussion & Limitations
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 iterative method like that needed with MCMC methods, and with the advantage that when training succeeds, the policy generates great diversity of samples near the main modes of the target distribution without being slowed by issues of mixing between modes.
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 local optimization in order to refine the generated samples and approach the local maxima of reward while keeping the batches of candidates diverse.
Limitations.
Negative Social Impact. The authors do not foresee negative social impacts of this work specifically.
10

---

Acknowledgments and Disclosure of Funding
This research enabled in part by computational resources provided by Calcul Québec (www. cal culcblquebec. 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 grateful 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 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 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.
References
Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy Colwell. Model-based reinforcement learning for biological sequence design. In International Conference on Learning Representations, 2020.
Philip Bachman and Doina Precup. Data generation sequential decision making, Advances in Neural Information Processing Systems, 28:3249-3257, 2015.
Guy W Bemis and Mark A Murcko. The properties of known drugs. 1. molecular frameworks. Journal of medicinal chemistry, 39(15):887-2881, 1996.
Emmanuel Bengio, Joelle Pineau, and Doina Precup. Interference and generalization in temporal difference learning. In International Conference on Machine Learning, pages 767-777. PMLR, 2020.
G Richard Bickerton, Gaia V Paolini, Jérémy Besnard, Sorel Muresan, and Andrew I. Hopkins. Quantifying the chemical beauty of drugs. Nature chemistry, 4(2):90-98, 2012.
Nathan Brown, Ben McKay, François Gilardoni, and Johann Gasteiger. A graph-based genetic algorithm and its application to the multiobjective evolution of median molecules. Journal of chemical information and computer sciences, 44(3):1079-1085 2004.
Lars Buesing, Nicolas Heess, and Theophane Weber. Approximate inference in discrete distributions with monte carlo tree search and value functions, 2019.
Nicola De Cao and Thomas Kipf. Molgan: An implicit generative model for small molecular graphs. 2018.
Yize Chen and Baosen Zhang. Learning to solve network flow problems via neural decoding. arXiv preprint arXiv:2002.04400 2020.
Benjamin E Childs, James H Brodeur, and Levente Kocsis. Transpositions and move groups in monte carlo tree search. In 2008 IEEE Symposium On Computational Intelligence and Games, pages 389-395. IEEE, 2008
Hanjun Dai, Rishabh Singh, Bo Dai, Charles Sutton, and Dale Schuurmans. Learning discrete energy-based models via auxiiiary-variable local exploration. In Neural Information Processing Systems (NeurIPS), 2020.
11

---

Zihang Dai, Amjad Almahairi, Philip Bachman, Eduard Hovy, and Aaron Courville. Calibrating energy-based generative adversarial networks. arXiv preprint arXiv:1702.01691, 2017.
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry, 2017.
Sai Krishna Gottipati, Boris Sattarov, Sufeng Niu, Yashaswi Pathak, Haoran Wei, Shengchao Liu, Karam M. J. Thomas, Simon Blackburn, Connor W. Coley, Jian Tang, Sarath Chandar, and Yoshua Bengio. Learning to navigate the synthetically accessible chemical space using reinforcement learning, 2020.
Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris J. Maddison. Oops took a gradient: Scalable sampling for discrete distributions, 2021.
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep mergy--based policies. In International Conference on Machine Learning, pages 1352-1361. PMLR, 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-poolicy maximum entropy deep reinforcement learning with stochastic actor, 2018.
Jan Jensen. A graph-based genetic algorithm and generative model/monte carlo tree search for the exploration of chemical space. Chemical science, 10(12):3567-3572, 2019.
Wengong Jin, Regina Barzilay, and Tommi Jaakkola. Chapter 11. junction tree variational autoencoder
for molecular graph generation. Drug Discovery, page 228-249. 249, 2020. ISSN 2041-3211, doi: 10. 1039/9781188016888 -00228. URL URL http:// áxx doi org/10. /0. 10009/978188006841-000000
Andreas Kirsch, Joost van Amersfoort, and Yarin Gal. Batchbald: Efficient and diverse batch acquisition for deep bayesian active learning, 2019.
Ashutosh Kumar, A Voet, and KYJ Zhang. Fragment based drug design: from experimental to
computational approaches. Current medicinal chemistry, 19(30):5128-5144 2012.
Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine. Implicit under- under-parameterizatti inhibits data-efficient deep reinforcement learning, 2020.
Greg Landrum. Rdkit: Open-source cheminformatics. URL http: rdkit. org.
Youzhi Luo, Keqiang Yan, and Shuiwang Ji. Graphdf: A discrete flow model for molecular graph generation, 2021.
Charlie Nash and Conor Durkan. Autoregressive energy machines, 2019.
Alireza Nazemi and Farahnaz Omidi. A capable neural network model for solving the maximum flow problem. Journal of Computational and Applied Mathematics, 236(14)) (1:6144:448-35133 2012.
Diana M. Negoescu, Peter Frazier, and Warren B. Powell. The knowledge-gradient algorithm for sequencing experiments in drug discovery. 23(3):346-333, 2011. ISSN 1526-5528. doi: 10. 10.1287/ijo. 1100.0447.. URL https: :ttpi:::/dooiorreo
1tt.o.11o//oocoo 1100:/10oo..100.......... 1100. 0417.
Andrew Y Ng, Stuart J Russell, et al. Algorithms for inverse reinforcement learning. In Ieml, volume 1, page 2, 2000.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein. Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Veto,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neurul Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019. URL http: /ttp::////p///e/pe/ea/eeeee neurips. ceturip..ccppppr.... 2015-pytorch an- aan-mmerativv--eee high wwwiggh-prffrrrmmmeec www.rformnnneeeeee-eeeeeeeee learning--nnnne library. pdf.
12

---

Abdurrahman Pektaş and Tankut Acarman. Deep learning to detect botnet via network flow summaries. Neural Computing and Applications. 31(11):8021-8333 2019.
Mariya Popova, Mykhailo Shvets, Junier Oliva, and Olexandr Isayev. Molecularrnn: Generating realistic molecular graphs with optimized properties, 2019.
RK Rahul, T Anjali, Vijay Krishna Menon, and KP Soman. Deep learning for network flow analysis and malware classification. In International Symposium on Security in Computing and Communication, pages 226-235. Springer, 2017.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
Ari Seff, Wenda Zhou, Farhan Damani, Abigail Doyle, and Ryan P Adams. Discrete object generation with reversible inductive construction. arXiv preprint (rXiv:1907008266 2019.
Marwin H. S. Segler, Thierry Kogej, Christian Tyrchan, and Mark P. Waller. Generating focussed molecule libraries for drug discovery with recurrent neural networks, 2017.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang, Graphaf: flow-based autoregressive model for molecular graph generation, 2020.
Yunsheng Shi, Zhengjie Huang. Shikun Feng, Hui Zhong, Wenjin Wang, and Yu Sun. Masked label prediction: Unified message passing model for semi-supervised classification, 2021.
Niranjan Srinivas, Andreas Krause, S. Kakade, and M. Seeger. Gaussian process optimization in the bandit setting: No regret and experimental design. In ICML, 2010.
Teague Sterling and John J Irwin. Zine 15-ligand discovery for everyone. Journal of chemical information and modeling, (:(11):2324-22 2015.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Kevin Swersky, Yulia Rubanova, David Dohan, and Kevin Murphy. Amortized bayesian optimization over discrete spaces. In Conference Uncertainty in Artificial Intelligence, pages 769-778. PMLR, 2020.
Oleg Trott and Arthur Olson. Autodock vina: improving the speed and accuracy of docking with new scoring function, efficient optimization, and multithreading. Journal of computational chemistry, (1 2010.
C. K. Williams and C. Rasmussen. Gaussian processes for regression. In Neural Information Processing Systems (NeurIPS), 1995.
Yutong Xie, Chence Shi, Hao Zhou, Yuwei Yang, Weinan Zhang, Yong Yu, and Lei (MARS): Markov molecular sampling for multi-objective drug discovery. In International Conference on Learning Representations, 2021. URL /ttttss://ppereeeeeee http:///onoooooooooooooaamrm////
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] We provide proofs (see Section 2 and Appendix) for our theoretical claims and fair empirical results with our proposed methods and baselines (see Section 4).
(b) Did you describe the limitations of your work? [Yes] GelowNet is limited by its use of bootstrapping, which known to be challenging in Deep RL (see Section 5).
(c) Did you discuss any potential negative societal impacts of your work? [Yes] Our theoretical work is fairly agnostic to applications, and aims compete with existing MCMCC-based methods (see Sections 1, 4, and 5). Our empirical work sitaaes itself in the context of automated drug-discovery.
13

---

(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] All complete proofs are available in the Appendix.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits. hyperparameters, how they were chosen)? [Yes] Hyperparameters and architectural choices are reported in the Appendix, and verifiable in the provided code.
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [No] We omit error bars for clarity, but we report standard deviations in the Appendix verify the significance of our results.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See A.3.
4. If you are using existing assets (e.g., code, data, models) or curating/relcasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you'le 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals. applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "noh7ArVqa4r", "reviewer": "Reviewer_UVLM", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea of the paper is very interesting but the paper is hard to follow as the paper borrows notations and concepts from RL but not in the exact same setting so it makes it a bit confusing for the readers. I personally feel like the paper could be made more intuitive and much easy to follow if it uses more direct and intuitive concepts and pushes less hard to connect with concepts in RL. \n\nThere are some minor comments and questions:\nComments:\nLine 71: The RL “agent” should have large support of what?\nLine 84: Missing citations of PPO and MCMC\nFigure 3: MARS looks to approximate better the proxy dataset distribution.\nFigure3: Didn’t understand the shifting of the density to the right.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper is a bit hard to follow. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8 hours"}, "extra_text": {}}, {"review_id": "hkUmH4sA2s5", "reviewer": "Reviewer_3Yk2", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed approach seems novel and is well described in the paper. The authors provide theoretical and practical details, as well as intuition behind the method. Experimental evaluation is sufficient, and results are demonstrated on a toy problem as well as on a practical task of molecular property optimization. Provided results in supplementary (Fig. 13 and Fig. 14) demonstrate huge advantage of the proposed approach.\n\nSome comments are provided below:\n1. More details on relation of the proposed approach and Bayesian networks would be profitable. Currently it seems as if model definition using a Bayesian network would yield a similar approach (e.g., flow-matching condition will be the same as total probability formulta).\n2. Experimental results should be provided on standard datasets. For example, optimization of logP and QED as provided in multiple papers, including GrammarVAE (https://arxiv.org/pdf/1703.01925.pdf). Such experiments are necessary to confirm that the proposed approach indeed yields good results against strong baselines.\n3. Diverse generation as mentioned in the title of the paper, is not built in the model and is a side-effect of the approach. I would recommend conducting experiments on the objective functions with high mode imbalance (e.g., 10k molecules with high reward in one mode and 5 molecules with high reward in another).\n\nUpdate: The authors answered my questions, I keep my original score.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "PoecrPjgff", "reviewer": "Reviewer_asFX", "summary": "This paper propose GFlowNet, a generative method that can turn a given positive reward into a generative policy that samples with a probability proportional to the return. \n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The most interesting part is that this novel generative method is based on flow network and local flow-conditions. And then the authors cast the flow consistency equations into an objective just like casting the Bellman equations into TD objectives. Compared with other methods, GFlowNet is designed especially for diverse candidate generation.\n\nOverall, I vote for accepting. This current submission is overall well written, clearly illustrated and appropriately structured. The motivations are well explained and experimental results are strong.\n\n\nQuestion:\n\n* Looking at Figure 4, it seems that MARS has the trend to have a higher average reward if the generative models are trained more than 10^6 molecules. \"For GFLowNet and for MARS, the more molecules are visited, the better they become, with slow converge towards the proxy's max reward\", which method converges slower? Are there any theoretical/experimental results?\n\n* In equation (9), the hyperparameter $\\epsilon$ is added to avoid numerical issues. It would be better if the authors can provide more details on why $\\epsilon$  can trade off how much pressure on large versus small flows. And add a brief description about how to setting this hyperparameter for different cases.\n\nUPDATE:\nI appreciate the authors' response to my questions and concerns.  I will keep my original score.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2.5 hours"}, "extra_text": {}}, {"review_id": "OU08eKZrXDF", "reviewer": "Reviewer_1AuD", "summary": "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).\n\nIn 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.\n\n\n**More specific details about GFlowNet:**   \nIn 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). \n\nThe 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.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# 1 Summary of My Review\nI found the approach presented in this paper very interesting and it appears to achieve impressive results. However, on a more negative note, some of the paper was hard to follow and the clarity could be improved in places (see main review below). I believe this would be easy enough to fix, and therefore, I have gone for a higher \"overall score\" and think this paper should be accepted. I have gone with a lower confidence score as I'm less knowledgable about how this idea relates to recent work in RL.\n \n# 2 Main Review\n## 2.1 Originality\nIn the molecule domain, the idea of using RL-based ideas to sample from the distribution of all molecules (weighted by their scores), rather than just find the top molecule, seems novel. Obtaining a set of diverse **and** well-performing molecules is an important problem and the proposed approach is an interesting way to do so.\n\nUnfortunately, I am less familiar with the recent, more general RL literature in this domain so feel less qualified to discuss the paper's originality in this respect, but the approach is fairly general and could hopefully also be of interest here.\n\nThe idea of visited states being equivalent (and so the state-action space being best represented with a DAG rather than a tree) has been considered in the context of Monte Carlo Tree Search, which the authors might briefly want to mention, e.g see the discussion about \"transpositions\" in Section 5.2.4 of:   \n> Browne, C. B., Powley, E., Whitehouse, D., Lucas, S. M., Cowling, P. I., Rohlfshagen, P., Tavener, S., Perez, D., Samothrakis, S. and Colton, S. (2012) ‘A Survey of Monte Carlo Tree Search Methods’, IEEE Transactions on Computational Intelligence in AI and Games, 4(1), pp. 1–43.\n\n## 2.2 Quality\na. As far as I can tell the paper is technically sound (although I did not look through the proofs in the appendix). The approach to make the method computationally tractable (leading to Eq 9) seems reasonable although I would have been interested to hear how dependent the results presented later are on the value of $\\epsilon$ chosen.\n\nb. One other currently popular way to sample from the distribution over molecules is using autoencoders with autoregressive decoders, for instance the junction tree VAE of Jin et al. (2020) for which the fragment-based scheme in this paper is said to be based upon. This seems like an obvious baseline to compare to in the experiments but is omitted.\n\nc. I found some of the experimental evaluation for the molecule task a little confusing: the authors invent a method for sampling from a distribution, yet evaluate it mostly in terms of how well it discovers the modes of the distribution (rather than how well it samples from the distribution as a whole). I realize that this is useful when the distribution is defined by a reward function, but what about when sampling more general distributions? (related to question 3.a. below)\n\n## 2.3 Clarity\n\na. One area for improvement is perhaps the paper's clarity. For instance, section 2 is rather dense and I had to read it twice to understand what was going on. Perhaps adding more signposting to the beginning of this section would allow the reader to follow its various parts.\n\nI also got a bit confused by the notation in this section. For instance $\\tau$ suddenly appears in section 2.2 to denote the trajectory, whereas earlier on in this section I believe $\\vec{a}$ was used instead (eg line 123). (Or do these variables mean different things....?)    \n\nLikewise, I got a bit confused with Equation 10. I think the second part of this should be $L_{\\theta^\\ast}(\\tau) = 0 \\quad \\forall \\tau \\sim P(\\tau)$, i.e. with $\\tau$ in place of $\\theta$? (On a more minor note I was also uncertain as to why $L_{\\theta^\\ast}(\\tau) $ had to equal zero, the approach can also deal with loss functions with any minimum right?)\n\nb. In section 4.1 (ie the toy experiment on the grid domain) I didn't understand why the states $s$ could not live in the same space as $x$ (defining the reward)? To convert from $s$ to $x$ you had to add the vector $-\\mathbf{1}$. This just seemed to add unnecessary complexity to the setup and made it harder to understand?\n\nc. I've only briefly glanced over the code, which hopefully should aid reproducibility. One thing that could be added here though is a description of which scripts to run (including runtime arguments) to actually reproduce the experiments. \n\n\n## 2.4 Significance\nThis paper presents interesting theoretical results as well as strong empirical performance. In the molecule optimization task it is nice to see the focus on obtaining a _diverse set_ of top molecules rather than just the top molecule score (although it would be nice if the top-1 score was also reported somewhere for completeness). The authors consider more interesting (and difficult) molecular optimization tasks (based on docking) than much of the previous work, and it is exciting to see improvements wrt recent (and reasonable) baselines. \n\nI think the ideas that GFlowNet is built around are interesting, particularly of how to sample from a distribution of discrete objects when the transition function (i.e. \"action\" using the paper's terminology) is irreversible. Hopefully this will inspire further work in this direction.\n\n\n# 3 Clarifying Questions\n\na. My understanding is that (if trained/run for long enough) the MCMC baseline and GFlowNet should ultimately sample from the same distribution, i.e. that defined by:\n$\\pi(x) \\propto R(x)$. However, in figure 3 it appears that GFlowNet's distribution is different (the mode is higher and the distribution has wider tails). Does this just mean the MCMC method has not converged? Is there any way to compare against the \"true\" distribution (for instance running the MCMC method for much longer, or using very many random trajectories and importance sampling)?\n\nb I am currently having a bit of trouble understanding figure 4. On lines 330-336 the authors state that when sampling 300k molecules using random trajectories that they obtain over 100 molecules with scores above 8 (233 such molecules to be precise). Looking at figure 4 the top-100 average for MARS is below 8 after visiting 300k molecules. Does this mean that it is worse than random search? Is this expected?  \n\nc For the MCMC methods (e.g. MARS) did the authors consider thinning?\n\nd A related question. For GFlowNet I assume every state is valid? If so I was wondering if the authors ever considered counting the reward for every state visited by GFlowNet rather than just the terminal ones. In some ways can one think of the \"stop action\" as a learnable thinning function?\n\ne I apologize if I missed this but what is the initial state for the molecule tasks?\n\n\n# 4 Very Minor Comments\n\ni. Lines 105, 138 and elsewhere on page 3 the word \"surjective\" is used. I think \"non-injective\" is meant here instead...? \n\nii. I found the figure on the right below line 265 helpful for understanding the toy problem (so thank you to the authors for this). However, its axis is too small and unreadable -- this should be made larger. Similarly, Figure 7 in the appendix could be bigger (especially given that there are no space constraints in that section).\n\niii. There are some weird black horizontal-ish lines on Figure 3 (at least on the two pdf viewers that I tried). Did these mean anything or were they just some left over plotting artifacts?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "# 5 Limitations\n\na. The authors briefly talk about the limitations of the approach in section 5. The main limitation they draw attention to is the challenge of moving closer to the local maxima of the reward function in the latter stages of optimization. To resolve this they discuss combining their method with local optimization techniques; however, I wonder whether the temperature approach they discuss in the earlier part of their paper (combined with some annealing scheme) could also be used here?\n\nb. One limitation the authors do not mention is how the method scales in terms of the size of the state and action space. The loss function requires for every current state the sum over all previous states and actions that may have led to the current state (see term 1 of Eq.9). I assume this may become intractable for very large state-action spaces (and the flows one is trying to model become very small). Can one approximate the sum using a subset? Also what about continuous state/action spaces?\n\n\n\n\n# 6 Societal impact\nThe authors state that they foresee no negative social impacts of their work (line 379). While I do not believe this work has the potential for significant negative social impact (and I'm not quite sure if/how I'm meant to review this aspect of their work), the authors could always mention the social impact of increased automation, or the risks from the dual use of their method, etc.\n", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: After an extensive back-and-forth discussion, all reviewers felt positively about the paper and its contribution, and lean towards acceptance.

However, there were issues regarding clarity in the submitted draft — this has been largely addressed through clarifications from the authors, but it is essential that the discussion here is incorporated into the final version of the manuscript (and in particular, the responses to reviewer 1AuD).
Decision: Accept (Poster)