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)∝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.
Figures & tables
Figure 1: Mixture of GFlowNets (MoG) and MoG exchange (MoGX) at a glance. A community of K GFlowNets is trained in parallel against a target R(x) , coupled to one jointly trained soft classifier fψ that softly partitions the state space. Each GFlowNet trains on the shaped reward (blue arrows). The classifier is the only object whose gradients cross rank boundaries. MoGX routes part of the on-policy trajectory to other ranks based on the classifier. Post-training, the Z -weighted mixture provably recovers p(x)∝R(x) .
Figure 2: Modes found at a fixed total-trajectory budget vs community size K (log K axis; right axis: % of all modes), off-policy regime (replay mix ratio 0.5 ). MoG (red) and MoGX (purple) scale with K while DDP (blue), NoComm (grey), and EGFN (green) tend to decline; MoGX gains most on the structured rewards.
Figure 3: As fig. 2 , pure on-policy regime . MoG and MoGX still separate from DDP and NoComm, and the MoGX–MoG gap is wider without replay.
Method
K=1
2
4
8
16
32
64
NoComm
702 ± 4
729 ± 6
726 ± 7
728 ± 7
732 ± 3
758 ± 23
713 ± 9
DDP
702 ± 4
715 ± 3
743 ± 2
784 ± 1
823 ± 2
950 ± 4
1249 ± 11
EGFN
702 ± 4
712 ± 3
748 ± 23
853 ± 14
876 ± 70
803 ± 6
975 ± 3
MoG
702 ± 4
721 ± 4
749 ± 7
800 ± 3
818 ± 6
923 ± 8
976 ± 7
MoGX
702 ± 4
762 ± 30
804 ± 15
841 ± 12
998 ± 12
1215 ± 49
1477 ± 12
Table 1: Mean per-iteration wall-clock time (ms) by method and K , pure on-policy, original reward, over 4 seeds and all steps. NoComm stays flat, MoG grows with the non-blocking classifier all-reduce, DDP with the ∇θ all-reduce, and MoGX adds the exchange all-to-all. At K=1 all methods are the same loop.
Figure 4: Mixture correctness ( eq. 2 ) for a K=4 MoGX checkpoint on HyperGrid original . Bars are empirical per-mode marginals from 106 on-policy trajectories; the dotted red line is the optimal policy π⋆ (uniform, 7.26×10−6 per mode). Each rank covers a distinct sparse subset of modes; the Z-weighted mixture recovers π⋆ .
Figure 5: Modes found at a fixed total-trajectory budget vs K on the ChromaX ideotype environment ( B=10 , β=4 ; 1024 optima), mean ±1 s.d. over 4 seeds. MoGX (purple) leads at every K , MoG (red) tracks NoComm (grey), and DDP (blue) is generally last. Dropping replay costs MoGX, MoG and NoComm 5−7% of their K=64 coverage and gains DDP 7% ; no ordering changes.
Appendix figures & tables29 assets
Supplementary material from the paper’s appendix.
Appendix
Component / hyperparameter
Value
Policy network
Architecture
MLP, 2 hidden layers, width 256
Forward/backward policy
shared parameters (tied PF / PB )
GFlowNet loss
trajectory balance (TB)
Sampling temperature / ϵ
T=1 , ϵ=0 (on-policy)
Optimization
Appendix
Table 2: Shared per-agent hyperparameters across all four population configurations. The four methods differ only in inter-agent coordination ( section 3 ); all other settings are identical, which is what makes the iso-budget comparisons in section 4 meaningful. Reward-specific grid heights H are given in section 4 . The EGFN rows apply to that baseline only.
Figure 6: Modes discovered vs total trajectories on the original reward, replay regime ( ρ=0.5 ). MoGX (purple) and MoG (red) discover modes faster than DDP (blue) and NoComm (grey).
Figure 7: Modes discovered vs total trajectories on the cosine reward, replay regime ( ρ=0.5 ).
Figure 8: Modes discovered vs total trajectories on the sparse reward, replay regime ( ρ=0.5 ).
Figure 9: Modes discovered vs total trajectories on the bitwise_xor reward, replay regime ( ρ=0.5 ). The MoGX–MoG gap widens with structural complexity.
Figure 10: Modes discovered vs total trajectories on the multiplicative_coprime reward, replay regime ( ρ=0.5 ). The MoGX–MoG gap widens with structural complexity.
Figure 11: Modes discovered vs total trajectories on the conditional_multiscale reward, replay regime ( ρ=0.5 ). The MoGX–MoG gap widens with structural complexity.
Figure 12: Modes discovered vs total trajectories on the original reward, pure on-policy regime ( ρ=0 , replay buffer record-only).
Figure 13: Modes discovered vs total trajectories on the cosine reward, pure on-policy regime ( ρ=0 , replay buffer record-only).
Figure 14: Modes discovered vs total trajectories on the sparse reward, pure on-policy regime ( ρ=0 , replay buffer record-only).
Figure 15: Modes discovered vs total trajectories on the bitwise_xor reward, pure on-policy regime ( ρ=0 , replay buffer record-only).
Figure 16: Modes discovered vs total trajectories on the multiplicative_coprime reward, pure on-policy regime ( ρ=0 , replay buffer record-only).
Figure 17: Modes discovered vs total trajectories on the conditional_multiscale reward, pure on-policy regime ( ρ=0 , replay buffer record-only).
Figure 18: Modes discovered vs total trajectories on the ChromaX ideotype environment ( section 4.1 ; B=10 , β=4 , MR=1024 optima), replay regime ( ρ=0.5 ), one panel per community size K , mean ± SEM over 4 seeds. MoGX (purple) pulls ahead of the other three methods from K=4 on, while MoG (red) and NoComm (grey) stay close together throughout.
Figure 19: Modes discovered vs total trajectories on the ChromaX ideotype environment ( section 4.1 ; B=10 , β=4 , MR=1024 optima), pure on-policy regime ( ρ=0 , replay buffer record-only), one panel per community size K , mean ± SEM over 4 seeds. The picture is the one from fig. 18 : MoGX (purple) separates from the other three methods from K=4 on, and MoG (red) stays on top of NoComm (grey).
Figure 20: Modes discovered at fixed total-trajectory budget vs community size K , SubTB , pure on-policy regime, mean ±1 s.d. over 4 seeds.
Figure 21: Same as fig. 20 for DB , with the tilt carried as a flow potential in MoG and MoGX.
Method
K=1
2
4
8
16
32
64
NoComm
702 ± 4
729 ± 6
726 ± 7
728 ± 7
732 ± 3
758 ± 23
713 ± 9
DDP
702 ± 4
715 ± 3
743 ± 2
784 ± 1
823 ± 2
950 ± 4
1249 ± 11
EGFN
702 ± 4
712 ± 3
748 ± 23
853 ± 14
876 ± 70
803 ± 6
975 ± 3
MoG
702 ± 4
721 ± 4
749 ± 7
800 ± 3
818 ± 6
923 ± 8
976 ± 7
MoGX
702 ± 4
762 ± 30
804 ± 15
841 ± 12
998 ± 12
1215 ± 49
1477 ± 12
Appendix
Table 3: Same per-method per-iteration wall-clock total as table 1 , reported also here for reference.
Method
K=1
2
4
8
16
32
64
NoComm
2936 ± 45
3103 ± 167
3278 ± 278
2971 ± 65
3303 ± 122
2999 ± 71
2969 ± 73
DDP
2936 ± 45
3106 ± 142
3127 ± 94
3181 ± 123
3451 ± 224
3988 ± 204
4727 ± 182
MoG
2936 ± 45
2982 ± 68
3099 ± 52
3197 ± 26
3274 ± 37
3628 ± 40
3994 ± 31
MoGX
2936 ± 45
2919 ± 24
3092 ± 35
3218 ± 15
3643 ± 64
4059 ± 34
4840 ± 17
Appendix
Table 4: Per-method per-iteration wall-clock total, similar to table 1 , but restricted to the cosine reward. Aggregation, units, and method definitions match table 1 . At K=1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
Method
K=1
2
4
8
16
32
64
NoComm
494 ± 3
498 ± 2
479 ± 5
479 ± 5
464 ± 3
474 ± 5
483 ± 4
DDP
494 ± 3
529 ± 3
560 ± 3
606 ± 3
905 ± 21
1475 ± 48
2780 ± 120
MoG
494 ± 3
501 ± 9
522 ± 6
558 ± 3
526 ± 9
488 ± 7
476 ± 13
MoGX
494 ± 3
556 ± 8
601 ± 8
701 ± 4
1519 ± 25
2640 ± 46
4535 ± 53
Appendix
Table 5: Per-method per-iteration wall-clock total, similar to table 1 , but restricted to the sparse reward. Aggregation, units, and method definitions match table 1 . At K=1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
Method
K=1
2
4
8
16
32
64
NoComm
1083 ± 22
1151 ± 48
1061 ± 26
1066 ± 22
1094 ± 47
1067 ± 9
937 ± 23
DDP
1083 ± 22
1206 ± 19
1402 ± 43
1636 ± 25
1829 ± 29
1904 ± 41
2530 ± 81
MoG
1083 ± 22
1208 ± 39
1363 ± 40
1567 ± 22
1543 ± 15
1578 ± 16
1470 ± 44
MoGX
1083 ± 22
1199 ± 59
1388 ± 89
1619 ± 130
1901 ± 147
2291 ± 171
2923 ± 215
Appendix
Table 6: Per-method per-iteration wall-clock total, similar to table 1 , but restricted to the bitwise_xor reward. Aggregation, units, and method definitions match table 1 . At K=1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
Method
K=1
2
4
8
16
32
64
NoComm
1162 ± 25
1320 ± 98
1238 ± 60
1328 ± 106
1319 ± 43
1216 ± 63
1102 ± 53
DDP
1162 ± 25
1348 ± 40
1495 ± 28
1725 ± 30
1906 ± 33
1945 ± 15
2315 ± 40
MoG
1162 ± 25
1244 ± 41
1415 ± 21
1598 ± 27
1615 ± 17
1789 ± 30
1671 ± 20
MoGX
1162 ± 25
1325 ± 17
1505 ± 11
1774 ± 20
2086 ± 74
2300 ± 26
2795 ± 48
Appendix
Table 7: Per-method per-iteration wall-clock total, similar to table 1 , but restricted to the multiplicative_coprime reward. Aggregation, units, and method definitions match table 1 . At K=1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
Method
K=1
2
4
8
16
32
64
NoComm
1114 ± 24
1272 ± 45
1110 ± 26
1220 ± 33
1174 ± 39
1171 ± 65
1087 ± 45
DDP
1114 ± 24
1241 ± 31
1433 ± 36
1597 ± 30
1736 ± 20
1758 ± 33
2329 ± 136
MoG
1114 ± 24
1180 ± 16
1323 ± 37
1518 ± 37
1580 ± 12
1674 ± 17
1672 ± 38
MoGX
1114 ± 24
1194 ± 72
1364 ± 99
1523 ± 112
1841 ± 165
2020 ± 171
2697 ± 265
Appendix
Table 8: Per-method per-iteration wall-clock total, similar to table 1 , but restricted to the conditional_multiscale reward. Aggregation, units, and method definitions match table 1 . At K=1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
K
Sample
Loss
Bwd
Exch. †
∇θ‡
fψ⋆
1
246 ± 1
224 ± 1
232 ± 2
–
–
–
2
245 ± 2
224 ± 2
229 ± 3
84 ± 45
24 ± 3
14 ± 4
4
245 ± 3
227 ± 3
233 ± 4
92 ± 44
55 ± 2
28 ± 14
8
248 ± 4
230 ± 4
235 ± 4
101 ± 29
90 ± 3
54 ± 22
16
244 ± 3
233 ± 4
238 ± 4
263 ± 42
129 ± 4
60 ± 22
32
244 ± 3
232 ± 6
238 ± 7
514 ± 102
256 ± 7
108 ± 37
Appendix
Table 9: Per-iteration timing phase breakdown, restricted to the original reward.
K
Sample
Loss
Bwd
Exch. †
∇θ‡
fψ⋆
1
914 ± 7
1051 ± 26
971 ± 20
–
–
–
2
915 ± 11
1052 ± 27
971 ± 22
88 ± 56
231 ± 153
21 ± 9
4
912 ± 14
1087 ± 50
993 ± 36
312 ± 106
210 ± 56
53 ± 22
8
909 ± 12
1042 ± 17
961 ± 15
362 ± 45
245 ± 42
155 ± 71
16
901 ± 15
1095 ± 39
995 ± 32
773 ± 217
464 ± 129
233 ± 98
32
895 ± 13
1056 ± 26
965 ± 20
1220 ± 60
983 ± 64
403 ± 161
Appendix
Table 10: Per-iteration timing phase breakdown, restricted to the cosine reward.
K
Sample
Loss
Bwd
Exch. †
∇θ‡
fψ⋆
1
170 ± 1
167 ± 1
157 ± 1
–
–
–
2
172 ± 2
172 ± 2
161 ± 2
22 ± 9
21 ± 3
9 ± 1
4
167 ± 1
170 ± 2
158 ± 2
79 ± 26
61 ± 8
23 ± 7
8
167 ± 2
169 ± 3
157 ± 3
188 ± 34
103 ± 7
40 ± 13
16
163 ± 2
165 ± 3
153 ± 3
1032 ± 50
416 ± 25
20 ± 6
32
162 ± 2
164 ± 2
152 ± 2
2164 ± 33
986 ± 53
6 ± 0
Appendix
Table 11: Per-iteration timing phase breakdown, restricted to the sparse reward.
K
Sample
Loss
Bwd
Exch. †
∇θ‡
fψ⋆
1
435 ± 7
325 ± 10
323 ± 10
–
–
–
2
423 ± 7
331 ± 11
329 ± 11
177 ± 32
142 ± 8
54 ± 19
4
410 ± 8
312 ± 12
308 ± 12
414 ± 58
380 ± 39
149 ± 59
8
394 ± 9
318 ± 11
313 ± 11
563 ± 109
608 ± 31
307 ± 113
16
373 ± 9
307 ± 8
304 ± 9
1006 ± 134
832 ± 19
295 ± 109
32
346 ± 11
312 ± 12
308 ± 11
1455 ± 131
964 ± 36
279 ± 104
Appendix
Table 12: Per-iteration timing phase breakdown, restricted to the bitwise_xor reward.
K
Sample
Loss
Bwd
Exch. †
∇θ‡
fψ⋆
1
463 ± 5
351 ± 14
349 ± 14
–
–
–
2
457 ± 5
366 ± 17
363 ± 17
165 ± 5
213 ± 54
58 ± 23
4
433 ± 9
354 ± 12
351 ± 13
402 ± 45
380 ± 40
158 ± 59
8
422 ± 9
365 ± 18
361 ± 19
706 ± 23
606 ± 31
260 ± 95
16
402 ± 8
364 ± 16
361 ± 17
1031 ± 97
829 ± 39
279 ± 106
32
380 ± 6
351 ± 16
348 ± 16
1270 ± 39
816 ± 49
424 ± 156
Appendix
Table 13: Per-iteration timing phase breakdown, restricted to the multiplicative_coprime reward.
K
Sample
Loss
Bwd
Exch. †
∇θ‡
fψ⋆
1
413 ± 7
351 ± 12
350 ± 12
–
–
–
2
416 ± 10
363 ± 15
361 ± 16
143 ± 25
123 ± 17
31 ± 13
4
404 ± 10
346 ± 14
343 ± 14
352 ± 95
306 ± 29
102 ± 47
8
386 ± 8
358 ± 11
354 ± 12
487 ± 78
532 ± 12
223 ± 83
16
379 ± 11
353 ± 10
350 ± 11
788 ± 133
697 ± 25
260 ± 100
32
362 ± 6
364 ± 14
360 ± 14
967 ± 144
710 ± 75
303 ± 115
Appendix
Table 14: Per-iteration timing phase breakdown, restricted to the conditional_multiscale reward.
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 (GFlowNets) learn to sample diverse candidates in proportion to a reward function, making them well-suited for scientific discovery, where exploring multiple promising solutions is crucial. Further extending GFlowNets to multi-objective settings has attracted growing interest as real-world applications often involve multiple, conflicting objectives. However, existing approaches require joint training for each combination of objectives, meaning that any change in the objective set necessitates retraining from scratch. We propose a framework that composes pre-trained GFlowNets at inference time, enabling rapid adaptation without fine-tuning or retraining. Importantly, our framework is flexible, capable of handling diverse reward combinations ranging from linear scalarization to complex nonlinear operators, which are often handled separately in previous literature. We prove that our method exactly recovers the target distribution for linear scalarization, and quantify the approximation quality for nonlinear operators through a distortion factor. Experiments on a synthetic 2D grid and real-world molecule generation tasks demonstrate that our approach achieves performance comparable to baselines.
Seokwon Yoon, Youngbin Choi, Seunghyuk Cho +3
Department of Computer Science & Engineering, POSTECH, South Korea · Graduate School of Artificial Intelligence, POSTECH, South Korea
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