Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, Yoshua Bengio
NeurIPS 2021
cs.LG
2021-06-08
GFlowNet treats generation as a flow on a DAG so completed objects are sampled in proportion to reward, recovering far more high-scoring molecular modes than MCMC or PPO.
Standard RL piles probability on the single highest-return trajectory. Molecule design wants something else: given a positive reward function, sample a batch of candidates in proportion to that reward, covering many modes rather than camping on one local peak.
MCMC can sample an unnormalized density, but mixing dies when modes sit across low-reward deserts. Tree-style generators are worse still. The same molecule can be assembled along many action sequences, so a tree value estimate weights objects by the number of paths, exponentially favoring larger graphs. Assay rounds that test tens of thousands of candidates need diverse batches at linear cost, not another single-point optimum.
GFlowNet treats states as nodes of a DAG and actions as deterministic edges. The source is the empty object; sinks are completed candidates whose outflow is pinned to a positive reward R(x). At every interior node, inflow must equal outflow. The policy picks actions in proportion to edge flow. If the flow equations hold, the probability of a terminal state is R(x)/Z, where Z is the sum of all terminal rewards, independent of how many construction paths exist.
The training loss is a Temporal-Difference analogue: on every visited node, match predicted log-inflow to log-outflow. Flows span many orders of magnitude, so the network emits log-flow and the loss uses log-sum-exp, keeping root-scale numbers from washing out the leaves. The objective does not depend on the sampling policy, so training can be off-policy as long as exploration covers the state space. The paper mixes 95% of the current policy with 5% uniform actions.
In the molecule domain, actions attach a fragment onto a junction tree under chemical validity. Reward comes from a pretrained MPNN proxy that predicts docking score against soluble epoxide hydrolase, rescaled into a positive range.
On a hypercube grid whose reward lives near 2^n corners, dropping the valley reward R0 from 0.1 to 0.001 makes MCMC need exponentially more samples to visit every mode. GFlowNet's L1 error barely moves, and it sweeps the modes faster than PPO with heavy entropy regularization.
The molecule MDP has about 10^16 states and 100–2000 actions per step. Training sees up to 10^6 molecules. The proxy training set has only 233 examples scoring above 8. GFlowNet finds 2339 unique molecules above 8, of which 39 were already in that set. Counted as Bemis-Murcko scaffolds with reward >8, it finds more than 1500 modes against fewer than 100 for MARS. Mean pairwise Tanimoto of the top 1000 is 0.44 for GFlowNet, 0.62 for PPO, and 0.59 for MARS (lower is more diverse). In multi-round active learning against real docking, GFlowNet's top-k scores stay above MARS and random acquisition.
| Method | Metric | Result |
| GFlowNet | scaffolds with R>8 | >1500 |
| MARS | scaffolds with R>8 | <100 |
| GFlowNet | top-1000 Tanimoto | 0.44 |
| PPO | top-1000 Tanimoto | 0.62 |
This is a way to turn an energy function into a single-pass generator, without running MCMC for every draw. When a round of experiments must submit a large batch and the proxy is itself noisy, covering modes beats chasing one maximum. The molecule numbers show flow matching works on a combinatorial DAG, not just a toy grid.
The objective is sampling in proportion to reward, not finding the global maximizer. If a task only cares about a single best object, PPO is the more direct tool.
The authors flag the usual TD issue: bootstrapping can make optimization brittle and cap how well the flows fit. Samples land near modes rather than on local maxima, so a later local search would still be needed. The molecule reward is a docking proxy with no ADMET or synthesizability. GFlowNet matches the proxy's induced distribution; if the proxy is wrong, so is the sampler. The hypergrid is tiny (4-D, side 8). How flow error compounds in much larger combinatorial spaces is left open.