arXiv preprint ·
Mixture of GFlowNets
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
Abstract
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.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning
Federated Coreset Subsampling for Distributed Memory Banks
ATTUNE: Agentic Throughput Tuning for LLM Training
Distributed $(Δ+1)$-Coloring in $O(\sqrtΔ+\log^* n)$ Rounds
DisaggKD: Efficient Knowledge Distillation through Teacher-Student Disaggregation and a Fused Loss
Prorsa: Unleashing the Potential of Prefill-Decode Disaggregation in Agentic Reinforcement Learning
Get the few papers worth your time in Machine Learning every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.