Trajectory balance: Improved credit assignment in GFlowNets

Nikolay Malkin; Moksh Jain; Emmanuel Bengio; Chen Sun; Yoshua Bengio

Source-linked AI summary

Trajectory balance: Improved credit assignment in GFlowNets

Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, Yoshua Bengio

arXiv:2201.13259v3Jan 31, 2022cs.LGstat.ML

TL;DR

Flow matching and detailed balance can propagate credit inefficiently across long GFlowNet action sequences. The paper introduces trajectory balance, proves its correctness for target-distribution sampling, and finds faster, more robust, and more diverse learning across experiments.

Problem
Flow matching and detailed balance objectives are local in action sequences and may require many iterations to propagate credit to early actions.
Method
Trajectory balance optimizes sampled complete trajectories and proves that global minimization yields a policy sampling proportionally to the unnormalized reward.
Results
Trajectory balance accelerated convergence, discovered more reward modes, improved diversity and reward divergence metrics, and was more robust to long sequences and large action spaces.
Takeaways & Limitations
Trajectory balance is presented as a superior training objective and the default choice for future GFlowNet work.
Takeaways & Limitations
The theoretical target assumes a nontrivial nonnegative reward function on terminal states, while optimizer momentum effects remain for future empirical investigation.

Abstract

from arXiv · show

Generative flow networks (GFlowNets) are a method for learning a stochastic policy for generating compositional objects, such as graphs or strings, from a given unnormalized density by sequences of actions, where many possible action sequences may lead to the same object. We find previously proposed learning objectives for GFlowNets, flow matching and detailed balance, which are analogous to temporal difference learning, to be prone to inefficient credit propagation across long action sequences. We thus propose a new learning objective for GFlowNets, trajectory balance, as a more efficient alternative to previously used objectives. We prove that any global minimizer of the trajectory balance objective can define a policy that samples exactly from the target distribution. In experiments on four distinct domains, we empirically demonstrate the benefits of the trajectory balance objective for GFlowNet convergence, diversity of generated samples, and robustness to long action sequences and large action spaces.

1 Introduction

GFlowNets learn stochastic policies for generating compositional objects so terminating states are sampled according to an unnormalized reward while preserving diversity across multiple action sequences. The paper introduces trajectory balance to improve credit assignment and evaluates it across four domains.

1 Introduction: GFlowNets generate compositional objects through stochastic action sequences that incrementally modify partial states.Examples include adding nodes or edges to graphs, starting from a universal initial state.
1 Introduction: GFlowNets target reward-proportional sampling rather than concentration on a maximal-reward state, including when multiple paths reach one terminal state.The likelihood of a state sums over potentially exponentially many action sequences.
1 Introduction: Flow matching and detailed balance use local flow constraints analogous to temporal-difference learning, which can make credit propagation inefficient across long sequences.Flow matching enforces conservation of incoming and outgoing flows at states.
1 Introduction: Trajectory balance is computed on sampled full trajectories and is proved to yield reward-proportional sampling at global minima.The objective is proposed as an alternative to flow matching and detailed balance.
1 Introduction: Experiments cover hypergrids, molecule synthesis, and sequence generation, including settings with wide troughs, large action spaces, and long sequences.The paper reports convergence, diversity, computational speed-up, and robustness benefits for trajectory balance.

2 Preliminaries

The preliminaries define GFlowNets on a directed acyclic state-action graph and explain how flow-based policies achieve reward-proportional terminal sampling. They motivate local flow-matching and detailed-balance objectives because directly computing terminal flow can be impractical.

2.1 Markovian flows: A GFlowNet state space is a directed acyclic graph with one initial source, terminal sinks, states, actions, children, and parents.Complete trajectories move from the initial state to a terminal state.
2.1 Markovian flows: A trajectory flow assigns nonnegative flow to complete paths, while state flow aggregates the flow of trajectories passing through a state.Flow values can be scaled because probabilities are obtained by normalization.
2.1 Markovian flows: A Markovian flow induces forward and backward policies over children and parents, with detailed balance relating these local transitions.The flow can be characterized through edge flows satisfying flow-matching constraints.
2.2 GFlowNets: GFlowNets seek a flow whose terminal flow equals the reward, so the forward policy terminates at x with likelihood proportional to R(x).Global minimization of a suitable objective guarantees the target-distribution property.
2.2 GFlowNets: Directly optimizing terminal flow is generally impractical because it sums over all trajectories leading to a terminal state.Flow matching and detailed balance were introduced as local alternatives optimized from sampled trajectories.
2.2 GFlowNets: Global minima of both flow matching and detailed balance recover the target distribution under full-support training policies.Detailed balance uses neural estimates of state flows and forward and backward action distributions.
2.2 GFlowNets: For autoregressive sequence generation on tree-structured graphs, the backward policy is trivial and detailed balance reduces to flow matching.The resulting objective is connected to Soft Q-Learning with specified entropy-regularization settings.

3 Trajectory balance

Trajectory balance converts a global path-consistency constraint into a loss evaluated on sampled complete trajectories. Under positive rewards, full-support sampling, and global minimization, it recovers a Markovian flow whose forward policy samples proportionally to reward.

3 Trajectory balance: Trajectory balance follows from an algebraic constraint relating the initial flow, forward-policy probability, backward-policy probability, and terminal reward along each complete trajectory.The constraint is a special case of a broader constraint that also includes detailed balance.
3 Trajectory balance: The trajectory-balance model outputs forward and backward policies plus a global scalar Zθ that estimates the initial flow.Zθ and the forward policy uniquely determine an implicit Markovian flow.
3 Trajectory balance: The trajectory loss is optimized with stochastic gradients on trajectories sampled from the forward policy or a tempered training policy.The full batch-size-one procedure samples a trajectory and updates the parameters until convergence.
3 Trajectory balance: At global optimum, the trajectory-balance loss is zero on every complete trajectory when the policies and normalizer correspond to a flow satisfying the target constraint.The proposition also states the converse when the loss is zero for all complete trajectories.
3 Trajectory balance: With full-support training and global minimization, trajectory balance yields a forward policy that samples proportionally to the reward.The proof assumes a positive reward function; smoothing can relax positivity.
3 Trajectory balance: On autoregressive trees, the backward policy is identically one, while logarithmic parameterization turns trajectory products into sums.The authors also use a higher learning rate for Z than for policy parameters.
3 Trajectory balance: The target constraint may admit multiple flows on graphs whose underlying undirected structure contains cycles, although fixing a backward policy selects a unique compatible flow.A uniform backward policy is suggested when modeling parent distributions is difficult.

4 Related work

The related-work discussion connects GFlowNets to reinforcement learning, MCMC, and variational inference while distinguishing their treatment of compositional structure and multimodal sampling. It also places trajectory balance within broader optimization and inference perspectives.

Reinforcement learning: On tree-structured autoregressive graphs, GFlowNets are equivalent to reinforcement-learning methods with suitable entropy regularization or soft Q-learning.The discussion contrasts this equivalence with failures in general DAGs that GFlowNets can handle.
Local exploration vs. amortized generalization to unseen modes: Unlike local MCMC exploration, GFlowNets treat compositional structure as a learning problem that can generalize to unseen modes.This connection is relevant when exact sampling is intractable.
Variational inference: Trajectory balance is connected to variational inference, including squared log-ratio losses used as control variates for evidence lower-bound optimization.The paper also relates on-policy trajectory-balance optimization to minimizing a KL divergence.

5 Experiments

Experiments compare trajectory balance with prior GFlowNet objectives and baselines across hypergrid, molecule, bit-sequence, and peptide-generation tasks. TB generally converges faster, fits reward-proportional sampling better, preserves greater diversity, and remains robust to long trajectories and large action spaces.

Hypergrid environment: TB and DB generally converge faster than FM on hypergrid environments, with TB slightly faster than DB in the 4-D setting.All models and settings reach the same stable minimum in the reported experiments.
Hypergrid environment: Learning the backward policy substantially accelerates convergence on the 64 × 64 grid compared with fixing it to a uniform distribution.The jointly learned policies may approach one of multiple optimal solutions more quickly.
Small drug molecule synthesis: TB produces higher reward–sampling correlations than FM and DB on held-out molecule states.The correlation compares log-reward with the log-likelihood that the learned policy terminates at each molecule.
Small drug molecule synthesis: TB consistently generates more diverse molecules than FM, measured by lower average pairwise Tanimoto similarity among the 1000 highest-reward samples.TB also achieves up to 5× runtime speedup over FM because parent enumeration is unnecessary.
Bit sequences: TB learns the highest reward correlations across action spaces and discovers more bit-sequence modes faster than competing methods.FM degrades with larger action spaces and improves mainly when trajectories become shorter, whereas TB is reported as robust to both factors.
Anti-Microbial Peptide generation: TB outperforms all baselines on both performance and diversity metrics for antimicrobial peptide generation.The metrics are mean reward and average pairwise edit distance among generated sequences.

6 Discussion and conclusion

The paper concludes that trajectory balance improves credit assignment and outperforms flow matching and detailed balance across the studied settings, while introducing a variance trade-off for long trajectories.

Trajectory balance provides faster and better training than flow matching and detailed balance by addressing their slow credit propagation.FM and DB are local in the action sequence, so updates may require many iterations to reach early actions.
Global minimization of TB yields a GFlowNet policy that samples from the target distribution specified by the unnormalized reward function.
TB discovers more energy-function modes faster and is more robust to exponential state-space growth from long sequences and large action spaces.The reported robustness reflects both trajectory length and action-space size.
TB has a potential limitation because sampling long trajectories can increase stochastic-gradient variance in difficult environments.Subsequent work studied interpolations between local and trajectory-level objectives to address this issue.
Across a broad set of experiments, the authors identify trajectory balance as a superior training objective and recommend it as the default choice for future GFlowNet work.

Checklist

The checklist records that the paper is primarily theoretical and algorithmic, reports reproducibility and ethics information, and identifies application-specific limitations for future work.

The work is mainly about theory and algorithms, with limitations and negative societal impacts left for future applications to consider.Molecule design for drug discovery is given as an example of an application requiring such consideration.
The paper states that complete proofs, training details, error bars, and compute information are included or described in the appendix.
Code and data are reported as reproducible for the grid and molecule environments but not for the other environments.
Existing assets are cited, and references to the molecule and AMP sequence data are provided.
The checklist reports no new data collection, no relevant personally identifiable or offensive content, and no human-subject procedures.

A.1 Proof of Proposition 1

The proof establishes that trajectory balance is equivalent to the desired flow constraints at zero loss, while extensions generalize the constraint to partial and mixed-direction trajectories.

Conversely, a Markovian flow satisfying the flow constraints makes the trajectory balance loss zero on every complete trajectory.The proof substitutes the reward for terminal flow in the trajectory balance constraint.
If trajectory balance loss is zero for every complete trajectory, the corresponding Markovian flow satisfies the flow constraints and samples proportionally to the reward.
Subtrajectory balance: The trajectory balance constraint can be generalized from complete trajectories to partial trajectories and used for modular or hierarchical GFlowNets.
Subtrajectory balance: Subtrajectory balance contains trajectory balance as its full-trajectory case and detailed balance as its one-edge case.The associated objective can estimate state flows only at selected nonterminal hub states.
Non-forward trajectories: Non-forward trajectory balance supports paths that move backward from one terminal state and then forward to another, with likelihood ratios determined by endpoint rewards.The construction does not require the flow at the shared intermediate state or an initial-state flow output.

A.3 GFlowNets and variational methods

The appendix connects trajectory balance with variational sequential sampling: its gradient matches a KL-based Reinforce gradient up to a constant, with lower variance near the optimum.

The paper connects the trajectory balance loss to a naïve variational approach for fitting sequential samplers.
With a fixed backward policy, the forward policy can be optimized by minimizing D_KL(P_F(τ)∥R(x)P_B(τ|x)) through sampled trajectories.The Reinforce estimator samples trajectories from the forward policy and uses the score-function gradient.
The expected trajectory balance gradient equals the expected Reinforce gradient up to a constant.
Near the optimum, where TB and KL are zero, the trajectory balance gradient estimator has lower variance.
The variance comparison is guaranteed in all directional derivatives when the relevant parenthesized term is always positive, including the neighborhood where P_F(τ)=R(x)P_B(τ|x).

B.1 Hypergrid

The hypergrid experiments compare TB and DB with uniform or learned backward policies across reward scales and grid geometries. Training uses reproduced FM and baseline implementations alongside specified model and optimization settings.

Implementation: The GFlowNet policy uses an MLP with 2 hidden layers of 256 units, trained with Adam and up to 106 sampled trajectories.PF and PB use learning rate 0.001, while Zθ uses 0.1; the batch size is 16.
Implementation: FM and non-GFlowNet baselines use published code, while TB and DB use learning rate 10−3 for flow and policy models and 10−1 for log Z.The 10−3 rate was the largest tested rate that did not cause rapid mode collapse.
Experimental design: TB and DB are evaluated with uniform or learned PB across R0 = 10−1, 10−2, 10−3 and grid settings (H, d) = (8, 4), (64, 2).The comparison spans two objectives, two backward-policy choices, three reward scales, and two grid geometries.

B.2 Molecule synthesis

The molecule-synthesis experiments use the provided dataset and proxy model while testing TB against FM and other baselines under varied reward exponents and learning rates. They report up to 5× runtime improvement for TB relative to FM, with the benefit depending on model evaluations, trajectory length, and hardware.

Molecule synthesis: The molecule task uses the dataset and proxy model from [3], with binding scores computed using AutoDock and a fixed uniform backward policy for TB.The GFlowNet uses the same architecture and hyperparameters as the comparison setup except β and learning rate.
Robustness: Hyperparameter robustness is tested over β = {4, 8, 10, 16} and learning rates {5 × 10−5, 10−4, 5 × 10−4, 10−3}.The training policy uses a 0.1 probability of stopping at a uniformly chosen length between 3 and 8.
Runtime: 5× runtime improvement is observed for TB relative to FM.The reported factors include model-evaluation requirements, average trajectory length, and the ratio of CPU and GPU load.
Runtime: FM evaluates models for the parents of sampled states, whereas TB and DB require one forward and backward policy evaluation per state.Shorter trajectories can further reduce evaluations when models terminate earlier.
Runtime: TB’s computation-time benefit remains present on newer hardware but appears smaller with identical batch-size settings.Molecule experiments used a Tesla K80 GPU, while grid experiments used CPU and were mainly controlled by trajectory length.
Experimental design: The experiments evaluate TB against FM, A2C, MARS, and other methods for autoregressive generation using PyTorch implementations.The broader evaluation includes molecule synthesis and other tasks, with MARS adapted from released code.