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