Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

Generative Flow Networks (GFlowNets) enable fine-tuning large language models to approximate reward-proportional posteriors, but they remain prone to mode collapse, manifesting as prefix collapse and length bias. We attribute this to two factors: (i) weak credit assignment to early prefixes, and (ii) biased replay that induces a shifted, non-representative training flow distribution. We propose Rooted absorbed prefix Trajectory Balance (RapTB), an objective that anchors subtrajectory supervision at the root and propagates terminal rewards to intermediate prefixes via absorbed suffix-based backups, providing dense prefix-level learning signals. To mitigate replay-induced distribution shift, we further introduce SubM, a submodular replay refresh strategy that promotes both high reward and diversity. Empirically, on tasks such as molecule generation with LLM using SMILES strings, RapTB combined with SubM consistently improves optimization performance and molecular diversity while preserving high validity. The code is released on https://github.com/ComDec/ChemGFN.

AI systems can be trained to generate diverse, high-quality molecules or text by learning to sample many good solutions rather than fixating on a single best one. However, current methods often fail at this: they tend to collapse, producing outputs that all look nearly identical or that are systematically too short or too long. We traced these failures to two causes — the learning signal only arrives at the very end of generation (making it hard for the model to learn what early decisions were good), and the training examples the model revisits become increasingly narrow and repetitive. To address this, we developed RapTB, a training objective that propagates reward information back to earlier steps in the generation process, giving the model denser and more informative feedback. We also introduced SubM, a strategy that keeps the model's training pool diverse by balancing quality, variety, and length coverage when selecting which past examples to revisit. Together, RapTB and SubM consistently improve both the quality and diversity of generated molecules and expressions across multiple tasks and model scales up to 32 billion parameters, offering a more reliable foundation for AI-driven scientific discovery such as drug design.