[Paper Review] Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation

Emmanuel Bengio, Moksh Jain+3 more|arXiv (Cornell University)|Jun 8, 2021

Machine Learning in Materials ScienceMaterials Science49 citations

TL;DR

GFlowNet learns a flow-based generative policy to sample candidates with probability proportional to a positive reward, enabling non-iterative, diverse high-reward outputs (e.g., 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.

Motivation & Objective

• Motivate sampling a distribution of high-reward candidates rather than a single top candidate.
• Introduce a flow-network view of the generative process to handle many action sequences leading to the same final object.
• Provide a learning objective that enforces flow consistency and yields a policy proportional to the reward.
• Demonstrate scalability and diversity gains in synthetic domains and molecule-generation tasks.
• Show off-policy/offline convergence properties of the training objective.

Proposed method

• Model the generation process as a flow network with a single source and terminal sinks corresponding to candidate objects.
• Define action probabilities as pi(a|s) = F(s,a)/F(s), where F denotes flow through edges and F satisfies flow-consistency equations.
• Use a flow-matching objective on log-flows to stabilize training: minimize the squared difference between log of incoming and outgoing flows (Eq. 12).
• Prove that a global minimum of the flow objective yields a policy with pi(x) proportional to R(x) for terminal states.
• Demonstrate that the method is off-policy: training can use trajectories from a broad-support policy while converging to the desired flow and distribution (Proposition 3).
• Apply to synthetic DAG/MDP domains and large-scale molecule generation with a docking-based reward proxy.

Experimental results

Research questions

• RQ1 Can a flow-network-based policy produce samples with probability proportional to a given reward R(x) even when multiple action sequences lead to the same final state?
• RQ2 Does the flow-matching objective recover all high-reward modes and outperform traditional MCMC/RL baselines in terms of diversity and sample efficiency?
• RQ3 Is the method scalable to large combinatorial spaces like molecule generation and robust to offline/off-policy training?
• RQ4 How does the generated distribution compare to the normalized reward distribution in practice on synthetic and molecular domains?

Key findings

• GFlowNet samples terminal states with probability proportional to R(x), even when C is non-injective (many paths to the same state).
• The log-flow matching objective yields a global minimum that corresponds to the desired distribution and enables off-policy/offline training.
• In a hypergrid domain, GFlowNet converges to pi(x) ∝ R(x) with fewer samples than MCMC and is robust to mode separation.
• In large-scale molecule generation, GFlowNet produces higher-reward and more diverse molecules faster than MARS, PPO, and JT-VAE+BO baselines.
• GFlowNet finds and covers many modes (e.g., >1500 modes with R>8) and achieves lower average pairwise similarity among top samples than baselines, indicating higher diversity.
• The approach scales to molecule graphs with sizable action spaces and yields more unique high-scoring molecules within the same compute budget.