Reward-proportional sampling

Overview

The goal is a policy whose probability of ending at a finished object x is proportional to a positive reward R(x), so that a generator yields many different good objects instead of one high-return sequence 2106.04399 §1 §2 eq. 1. The shared invariant is flow consistency: at every state, the flow in equals the flow out, and the policy picks each action in proportion to its edge flow 2106.04399 §2.1 eq. 4 eq. 5. Trajectory balance enforces the same target with one constraint per complete trajectory instead of one per state 2201.13259 §3 eq. 14.

The difficulty is that several action sequences can build the same object. A tree-based rule that gives each state the sum of rewards below it is correct only when action sequences and states are in one-to-one correspondence. With n(x) build orders for an object, it over-weights that object by about n(x), which the paper calls exponentially bad for molecules 2106.04399 §2 §2.1.

A second difficulty, raised by the trajectory-balance paper, is credit assignment. Local consistency checks must pass a late, high reward back through many steps before the early choices are corrected, which the authors argue is slow 2201.13259 §6.

A third card, a position essay, proposes the same sampling target for discovering theorems. It defines no method and reports no results 2403.04571 §8. It is listed under Variants so the framing is not lost, but it is not an algorithm.

Variants

Flow-matching generator (GFlowNet)

Treats the construction process as a flow network over a DAG of partial objects: one source with inflow Z, one sink per finished object with outflow R(x), and an edge flow for each allowed action 2106.04399 §2.1. The policy samples an action with probability F(s, a) / F(s) [arxiv:2106.04399 eq. 5]. Training fits a flow predictor so that inflow and outflow match at every state, which needs no MCMC 2106.04399 §2.2 eq. 11.

The loss matches log-flows. It exponentiates the predicted log-flows, sums the parent flows and the child flows plus reward, adds a small epsilon inside each log, and penalises the squared difference of the two log sides [arxiv:2106.04399 §2.2] eq. 12. The paper says epsilon should be close to the smallest reward value 2106.04399 §2.2.

Trajectories may come from any exploratory policy whose support covers the optimal policy's support, and a global minimiser of the loss then recovers the target distribution 2106.04399 §2.2. The proof assumes the predictor can represent a consistent flow 2106.04399 app. A.1. The molecule runs mix the policy with a uniform random action at probability 0.05 [arxiv:2106.04399 §4.2]. Terminal transitions are weighted by lambda_T = 10 [arxiv:2106.04399 app. A.4].

Reported results, author-reported:

• Molecules, one round. With a docking-score proxy and up to 8 building blocks, top-10 reward at 10^5 samples is 8.36 for GFlowNet, 8.06 for PPO and 8.05 for MARS 2106.04399 app. A.7 fig. 4. At 10^6 samples the gap is 8.45 against 8.25 for PPO and 8.31 for MARS 2106.04399 app. A.7.
• Molecule diversity. Mean pairwise Tanimoto similarity of the top-1000 molecules is 0.44 for GFlowNet, 0.62 for PPO and 0.59 for MARS; lower means more diverse 2106.04399 §4.2 app. A.7.
• Hypergrid. In a 4-dimensional grid with three values of R0, GFlowNet reaches low L1 error to the target distribution, while MCMC needs exponentially more samples as R0 falls 2106.04399 §4.1 fig. 2.
• Batch active learning. Inside a loop that refits a proxy and queries batches, GFlowNet reaches a molecule top-10 of 8.83 after 1800 docking evaluations, against 8.27 for MARS. PPO diverged and reports no numbers 2106.04399 app. A.5.2 §4.3 alg. 1. Card: GFlowNet.

Trajectory balance

Trajectory balance (TB) takes the whole sampled path as its unit. It learns a scalar log Z for the initial flow, and it penalises the squared log-ratio between log Z plus the summed forward log-probabilities and log R(x) plus the summed backward log-probabilities 2201.13259 §3 eq. 14. If this loss is zero on every complete trajectory, the forward policy samples in proportion to reward 2201.13259 app. A.1. The backward policy can be learned or fixed to uniform over parents 2201.13259 §3.1. The molecule runs fix it to uniform 2201.13259 app. B.2.

Each step samples one trajectory, takes one gradient step on its loss, and stops when the running loss converges 2201.13259 alg. 1. The authors use a higher learning rate for log Z than for the policy, and explain it as a moving average of the gap between the two sides 2201.13259 §3, footnote. The paper's values differ by task: 1e-1 on the hypergrid [arxiv:2201.13259 app. B.1], 1e-3 on bit sequences app. B.3 and 1e-2 on AMP app. B.4.

The paper also gives a sub-trajectory variant that applies the same constraint to segments between chosen states, so that not every state needs a flow estimate 2201.13259 app. A.2.

Reported results, author-reported and on different tasks from the flow-matching card:

• AMP sequences, top 100 of 2048 samples. Reward and mean pairwise edit distance are 0.85 and 18.35 for TB, 0.78 and 12.61 for flow matching or detailed balance, 0.80 and 8.36 for SAC, and 0.75 and 12.56 for MCMC 2201.13259 tab. 1.
• Molecules. On a held-out set, the Pearson correlation between log reward and log sampling probability is higher for TB than for flow matching and detailed balance. The authors report a runtime speed-up of up to 5x over flow matching 2201.13259 §5.2 fig. 3 app. B.2.
• Hypergrid. Detailed balance and TB converge faster than flow matching, with a slight edge for TB on the 4-dimensional grid 2201.13259 §5.1 fig. 2. The paper gives no table of these values.
• Bit sequences. TB has the highest Spearman correlation between sampling probability and reward for every tested word size k 2201.13259 §5.3.1 fig. 4.

Card: Trajectory balance.

Conjecture sampling for theorem discovery (position essay)

A position essay proposes sampling statements so that each is drawn in proportion to a reward for its usefulness, names GFlowNet-style samplers as an option, and prefers this to reinforcement learning, which targets one best statement 2403.04571 §8. It proposes a usefulness measure based on how much a set of theorems compresses the provable statements [arxiv:2403.04571 §6]. The essay defines no reward formula, no number of derivation steps and no prover, and it reports no experiments 2403.04571 §6 §8 §9. Card: AI Mathematician essay. Use it for the framing of the objective. It gives no algorithm to run.

When to use which

Constraints are taken from what the papers report. Where a paper says nothing, the table says "not stated".

Implementation sketch

Python-flavoured pseudocode of the common core. Lines follow the papers' equations; log_flow(s, a) is the flow predictor's log output.

# Sample one object: pick each action with prob F(s, a) / F(s)  [eq. 5]
def sample_object(log_flow, env):
    s = env.initial_state()
    while not env.is_terminal(s):
        flows = [exp(log_flow(s, a)) for a in env.allowed(s)]
        a = sample_index(flows)  # proportional to the flows
        s = env.step(s, env.allowed(s)[a])
    return s

# Flow-matching loss at one state s, on the log scale  [eq. 12]
def state_loss(s, log_flow, eps, reward, lam_terminal):
    # Log-inflow: sum over every parent edge (p, a) with T(p, a) = s  [§2.1]
    parent_flows = [exp(log_flow(p, a)) for p, a in parents(s)]
    log_in = log(eps + sum(parent_flows))
    # Log-outflow: reward (0 if interior) plus flows of allowed children
    child_flows = [exp(log_flow(s, a)) for a in env.allowed(s)]
    log_out = log(eps + reward(s) + sum(child_flows))
    # Squared log-ratio; terminal transitions weighted by lambda_T  [app. A.4]
    loss = (log_in - log_out) ** 2
    return lam_terminal * loss if env.is_terminal(s) else loss

# Trajectory balance loss for one complete trajectory  [eq. 14]
def tb_loss(states, reward, log_z, p_fwd, p_bwd):
    # Sum the forward log-probabilities along the trajectory  [eq. 14]
    log_fwd = sum(log(p_fwd(nxt, given=prev)) for prev, nxt in pairs(states))
    # Sum the backward log-probabilities along the same trajectory  [eq. 14]
    log_bwd = sum(log(p_bwd(prev, given=nxt)) for prev, nxt in pairs(states))
    # Squared log-ratio; the reward must be positive  [§3]
    return (log_z + log_fwd - log(reward) - log_bwd) ** 2

Notes on the sketches:

• The paper does not say how the source state s_0 is handled in the flow-matching sketch, so the sketch leaves it out.
• The exp and log calls can overflow or lose precision with large predicted log-flows. A log-sum-exp form is a safer implementation choice; the paper writes the direct form 2106.04399 §2.2 eq. 12.
• Epsilon is ambiguous. The paper gives 2.5e-5 2106.04399 app. A.4, and the released molecule config uses a far smaller value (about 2.4e-8) gflownet.py. The paper does not say which one produced its numbers.
• The paper's eq. 12 sums over the states of a trajectory, but the released code averages the terminal and non-terminal losses over the batch gflownet.py gflownet.py. The learning rate and lambda_T do not transfer directly between the two.
• The released molecule config uses a reward floor of 0.1, where the paper gives 0.01 gflownet.py 2106.04399 app. A.4.
• The code masks the stop action for the first two blocks, which the paper does not state gflownet.py.
• The code's uniform exploration draws over the whole logit vector, including the stop index, while the paper says uniform over allowed actions gflownet.py 2106.04399 §4.2.
• The code marks a molecule terminal when stop is chosen, even if stems are still open gflownet.py. The paper defines terminal as having no children 2106.04399 §2.1.
• Hypergrid learning rate, minibatch and steps are not stated in the paper; the code's defaults are a learning rate of 1e-4, minibatch 16 and 20000 steps toy_grid_dag.py.
• For trajectory balance, the paper links a non-default trajectory_balance branch for its molecule code 2201.13259 §5.2. Check which branch you are reading before you copy any settings.
• The molecule settings differ between the paper and that branch: the paper sweeps the reward exponent over {4, 8, 10, 16} and the learning rate over {5e-5, 1e-4, 5e-4, 1e-3} without stating the values used in its figures 2201.13259 app. B.2.

Contradictions

The two learning-objective papers do not test each other's objective on the same task, so their numbers are not directly comparable.

• The TB paper reports that TB and detailed balance converge faster than flow matching on its hypergrid 2201.13259 §5.1 fig. 2, and reports a runtime speed-up over flow matching on molecules 2201.13259 §5.2. The flow-matching paper's card records no comparison with TB, and its own molecule results are reported against PPO and MARS [arxiv:2106.04399 app. A.7].
• Diversity is measured differently. The flow-matching paper reports mean pairwise Tanimoto similarity of the top-1000 molecules, with GFlowNet at 0.44 2106.04399 §4.2 app. A.7. The TB paper reports the same kind of similarity only as a curve, with lower values for TB than for flow matching 2201.13259 §5.2 fig. 3. No absolute values can be compared.
• The flow-matching paper's mode counts use different definitions in the main text and the appendix (scaffolds above reward 8, against Tanimoto-separated modes above reward 7), so its numbers should be quoted with their definitions 2106.04399 §4.2 app. A.7.
• The TB paper's flow-matching baseline uses the learning rate and leaf-loss coefficient lambda_T copied from an earlier paper, and lambda_T is not defined in the TB paper itself 2201.13259 app. B.3 app. B.4. Its flow-matching numbers therefore depend on that baseline's settings.

Papers

• Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation: introduces GFlowNet, a flow-matching trainer that samples in proportion to a positive reward over a DAG of construction steps.
• Trajectory balance: Improved credit assignment in GFlowNets: trains a GFlowNet with one squared log-ratio loss per complete trajectory, with a learned log Z.
• Machine learning and information theory concepts towards an AI Mathematician: a position essay that proposes reward-proportional conjecture sampling for theorem discovery, with no method or experiments.