# f-Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 41 references that are in the map.
- First posted on arXiv 14 May 2026.

## Builds on (most cited first, 30 of 33)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [PubMedQA: A Dataset for Biomedical Research Question Answering](https://mlsanity.com/paper/1909.06146/) (2019)
- [Kimi k1.5: Scaling Reinforcement Learning with LLMs](https://mlsanity.com/paper/2501.12599/) (2025)
- [Kimi K2: Open Agentic Intelligence](https://mlsanity.com/paper/2507.20534/) (2025)
- [f-GAN: Training Generative Neural Samplers using Variational Divergence Minimization](https://mlsanity.com/paper/1606.00709/) (2016)
- [Policy Distillation](https://mlsanity.com/paper/1511.06295/) (2015)
- [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)
- [MiniMax-M1: Scaling Test-Time Compute Efficiently with Lightning Attention](https://mlsanity.com/paper/2506.13585/) (2025)
- [Trajectory balance: Improved credit assignment in GFlowNets](https://mlsanity.com/paper/2201.13259/) (2022)
- [A Divergence Minimization Perspective on Imitation Learning Methods](https://mlsanity.com/paper/1911.02256/) (2019)
- [Imitation Learning as f-Divergence Minimization](https://mlsanity.com/paper/1905.12888/) (2019)
- [Aligning Language Models with Preferences through f-divergence Minimization](https://mlsanity.com/paper/2302.08215/) (2023)
- [GFlowNets and variational inference](https://mlsanity.com/paper/2210.00580/) (2022)
- [A theory of continuous generative flow networks](https://mlsanity.com/paper/2301.12594/) (2023)
- [Hierarchical Variational Models](https://mlsanity.com/paper/1511.02386/) (2015)
- [Solving high-dimensional Hamilton-Jacobi-Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space](https://mlsanity.com/paper/2005.05409/) (2020)
- [Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimization](https://mlsanity.com/paper/2407.13399/) (2024)
- [VarGrad: A Low-Variance Gradient Estimator for Variational Inference](https://mlsanity.com/paper/2010.10436/) (2020)
- [Amortizing intractable inference in diffusion models for vision, language, and control](https://mlsanity.com/paper/2405.20971/) (2024)
- [Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond](https://mlsanity.com/paper/2403.06279/) (2024)
- [Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Training](https://mlsanity.com/paper/2503.18929/) (2025)
- [On a few pitfalls in KL divergence gradient estimation for RL](https://mlsanity.com/paper/2506.09477/) (2025)
- [SynFlowNet: Design of Diverse and Novel Molecules with Synthesis Constraints](https://mlsanity.com/paper/2405.01155/) (2024)
- [Variational Inference with Tail-adaptive f-Divergence](https://mlsanity.com/paper/1810.11943/) (2018)
- [f-PO: Generalizing Preference Optimization with f-divergence Minimization](https://mlsanity.com/paper/2410.21662/) (2024)
- [From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training](https://mlsanity.com/paper/2501.06148/) (2025)

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

- [R2 Flow: Recursive Self-Improvement via Recursive Skill Evolution](https://mlsanity.com/paper/2609.33867/) (2026)
- [One Proposal for Every Margin: Zero-Shot Amortized Sequential Importance Sampling for Binary Matrices](https://mlsanity.com/paper/2609.35514/) (2026)
- [EvoSteer: Online Self-Evolving Graph Orchestration via Reference-Anchored Credit Assignment](https://mlsanity.com/paper/2609.38661/) (2026)
- [Towards Understanding the Training of GFlowNets: A Theoretical Perspective](https://mlsanity.com/paper/2610.13711/) (2026)

## Abstract (arXiv)

In GFlowNets and variational inference, it has been shown that the mean square error between target and model log probabilities is an effective, low variance, surrogate loss for training generative models. This loss has the property that when evaluated on-policy its gradients correspond to those of the KL divergence, while off-policy it remains a valid loss with the same global minimizer. In this work, we demonstrate that this construction can be extended to the whole family of f-divergences, leading to a family of losses whose on-policy gradients are that of the corresponding f-divergence, but retain the same global minimizer off-policy. Specifically, we show that the on-policy gradients lead to a one to one correspondence between translation invariant loss functions on the target and model log probabilities, and f-divergences. This equivalence allows us to design new surrogate loss functions for tuning a wide class of generative models that inherit the properties of the corresponding f-divergence, such as being more mode covering, whilst being applicable to off-policy data. We apply our losses on a range of tasks, including classic synthetic examples, SynFlowNets for molecule discovery, and asynchronous large language model (LLM) tuning, demonstrating that our models retain their predicted properties on- and off-policy in a wide class of generative models.