Avoid What You Know: Divergent Trajectory Balance for GFlowNets

Pedro Dall'Antonia; Tiago da Silva; Daniel Csillag; Salem Lahlou; Diego Mesquita

cs.LGFeb 19, 2026

Avoid What You Know: Divergent Trajectory Balance for GFlowNets

Authors: Pedro Dall'Antonia, Tiago da Silva, Daniel Csillag, Salem Lahlou, Diego Mesquita

Organizations: School of Applied Mathmatics, Getulio Vargas Foundation · MBZUAI

Abstract

Generative Flow Networks (GFlowNets) are a flexible family of amortized samplers trained to generate discrete and compositional objects with probability proportional to a reward function. However, learning efficiency is constrained by the model's ability to rapidly explore diverse high-probability regions during training. To mitigate this issue, recent works have focused on incentivizing the exploration of unvisited and valuable states via curiosity-driven search and self-supervised random network distillation, which tend to waste samples on already well-approximated regions of the state space. In this context, we propose Adaptive Complementary Exploration (ACE), a principled algorithm for the effective exploration of novel and high-probability regions when learning GFlowNets. To achieve this, ACE introduces an exploration GFlowNet explicitly trained to search for high-reward states in regions underexplored by the canonical GFlowNet, which learns to sample from the target distribution. Through extensive experiments, we show that ACE significantly improves upon prior work in terms of approximation accuracy to the target distribution and discovery rate of diverse high-reward states.

Figures & tables

Figure 1: When learning the exploratory policy based on a combination of intrinsic and extrinsic rewards—see Equations 6 and 7 —, the model may overemphasize well-learned states (bottom row). In contrast, ACE avoids sampling trajectories from over-explored regions of the state space by design (top row), which improves mode discovery and accelerates learning convergence (rightmost panel) to the target. This figure shows the marginal distribution of forward and exploratory policies at different training points (marked as dashed vertical lines).

Panel of Figure 2 — Figure 2: A GFlowNet trained on the Rings distribution (left) via $\epsilon$ -greedy exploration may overdraw samples from a well-approximated region (polygon), misrepresenting other high-probability regions. The TB residual on the rightmost panel for the inner (top) and outer (bottom) rings shows ACE avoids this issue.

Panel of Figure 3 — Figure 3: ACE significantly accelerates mode-discovery for autoregressive sequence generation with GFlowNets. Each plot shows the average reward of the unique 200 highest-valued discovered states as a function of the number of trajectories sampled throughout training.

Panel of Figure 4 — Figure 4: ACE finds diverse and high-reward states faster than prior approaches for improved GFlowNet exploration for the bag generation (left) and quadratic knapsack (right) problems. In both (a) and (b), $K$ denotes the number of available items for selection.

Panel of Figure 5 — Figure 5: ACE results in faster learning convergence than prior art for GFlowNet exploration for the Lazy Random Walk task. Please consult Figure 6 below for the legend.

Figure 6: ACE achieves the best goodness-of-fit to the target distribution in the grid world task described in Equation 13 .

Panel of Figure 7 — Figure 7: Peptide embeddings colored by predicted micro-organism group. Each dot represents the 2D projection of the k-mer embedding of the AMPs in Figure 8 . As we can see, ACE generates a substantially more diverse set of AMPs with likely antimicrobial activity than alternative methods. We showcase a subset of ACE’s generated peptides in in the leftmost panel.

Panel of Figure 8 — Figure 8: ACE identifies a larger number of high-fitness AMPs than prior methods. The plot shows the number of unique AMPs found during training with at least $95\%$ probability of exhibiting antimicrobial activity as a function of the number of trajectories.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 9 — Figure 9: Lazy Random Walk target distributions.

Figure 10: Sensitivity to DTB hyperparameters. TV distance vs. sampled trajectories on Rings for a sweep over

\alpha\in\{0.1,0.2,0.4,0.7\}

(colors) and

\beta\in\{0.25,0.5,1\}

(panels). Solid lines denote the mean across seeds and shaded regions show

\pm 1

standard deviation. We observe a broad stable regime for

\beta=0.25

, whereas larger

\beta

makes training more sensitive to

\alpha

and can induce failure modes for

\alpha\geq 0.4

\beta=1

(TV plateau), together with markedly increased variance.

Generative Flow Networks

Generative Flow Networks learn to sample diverse objects with probabilities proportional to their rewards.

See all 56 papers

Learning a Mixture of GFlowNets
Oct 6, 2026Tiago da Silva, Amauri H. Souza, Salem LahlouGenerative Flow Networks
MBZUAI, UAE · Avra, Brazil & Federal Institute of Cear´a, Brazil
Partial GFlowNet: Accelerating Convergence in Large State Spaces via Strategic Partitioning
Feb 12, 2026Xuan Yu, Xu Wang, Rui Zhu +2Reinforcement Learning Generative Flow Networks
University of Science and Technology of China (USTC), Hefei, China · Suzhou Institute for Advanced Research, USTC, Suzhou, China · State Key Laboratory of Precision and Intelligent Chemistry, USTC, Hefei, China
Stable GFlowNets with Probabilistic Guarantees
May 3, 2026Zengxiang Lei, Ananth Shreekumar, Jonathan Rosenthal +6Generative Flow Networks Flow-Based Generative Modeling
Lyles School of Civil and Construction Engineering, Purdue University, IN, USA · Department of Computer Science, Purdue University, IN, USA · Computer Science and Engineering, University of California, Santa Cruz, USA

Oct 6, 2026J/K move · Enter open · S save

Learning a Mixture of GFlowNets

Tiago da Silva, Amauri H. Souza, Salem Lahlou

MBZUAI, UAE · Avra, Brazil & Federal Institute of Cear´a, Brazil

Generative Flow Networks

Open paper AskPDF arXiv

Learning an ensemble of GFlowNets to sample from a discrete target distribution has become a common approach for achieving better state space exploration and convergence than that of a monolithic sampler. However, these methods often add a substantial runtime overhead to the base model, and their conceptual connection remains elusive. To address this, we first propose a general-purpose theoretical framework for describing a mixture of GFlowNets, which we specialize into continuously (CI) and discretely indexed (DI) collections. On the one hand, we show CI GFlowNets can be interpreted through the lens of a random features expansion, provably boosting the sampler's expressivity in graph-structured tasks and reducing learning instability via spectral shifting. On the other hand, we demonstrate DI GFlowNets encompass prior approaches for GFlowNet training and provide the foundation for the newly proposed Stratum-Conditioned (SC) GFlowNets. This method, which is inspired by the Doob's h-transform of Markov chains, decomposes the state space according to a prescribed modular function and restricts each component to sample from a distinct subset of it. Importantly, SC GFlowNets support centralized and component-wise embarrassingly parallel training, and we show both of them significantly speed up learning convergence and mode coverage without introducing any non-negligible extra computation.

Abstract

Generative Flow Networks

Explore similar work

Learning a Mixture of GFlowNets

Partial GFlowNet: Accelerating Convergence in Large State Spaces via Strategic Partitioning

Stable GFlowNets with Probabilistic Guarantees