Loss-Guided Auxiliary Agents for Overcoming Mode Collapse in GFlowNets

2025-05-21

AI Summary

GFlowNets suffer from mode collapse in multimodal generation, resulting in low exploration efficiency and insufficient sample diversity. To address this, we propose Loss-Guided Auxiliary GFlowNets (LG-GFN), the first method to directly align exploration objectives with the primary model’s training loss—replacing heuristic novelty signals. LG-GFN employs backward flow modeling and loss-sensitive trajectory prioritization to steer the auxiliary agent toward high-error states, while a dual-agent cooperative training framework enables targeted diversity enhancement. Evaluated on grid-world environments, structured sequence generation, and Bayesian network structure learning, LG-GFN significantly improves exploration efficiency: in sequence tasks, it discovers 40× more valid modes, reduces exploration error by 99%, and achieves substantially higher sample diversity than all baselines.

Technology Category

• Machine Learning: Multimodal Learning
• Natural Language Processing: Generation
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

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 may rely on heuristic novelty signals. We propose Loss-Guided GFlowNets (LGGFN), a novel approach where an auxiliary GFlowNet's exploration is directly driven by the main GFlowNet's training loss. By prioritizing trajectories where the main model exhibits 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 various benchmarks including grid environments, structured sequence generation, and Bayesian structure learning, LGGFN consistently enhances exploration efficiency and sample diversity compared to baselines. 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%.

Problem

• Overcoming mode collapse in GFlowNets during training
• Enhancing exploration efficiency in diverse reward regions
• Improving sample diversity in generative flow networks

Innovation

• Auxiliary GFlowNet driven by main model's loss
• Prioritizes high-loss trajectories for targeted exploration
• Enhances discovery of diverse high-reward samples