PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, Yoshua Bengio

arXiv:2106.04399v2cs.LG

TL;DR

The paper studies how to generate diverse high-reward objects when the target distribution should be proportional to a positive reward, rather than concentrating on one optimum. It introduces GFlowNet, which learns flow-matching conditions on sequential state-transition networks, including cases with multiple trajectories per final state. GFlowNet recovers reward-proportional distributions and demonstrates stronger diversity and mode discovery than PPO and MCMC baselines in synthetic and molecule-generation experiments.

  • Problem

    The problem is learning a stochastic policy whose object-generation probabilities are proportional to positive rewards while preserving diverse high-reward solutions, including when multiple trajectories yield one object.

  • Method

    GFlowNet models sequential generation as a flow network and learns local incoming–outgoing flow matching so action probabilities proportional to learned flows produce the target terminal distribution.

  • Results

    GFlowNet recovers π(x) ∝R(x), recovers all synthetic modes faster than MCMC and PPO, and finds 2339 unique molecules scoring above 8 with mean pairwise Tanimoto similarity 0.44 ± 0.01.

  • Takeaways & Limitations

    GFlowNet provides fast, diverse candidate generation for multimodal reward landscapes and molecule-design settings where batch diversity is important.

  • Takeaways & Limitations

    Bootstrapping can create optimization challenges and limit performance, while combining GFlowNet with local optimization remains future work.

Abstract

from arXiv · show

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.

1 Introduction

The paper targets diverse high-reward generation rather than concentrating probability on one maximizing trajectory, especially for expensive black-box evaluations. GFlowNet addresses this by learning flow consistency over sequential state transitions so terminal sampling follows reward-proportional probabilities.

  • The objective is to sample trajectories with probability proportional to a positive reward, rather than select only the highest-return action sequence.
  • Diversity matters in batch black-box optimization because expensive or uncertain oracles evaluate many candidates per round, including in drug discovery.
  • A proxy model trained on previously collected oracle evaluations can serve as the reward function for generating the next batch of candidates.
  • GFlowNet represents action probabilities as flows on a state-transition network and trains local incoming–outgoing flow matching toward terminal rewards.
  • The method is designed for settings where multiple action sequences can produce the same final object, a case that complicates tree-based trajectory modeling.
  • The paper demonstrates GFlowNet on multimodal synthetic rewards and large-scale molecule synthesis, comparing it with PPO and MCMC methods.

2 Approximating Flow Network generative models with a TD-like objective

GFlowNet represents sequential generation as flow on a possibly non-injective DAG and learns local flow consistency so terminal states are sampled proportionally to reward. A log-scale flow-matching objective addresses numerical imbalance while preserving the correct global optimum.

  • 2.2 Objective Functions for GFlowNet: The flow-matching objective is TD-like: it penalizes discrepancies between incoming and outgoing flows along sampled trajectories, and any global optimum recovers valid flows.The proof constructs a zero-loss solution when the target flow is representable.
  • 2.1 Flow Networks: When multiple action sequences lead to the same object, tree-based policies can bias generation toward objects with more trajectories and miss reward modes.This non-injective setting forms a DAG, as in molecular graph generation, rather than a tree.
  • 2.1 Flow Networks: GFlowNet defines a policy from positive edge flows, with each state’s incoming flow matching its outgoing flow and terminal out-flow equal to its reward.The root flow equals the total terminal reward, denoted the partition function.
  • 2.1 Flow Networks: Under flow consistency, the induced policy visits state s with probability F(s)/F(s0) and samples each terminal x with probability R(x)/Z.The construction works whether the action-to-state mapping is bijective or non-injective.
  • 2.2 Objective Functions for GFlowNet: Raw flows create numerical problems because early states can have exponentially larger flows than late states in high-dimensional spaces.The method therefore matches logarithms of incoming and outgoing flows, with epsilon stabilizing tiny values while leaving the global minimum unchanged.

3 Related Work

GFlowNet extends reward-proportional generative modeling to settings with multiple paths per state, where earlier approaches may be inadequate. Unlike iterative MCMC, it amortizes sampling through training and targets deterministic generative settings.

  • 3 Related Work: GFlowNet generalizes reward-proportional policy learning beyond bijective action-to-state mappings, allowing many paths to lead to one state.This contrasts with the earlier Buesing et al. objective, which only makes sense in the bijective case.
  • 3 Related Work: MCMC methods can target unnormalized positive functions, but separated modes can make convergence extremely slow.The paper contrasts this iterative mixing challenge with GFlowNet’s non-iterative generation after training.
  • 3 Related Work: Prior molecule-generation work spans generative models, MCMC, reinforcement learning, and evolutionary methods, with some methods relying only on positive examples.Such methods do not exploit negative examples and continuous measurement differences as fully as reward-weighted generation.
  • 3 Related Work: Soft Q-Learning differs from GFlowNet by using only the trajectory’s parent, producing a policy proportional to trajectory rewards rather than object rewards.GFlowNet includes all parents of a state in the in-flow.
  • 3 Related Work: Related network-flow deep-learning work mainly addresses maximum-flow problems or classification, while MCTS connects through its treatment of DAG transpositions.These links concern neighboring uses of flow networks or DAG structure rather than the proposed generative framework.

4 Empirical Results

GFlowNet recovers high-reward modes and diverse candidates faster than MCMC and RL baselines across synthetic, molecule-generation, and active-learning experiments. Its advantages include robustness to separated modes, higher reward, and greater molecular or batch diversity.

  • 4.1 A (hyper-)grid domain: GFlowNet converges to a distribution proportional to reward and recovers all modes faster than MCMC and PPO in the hyper-grid domain.The comparison uses an artificial domain where the partition function is exactly computable.
  • 4.2 Generating small molecules: In molecule generation, GFlowNet finds higher-reward and more diverse molecules faster than MARS, PPO, and JT-VAE with Bayesian optimization.The environment contains up to 10^16 states and between 100 and 2000 actions depending on the state; rewards come from a pretrained docking proxy.
  • 4.1 A (hyper-)grid domain: As R0 decreases, GFlowNet remains robust while MCMC requires exponentially more samples to reach comparable L1 error.MCMC also takes much longer to visit each mode as the low-reward region separating modes becomes harder to explore.
  • 4.1 A (hyper-)grid domain: PPO finds all modes more slowly than GFlowNet and requires unusually large entropy regularization, although it is robust to R0.The entropy coefficient used to discover all modes was 0.5, much higher than usual values below 1.
  • 4.2 Generating small molecules: GFlowNet finds 2339 unique molecules scoring above 8, with mean top-1000 Tanimoto similarity 0.44 ± 0.01 and more than 1500 modes above that threshold.The corresponding mode counts are below 100 for MARS, while PPO and MARS have higher similarity values of 0.62 ± 0.03 and 0.59 ± 0.02.

5 Discussion & Limitations

The paper presents GFlowNet as a TD-like flow-learning method that generates diverse terminal samples proportional to reward, while noting bootstrapping challenges and a need for local refinement.

  • Discussion: GFlowNet learns state and state-action flows so its policy samples terminal states in proportion to their rewards.The method is framed as an alternative to iterative energy-based generation.
  • Discussion: Successful training yields diverse samples near the main modes of the target distribution without MCMC mixing delays.
  • Limitations: Bootstrapping, as in TD-based methods, may create optimization challenges and limit performance.
  • Limitations: Future work should combine generative sampling with local optimization to refine samples while preserving batch diversity.
  • Limitations: The authors do not foresee negative social impacts specifically from this work.

A Appendix

The appendix reports the software libraries, computational resources, and implementation availability used for the experiments.

  • Implementation: The implementations use PyTorch, reimplemented RL and other baselines, AutoDock Vina for binding-energy estimation, and RDKit for chemistry routines.
  • Computational resources: 26 GPU days were estimated for the molecule experiments, using a cluster with NVIDIA V100 GPUs.
  • Computational resources: The grid experiments were estimated to require 8 single-core CPU days.
  • Implementation: All implementations are publicly available in the paper’s GitHub repository.

A.1 Proofs

The proofs establish how flow matching induces reward-proportional sampling, including the non-injective case where multiple action sequences reach the same terminal state, and characterize global-loss optima.

  • Proposition 1: The state-value flow is defined as the sum of rewards over all terminal states reachable from a state.
  • Proposition 1: Under the flow-based policy, a state’s visitation probability is determined by its value flow relative to the root flow.
  • Proposition 2: When multiple action sequences reach one state, visitation probability sums the probabilities of all mutually exclusive paths reaching it.
  • Proposition 1: For terminal states in the bijective case, the policy probability equals the terminal reward divided by the total reward.
  • Proposition 2: For a non-injective mapping with n(x) paths to terminal state x, the induced terminal probability is proportional to n(x)R(x).
  • Proposition 3: A global optimum of the expected training loss provides the correct flows when the exploratory policy covers reward-positive trajectories and the predictor family can represent the consistent flow.
  • Flow solutions: Flow-matching equations can admit infinitely many valid solutions when multiple trajectories merge into the same terminal state.

A.2 Action-value function equivalence

This appendix relates learned flows to action-value functions in reinforcement learning, obtaining a direct equivalence for bijective tree-structured environments while leaving the general non-injective case unresolved.

  • Motivation: The appendix asks whether a learned flow can correspond to a real RL action-value function under some policy.
  • Bijective case: The proposed transformed reward and scaling satisfy Qµ(s, a; ˆR) = F(s, a; R)f(s) for the uniform policy.
  • Bijective case: The equivalence relies on deterministic transitions and a tree structure between action sequences and states.
  • Bijective case: In the bijective case, the learned flow corresponds to the action-value function of the uniform policy after a state-dependent scaling.
  • Non-injective case: In the non-injective case, infinitely many valid flows exist, so a comparably simple flow–action-value equivalence is unclear.
  • Non-injective case: The appendix conjectures that more general mappings between flows and action-value functions may exist.

A.3 Molecule domain details

The molecule domain builds compounds by attaching predefined building blocks at highlighted stem atoms, with rewards derived from a normalized binding-affinity proxy. The setup also tests how easily GFlowNet optimizes alternative molecular objectives.

  • Molecule construction: The agent selects from 72 predefined building blocks, duplicated across symmetry groups to yield 105 actions per stem, and may generate up to 8 blocks.Stems are atoms where new blocks can be attached.
  • Molecule construction: Figure 8 depicts the building-block library and highlights each stem, the atom connecting a block to the rest of the molecule.
  • Reward: The reward uses the opposite of AutoDock binding affinity, then subtracts its mean and divides by its standard deviation.The proxy is trained on 300k randomly generated molecules with computed binding affinities.
  • Alternative objectives: GFlowNet quickly finds molecules with biologically uninteresting logP above 20 and reaches the action space’s maximum QED of 0.948.For reference, ibuprofen’s logP is between 3.5 and 4.

A.4 Molecule domain implementation details

The implementation uses separate message-passing neural networks for reward prediction and flow prediction, with a terminal-transition weighting for stability. Proxy training combines mostly random molecules with a smaller set from prior reinforcement-learning runs.

  • Proxy model: The reward proxy is an MPNN operating on RDKit atom graphs, using atom features and 12 recurrent graph-convolution steps.
  • Training data: The non-active-learning proxy dataset contains 300k molecules, comprising 80% random trajectories and 20% molecules from previous reinforcement-learning runs.The additional 20% contains slightly higher-scoring molecules.
  • Flow model: The flow predictor is an MPNN over a tree-structured block graph, with learned block and bond embeddings and 10 recurrent convolution steps.It produces predictions for each stem where a new block can be attached.
  • Optimization: Terminal-transition losses are multiplied by λT > 1 to prioritize endpoint flow predictions and improve stability through bootstrapping.The paper relates this weighting to reward prediction as an auxiliary reinforcement-learning task.

A.5 Multi-Round Experiments

The multi-round procedure alternates between fitting a proxy, training a generative policy against its powered predictions, querying the oracle, and updating the dataset. The inverse-temperature parameter β controls reward concentration.

  • Multi-round active learning: The generative policy’s exploratory behavior matters because the proxy is trained on samples proposed by the generative model and scored by the oracle.
  • Multi-round active learning: β acts as an inverse temperature that makes reward modes more or less peaked.
  • Multi-round active learning: Algorithm 1 repeatedly fits the proxy on the current dataset, trains the policy with r(x) = M(x)^β, samples a batch, evaluates it with the oracle, and appends the results.

A.5.1 Hyper-grid

The hyper-grid experiment uses a Gaussian-process proxy with batches of 16 samples per round and compares mode discovery across methods. GFlowNet and MCMC discover more modes than PPO under the reported setting.

  • Experimental setup: The experiment uses a BoTorch Gaussian-process proxy and generates query batches of 16 samples after each round.Generative-model hyperparameters come from the best-performing single-round experiments.
  • Results: GFlowNet discovered 10 modes within 5 rounds, MCMC discovered 10 within 10 rounds, and PPO discovered 8 modes by the end.The initial dataset contained 4 modes, and the PPO result uses R0 = 10^-1.

A.5.2 Molecules

GFlowNet learns reward-proportional molecular distributions while producing substantially more diverse high-reward molecules than the compared methods. Additional diagnostics show approximate flow-reward agreement, though high rewards are underestimated and offline generalization depends on the training data.

  • Learned distributions: In the hypergrid diagnostic, the learned distribution πθ(x) matches p(x) ∝ R(x) almost perfectly, while sampled paths to a mode are spread but not uniform.The experiment uses n = 2, H = 8 and 30k sampled points; the flow objective does not enforce uniformity over paths.
  • Offline and loss behavior: Offline training can recover the target distribution from fixed trajectory datasets, with better matching when the dataset is generated by uniformly sampled grid points and reverse parent sampling.Many points are intentionally omitted as a generalization test, and the reported plots are illustrative rather than hyperparameter-tuned.
  • Molecule quality and diversity: GFlowNet produces significantly better and more diverse molecules than PPO and MARS across the molecule experiments.Its top-1000 molecules have mean pairwise Tanimoto similarity 0.44±0.01, compared with 0.62±0.03 for PPO and 0.59±0.02 for MARS.
  • Molecule quality and diversity: For reward R > 7, GFlowNet discovers over 500 Tanimoto-separated modes, whereas MARS discovers fewer than 100.Modes use a Tanimoto similarity threshold S = 0.7 and reward thresholds of 7 or 8.
  • Molecule quality and diversity: GFlowNet finds more diverse Bemis-Murcko scaffolds above reward thresholds T = 7.5 and T = 8.The scaffold counts provide a second approximation to the number of high-reward modes.
  • Learned distributions: The molecule-domain diagnostics suggest that the learned distribution approximately matches p(x) ∝ R(x)^β, but direct verification is difficult because many trajectories can lead to one molecule.A log-log fit between R(x)^β and predicted leaf inflow has slope 0.58 and r = 0.69 for β = 10, indicating underestimation of high rewards.
Loading 2106.04399v2…