Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

TL;DR

This paper introduces RapTB and SubM, novel methods to improve GFlowNet training by addressing mode collapse, enhancing reward propagation, and promoting diversity, leading to better performance in molecule generation tasks.

Contribution

The paper proposes a rooted prefix trajectory balance objective and a submodular replay strategy to enhance GFlowNet training, reducing mode collapse and improving diversity.

Findings

Improved molecular diversity in generated molecules.

Enhanced reward optimization performance.

Maintained high validity of generated molecules.

Abstract

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.