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

In plain words

Imagine you want a machine that builds objects one piece at a time, such as a molecule built from blocks, and you want it to produce many different good objects, not just the single best one. A standard reinforcement-learning trainer tends to find one winning recipe and repeat it. Slow sampling methods (MCMC) can cover more ground but take a long time to wander between good regions.

This paper treats the building process like water flowing through a network of pipes. Each partial object is a junction, each possible next piece is a pipe, and finished objects are drains whose size is set by how good they are. If the flow entering and leaving every junction balances, then a builder that picks its next piece in proportion to the flow down each pipe ends up finishing good objects as often as they deserve. Because many build orders can lead to the same object, the method counts all of those paths. The training rule nudges the flows toward balance, and it works on a logarithmic scale so that very large and very small numbers stay manageable. The result is a fast generator that yields diverse good samples, and it can also learn from data gathered by other policies.

TL;DR

• Learns a policy that builds a discrete object step by step, so that the probability of ending at object x is proportional to a positive reward R(x) §1 §2.
• Training uses a flow-matching loss over the construction DAG, with flows kept on a log scale, and needs no MCMC §2.2 eq. 12.
• Author-reported: on a molecule task (up to 8 building blocks, sEH docking proxy), it finds higher top-k rewards and more distinct scaffolds than MARS and PPO §4.2 app. A.7.
• The official code is a reduced snapshot of the internal repository; several defaults differ from the paper (see Paper vs code).

Problem

Standard return maximisation converges to one high-return sequence of actions, so it under-samples the other good modes of a reward function §1. The paper wants a policy whose sample distribution is proportional to a given positive reward, which is useful for black-box optimisation with few, large batches of queries, such as diverse molecule design §1.

Existing ways to do this have a problem when several action sequences build the same object. MCMC samplers reach the target distribution but mix slowly between modes §3. Tree-based methods that assign each state the sum of rewards below it are correct only when action sequences and states are in one-to-one correspondence. When one object has n(x) build orders, they over-weight it by about n(x), which becomes exponentially bad for molecules §2 §2.1 (Proposition 1, §2). The paper's method is designed for the general, non-injective case, where the construction graph is a directed acyclic graph rather than a tree §2.

Method

Setting. A deterministic MDP builds x from a sequence of actions. Terminal states x have reward R(x) > 0, and interior states have reward 0. The paper asks for a policy with pi(x) approximately R(x)/Z, where Z is the sum of rewards over all terminal states §2 eq. 1.

Flow network. Treat the MDP as a flow network: a single source s_0 with inflow Z, one sink per terminal state with outflow R(x), and an edge flow F(s, a) for each allowed action §2.1. Flow consistency at every node requires that inflow equals outflow, where the inflow is the sum over all parent edges (s, a) with T(s, a) = s' and the outflow is R(s') plus the sum over the allowed actions a' of F(s', a') §2.1 eq. 4.

Policy from flows. Sample actions with probability F(s, a) / F(s), where F(s) is the total outgoing flow of s §2.1 eq. 5. Proposition 2 shows that, under consistent flows, the visit probability of each state equals F(s)/F(s_0), and that pi(x) = R(x)/sum over x' of R(x') §2.1.

Training objective. Approximate the flows with a predictor F_theta and penalise violations of flow consistency, in the same way TD methods penalise violated Bellman equations §2.2 eq. 11. Because flows near the root are huge and flows near leaves are tiny, the paper matches log-flows instead. It exponentiates the predicted log-flows, sums them, adds a small epsilon inside each log, and squares the difference of the two log sides §2.2 eq. 12. Epsilon both keeps the log finite and sets how much large flows are emphasised over small ones; the paper says it is set close to the smallest reward value §2.2.

Off-policy property. Proposition 3 says that trajectories may come from any exploratory policy whose support covers the optimal policy's support. A global minimiser of the expected loss then recovers the correct flows and the target distribution §2.2. The proof assumes the predictor family is rich enough to represent a consistent flow §2.2 app. A.1.

Sampling used in the molecule experiments. Training trajectories come from a mixture: with probability 0.95 the policy from the flows, and with probability 0.05 a uniform random action over the allowed actions §4.2. The molecule environment adds blocks to a junction tree of 72 predefined building blocks, with 105 actions per stem, up to 8 blocks, and an explicit stop action §4.2 app. A.3.

Reward used for molecules. A proxy MPNN predicts a docking score; the reward is a renormalised version of it, with values scaled to be positive §4.2 app. A.3. The paper then maps the reward as R_hat(x) = (R(x)/T)^beta with T = 8 and beta = 10, and clips R below R_min app. A.4.

# Flow-matching loss on the log scale  [§2.2] [eq. 12]
# Input:  transitions of sampled trajectories; each transition has parent
#         edges (s, a) with T(s, a) = s_next, a reward R(s_next) (0 for
#         interior states) and a flag marking terminal states
# Output: scalar loss

def flow_matching_loss(transitions, log_flow, eps, leaf_weight):
    total = 0.0
    for s_next in states_reached(transitions):
        # Log-inflow: sum over all parent edges, not only the one in
        # the trajectory  [§2.2]
        parent_flows = [exp(log_flow(s, a)) for s, a in parents(s_next)]
        log_inflow = log(eps + sum(parent_flows))
        # Log-outflow: reward plus sum over allowed children; terminal
        # states have no children  [§2.1] [§2.2]
        child_flows = [exp(log_flow(s_next, a)) for a in allowed(s_next)]
        log_outflow = log(eps + reward(s_next) + sum(child_flows))
        # Squared log-ratio; terminal transitions weighted by lambda_T
        # [§2.2] [app. A.4]
        term = (log_inflow - log_outflow) ** 2
        if is_terminal(s_next):
            total = total + leaf_weight * term
        else:
            total = total + term
    return total

# Reward used for training  [app. A.4]
# Input:  score = proxy output for a finished molecule
# Output: positive reward

def molecule_reward(score, r_min, T=8, beta=10):
    # Clip below R_min, then scale by T and raise to beta  [app. A.4]
    clipped = max(r_min, score)
    return (clipped / T) ** beta

Hypergrid (toy) domain. An n-dimensional grid of side H; the agent increases one coordinate at a time, and a stop action ends the episode. Many paths reach each cell, so the MDP is a DAG §4.1. The reward is R(x) = R0 + R1 * prod_i I(0.25 < |x_i/H - 0.5|) + R2 * prod_i I(0.3 < |x_i/H - 0.5| < 0.4), with R1 = 1/2 and R2 = 2 and R0 varied over {0.1, 0.01, 0.001}; the paper uses n = 4, H = 8 §4.1. The reference algorithms are MCMC (Metropolis-Hastings) and PPO (entropy coefficient 0.5, far above the usual value), and SAC was also run §4.1.

Multi-round (active learning) loop. Fit a proxy on the data so far, train the generator on M(x)^beta as the target, sample a query batch, label it with the oracle, add it to the dataset, and repeat §4.3 alg. 1. For the hypergrid the proxy is a Gaussian Process; for molecules it is an MPNN §4.3 app. A.5.1 app. A.5.2.

Hyperparameters

Values are from the paper unless marked "code". Where the paper gives a range or says "best values", the table says so.

Code defaults that the paper does not state (hypergrid script, grid/toy_grid_dag.py): learning rate 1e-4, minibatch 16, 256 hidden units, 2 layers, 20000 training steps, horizon 8, 2 dimensions toy_grid_dag.py. The paper's grid experiments use n = 4 §4.1, so the flag must be set to 4 to match.

Results

Author-reported. Each row gives the setup and the source.

Hypergrid, single round §4.1 fig. 2. Setup: n = 4, H = 8, three values of R0, 10 runs. Metric: empirical L1 error between the sampling distribution and p(x) = R(x)/Z, and the number of modes visited. The paper reports only plots here. Claims: GFlowNet reaches low L1 error for all R0, the MCMC sampler needs exponentially more samples as R0 falls, and PPO finds all modes more slowly than GFlowNet §4.1 fig. 2.

Molecules, reward at 10^5 samples fig. 4 app. A.7. Setup: sEH docking proxy, up to 8 blocks, 3 runs. Top-k reward, mean +/- std:

Molecules, reward at 10^6 samples app. A.7:

Diversity of molecules §4.2 app. A.7. Mean pairwise Tanimoto similarity of the top-1000 molecules: GFlowNet 0.44 +/- 0.01, PPO 0.62 +/- 0.03, MARS 0.59 +/- 0.02 (3 runs). Lower means more diverse. A random agent gives 0.231.

Modes §4.2 fig. 5 fig. 14. The main text reports more than 1500 Bemis-Murcko scaffolds with reward above 8 for GFlowNet, against fewer than 100 for MARS. The appendix counts Tanimoto-separated modes (threshold 0.7) and reports more than 500 modes above reward 7 for GFlowNet, against fewer than 100 for MARS (see Pitfalls).

Multi-round, molecules app. A.5.2. After 1800 docking evaluations: GFlowNet top-10 8.83 +/- 0.15, top-100 7.76 +/- 0.11; MARS top-10 8.27 +/- 0.20, top-100 7.08 +/- 0.13. The top-k curve over rounds is fig. 7. PPO training was unstable and diverged, so no numbers are reported §4.3.

Multi-round, hypergrid §4.3 fig. 6. Top-k return with k = 10 in the 4-D hypergrid; the paper reports that GFlowNet outperforms the baselines. Mean pairwise L2 distance between top-k points at the last round: GFlowNet 0.83 +/- 0.03, MCMC 0.61 +/- 0.01, PPO 0.51 +/- 0.02. The hypergrid multi-round run found 10 of the modes within 5 rounds for GFlowNet, 10 within 10 rounds for MCMC, and 8 for PPO app. A.5.1.

Calibration checks app. A.7 fig. 16 fig. 17. Log-log regression of leaf inflow against R(x)^beta gives slope 0.58 and r = 0.69; the paper reads the slope below 1 as underestimating high rewards.

Paper vs code

Code: bengioe/gflownet at commit 3ba65d243bb2e4fe77b3b93dc5f1634fd3ea5ea6, MIT license, owned by the first author LICENSE. Its README states this is a "minimum working version" extracted from an internal repository and that the original commit history is lost README.md.

• Epsilon. Paper: 2.5e-5 app. A.4. Code argparse default is 2.5e-5 gflownet.py, but the molecule good_config sets log_reg_c = (0.1/8)**4, about 2.4e-8 gflownet.py. The paper does not say which one produced its numbers.
• R_min. Paper: 0.01 app. A.4. Code good_config: 0.1 gflownet.py and argparse default 0.1 gflownet.py.
• Loss aggregation. Paper's eq. 12 sums over the states of a trajectory §2.2 eq. 12. The code averages the terminal and non-terminal losses separately over the batch and computes term_loss * leaf_coef + flow_loss gflownet.py gflownet.py. The scale differs from the paper's sum, so learning rates will not transfer directly.
• Stop action. The code masks the stop action for the first two blocks (min_blocks = 2) gflownet.py gflownet.py. The paper does not state this.
• Uniform exploration. Paper: uniform over the allowed actions §4.2. The code draws uniformly over the whole logit vector, which includes the stop index, and applies no masking in that line gflownet.py. Whether invalid stems are masked upstream is not visible in this function.
• Reward normalisation. Paper: the docking energy is negated and standardised app. A.3. In the training path the code uses the proxy output directly, and the target_norm = (-8.6, 1.10) constants only apply to the docking-score path that training does not call gflownet.py gflownet.py gflownet.py. The proxy in the repo is trained on the standardised data, so the code's reward is (max(R_min, proxy)/8)**10.
• Replay. The paper does not mention replay. The code default is replay_mode='online' gflownet.py, but with sample_prob = 1 the replay buffer is never read to build batches gflownet.py.
• Multi-round molecule script. Paper: reward beta 10 and up to 8 blocks for the flow model app. A.4 and, for molecules, all other hyperparameters are said to match the single-round setting app. A.5.2. The active-learning script defaults to reward_exp = 8 and max_blocks = 6 gflownet_activelearning.py gflownet_activelearning.py.
• Multi-round hypergrid script. Paper: top-k with k = 10 §4.3 and batch 16 app. A.5.1. Code defaults: reward_topk = 5 toy_grid_dag_al.py, num_samples = 8 toy_grid_dag_al.py, method ppo toy_grid_dag_al.py, and reward function corners_floor_B toy_grid_dag_al.py. The exact run configuration is not in the repo.
• Hypergrid reward. Matches the paper's R1 = 0.5 and R2 = 2 with thresholds on |x| in [-1, 1], which equal the paper's |x_i/H - 0.5| thresholds toy_grid_dag.py. The code adds R0 = 0.1 by default toy_grid_dag.py. The code's flow loss for the grid has no epsilon toy_grid_dag.py.
• Proxy and dataset. The paper says the dataset is on the authors' GitHub repository app. A.4. The repo contains mols/data/docked_mols.h5.gz and the pretrained proxy . Its size and exact contents are not stated in the paper.

Pitfalls

• Epsilon is ambiguous. The paper says epsilon is "close to the smallest value R can take" §2.2 and gives 2.5e-5 app. A.4. The code uses about 2.4e-8 for the molecule config. Pick one and document it.
• Numerical stability. The loss needs logs of sums of exponentials. The paper writes exp directly; the code uses torch.exp on log-flow outputs and then log(exp_sum + log_reg_c) gflownet.py. Large predicted log-flows can overflow. A log-sum-exp implementation is safer. This is an implementation choice, not stated in the paper.
• Loss scale and batch size. Eq. 12 is a sum over states, while the code averages per batch. Learning rates and the terminal weight lambda_T = 10 depend on which convention is used.
• Terminal molecules with free stems. When the stop action is chosen, the code marks the molecule terminal even if it still has open stems, and no children are added for it gflownet.py. The paper's terminal definition is "no children" §2.1. These are different states in the DAG; follow the code if you want the same behaviour.
• Exploration mixture. The 0.05 uniform mixture is what keeps the off-policy support condition of Proposition 3 §2.2 plausible. Remove it and the support assumption may fail on large spaces.
• Non-uniqueness. With several parents or several build orders, consistent flows are not unique app. A.1. The paper's claim that the policy is correct does not identify a unique flow, so do not expect the learned flows to match any particular split.
• Mode counting definitions differ. The main text counts Bemis-Murcko scaffolds above reward 8 §4.2. The appendix counts Tanimoto-separated modes above reward 7 with threshold 0.7 app. A.7. Quote the numbers with their definition.
• Bootstrapping. The paper names bootstrapping as a source of optimisation difficulty and limitation §5. Its leaf-weight lambda_T is described as a stabiliser app. A.4.
• Off-policy claim is partly conditional. Offline training on a fixed dataset works in the paper's examples, but these are illustrative. The paper says no hyperparameter search was done app. A.6.
• Hypergrid learning rate. The paper gives no hypergrid learning rate or batch size in the main text. The code defaults above are the only values in the repo.
• Cost. Molecule experiments took about 26 GPU days in total app. A. Each step needs a forward pass for both parents and children of every sampled state, so the cost per step is larger than a single-policy update.

Lineage

• Builds on: tree-structured reward-proportional sampling of Buesing et al. (2019), which the paper shows is correct only for injective action-to-state maps §2 §3. It also relates to Soft Q-Learning (control as inference), which induces a different policy over trajectories rather than over states §3.
• Compared against: MCMC (MARS), PPO, SAC, JT-VAE with Bayesian optimisation, and a random agent §4 app. A.4 app. A.6.
• Superseded by: not stated in the paper. The code README points to newer GFlowNet implementations, which are not cited here README.md.