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

1 Introduction . 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 . 2 Approximating Flow Network generative models with a TD-like objective . 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 ) . 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 ( 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 can not 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 ) .

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).  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.