Source-linked AI summary
Generative Augmented Flow Networks
Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, Yoshua Bengio
TL;DR
GFlowNets sample objects in proportion to reward, but terminal-only feedback can limit their applicability, especially in sparse-reward settings. GAFlowNet introduces intermediate intrinsic rewards as augmented flows, and the paper reports improvements across sparse-reward tasks and scalability to molecule generation.
Problem
GFlowNets learn from terminal rewards alone, which can limit applicability when sparse rewards provide little intermediate feedback.
Source: §1 · p. 2Method
GAFlowNet represents intermediate intrinsic rewards as augmented flows to encourage exploration in sparse-reward tasks.
Source: §1 · p. 2Results
GAFlowNet significantly outperforms strong baselines in diversity, convergence, and performance under very sparse rewards, and scales to complex molecular graph generation.
Source: §6 · pp. 11–12Takeaways & Limitations
The experiments support applying GAFlowNet to sparse-reward exploration and complex tasks such as molecular graph generation.
Source: §6 · pp. 11–12Edge-based augmentation alone can remain limited to local exploration and fail to discover all modes.
Source: §4.2 · p. 7
Abstract
from arXiv · showhide
The Generative Flow Network is a probabilistic framework where an agent learns a stochastic policy for object generation, such that the probability of generating an object is proportional to a given reward function. Its effectiveness has been shown in discovering high-quality and diverse solutions, compared to reward-maximizing reinforcement learning-based methods. Nonetheless, GFlowNets only learn from rewards of the terminal states, which can limit its applicability. Indeed, intermediate rewards play a critical role in learning, for example from intrinsic motivation to provide intermediate feedback even in particularly challenging sparse reward tasks. Inspired by this, we propose Generative Augmented Flow Networks (GAFlowNets), a novel learning framework to incorporate intermediate rewards into GFlowNets. We specify intermediate rewards by intrinsic motivation to tackle the exploration problem in sparse reward environments. GAFlowNets can leverage edge-based and state-based intrinsic rewards in a joint way to improve exploration. Based on extensive experiments on the GridWorld task, we demonstrate the effectiveness and efficiency of GAFlowNet in terms of convergence, performance, and diversity of solutions. We further show that GAFlowNet is scalable to a more complex and large-scale molecule generation domain, where it achieves consistent and significant performance improvement.
1 Introduction
GFlowNets sample objects in proportion to reward, preserving diversity rather than only maximizing reward, but their terminal-reward-only training can limit learning when feedback is sparse. GAFlowNet adds intrinsic intermediate rewards through augmented flows and combines edge- and state-based signals; the paper reports benefits across GridWorld and molecule generation.
1 Introduction
- GFlowNets sample composite objects with probability proportional to their return, modeling diverse modes rather than only seeking the reward maximum.This is useful where many high-reward candidates are desired, including molecule generation with uncertain and imperfect simulated rewards.Source: §1 · p. 1
- GFlowNets learn only from terminal rewards, while intermediate intrinsic feedback can support exploration even when extrinsic rewards are sparse.The paper motivates this feedback as a way to address limitations of terminal-only learning.Source: §1 · p. 2
- The proposed framework combines edge-based intrinsic rewards at transitions with state-based augmented flows to improve exploration and learning efficiency.Edge-based rewards improve learning efficiency but explore locally; state-based rewards explore more broadly but can converge more slowly and with greater empirical bias.Source: §1 · p. 2
- GAFlowNet represents intermediate intrinsic rewards as augmented flows, and its augmented objective asymptotically yields an unbiased solution to the original formulation.The framework targets exploration in sparse-reward tasks, where GFlowNet training may become trapped in a few modes.Source: §1 · p. 2
- GAFlowNet achieves improvements in diversity, convergence, and performance on GridWorld and substantial improvements over strong baselines on larger-scale molecule generation.The framework is also reported to apply to different types of GFlowNets.Source: §1 · p. 2
2 Background
GFlowNets represent object construction as trajectories through a directed acyclic graph and use flows to define policies that sample terminal objects in proportion to reward. The background then describes consistency and several training criteria for learning those flows.
2 Background
- GFlowNets construct discrete objects through state transitions in a directed acyclic graph, with trajectories running from an initial state to a terminal state.States encode partial constructions, and state-action pairs correspond to graph edges.Source: §2 · pp. 2–3
- Trajectory flow induces a probability measure over trajectories, while state and edge flows aggregate trajectory flow through a state or transition.The corresponding forward policy is defined from edge flow relative to state flow.Source: §2 · p. 3
2.1 GFlowNets training criterion
- A consistent flow equates incoming and outgoing flow at each internal state; when terminal flow equals reward, the resulting forward policy samples objects proportionally to reward.This is the flow-matching constraint underlying the target distribution.Source: §2.1 · p. 3
- Flow matching trains an edge-flow model by balancing incoming and outgoing flow, with terminal incoming flow matched to reward.The criterion samples training trajectories from a full-support policy; at a global minimum, that policy yields the target distribution.Source: §2.1 · p. 3
- Detailed balance uses state flow and forward and backward policies to avoid summing over a state's parents or children; trajectory balance is proposed for faster credit assignment over longer trajectories.Both detailed balance and trajectory balance are presented as alternative training criteria.Source: §2.1 · p. 3
3 Related Work
Prior work presents GFlowNets as a general-purpose probabilistic inference framework, but terminal-only training can provide weak signals for intermediate states when rewards are sparse. Intrinsic motivation is established in RL, yet had not been explored with GFlowNets in this framework.
3 Related Work
- Earlier GFlowNet work focuses on improving, understanding, and applying the framework, while relying on terminal rewards that can leave intermediate states without a good training signal under sparse rewards.Sparse rewards are described as being significantly non-zero in only a tiny fraction of terminal states.Source: §3 · pp. 3–4
- Entropy-regularized RL can improve diversity only on tree-structured DAGs and can fail on general DAGs, where terminals may have different numbers of generating trajectories.RL otherwise learns a reward-maximizing policy rather than sampling in proportion to reward.Source: §3 · p. 4
- Intrinsic motivation has been studied to improve RL exploration, but had not been incorporated into GFlowNets because their existing mathematical framework permits only terminal rewards.This limitation motivates adding intrinsic intermediate rewards to GFlowNets.Source: §3 · p. 4
4 Generative Augmented Flow Networks
GAFlowNet adds intrinsic intermediate rewards to GFlowNets to address sparse-reward exploration, where standard training can become trapped in only some modes. The section develops edge-based, state-based, and joint augmentations, with the joint method combining their strengths.
4 Generative Augmented Flow Networks
- GAFlowNet represents intrinsic intermediate rewards as augmented flows, targeting the sparse-reward exploration challenge that can trap GFlowNets in a subset of modes.The task's only positive rewards are at three goals; other states provide zero reward.Source: §4 · p. 4
4.1 Edge-based intermediate reward augmentation
- Edge-based augmentation adds intrinsic reward to transitions as extra flow, modeled through pseudo-exits that attract probability toward their ancestors.The pseudo-exit is virtual and is never taken.Source: §4.1 · pp. 4–5
- Edge-based GAFlowNet discovers more modes, learns more efficiently, and achieves smaller empirical L1 error than standard GFlowNet, but still misses some modes.Intrinsic rewards use RND novelty estimates that generally shrink for frequently seen or similar states.Source: §4.1 · pp. 5–6, Figure 3 · p. 6, §4.2 · p. 7
4.2 State-based intermediate reward augmentation
- State-based augmentation uses intermediate rewards in the trajectory return and can discover all modes, but changes the target distribution and converges more slowly.The authors attribute the slower convergence to trajectory-length sensitivity and the altered target distribution.Source: Figure 3 · p. 6, §4.2 · p. 7
4.3 Joint intermediate reward augmentation
- The joint method adds intrinsic reward to the terminal reward and uses edge-based intrinsic rewards for internal states, combining broader exploration with more efficient learning.It uses RND intrinsic rewards, which decrease as the agent gains knowledge about a state.Source: §4.3 · p. 7
- The joint augmentation is theoretically justified to yield an unbiased solution asymptotically as state-based intrinsic rewards converge to zero.The theorem states this result under its stated assumptions, including zero training loss and positive augmented terminal rewards.Source: §4.3 · pp. 7–8
- The joint method discovers all modes, reaches the smallest empirical L1 error, and is more efficient than the state-based and edge-based formulations.Figure 3 compares mode count and empirical L1 error.Source: Figure 3 · p. 6, §4.3 · p. 8
5 Experiments
The experiments compare GAFlowNet with baselines on sparse-reward GridWorld and molecule-generation tasks, examining performance, diversity, augmentation choices, and scalability.
5.1 GridWorld
- GridWorld uses sparse rewards and increasing size H to test longer-horizon planning and learning from sparse signals against MCMC, PPO, PPO-RND, and GFlowNet baselines.Each algorithm is run with five random seeds, and mean and standard deviation are reported.Source: §5.1 · p. 8, Figure 4 · p. 8
5.1.1 Performance Comparison
- GAFlowNet converges fastest with the smallest empirical L1 error and discovers all modes efficiently in GridWorld, while MCMC and PPO fail to converge.PPO-RND improves over PPO but still underperforms GFlowNets by a large margin.Source: §5.1.1 · p. 9, Figure 5 · p. 9
5.1.2 Ablation Study
- Joint state- and edge-based intrinsic rewards achieve the smallest L1 error and improve diversity, unlike the slower state-only and insufficiently exploratory edge-only variants.Terminal-state-only augmentation is also less efficient, supporting the use of both edge-based and terminal state-based intrinsic rewards.Source: §5.1.2 · p. 10, Figure 6 · p. 10
- An intermediate intrinsic-reward coefficient gives the best trade-off; too small a coefficient does not improve performance, while too large a coefficient slows convergence.The coefficient scales intrinsic rewards.Source: §5.1.2 · p. 10, Figure 6 · p. 10
5.1.3 Versatility
- GAFlowNet also significantly improves training convergence when built on detailed-balance and flow-matching objectives.Figure 7 compares empirical L1 error across increasing GridWorld sizes for the two objectives.Source: §5.1.3 · p. 10, Figure 7 · p. 10
5.2.1 Experimental Setup
- The molecule task requires sequential graph construction under chemical-validity constraints, using a sparse reward for meeting the target score in large state and action spaces.The evaluation considers top-10 molecule reward, modes with R > 7.5, and pairwise Tanimoto similarity.Source: §5.2.1 · pp. 10–11, §5.2.2 · p. 11, Figure 8 · p. 11
5.2.2 Performance Comparison
- GAFlowNet significantly outperforms molecule-generation baselines in both solution performance and diversity.PPO variants find higher-quality solutions but produce highly similar samples, while unaugmented GFlowNet finds more diverse molecules but has lower solution quality.Source: §5.2.2 · p. 11, Figure 8 · p. 11
6 Conclusion
The conclusion presents GAFlowNet as a framework for incorporating intermediate rewards into GFlowNets, using intrinsic motivation to address exploration in sparse-reward tasks. It reports improvements over strong baselines and scalability to molecule generation.
6 Conclusion
- GAFlowNet incorporates intermediate rewards into GFlowNets, using intrinsic motivation to address exploration in sparse-reward tasks where GFlowNets can get trapped in a few modes.Source: §6 · p. 11
- In sparse-reward settings, GAFlowNet significantly outperforms strong baselines in diversity, convergence, and performance.Source: §6 · p. 11
- GAFlowNet is scalable to complex tasks such as molecular graph generation; Figure 9 visualizes five candidates, with top-10 results in Appendix B.6.Source: §6 · p. 12, Figure 9 · p. 12
A Proof of Theorem 1
The theorem states limiting-distribution results under conditions on the GAFlowNet loss, augmented rewards, and intrinsic rewards. Its state-based case yields an unbiased sample distribution when state-based intrinsic rewards vanish.
A Proof of Theorem 1
- Under zero GAFlowNet trajectory loss and positive augmented rewards, Theorem 1 states limiting results as intrinsic rewards converge to zero.Source: A Proof of Theorem 1 · p. 15
- When state-based intrinsic rewards converge to zero, the theorem states that P(x) is an unbiased sample distribution.Source: A Proof of Theorem 1 · p. 15
- The proof invokes the law of total probability in deriving the edge-based result.Source: A Proof of Theorem 1 · p. 15
B.1.1 Task
The appendix describes sparse-reward molecule generation and implementation choices for GridWorld and molecule-generation baselines. It also reports GridWorld comparisons across increasing task sizes.
B.1.1 Task
- Molecule generation uses a sparse reward: the agent receives the original non-zero reward only when the normalized score reaches 7.0.The agent attaches blocks from a vocabulary of 105 basic building blocks.Source: §B.1.1 · p. 15
- Figure 10 compares GAFlowNets and baselines across GridWorld sizes H ∈ {8, 16, 32, 64, 128}, reporting empirical L1 error and discovered modes.Source: Figure 10 · p. 16
B.1.2 Baseline
- GridWorld uses feedforward models with two 256-unit hidden layers and LeakyReLU, with RND and GFlowNet trained jointly using Adam.The learning rates are 0.001 for policy models and 0.1 for Z.Source: §B.1.2 · p. 16
- For molecule generation, all models use MPNNs and GAFlowNet builds on flow matching; its intrinsic-reward coefficient α is tuned by grid search and set to 0.1 for this task.Source: §B.1.2 · pp. 16–17
B.2 Full results in GridWorld
Across GridWorld horizons, GAFlowNet outperforms baselines on empirical L1 error and discovered modes, and also surpasses GFlowNet in per-update performance and efficiency.
B.2 Full results in GridWorld
- Across H ∈ {8, 16, 32, 64, 128}, GAFlowNet significantly outperforms baselines in empirical L1 error and number of modes found.The full GridWorld comparison reports both measures across increasing task sizes.Source: §B.2 · p. 17
B.3 Performance comparison
- For top-5 solutions among 16 parallel rollouts after each update, GAFlowNet significantly outperforms baselines and is more efficient than GFlowNet.The comparison covers GridWorld horizons H ∈ {8, 16, 32, 64, 128}.Source: §B.3 · p. 17, Figure 11 · p. 17
B.4 Additional Ablation Study of GAFlowNet
The ablation compares intrinsic-reward forms and a constant-reward baseline, finding robustness across forms and faster convergence with RND.
B.4 Additional Ablation Study of GAFlowNet
- GAFlowNet is not sensitive to the tested intrinsic-reward forms, while RND enables the fastest convergence in these simulations.The comparison uses fine-tuned intrinsic-reward coefficients.Source: §B.4 · p. 17, Figure 12 · p. 17
- Novelty-based intrinsic rewards are more efficient than constant intrinsic rewards, which mimic ϵ-greedy exploration.The paper attributes this comparison to novelty-based methods being more efficient than blindly wandering in the maze.Source: §B.4 · p. 17, Figure 12 · p. 17
- The ablation tests ICM, NovelD, RND, and constant intrinsic rewards as alternatives for GAFlowNet.The constant-reward baseline is included to mimic ϵ-greedy exploration.Source: §B.4 · p. 17, Figure 12 · p. 17
B.5 Additional Performance Comparison on the Molecule Generation Task
On molecule generation, GAFlowNet improves over previous baselines on additional performance measures, and its displayed top-10 molecules are more diverse and higher quality.
B.5 Additional Performance Comparison on the Molecule Generation Task
- GAFlowNet achieves consistent and significant performance improvement over previous baselines on top-100 average reward and modes with R > 8.0.Figure 13 reports the average reward of the top-100 molecules and the number of modes with R > 8.0.Source: §B.5 · p. 18, Figure 13 · p. 18
B.6 Full visualization of top-10 molecules
- Among the displayed top-10 molecules, GAFlowNet discovers more diverse and higher-quality solutions than GFlowNet.Figure 14 visualizes the top-10 molecules generated by both methods.Source: §B.6 · p. 18, Figure 14 · p. 18