Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

TL;DR

This work proposes 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.

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.

Figures (5)

• Figure 1: Training objectives for LLM-GFlowNets. TB uses only terminal reward $\log R(s_\tau)$ ($O(1)$). SubTB adds $O(N^2)$ windowed consistency constraints. RapTB replaces prefix stop-rewards with suffix-absorbed targets $u_j$ and applies $O(N)$ rooted prefix constraints. ($N$: trajectory length). The right column shows examples of generated molecules trained with different losses. QED (Quantitative Estimate of Drug-likeness) is a comprehensive metric for measuring a molecule's drug-likeness and is used as the task reward during training.
• Figure 2: Length-stratified analysis on SMILES ($L_{\max}=10$). (a) Distribution of valid generation lengths. (b) Mean Score and FPDiv conditioned on length.
• Figure 3: Prefix-collapse diagnostics on SMILES ($L_{\max}=10$). Metrics vs. prefix length $k$ computed on correct samples: prefix survival (fraction of samples reaching length $k$), prefix entropy (diversity of prefix), and top-1 mass (frequency of the most common prefix).
• Figure 4: EBNF grammar used for constrained SMILES decoding.
• Figure 5: EBNF grammar used for constrained Expr24 decoding.