cs.LGOct 12, 2026

Mixture of GFlowNets

Authors: Omar G. Younis, Joseph D. Viviano, Chirayu Anant Haryan, MD Vasimuddin, Sanchit Misra, Yoshua Bengio

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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Generative Flow Networks (GFlowNets) learn a policy that samples compositional objects in proportion to an unnormalized reward R(x)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 KK 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)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=64K{=}64, MoG and MoGX keep finding new modes as KK grows, while data-parallel training saturates. On a genomic ideotype-design task, MoGX covers 3939% of the 10241024 optima at K=64K{=}64, against 31%31\% for independent agents and 20%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 KK GFlowNets is trained in parallel against a target R(x)R(x) , coupled to one jointly trained soft classifier fψf_{\psi} 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 ZZ -weighted mixture provably recovers p(x) ⁣∝ ⁣R(x)p(x)\!\propto\!R(x) .
Figure 2: Modes found at a fixed total-trajectory budget vs community size KK (log KK axis; right axis: % of all modes), off-policy regime (replay mix ratio 0.50.5 ). MoG (red) and MoGX (purple) scale with KK 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.
MethodK=1K{=}1224488161632326464
NoComm702 ±\pm 4729 ±\pm 6726 ±\pm 7728 ±\pm 7732 ±\pm 3758 ±\pm 23713 ±\pm 9
DDP702 ±\pm 4715 ±\pm 3743 ±\pm 2784 ±\pm 1823 ±\pm 2950 ±\pm 41249 ±\pm 11
EGFN702 ±\pm 4712 ±\pm 3748 ±\pm 23853 ±\pm 14876 ±\pm 70803 ±\pm 6975 ±\pm 3
MoG702 ±\pm 4721 ±\pm 4749 ±\pm 7800 ±\pm 3818 ±\pm 6923 ±\pm 8976 ±\pm 7
MoGX702 ±\pm 4762 ±\pm 30804 ±\pm 15841 ±\pm 12998 ±\pm 121215 ±\pm 491477 ±\pm 12
Table 1: Mean per-iteration wall-clock time (ms) by method and KK , 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 ∇θ\nabla_{\theta} all-reduce, and MoGX adds the exchange all-to-all. At K=1K{=}1 all methods are the same loop.
Figure 4: Mixture correctness ( eq. 2 ) for a K=4K{=}4 MoGX checkpoint on HyperGrid original . Bars are empirical per-mode marginals from 10610^{6} on-policy trajectories; the dotted red line is the optimal policy π⋆\pi^{\star} (uniform, 7.26 ⁣× ⁣10−67.26\!\times\!10^{-6} per mode). Each rank covers a distinct sparse subset of modes; the Z-weighted mixture recovers π⋆\pi^{\star} .
Figure 5: Modes found at a fixed total-trajectory budget vs KK on the ChromaX ideotype environment ( B=10B{=}10 , β=4\beta{=}4 ; 10241024 optima), mean ±1\pm 1 s.d. over 44 seeds. MoGX (purple) leads at every KK , MoG (red) tracks NoComm (grey), and DDP (blue) is generally last. Dropping replay costs MoGX, MoG and NoComm 5−7%5{-}7\% of their K=64K{=}64 coverage and gains DDP 7%7\% ; no ordering changes.
Appendix figures & tables29 assets

Supplementary material from the paper’s appendix.

Appendix

Component / hyperparameterValue
Policy network
ArchitectureMLP, 2 hidden layers, width 256
Forward/backward policyshared parameters (tied PFP_{F} / PBP_{B} )
GFlowNet losstrajectory balance (TB)
Sampling temperature / ϵ\epsilonT=1T{=}1 , ϵ=0\epsilon{=}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 HH 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\rho{=}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\rho{=}0.5 ).
Figure 8: Modes discovered vs total trajectories on the sparse reward, replay regime ( ρ=0.5\rho{=}0.5 ).
Figure 9: Modes discovered vs total trajectories on the bitwise_xor reward, replay regime ( ρ=0.5\rho{=}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\rho{=}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\rho{=}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\rho{=}0 , replay buffer record-only).
Figure 13: Modes discovered vs total trajectories on the cosine reward, pure on-policy regime ( ρ=0\rho{=}0 , replay buffer record-only).
Figure 14: Modes discovered vs total trajectories on the sparse reward, pure on-policy regime ( ρ=0\rho{=}0 , replay buffer record-only).
Figure 15: Modes discovered vs total trajectories on the bitwise_xor reward, pure on-policy regime ( ρ=0\rho{=}0 , replay buffer record-only).
Figure 16: Modes discovered vs total trajectories on the multiplicative_coprime reward, pure on-policy regime ( ρ=0\rho{=}0 , replay buffer record-only).
Figure 17: Modes discovered vs total trajectories on the conditional_multiscale reward, pure on-policy regime ( ρ=0\rho{=}0 , replay buffer record-only).
Figure 18: Modes discovered vs total trajectories on the ChromaX ideotype environment ( section 4.1 ; B=10B{=}10 , β=4\beta{=}4 , MR=1024M_{R}{=}1024 optima), replay regime ( ρ=0.5\rho{=}0.5 ), one panel per community size KK , mean ±\pm SEM over 44 seeds. MoGX (purple) pulls ahead of the other three methods from K=4K{=}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=10B{=}10 , β=4\beta{=}4 , MR=1024M_{R}{=}1024 optima), pure on-policy regime ( ρ=0\rho{=}0 , replay buffer record-only), one panel per community size KK , mean ±\pm SEM over 44 seeds. The picture is the one from fig. 18 : MoGX (purple) separates from the other three methods from K=4K{=}4 on, and MoG (red) stays on top of NoComm (grey).
Figure 20: Modes discovered at fixed total-trajectory budget vs community size KK , SubTB , pure on-policy regime, mean ±1\pm 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.
MethodK=1K{=}1224488161632326464
NoComm702 ±\pm 4729 ±\pm 6726 ±\pm 7728 ±\pm 7732 ±\pm 3758 ±\pm 23713 ±\pm 9
DDP702 ±\pm 4715 ±\pm 3743 ±\pm 2784 ±\pm 1823 ±\pm 2950 ±\pm 41249 ±\pm 11
EGFN702 ±\pm 4712 ±\pm 3748 ±\pm 23853 ±\pm 14876 ±\pm 70803 ±\pm 6975 ±\pm 3
MoG702 ±\pm 4721 ±\pm 4749 ±\pm 7800 ±\pm 3818 ±\pm 6923 ±\pm 8976 ±\pm 7
MoGX702 ±\pm 4762 ±\pm 30804 ±\pm 15841 ±\pm 12998 ±\pm 121215 ±\pm 491477 ±\pm 12
Appendix
Table 3: Same per-method per-iteration wall-clock total as table 1 , reported also here for reference.
MethodK=1K{=}1224488161632326464
NoComm2936 ±\pm 453103 ±\pm 1673278 ±\pm 2782971 ±\pm 653303 ±\pm 1222999 ±\pm 712969 ±\pm 73
DDP2936 ±\pm 453106 ±\pm 1423127 ±\pm 943181 ±\pm 1233451 ±\pm 2243988 ±\pm 2044727 ±\pm 182
MoG2936 ±\pm 452982 ±\pm 683099 ±\pm 523197 ±\pm 263274 ±\pm 373628 ±\pm 403994 ±\pm 31
MoGX2936 ±\pm 452919 ±\pm 243092 ±\pm 353218 ±\pm 153643 ±\pm 644059 ±\pm 344840 ±\pm 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=1K{=}1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
MethodK=1K{=}1224488161632326464
NoComm494 ±\pm 3498 ±\pm 2479 ±\pm 5479 ±\pm 5464 ±\pm 3474 ±\pm 5483 ±\pm 4
DDP494 ±\pm 3529 ±\pm 3560 ±\pm 3606 ±\pm 3905 ±\pm 211475 ±\pm 482780 ±\pm 120
MoG494 ±\pm 3501 ±\pm 9522 ±\pm 6558 ±\pm 3526 ±\pm 9488 ±\pm 7476 ±\pm 13
MoGX494 ±\pm 3556 ±\pm 8601 ±\pm 8701 ±\pm 41519 ±\pm 252640 ±\pm 464535 ±\pm 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=1K{=}1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
MethodK=1K{=}1224488161632326464
NoComm1083 ±\pm 221151 ±\pm 481061 ±\pm 261066 ±\pm 221094 ±\pm 471067 ±\pm 9937 ±\pm 23
DDP1083 ±\pm 221206 ±\pm 191402 ±\pm 431636 ±\pm 251829 ±\pm 291904 ±\pm 412530 ±\pm 81
MoG1083 ±\pm 221208 ±\pm 391363 ±\pm 401567 ±\pm 221543 ±\pm 151578 ±\pm 161470 ±\pm 44
MoGX1083 ±\pm 221199 ±\pm 591388 ±\pm 891619 ±\pm 1301901 ±\pm 1472291 ±\pm 1712923 ±\pm 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=1K{=}1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
MethodK=1K{=}1224488161632326464
NoComm1162 ±\pm 251320 ±\pm 981238 ±\pm 601328 ±\pm 1061319 ±\pm 431216 ±\pm 631102 ±\pm 53
DDP1162 ±\pm 251348 ±\pm 401495 ±\pm 281725 ±\pm 301906 ±\pm 331945 ±\pm 152315 ±\pm 40
MoG1162 ±\pm 251244 ±\pm 411415 ±\pm 211598 ±\pm 271615 ±\pm 171789 ±\pm 301671 ±\pm 20
MoGX1162 ±\pm 251325 ±\pm 171505 ±\pm 111774 ±\pm 202086 ±\pm 742300 ±\pm 262795 ±\pm 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=1K{=}1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
MethodK=1K{=}1224488161632326464
NoComm1114 ±\pm 241272 ±\pm 451110 ±\pm 261220 ±\pm 331174 ±\pm 391171 ±\pm 651087 ±\pm 45
DDP1114 ±\pm 241241 ±\pm 311433 ±\pm 361597 ±\pm 301736 ±\pm 201758 ±\pm 332329 ±\pm 136
MoG1114 ±\pm 241180 ±\pm 161323 ±\pm 371518 ±\pm 371580 ±\pm 121674 ±\pm 171672 ±\pm 38
MoGX1114 ±\pm 241194 ±\pm 721364 ±\pm 991523 ±\pm 1121841 ±\pm 1652020 ±\pm 1712697 ±\pm 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=1K{=}1 all four methods collapse to the same single-agent loop and share the pooled K=1 baseline.
KKSampleLossBwdExch. †∇θ‡\nabla_{\theta}^{\ddagger}fψ⋆f_{\psi}^{\star}
1246 ±\pm 1224 ±\pm 1232 ±\pm 2–––
2245 ±\pm 2224 ±\pm 2229 ±\pm 384 ±\pm 4524 ±\pm 314 ±\pm 4
4245 ±\pm 3227 ±\pm 3233 ±\pm 492 ±\pm 4455 ±\pm 228 ±\pm 14
8248 ±\pm 4230 ±\pm 4235 ±\pm 4101 ±\pm 2990 ±\pm 354 ±\pm 22
16244 ±\pm 3233 ±\pm 4238 ±\pm 4263 ±\pm 42129 ±\pm 460 ±\pm 22
32244 ±\pm 3232 ±\pm 6238 ±\pm 7514 ±\pm 102256 ±\pm 7108 ±\pm 37
Appendix
Table 9: Per-iteration timing phase breakdown, restricted to the original reward.
KKSampleLossBwdExch. †∇θ‡\nabla_{\theta}^{\ddagger}fψ⋆f_{\psi}^{\star}
1914 ±\pm 71051 ±\pm 26971 ±\pm 20–––
2915 ±\pm 111052 ±\pm 27971 ±\pm 2288 ±\pm 56231 ±\pm 15321 ±\pm 9
4912 ±\pm 141087 ±\pm 50993 ±\pm 36312 ±\pm 106210 ±\pm 5653 ±\pm 22
8909 ±\pm 121042 ±\pm 17961 ±\pm 15362 ±\pm 45245 ±\pm 42155 ±\pm 71
16901 ±\pm 151095 ±\pm 39995 ±\pm 32773 ±\pm 217464 ±\pm 129233 ±\pm 98
32895 ±\pm 131056 ±\pm 26965 ±\pm 201220 ±\pm 60983 ±\pm 64403 ±\pm 161
Appendix
Table 10: Per-iteration timing phase breakdown, restricted to the cosine reward.
KKSampleLossBwdExch. †∇θ‡\nabla_{\theta}^{\ddagger}fψ⋆f_{\psi}^{\star}
1170 ±\pm 1167 ±\pm 1157 ±\pm 1–––
2172 ±\pm 2172 ±\pm 2161 ±\pm 222 ±\pm 921 ±\pm 39 ±\pm 1
4167 ±\pm 1170 ±\pm 2158 ±\pm 279 ±\pm 2661 ±\pm 823 ±\pm 7
8167 ±\pm 2169 ±\pm 3157 ±\pm 3188 ±\pm 34103 ±\pm 740 ±\pm 13
16163 ±\pm 2165 ±\pm 3153 ±\pm 31032 ±\pm 50416 ±\pm 2520 ±\pm 6
32162 ±\pm 2164 ±\pm 2152 ±\pm 22164 ±\pm 33986 ±\pm 536 ±\pm 0
Appendix
Table 11: Per-iteration timing phase breakdown, restricted to the sparse reward.
KKSampleLossBwdExch. †∇θ‡\nabla_{\theta}^{\ddagger}fψ⋆f_{\psi}^{\star}
1435 ±\pm 7325 ±\pm 10323 ±\pm 10–––
2423 ±\pm 7331 ±\pm 11329 ±\pm 11177 ±\pm 32142 ±\pm 854 ±\pm 19
4410 ±\pm 8312 ±\pm 12308 ±\pm 12414 ±\pm 58380 ±\pm 39149 ±\pm 59
8394 ±\pm 9318 ±\pm 11313 ±\pm 11563 ±\pm 109608 ±\pm 31307 ±\pm 113
16373 ±\pm 9307 ±\pm 8304 ±\pm 91006 ±\pm 134832 ±\pm 19295 ±\pm 109
32346 ±\pm 11312 ±\pm 12308 ±\pm 111455 ±\pm 131964 ±\pm 36279 ±\pm 104
Appendix
Table 12: Per-iteration timing phase breakdown, restricted to the bitwise_xor reward.
KKSampleLossBwdExch. †∇θ‡\nabla_{\theta}^{\ddagger}fψ⋆f_{\psi}^{\star}
1463 ±\pm 5351 ±\pm 14349 ±\pm 14–––
2457 ±\pm 5366 ±\pm 17363 ±\pm 17165 ±\pm 5213 ±\pm 5458 ±\pm 23
4433 ±\pm 9354 ±\pm 12351 ±\pm 13402 ±\pm 45380 ±\pm 40158 ±\pm 59
8422 ±\pm 9365 ±\pm 18361 ±\pm 19706 ±\pm 23606 ±\pm 31260 ±\pm 95
16402 ±\pm 8364 ±\pm 16361 ±\pm 171031 ±\pm 97829 ±\pm 39279 ±\pm 106
32380 ±\pm 6351 ±\pm 16348 ±\pm 161270 ±\pm 39816 ±\pm 49424 ±\pm 156
Appendix
Table 13: Per-iteration timing phase breakdown, restricted to the multiplicative_coprime reward.
KKSampleLossBwdExch. †∇θ‡\nabla_{\theta}^{\ddagger}fψ⋆f_{\psi}^{\star}
1413 ±\pm 7351 ±\pm 12350 ±\pm 12–––
2416 ±\pm 10363 ±\pm 15361 ±\pm 16143 ±\pm 25123 ±\pm 1731 ±\pm 13
4404 ±\pm 10346 ±\pm 14343 ±\pm 14352 ±\pm 95306 ±\pm 29102 ±\pm 47
8386 ±\pm 8358 ±\pm 11354 ±\pm 12487 ±\pm 78532 ±\pm 12223 ±\pm 83
16379 ±\pm 11353 ±\pm 10350 ±\pm 11788 ±\pm 133697 ±\pm 25260 ±\pm 100
32362 ±\pm 6364 ±\pm 14360 ±\pm 14967 ±\pm 144710 ±\pm 75303 ±\pm 115
Appendix
Table 14: Per-iteration timing phase breakdown, restricted to the conditional_multiscale reward.
More in this topic

Generative Flow Networks

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

See all 56 papers

Explore similar work

CardsList

  1. 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

  2. Routing by Reaching: Composition of Pre-trained GFlowNets for Multi-Objective Generation

    Feb 25, 2026Seokwon Yoon, Youngbin Choi, Seunghyuk Cho +3Generative Flow NetworksGenerative Modeling

    Department of Computer Science & Engineering, POSTECH, South Korea · Graduate School of Artificial Intelligence, POSTECH, South Korea

  3. Avoid What You Know: Divergent Trajectory Balance for GFlowNets

    Feb 19, 2026Pedro Dall'Antonia, Tiago da Silva, Daniel Csillag +2Generative Flow NetworksAdaptive Sampling

    School of Applied Mathmatics, Getulio Vargas Foundation · MBZUAI