Trajectory balance: Improved credit assignment in GFlowNets

In plain words

A GFlowNet is a model that builds objects one piece at a time, and it is trained so that a finished object turns up as often as its reward says it should, rather than only ever producing the single best object. Many different build orders can end in the same object, which makes training hard.

Older training rules for these models check each step on its own, locally. A mistake found late in a build then has to travel back through many steps before the early choices are corrected, and that is slow. Trajectory balance instead looks at a complete build at once. It compares how likely the whole build was under the model with the object's reward, and pushes the two toward agreement. The authors report that this speeds up training and gives more varied outputs.

TL;DR

Trajectory balance replaces step-by-step consistency checks with one loss per complete trajectory, using a learned scalar for the initial flow. Minimising it to zero makes the forward policy sample in proportion to reward. The authors report faster convergence and more diverse samples than flow matching and detailed balance on four tasks. Consider it if you train a GFlowNet, and read the pitfalls: the official branch the paper links differs from its default branch, and the paper's and code's settings do not fully match.

Problem

GFlowNets learn a policy that generates a compositional object x through a sequence of actions, so that P(x) is proportional to a reward R(x) §1 §2.2. The same object can be reached by many action sequences, so the probability of ending at x is a sum over exponentially many paths §1.

The two earlier local objectives, flow matching and detailed balance, enforce consistency one state or one edge at a time eq. 9 eq. 11. The authors argue that these objectives propagate credit slowly across long action sequences: when a new high-reward state appears far from the start, the updates must reach its parents before reaching the start of the sequence §6. Flow matching also needs the parents of every visited state, which the paper says is costly on large graphs app. B.2.

Method

Setting. A Markovian flow F on a directed acyclic graph of states, with a forward policy P_F, a backward policy P_B, and an initial flow Z = F(s_0) §2.1. A complete trajectory is a path from s_0 to a terminal state x. Its probability under the flow is F(τ)/Z, and the forward probability is the product of forward transition probabilities eq. 3 eq. 4. The reward-matching requirement is F(x) = R(x) eq. 8.

Trajectory balance constraint. For any complete trajectory τ ending in x, the flow satisfies Z times the product of P_F equals F(x) times the product of P_B eq. 13. Substituting R(x) for F(x) gives the form used for training eq. 14.

Loss. The authors define the per-trajectory loss as the squared logarithm of the ratio between (Z times the forward product) and (R(x) times the backward product) eq. 14. It is minimised along trajectories sampled from a training policy, which is usually the forward policy or a tempered version of it §3 eq. 15. Log Z, the forward logits and the backward logits are the trainable quantities. The backward policy P_B can be learned or fixed to uniform over parents §3.1.

# Trajectory balance loss for one trajectory  [alg. 1] [eq. 14]
# Input: states (s_0 ... s_n), reward R(x) > 0, log Z, P_F, P_B
# Output: scalar loss

def trajectory_balance_loss(states, reward, log_z, p_fwd, p_bwd):
    # Sum log P_F over the transitions s_{t-1} -> s_t  [eq. 14]
    log_forward = 0
    for prev, nxt in zip(states[:-1], states[1:]):
        log_forward = log_forward + log(p_fwd(nxt, given=prev))
    # Sum log P_B over the same transitions, read backwards  [eq. 14]
    log_backward = 0
    for prev, nxt in zip(states[:-1], states[1:]):
        log_backward = log_backward + log(p_bwd(prev, given=nxt))
    # Squared log-ratio of Z * prod P_F to R(x) * prod P_B  [eq. 14]
    residual = log_z + log_forward - log(reward) - log_backward
    return residual ** 2

# Training loop  [alg. 1]
# Repeats until the running loss has converged  [alg. 1]

def train(model, reward_fn, learning_rate):
    # Sample a trajectory from P_F, or a tempered version  [alg. 1] [§3]
    states = sample_trajectory(model.p_fwd)
    # One gradient step on the loss; the algorithm uses batch size 1  [alg. 1]
    loss = trajectory_balance_loss(
        states,
        reward_fn(states[-1]),
        model.log_z,
        model.p_fwd,
        model.p_bwd,
    )
    model.update(loss, learning_rate)

Training procedure. The algorithm samples one trajectory, computes the loss, and takes a gradient step with learning rate η. Its stopping check is convergence of the running loss alg. 1.

Correctness. Proposition 1 has two parts. (a) If the policies and Z come from a flow with reward matching, the loss is zero for every complete trajectory. (b) If the loss is zero for every complete trajectory, the induced flow matches the reward and the forward policy samples in proportion to reward. The proof is given in the appendix app. A.1. The main-text assumption is that R is positive, because the loss divides by R(x); smoothing can relax this §3.

Autoregressive case. When the graph is a tree, each state has one parent, so P_B is 1 and the loss simplifies eq. 16. The authors note that in this case the detailed balance objective reduces to flow matching, which they relate to soft Q-learning §2.2, remark. §4 also relates trajectory-level losses to RL objectives on subtrajectories.

Parametrisation. The authors parametrise Z in log space and output logits for P_F and P_B, so the products in the loss become sums §3, remark 2. They found a higher learning rate for log Z helpful, and explain it as a weighted moving average of the discrepancy between the two sides §3, footnote.

Extensions (appendix). Sub-trajectory balance applies the same constraint to segments between chosen states, so that not every state needs a flow estimate; detailed balance and trajectory balance are the two extremes app. A.2. The paper also gives a constraint for paths that mix backward and forward steps, which needs no initial flow app. A.2. The authors show that on-policy trajectory balance has the same expected gradient as a REINFORCE-style KL objective, with lower variance near the optimum app. A.3.

Experiments. Four tasks §1:

• Hypergrid §5.1: a synthetic D-dimensional grid with a reward that has peaks at the corners, plateaus and troughs. Models are an MLP with one-hot input; the policy and backward policy share all but the last layer fig. 1 §5.1.
• Molecule synthesis §5.2: junction-tree graphs built from a vocabulary of building blocks, with reward from a learned proxy of docking energy against a sEH target. Trajectory length is at most 8 §5.2.
• Bit sequences §5.3.1: fixed-length binary strings with reward based on edit distance to a hidden set of modes, generated in words of k bits.
• Peptides (AMP) §5.3.2: sequences over 20 amino acids with a classifier reward.

Hyperparameters

Values are as stated in the paper unless marked as from code.

Results

Author-reported. Setup as described in each section.

AMP generation (Transformer policies; 2048 samples per method; top 100 by reward; diversity is mean pairwise edit distance) §5.3.2 tab. 1:

Hypergrid §5.1 fig. 2: the L1 error between the sampled and true reward distributions falls for all models. Detailed balance and trajectory balance converge faster than flow matching, with a slight advantage for trajectory balance on the 4-D grid. Fixing P_B to uniform slows learning, especially on the 64×64 grid §5.1, "Effect of uniform P_B". PPO and MCMC curves lie outside the plotted range fig. 2. The paper gives no table of these values.

Molecule synthesis §5.2 fig. 3: the Pearson correlation between log reward and log sampling probability is higher for trajectory balance than for flow matching and detailed balance on a held-out set (left, centre) fig. 3. Over 16 hyperparameter settings and 3 seeds, trajectory balance fits better (centre) fig. 3. The mean pairwise Tanimoto similarity of the top-1000 samples is lower for trajectory balance than for flow matching throughout training, which the authors read as greater diversity (right) fig. 3. The paper reports a runtime speed-up of up to 5× over flow matching §5.2 app. B.2.

Bit sequences §5.3.1 fig. 4: trajectory balance has the highest Spearman correlation between sampling probability and reward across all tested k (left) fig. 4. Flow matching improves as k grows, that is, as trajectories get shorter; large action spaces hurt flow matching but not trajectory balance §5.3.1. For k=1, trajectory balance finds modes faster than the other methods (right) fig. 4.

Paper vs code

Which code. The paper links github.com/GFNOrg/gflownet on the trajectory_balance branch for the molecule experiments §5.2. That branch is the one code_commit refers to (commit 4d3a7d5, dated 2022-08-30, read from the GitHub branches API). The default branch, which we cloned to raw/code/2201.13259 at commit 3ba65d2, holds the flow-matching code instead, with grid/ and mols/ folders and an MIT LICENSE (copyright Emmanuel Bengio, 2022). The license of the trajectory_balance branch was not read separately. The code was read through WebFetch, not as a local checkout, so no line numbers are given: citations name file and function only.

Hypergrid and bit sequences. The paper links GitHub gists, not repositories, for the hypergrid §5.1 and bit-sequence app. B.3 code. These are not repositories and are not used here. The paper's checklist says code is released only for the grid and molecule environments [checklist, item 3a].

Molecule settings.

• Reward exponent: the paper sweeps β in {4, 8, 10, 16} app. B.2. The trajectory_balance branch defaults reward_exp to 10 gflownet.py (trajectory_balance branch).
• Learning rate: the paper's sweep is {5e-5, 1e-4, 5e-4, 1e-3} app. B.2. The branch default is 5e-4 gflownet.py (trajectory_balance branch). The paper does not state which value its figures use.
• Random actions: the paper's exploration rule is a 0.1 chance to stop at a sampled length k app. B.2. The branch has a random_action_prob option with default 0.05 gflownet.py (trajectory_balance branch). We did not confirm that this option implements the paper's stop-length rule.
• Default objective: the branch default is fm; trajectory balance runs only with --objective tb gflownet.py (trajectory_balance branch).
• Iterations and batch: the branch default is 250,000 iterations with batch 4; its good_config overrides iterations to 30,000 and keeps batch 4 gflownet.py (trajectory_balance branch). The paper does not state the molecule batch size app. B.2.

Loss. The branch's inline trajectory-balance branch computes the paper's loss: a learned log Z plus summed forward log-probabilities, minus the log of the reward and a summed uniform backward term, squared and averaged over the minibatch gflownet.py (trajectory_balance branch) eq. 14. The backward term is a per-step log(1/n), where n is not clarified in the code excerpt we read, so check it before reuse. The reward is (score / norm) ** reward_exp, which matches the paper's R(x)^β form app. B.4 §5.2.

Not checked. We did not read the grid code on the default branch (grid/toy_grid_dag.py) against the paper's hypergrid settings.

Pitfalls

• Wrong branch. The default branch of GFNOrg/gflownet is flow matching. Trajectory balance is only on the trajectory_balance branch gflownet.py.
• Uniform backward policy in the molecule experiment. It is fixed to uniform app. B.2. Learning P_B jointly is faster on the 64×64 grid §5.1.
• Two different "random action" settings. The paper's 0.1 stop-length rule app. B.2 and the code's random_action_prob gflownet.py are not the same mechanism.
• Log Z needs its own learning rate. The paper sets a higher one for log Z, noting that momentum complicates the moving-average view §3, footnote. Values differ by task: 1e-1 on hypergrid app. B.1, 1e-3 for bit sequences app. B.3, 1e-2 for AMP app. B.4.
• Positive reward required. The loss divides by R(x); zero or negative rewards need smoothing constants §3. The paper does not give the constants it used.
• Variance. The authors note that trajectory-level sampling gives higher gradient variance than detailed balance, and that this is a possible limit in hard environments §6 fig. 3 caption.
• Flow-matching numbers are not comparable. The FM baseline uses its own learning rate and leaf-loss coefficient λ_T app. B.3 app. B.4, copied from the earlier paper. λ_T is not defined in this paper.
• Tree-shaped graphs. In the autoregressive case P_B is 1 §3, remark 1, so the TB loss reduces to a sum over forward log-probabilities and log R.
• Checklist. The authors answer "No" for code of the environments other than grid and molecule [checklist, item 3a]. Bit-sequence and AMP results cannot be reproduced from released code.
• Equation numbers in this card are counted from the LaTeX source (arXiv v4) in order. Check them against the compiled PDF before publishing.

Lineage

• Builds on: flow matching and detailed balance for GFlowNets (Bengio et al., 2021, cited as bengio2021flow and bengio2021foundations in the paper) §2.2.
• Related: soft Q-learning and control as inference §4; trajectory-level RL losses on subtrajectories §4; variational sequential samplers app. A.3.
• Compared against: flow matching and detailed balance §5; A2C with entropy regularisation, SAC, MARS and MCMC §5.3 tab. 1.
• Extended by: sub-trajectory balance, cited as a work that appeared during review app. A.2 §6. The paper also cites two concurrent works on the variational view app. A.3.
• Superseded by: not stated in this paper.