Source-linked AI summary
Better Training of GFlowNets with Local Credit and Incomplete Trajectories
Ling Pan, Nikolay Malkin, Dinghuai Zhang, Yoshua Bengio
TL;DR
Previous GFlowNet objectives require complete trajectories and terminal rewards, limiting training from incomplete paths. FL-GFN uses energy information available at intermediate states to provide local credit, and the paper reports convergence guarantees and effectiveness across generation tasks.
Problem
Previous GFlowNet objectives learn from terminal rewards and require complete trajectories, leaving training from incomplete trajectories unsupported.
Source: §1 · p. 2Method
FL-GFN exploits per-state or per-transition energy information to provide local credit as partial reward factors accrue.
Source: §1 · p. 2, §4.1 · pp. 4–5Results
FL-GFN is shown to guarantee convergence while training from incomplete trajectories, and experiments report effectiveness on set and bit-sequence generation and scalability to molecule discovery.
Source: §1 · p. 2Takeaways & Limitations
The approach supports earlier credit during training and can also accelerate training when complete trajectories are available.
Source: §1 · p. 2, §6 · p. 9FL-GFN relies on per-state or per-transition energy information, as in cases where additive energy terms are available along the trajectory.
Source: §1 · p. 2
Abstract
from arXiv · showhide
Generative Flow Networks or GFlowNets are related to Monte-Carlo Markov chain methods (as they sample from a distribution specified by an energy function), reinforcement learning (as they learn a policy to sample composed objects through a sequence of steps), generative models (as they learn to represent and sample from a distribution) and amortized variational methods (as they can be used to learn to approximate and sample from an otherwise intractable posterior, given a prior and a likelihood). They are trained to generate an object $x$ through a sequence of steps with probability proportional to some reward function $R(x)$ (or $\exp(-\mathcal{E}(x))$ with $\mathcal{E}(x)$ denoting the energy function), given at the end of the generative trajectory. Like for other RL settings where the reward is only given at the end, the efficiency of training and credit assignment may suffer when those trajectories are longer. With previous GFlowNet work, no learning was possible from incomplete trajectories (lacking a terminal state and the computation of the associated reward). In this paper, we consider the case where the energy function can be applied not just to terminal states but also to intermediate states. This is for example achieved when the energy function is additive, with terms available along the trajectory. We show how to reparameterize the GFlowNet state flow function to take advantage of the partial reward already accrued at each state. This enables a training objective that can be applied to update parameters even with incomplete trajectories. Even when complete trajectories are available, being able to obtain more localized credit and gradients is found to speed up training convergence, as demonstrated across many simulations.
1. Introduction
The paper addresses the need to train GFlowNets without complete trajectories, proposing FL-GFN to use intermediate energy information for earlier credit. It reports convergence guarantees and experiments across several generation tasks.
1. Introduction
- Previous GFlowNet objectives learn from terminal rewards, so complete trajectories are needed and delayed feedback can make credit assignment inefficient.The paper identifies long trajectories and complex terminal objects as settings where complete-trajectory training may be infeasible.Source: §1 · p. 2
- Partial energy factors can be available along a trajectory when the energy is additive, including terms computed as variables in a sparse factor graph are specified.The paper motivates incomplete-trajectory learning through partial inferences that do not fully specify latent variables.Source: §1 · p. 2
- FL-GFN exploits per-state or per-transition energy values to provide local credit and gradients as reward factors accrue.The method is also reported to accelerate training when complete trajectories are available.Source: §1 · p. 2
- FL-GFN is shown to guarantee convergence to the reward-matched distribution while learning from incomplete trajectories without terminal states.The formulation makes credit information available early.Source: §1 · p. 2
- Experiments on set and bit-sequence generation demonstrate FL-GFN’s effectiveness, and the approach scales to molecule discovery.The introduction describes the molecule discovery task as more complex and challenging.Source: §1 · p. 2
2. Background
GFlowNets construct objects as trajectories through a directed acyclic graph and aim for terminal-state probabilities proportional to reward. The section introduces balance-based training criteria and their relationship to complete and partial trajectories.
2.1. GFlowNet Preliminaries
- A GFlowNet represents sequential construction as trajectories from an initial state through a directed acyclic graph to terminal states, with a forward policy targeting reward-proportional terminal sampling.The forward policy is parameterized as a neural network that maps a state to a distribution over its children.Source: §2.1 · pp. 2–3
2.2. Training Criteria for GFlowNets
- Detailed balance learns a state-flow function and backward policy, with terminal-state flow set to the reward and a constraint for each transition.For a given forward policy, the section states that the satisfying backward policy and state-flow function are unique.Source: §2.2 · p. 3
- Trajectory balance applies its constraint to complete trajectories, whereas subtrajectory balance applies it to partial trajectories whose endpoints need not be initial or final.Detailed balance and trajectory balance are special cases of the SubTB constraint at length-one partial trajectories and complete trajectories, respectively.Source: §2.2 · p. 3
- Satisfying SubTB constraints for partial trajectories of all lengths implies the detailed-balance and trajectory-balance constraints and is sufficient for reward-proportional sampling.Prior empirical testing reports that the SubTB objective reduces the gradient variance of the trajectory-balance objective.Source: §2.2 · p. 3
3. Related Work
Related work places GFlowNets among methods for sequential generation and sampling, while highlighting limitations in credit assignment and complete-trajectory training. It also contrasts their objectives with reinforcement learning and other GFlowNet applications.
3. Related Work
- Prior work applies GFlowNets to molecule discovery, sequence design, Bayesian structure learning, and theory, and has also investigated improving training efficiency.The section additionally notes work on intrinsic exploration rewards and stochastic environments.Source: §3 · p. 4
- Terminal-only trajectory rewards make it difficult to assign credit to actions, especially early ones, and complete-trajectory requirements can be infeasible for large composite objects.The related-work discussion identifies long trajectories and complex terminal objects as practical constraints.Source: §3 · p. 4
- Unlike GFlowNets’ reward-proportional sampling objective, reinforcement learning typically seeks reward maximization; terminal feedback can also hinder learning efficiency in long-horizon problems.The passage also notes that Soft Q-learning can be biased toward terminal states with more trajectories in general DAGs.Source: §3 · p. 4
4. Proposed Method
FL-GFN uses energy available at intermediate states to reparameterize state flow, aiming to learn from incomplete trajectories and provide more local credit. The method applies when terminal energy extends to intermediate states, including additive-energy settings and domains where the target-space energy can be evaluated on nonterminal states.
4. Proposed Method
- Previous DB, TB, and SubTB objectives rely on terminal energy and complete trajectories, a constraint that becomes critical for long trajectories or complex terminal objects.FL-GFN instead uses computable intermediate-state energy to support incomplete-trajectory learning and faster learning.Source: §4 · p. 4
- The method assumes terminal energy extends to all states, allowing state energy to be decomposed additively over transitions.The paper defines transition energy and trajectory energy differences to express the energy of a state along a trajectory.Source: §4 · p. 4
4.1. Sources of Partial Energies
- Purely additive energies cover set construction, where inserting an element contributes an energy term; the paper also applies the approach to parsing under probabilistic context-free grammars.The parsing example concerns posteriors over compositional latent variables.Source: §4.1 · pp. 4–5
- When states share the target object's space, the same energy function can be evaluated on nonterminal states, and transition energy measures stepwise marginal improvement in log-reward.This setting includes molecule generation when the reward function extends to nonterminal states.Source: §4.1 · p. 5
4.2. Forward-looking GFlowNets (FL-GFN)
- FL-GFN factors known intermediate-state energy out of terminal-energy flow, defining a forward-looking flow over future transition energies and a corresponding FL-DB constraint.The flow uses backward-policy reachability from intermediate states to terminal states.Source: §4.2 · p. 5
- The FL-DB constraint supplies intermediate energy signals and permits incomplete-trajectory learning when those signals contain enough information to generalize downstream.The implementation optimizes the transition-level FL-DB loss in log space; an alternative substitutes a learned flow model into DB.Source: §4.2 · p. 5
- If the FL-DB loss is zero for all transitions under the stated energy-extension assumption, the forward policy samples proportionally to reward.At a global minimum of expected FL-GFN loss, the method is theoretically justified to sample from the target reward distribution.Source: §4.2 · pp. 5–6
5. Experiments
The experiments compare FL-GFN with existing GFlowNet objectives across set and bit-sequence generation, test training from incomplete trajectories, and extend evaluation to molecule generation. Across these settings, FL-GFN improves training efficiency and solution quality, including when trained on incomplete trajectories.
5.1.1. EXPERIMENTAL SETUP
- The set-generation experiments use sequential construction, adding distinct elements with fixed per-element energy across small, medium, and large settings.The set-generation task terminates when the set contains exactly |S| elements.Source: §5.1.1 · p. 6
5.1.2. PERFORMANCE COMPARISON
- FL-GFN significantly outperforms DB, TB, and SubTB in set quality and discovers more high-reward modes faster, with larger gaps at larger scales.Figure 2 compares average reward of unique top-100 sets and the number of high-reward modes across small, medium, and large settings.Source: Figure 2 · p. 6, §5.1.2 · p. 7
5.1.3. APPLICABILITY TO OTHER OBJECTIVES
- FL-GFN (SubTB) improves average reward of unique top-100 sets over SubTB across set-generation problem sizes.The paper reports faster convergence and better solutions, especially for larger sets.Source: §5.1.3 · p. 7, Figure 3 · p. 7
5.1.4. LEARNING WITH INCOMPLETE TRAJECTORIES
- With incomplete set-generation trajectories, FL-GFN learns well at different scales and performs close to its counterpart trained on complete trajectories.Incomplete trajectory lengths are uniformly sampled from 1 to |S|−1; DB and SubTB do not receive informative learning signals without terminal states and rewards, and TB cannot be implemented in that setting.Source: §5.1.4 · pp. 7–8, Figure 4 · p. 7
5.2. Bit Sequence Generation
- In bit-sequence generation, FL-GFN learns more efficiently and discovers more modes than baselines, with a larger advantage for longer sequences.Intermediate transition energies are obtained by applying the energy function to intermediate states.Source: §5.2 · p. 8, Figure 5 · p. 8
5.3. Molecule Generation
- In molecule generation, FL-GFN trains faster and achieves consistent, significant performance improvement on the larger, more challenging task.The reported measures include top-100 molecule reward, Tanimoto similarity for diversity, and held-out Spearman correlation between log reward and log sampling probability.Source: Figure 6 · p. 8, §5.3 · p. 9
6. Conclusion
Forward-looking GFlowNets use computable per-state or per-transition energy to improve credit assignment and learn from incomplete trajectories. Experiments demonstrate their effectiveness on challenging tasks, including molecular graph generation.
6. Conclusion
- FL-GFN exploits computable per-state or per-transition energy while retaining sampling proportional to the target reward distribution.The formulation is described as speeding training through more efficient credit assignment.Source: §6 · p. 9
- FL-GFN can be trained on incomplete trajectories without access to terminal states, unlike previous GFlowNet objectives.Source: §6 · p. 9
- Experiments demonstrate effectiveness on complex and challenging tasks, including molecular graph generation.Source: §6 · p. 9
A. Proof of Proposition 4.2
Under Assumption 4.1, when L(s, s′) is zero for every transition, Proposition 4.2 establishes reward-proportional sampling. Its proof reduces the forward-looking formulation to the detailed-balance training theorem.
A. Proof of Proposition 4.2
- When Assumption 4.1 holds and L(s, s′) = 0 for all transitions, the forward policy samples proportionally to the reward function.Source: A. Proof of Proposition 4.2 · p. 12
- The proof invokes the detailed-balance theorem, which guarantees reward-proportional sampling when its constraint holds for all transitions.The cited theorem is attributed to Bengio et al. (2021b).Source: A. Proof of Proposition 4.2 · p. 12
- The proof algebraically defines a state flow from the forward-looking formulation and concludes that the resulting triple satisfies detailed balance.Source: A. Proof of Proposition 4.2 · p. 12
B. Experimental Details
Baseline methods follow the open-source implementations and the default hyperparameters and setup described in the paper. The authors state that the code will be released upon publication.
B. Experimental Details
- All baseline methods use open-source implementations with the default hyperparameters and setup described in the following sections.Source: B. Experimental Details · p. 12
- The authors state that the code will be released upon publication.Source: B. Experimental Details · p. 12
- The implementation details for the baseline methods are deferred to the following sections.Source: B. Experimental Details · p. 12
B.1. Set Generation
The set-generation experiments scale action-space and target-set sizes across small, medium, and large tasks, with multiple optimal solutions. Models use a two-layer feedforward network and parallel rollouts for training.
B.1. Set Generation
- The set-generation task uses action-space dimensions 30, 80, and 100 and target-set sizes 20, 60, and 80 for small, medium, and large scales, respectively.Source: §B.1 · p. 12
- Each element's predefined energy is randomly drawn from [−1, 1], with |S|/10 elements sharing the same energy to create multiple optimal solutions.Source: §B.1 · p. 12
- The GFlowNet is a two-hidden-layer feedforward network with 256 units per layer, trained using 16 parallel rollouts and Adam.DB, SubTB, and FL-DB use learning rate 0.001; TB uses learning rate 0.1 for its learnable parameter Z.Source: §B.1 · p. 12
B.2. Bit Sequence Generation
Bit-sequence experiments vary sequence length across three settings and use a Transformer GFlowNet with specified training configurations.
B.2. Bit Sequence Generation
- Sequence lengths are normal n = 120, long n = 140, and very long n = 160.The experiments study bit-sequence generation tasks at these increasing lengths.Source: §B.2 · p. 12
- The GFlowNet uses a Transformer with 3 hidden layers, 64 hidden units per layer, and 8 attention heads.Source: §B.2 · p. 12
- Training uses minibatches of 16, random-action probability 0.0005, reward exponent 3, and forward-policy sampling temperature 1.Source: §B.2 · p. 12
- Learning rates are 1 × 10^-4 for TB policy parameters, 1 × 10^-3 for Z, and 5 × 10^-3 for DB and FL-GFN.Source: §B.2 · p. 12
B.3. Molecule Discovery
Molecule discovery uses an atom-graph MPNN setup, with an appendix visualization of top-three molecules generated by different methods.
B.3. Molecule Discovery
- All baselines use an MPNN architecture because molecules are represented as atom graphs, with a reward proxy from Bengio et al. (2021a).Implementations use open-source code and follow Malkin et al. (2022a)'s default setup.Source: §B.3 · p. 13
- The molecule experiments set the reward exponent to 4 after grid search, random-action probability to 0.1, and learning rate to 5 × 10^-4.Source: §B.3 · p. 13
- Figure 7 visualizes the top three molecules generated by different methods.Source: Figure 7 · p. 13