Loss-Guided Auxiliary Agents for Overcoming Mode Collapse in GFlowNets
This addresses mode collapse for users of GFlowNets in applications like sequence generation and Bayesian learning, offering a novel method rather than an incremental improvement.
The paper tackled mode collapse in Generative Flow Networks (GFlowNets) by proposing Loss-Guided GFlowNets (LGGFN), which uses the main model's training loss to guide exploration, resulting in over 40 times more unique valid modes and a 99% reduction in exploration error on a sequence generation task.
Although Generative Flow Networks (GFlowNets) are designed to capture multiple modes of a reward function, they often suffer from mode collapse in practice, getting trapped in early-discovered modes and requiring prolonged training to find diverse solutions. Existing exploration techniques often rely on heuristic novelty signals. We propose Loss-Guided GFlowNets (LGGFN), a novel approach where an auxiliary GFlowNet's exploration is \textbf{directly driven by the main GFlowNet's training loss}. By prioritizing trajectories where the main model exhibits \textbf{high loss}, LGGFN focuses sampling on poorly understood regions of the state space. This targeted exploration significantly accelerates the discovery of diverse, high-reward samples. Empirically, across \textbf{diverse benchmarks} including grid environments, structured sequence generation, Bayesian structure learning, and biological sequence design, LGGFN consistently \textbf{outperforms} baselines in exploration efficiency and sample diversity. For instance, on a challenging sequence generation task, it discovered over 40 times more unique valid modes while simultaneously reducing the exploration error metric by approximately 99\%.