---
kind: concept
slug: reward-proportional-sampling
title: Reward-proportional sampling
summary: >-
  Generate discrete objects step by step so that each finished object is sampled with probability proportional to a positive reward, not only the single best one. Core issue: many build orders lead to one object. Variants: flow matching and trajectory balance (GFlowNets). Trade-off: flow training and tuning against simpler return-maximising RL.
tags: [sampling, diversity, generative-models, flow-networks]
papers: ["2106.04399", "2201.13259", "2403.04571"]
related: [active-learning, bayesian-optimization, goal-conditioned-reinforcement-learning, policy-gradient]
status: draft
updated_at: 2026-10-24
---

## 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 [arxiv: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 [arxiv: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 [arxiv: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 [arxiv: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 [arxiv: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
[arxiv: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 [arxiv: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 [arxiv: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 [arxiv: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 [arxiv:2106.04399 §2.2]. The proof assumes the predictor can
represent a consistent flow [arxiv: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 [arxiv: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 [arxiv: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 [arxiv: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 [arxiv: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
  [arxiv:2106.04399 app. A.5.2] [§4.3] [alg. 1]. Card:
  [GFlowNet](/papers/2106.04399).

### 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 [arxiv:2201.13259 §3] [eq. 14]. If this loss
is zero on every complete trajectory, the forward policy samples in proportion
to reward [arxiv:2201.13259 app. A.1]. The backward policy can be learned or
fixed to uniform over parents [arxiv:2201.13259 §3.1]. The molecule runs fix it
to uniform [arxiv:2201.13259 app. B.2].

Each step samples one trajectory, takes one gradient step on its loss, and stops
when the running loss converges [arxiv: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 [arxiv: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 [arxiv: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
  [arxiv: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 [arxiv: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 [arxiv: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 [arxiv:2201.13259 §5.3.1]
  [fig. 4].

Card: [Trajectory balance](/papers/2201.13259).

### 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 [arxiv: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 [arxiv:2403.04571 §6] [§8] [§9]. Card:
[AI Mathematician essay](/papers/2403.04571). 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".

| Situation | Use | Why |
|---|---|---|
| Many build orders lead to the same object, and you need samples in proportion to reward | Flow-matching generator (GFlowNet), or trajectory balance | Tree-based reward sums over-weight objects with many build orders [arxiv:2106.04399 §2]; TB targets the same distribution [arxiv:2201.13259 §3] |
| You only need the single best object | Return-maximising RL such as PPO | Return maximisation converges to one high-return sequence [arxiv:2106.04399 §1]; in the molecule runs, PPO's top-10 is close to GFlowNet's, but its diversity is lower [arxiv:2106.04399 app. A.7] |
| Late steps carry the reward, so credit must travel far back | Trajectory balance | The authors argue local consistency checks propagate credit slowly [arxiv:2201.13259 §6]; the paper's hypergrid and molecule results are consistent with this, but it is the authors' argument [§5] |
| Large action spaces, such as wide word sizes | Trajectory balance | Large action spaces hurt flow matching but not TB in the bit-sequence runs [arxiv:2201.13259 §5.3.1] |
| Many parents per state, or large graphs | Trajectory balance | Flow matching needs the parents of every visited state, which the paper calls costly on large graphs [arxiv:2201.13259 app. B.2] |
| Reward is positive, and you can afford many samples | Flow-matching generator, or MCMC if you can wait for mixing | MCMC reaches the target but mixes slowly [arxiv:2106.04399 §3]; the hypergrid shows its sample cost rising as R0 falls [arxiv:2106.04399 §4.1] |
| Reward can be zero or negative | Not covered by either paper as stated | TB divides by R(x), and the paper says smoothing is needed but does not give the constants it used [arxiv:2201.13259 §3] |
| Few, large batches of expensive queries (e.g. docking) | Flow-matching generator inside a batch loop | The paper's motivation for this setting [arxiv:2106.04399 §1]; the batch-loop result is one paper's, with 3 runs [arxiv:2106.04399 app. A.5.2] |
| You have logged trajectories from other policies | Flow-matching generator, off-policy | Off-policy training is allowed if support covers the optimal policy's [arxiv:2106.04399 §2.2]; the paper reports no hyperparameter search for it [app. A.6] |
| Tight compute budget | Not stated for either method | Molecule runs took about 26 GPU days on V100s [arxiv:2106.04399 app. A]; TB's compute is given per task, with the molecule run on a Tesla K80 [arxiv:2201.13259 app. B] |
| Memory and latency at inference | Not stated | Neither paper reports an inference cost or memory figure |
| Theorem or conjecture generation | Essay only, no method | The essay gives a framing and no experiments [arxiv:2403.04571 §8] |

## 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.

```pseudocode
# 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
```

```pseudocode
# 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
```

```pseudocode
# 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 [arxiv:2106.04399 §2.2] [eq. 12].
- Epsilon is ambiguous. The paper gives 2.5e-5 [arxiv:2106.04399 app. A.4],
  and the released molecule config uses a far smaller value (about 2.4e-8)
  [code: mols/gflownet.py#L578]. 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
  [code: mols/gflownet.py#L474] [code: mols/gflownet.py#L479-L482]. 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 [code: mols/gflownet.py#L577] [arxiv:2106.04399 app. A.4].
- The code masks the stop action for the first two blocks, which the paper
  does not state [code: mols/gflownet.py#L167-L168].
- The code's uniform exploration draws over the whole logit vector, including
  the stop index, while the paper says uniform over allowed actions
  [code: mols/gflownet.py#L175-L176] [arxiv:2106.04399 §4.2].
- The code marks a molecule terminal when stop is chosen, even if stems are
  still open [code: mols/gflownet.py#L180-L186]. The paper defines terminal
  as having no children [arxiv: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 [code: grid/toy_grid_dag.py#L27-L42].
- For trajectory balance, the paper links a non-default `trajectory_balance`
  branch for its molecule code [arxiv: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
  [arxiv: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 [arxiv:2201.13259 §5.1] [fig. 2], and reports a
  runtime speed-up over flow matching on molecules [arxiv: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 [arxiv: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
  [arxiv: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
  [arxiv: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 [arxiv: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](/papers/2106.04399): 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](/papers/2201.13259): 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](/papers/2403.04571): a position essay that proposes reward-proportional conjecture sampling for theorem discovery, with no method or experiments.