cs.LGFeb 28, 2026

Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

Authors: Xi Wang, Wenbo Lu, Shengjie Wang

Organizations: Courant Institute School of Mathematics, Computing, and Data Science, New York University

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. The code is released on https://github.com/ComDec/ChemGFN.

Figures & tables

Appendix figures & tables28 assets

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

CardsList
  1. GFlowPO: Generative Flow Network as a Language Model Prompt Optimizer

    Feb 3, 2026Junmo Cho, Suhan Kim, Sangjune An +5LLM Input Optimization

  2. GFlowRL: Scaling Distribution-Matching RL to Large Language Models

    Jul 15, 2026Xiaodong Liu, Michael Xu, Jack W. Stokes +3RL for Code GenerationReinforcement Learning

  3. Signal from Structure: Exploiting Submodular Upper Bounds in Generative Flow Networks

    Jan 28, 2026Alexandre Larouche, Audrey DurandGenerative Flow NetworksGraph Generation