# Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

arXiv: https://arxiv.org/abs/2603.00454 · HTML: https://mlsanity.com/paper/2603.00454/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 41 references that are in the map.
- First posted on arXiv 28 Feb 2026.

## Builds on (most cited first, 27 of 27)

- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation](https://mlsanity.com/paper/2106.04399/) (2021)
- [GFlowNet Foundations](https://mlsanity.com/paper/2111.09266/) (2021)
- [Trajectory balance: Improved credit assignment in GFlowNets](https://mlsanity.com/paper/2201.13259/) (2022)
- [GLISTER: Generalization based Data Subset Selection for Efficient and Robust Learning](https://mlsanity.com/paper/2012.10630/) (2020)
- [Biological Sequence Design with GFlowNets](https://mlsanity.com/paper/2203.04115/) (2022)
- [Learning GFlowNets from partial episodes for improved convergence and stability](https://mlsanity.com/paper/2209.12782/) (2022)
- [GFlowNets and variational inference](https://mlsanity.com/paper/2210.00580/) (2022)
- [Multi-Objective GFlowNets](https://mlsanity.com/paper/2210.12765/) (2022)
- [Amortizing intractable inference in large language models](https://mlsanity.com/paper/2310.04363/) (2023)
- [Towards Understanding and Improving GFlowNet Training](https://mlsanity.com/paper/2305.07170/) (2023)
- [Generative Flow Networks as Entropy-Regularized RL](https://mlsanity.com/paper/2310.12934/) (2023)
- [PRISM: A Rich Class of Parameterized Submodular Information Measures for Guided Subset Selection](https://mlsanity.com/paper/2103.00128/) (2021)
- [Local Search GFlowNets](https://mlsanity.com/paper/2310.02710/) (2023)
- [Discrete Probabilistic Inference as Control in Multi-path Environments](https://mlsanity.com/paper/2402.10309/) (2024)
- [GFlowNet-EM for learning compositional latent variable models](https://mlsanity.com/paper/2302.06576/) (2023)
- [Full-Atom Peptide Design based on Multi-modal Flow Matching](https://mlsanity.com/paper/2406.00735/) (2024)
- [CASCADE: Cumulative Agentic Skill Creation through Autonomous Development and Evolution](https://mlsanity.com/paper/2512.23880/) (2025)
- [Crystal-GFN: sampling crystals with desirable properties and constraints](https://mlsanity.com/paper/2310.04925/) (2023)
- [Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Training](https://mlsanity.com/paper/2503.18929/) (2025)
- [An Empirical Study of the Effectiveness of Using a Replay Buffer on Mode Discovery in GFlowNets](https://mlsanity.com/paper/2307.07674/) (2023)
- [Looking Backward: Retrospective Backward Synthesis for Goal-Conditioned GFlowNets](https://mlsanity.com/paper/2406.01150/) (2024)
- [Improving GFlowNets with Monte Carlo Tree Search](https://mlsanity.com/paper/2406.13655/) (2024)
- [GFlowNet Training by Policy Gradients](https://mlsanity.com/paper/2408.05885/) (2024)
- [Evolution Guided Generative Flow Networks](https://mlsanity.com/paper/2402.02186/) (2024)
- [Geometric Point Attention Transformer for 3D Shape Reassembly](https://mlsanity.com/paper/2411.17788/) (2024)
- [RTMol: Rethinking Molecule-text Alignment in a Round-trip View](https://mlsanity.com/paper/2511.12135/) (2025)

## Built on by (most cited first, 3 of 3)

- [SHAPE: Unifying Safety, Helpfulness and Pedagogy for Educational LLMs](https://mlsanity.com/paper/2604.22134/) (2026)
- [R2 Flow: Recursive Self-Improvement via Recursive Skill Evolution](https://mlsanity.com/paper/2609.33867/) (2026)
- [SNACK: A Sequential Notation Framework for Probabilistic Graph Generation](https://mlsanity.com/paper/2610.19781/) (2026)

## Abstract (arXiv)

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.