Mixture of GFlowNets

Omar G. Younis, Joseph D. Viviano, Chirayu Anant Haryan, MD Vasimuddin, Sanchit Misra +1

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…