Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

2026-02-27

AI Summary

This work addresses mode collapse in sequence generation with GFlowNets—manifesting as prefix collapse and length bias—stemming from insufficient credit assignment to early prefixes and replay-induced distributional bias. To mitigate these issues, the authors propose the RapTB objective, which anchors sub-trajectories at the root node and employs an absorbing-suffix backup mechanism to densely propagate terminal rewards back to intermediate prefixes, thereby delivering strong prefix-level supervision signals. Additionally, they introduce SubM, a replay strategy grounded in submodular functions, which enhances sample diversity while preserving high-reward trajectories, thus alleviating distributional shift. Evaluated on SMILES-based molecular generation, the proposed approach significantly improves optimization performance, generation diversity, and validity.

Technology Category

• Natural Language Processing: Generation
• Search and Optimization: Sampling/Simulation-based Search
• Planning, Routing, and Scheduling: Replanning and Plan Repair

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

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.

Problem

• mode collapse
• prefix collapse
• length bias
• GFlowNets
• reward-proportional posterior

Innovation

• GFlowNet
• Trajectory Balance
• submodular replay
• prefix collapse
• credit assignment