Hierarchical Cooperative GFlowNets

Pedro Dall'Antonia; Juan Belieni; Ana Carolina Erthal; Saullo H. G. de Oliveira; Eliezer de Souza da Silva; Diego Mesquita

cs.LGOct 25, 2026

Hierarchical Cooperative GFlowNets

Authors: Pedro Dall'Antonia, Juan Belieni, Ana Carolina Erthal, Saullo H. G. de Oliveira, Eliezer de Souza da Silva, Diego Mesquita

Organizations: School of Applied Mathematics (EMaP), Getulio Vargas Foundation, Rio de Janeiro, Brazil · Department of Informatics Engineering, University of Coimbra, Coimbra, Portugal

Abstract

Generative Flow Networks (GFlowNets) are trained to sample objects from compositional spaces with probability proportional to a reward function, but ineffective exploration can cause learned samplers to concentrate on a few high-reward regions. An intuitive remedy is to train multiple GFlowNets independently in parallel and combine their samples. However, this does not prevent multiple models from collapsing onto the same high-reward areas, leaving large regions of the state space unexplored. With this in mind, we propose Hierarchical Cooperative GFlowNets (HCG), an approach that dynamically partitions the search space to coordinate discovery without relying on hand-crafted boundaries. HCG imposes an ordering in which the highest-ranked model may explore the entire space, while lower-ranked models are progressively pushed toward regions not covered by higher-ranked ones. We provide a theoretical characterization of how specific learning objectives can shift probability mass away from designated areas, instantiating this mechanism within our architecture. A shared classifier coordinates the hierarchy by partitioning the search space, keeping the GFlowNets independently parameterized and enabling asynchronous training. Furthermore, we derive a normalizer-weighted aggregation rule for the ensemble and prove that it recovers the target reward-proportional distribution under the stated assumptions. Finally, experiments across multimodal distributions, sequence generation, combinatorial optimization, and biological sequence design show that HCG improves mode coverage and target approximation, translating increases in the number of coordinated samplers into broader, complementary exploration.

Figures & tables

Figure 1: Three GFlowNets on a two-ring reward with (top) and without (bottom) an exploration hierarchy. Left: samples, colored by generating model. Center: regions of $\mathcal{X}$ owned by each model. Right: zoom on the inner mode. Without hierarchical coordination, all models split the inner mode and the outer ring is never sampled. With hierarchy, $G_{0}$ owns the inner mode and a lower-ranked model is redirected to the outer ring.

Method	8 Gaussians	Rings	Moons	Spiral
TB	$0.8750\pm 0.0000$	$0.6143\pm 0.2650$	$0.5590\pm 0.1027$	$0.3446\pm 0.1217$
SA	$0.8241\pm 0.0939$	$0.6894\pm 0.1492$	$0.5071\pm 0.2526$	$0.3345\pm 0.1218$
ACE	$0.4946\pm 0.3282$	$0.5994\pm 0.3190$	$0.1795\pm 0.1855$	$\mathbf{0.1057\pm 0.0635}$
HCG (ours)	$\mathbf{0.3054\pm 0.1249}$	$\mathbf{0.3324\pm 0.2637}$	$\mathbf{0.1703\pm 0.1534}$	$0.1340\pm 0.0873$

Table 1: Final total variation distance (

\downarrow

) on DiffGrid (

81\times 81

grid), reported as mean

\pm

sample standard deviation over five seeds. Bold indicates the lowest mean in each column. HCG achieves the lowest distance on three of the four landscapes, and remains competitive on Spiral.

Figure 2: Total variation distance (

\downarrow

) to the target distribution over training on DiffGrid rewards (

81\times 81

grid). Lines show the mean over five seeds and shaded bands

\pm 1

std. On 8 Gaussians and Rings, HCG continues to improve throughout training and ends with the lowest TV distance. On Moons and Spiral, whose reward mass forms connected structures rather than isolated modes, HCG remains competitive with ACE, the strongest baseline.

Figure 3: Effect of the number of GFlowNets on the DiffGrid 8 Gaussians reward on an $81\times 81$ grid. Left: TV distance ( $\downarrow$ ) to the target distribution over training. Lines show the mean over five seeds and bands $\pm 1$ std. Right: final samples for ensembles of 1, 2, 4, and 8 GFlowNets. Each sample is colored by the GFlowNet that generated it, and the background hue shows the region the classifier assigns to each GFlowNet. While a single GFlowNet collapses onto one mode, each additional component is steered toward a mode not yet covered, so that the number of discovered modes grows with the ensemble and the TV distance decreases accordingly.

Figure 4: Top-200 mean reward (

\uparrow

) over two environments using two configurations: the knapsack task at

n=256

and

n=512

, and the multi-modal bit-sequence task at

M=32

and

M=64

. Results are the mean of 3 seeds and bands

\pm 1

std. HCG reaches the highest top-200 reward in all four conditions, and its margin over the strongest baseline widens as the problem grows. This is more pronounced in the bit-sequences task (

M=64

), where most baselines stall below

10^{7}

, ACE plateaus and HCG reaches more than

6\times 10^{7}

mean reward.

Figure 5: Cumulative count of unique high-scoring peptides discovered per method throughout training (log scale). Lines show the mean over 5 seeds, and bands span

\pm 1

std. HCG shows faster discovery and sustained exploration over training.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: DiffGrid reward landscapes on an $81\times 81$ grid ( $m=40$ ). Each cell shows the implemented reward divided by its maximum within that landscape, using a shared linear color scale. (a) Eight Gaussian peaks lie on a circle of radius $0.8m$ . (b) Two rings have radii $0.2m$ and $0.8m$ . (c) Two offset semicircular arcs are represented by 256 Gaussian anchors each. (d) A single spiral arm makes two turns, with 512 anchors spaced uniformly in arc length and outer radius $0.85m$ . All landscapes use width $\sigma=m/30=4/3$ and floor parameter $\lambda=10^{-6}$ .

Component	Setting
Sequences	19 amino acids + STOP ; $1\leq L\leq L_{\max}=10$
Batch size	4096 trajectories per rollout
Exploration $(\varepsilon)$	0.1
Seeds	$\{10,\dots,14\}$
Policy	Windowed autoregressive MLP (embedding 32, context 8, hidden 64) over $\|\mathcal{V}\|=20$ actions
Optimizer	AdamW; learning rates $(P_{F},\log Z)=(5\times 10^{-2},10^{-1})$

Appendix

Table 2: Default training setup for peptide generation.

Figure 7: Effect of the classifier update cadence on the top-200 mean reward. Curves show the mean over three seeds, with bands at

\pm 1

standard deviation. Every cadence improves on TB and

c=25

attains the highest mean in both families.

Figure 8: Cost of the classifier cadence. (a) Classifier gradient steps: with fixed steps the number of updates falls in proportion to the cadence, from 8000 at

c=1

to 80 at

c=100

, whereas the compute-matched family holds it at 8000. (b) Measured training wall-clock time: sparser fixed-step cadences save almost no time despite performing fewer updates, while concentrating the same updates into rarer, longer bursts grows the cost. Trajectory Balance is shown for reference.

	Classifier	Reward ( $\times 10^{6}$ , $\uparrow$ )		Distance to target ( $\downarrow$ )
$c$	updates	mean	$\Delta$ vs $c{=}1$	Pooled TV	FCS
TB	0	$1.59\pm 0.03$	–	$\mathbf{0.1677\pm 0.0006}$	$\mathbf{0.1657\pm 0.0008}$
1	8000	$5.26\pm 1.14$	–	$0.1748\pm 0.0044$	$0.1717\pm 0.0034$
10	800	$5.93\pm 1.64$	$+0.67\pm 0.54$	$0.1731\pm 0.0041$	$0.1705\pm 0.0041$
25	320	$\mathbf{8.01\pm 1.28}$	$+2.75\pm 1.61$	$0.1715\pm 0.0047$	$0.1691\pm 0.0045$
50	160	$6.51\pm 2.26$	$+1.25\pm 1.22$	$0.1728\pm 0.0046$	$0.1702\pm 0.0045$

Appendix

Table 3: Fixed-steps family. The classifier takes

s=4

gradient steps whenever it fires, so the total number of updates falls with the cadence

c

. Mean

\pm

standard deviation over three seeds (42–44);

\Delta

is the per-seed paired difference against the default

c=1

. Reward is the mean of the top-200 unique terminal states. The TB reference is trained with

\epsilon=0.1

, whereas HCG is on-policy (

\epsilon=0

	Classifier	Reward ( $\times 10^{6}$ , $\uparrow$ )		Distance to target ( $\downarrow$ )
$c$	updates	mean	$\Delta$ vs $c{=}1$	Pooled TV	FCS
TB	0	$1.59\pm 0.03$	–	$\mathbf{0.1677\pm 0.0006}$	$\mathbf{0.1657\pm 0.0008}$
1	8000	$5.26\pm 1.14$	–	$0.1748\pm 0.0044$	$0.1717\pm 0.0034$
10	8000	$5.40\pm 0.58$	$+0.13\pm 1.35$	$0.1778\pm 0.0071$	$0.1739\pm 0.0058$
25	8000	$\mathbf{6.14\pm 0.91}$	$+0.88\pm 0.86$	$0.1783\pm 0.0051$	$0.1737\pm 0.0025$
50	8000	$5.54\pm 1.05$	$+0.28\pm 0.49$	$0.1790\pm 0.0061$	$0.1735\pm 0.0037$

Appendix

Table 4: Compute-matched family. Steps per burst scale with the cadence,

s=4c

, holding the total at 8000 updates for every arm, so this family varies staleness and burstiness at a fixed budget. Columns, seeds, and

\Delta

as in Table 3 ;

c=1

is the shared reference point of both families.

Figure 9: TV distance (

\downarrow

) to the target distribution against wall-clock training time on 8 Gaussians (left) and Rings (right), truncated at TB’s finishing time so that all methods are compared over the same time budget. Mean

\pm 1

standard deviation over five seeds.

Figure 10: HCG versus independent GFlowNet ensembles on Rings. Exact terminal distributions on the $31\times 31$ grid after 2,000 updates per GFlowNet. Rows correspond to seeds 42, 43, and 44, from top to bottom. Columns show the reward-proportional target, the mixture of three independently trained TB GFlowNets, and the three-component HCG mixture. Both learned mixtures use weights $Z_{i}/\sum_{j}Z_{j}$ , and all nine panels share a linear probability color scale. Independent TB concentrates on the inner ring in all three runs, whereas HCG assigns substantial probability to both rings.

Generative Flow Networks

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

See all 57 papers

Feb 19, 2026

Avoid What You Know: Divergent Trajectory Balance for GFlowNets

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.

Pedro Dall'Antonia, Tiago da Silva, Daniel Csillag +2

School of Applied Mathmatics, Getulio Vargas Foundation · MBZUAI

Generative Flow Networks Adaptive Sampling RL Exploration

Ask PDF

Oct 12, 2026

Mixture of GFlowNets

Generative Flow Networks (GFlowNets) learn a policy that samples compositional objects in proportion to an unnormalized reward

R(x)

. This makes them useful wherever many diverse, high-reward candidates are needed and oracle queries are expensive. In practice, a single GFlowNet often fails to cover all modes of a complex distribution, and the obvious ways to scale compute saturate: larger batches and data-parallel gradient averaging draw every trajectory from one, possibly biased, policy, so extra rollouts rediscover known modes instead of finding new ones. We explore whether this mode-bias problem (mode collapse in the extreme) can be addressed with a population of asynchronously trained policies which divide the state space between them. Mixture of GFlowNets (MoG) replaces the single forward policy with a community of

K

independent GFlowNets coupled through one jointly trained soft classifier, which assigns each member its own region of the state space through reward shaping. Mixing the members under their learned partition functions provably recovers the target

p(x)\propto R(x)

. MoG exchange (MoGX) also lets the classifier route sampled trajectories: a trajectory the classifier assigns to another member is sent to that member and trained on there, with no importance sampling. We evaluate both methods by keeping the environment-interaction budget fixed while scaling the community size, turning parallel compute into mode coverage without extra oracle calls. On HyperGrid, across six reward landscapes and communities of up to

K{=}64

, MoG and MoGX keep finding new modes as

K

grows, while data-parallel training saturates. On a genomic ideotype-design task, MoGX covers

39%

of the

1024

optima at

K{=}64

, against

31\%

for independent agents and

20\%

for data-parallel training. MoG also has a lower per-iteration synchronization cost than data-parallel training.

Omar G. Younis, Joseph D. Viviano, Chirayu Anant Haryan +3

Mila – Québec AI Institute · Intel · Université de Montréal

Generative Flow Networks Mode Collapse

Ask PDF

Oct 6, 2026

Learning a Mixture of GFlowNets

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.

Tiago da Silva, Amauri H. Souza, Salem Lahlou

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

Generative Flow Networks

Ask PDF

Abstract

Generative Flow Networks

Explore similar work

Avoid What You Know: Divergent Trajectory Balance for GFlowNets

Mixture of GFlowNets

Learning a Mixture of GFlowNets