🤖 machine learning

GFlowNet Training by Policy Gradients

This paper proposes a novel GFlowNet training framework that bridges flow balance with expected reward optimization to derive new policy-based methods, featuring a coupled strategy for jointly training forward policies and designing backward policies, which is theoretically guaranteed and empirically shown to improve performance on both simulated and real-world datasets.

Original authors: Puhua Niu, Shili Wu, Mingzhou Fan, Xiaoning Qian

Published 2026-08-06

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Puhua Niu, Shili Wu, Mingzhou Fan, Xiaoning Qian

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where you need to find the perfect recipe, the most efficient delivery route, or a new drug molecule, but the number of possible combinations is so vast it would take longer than the age of the universe to check them all one by one. This is the challenge of "combinatorial explosion," a problem that plagues everything from biology to engineering. To solve this, scientists use a clever tool called a Generative Flow Network (GFlowNet). Think of a GFlowNet not as a rigid rulebook, but as a magical water system. It builds complex objects step-by-step, like a river carving a path through a canyon. The goal is to make sure the "water" (or probability) flows so that the river ends up in the most beautiful, high-reward valleys (the best solutions) exactly as often as those valleys deserve.

Traditionally, teaching this water system to flow correctly has been like trying to balance a giant, invisible scale. The old methods, known as "value-based" approaches, focus on checking if the water levels at every single junction match a specific equation. It's a bit like a plumber constantly measuring the pressure at every pipe to ensure nothing leaks. While this works, it can be slow and clumsy, especially when the landscape is full of isolated, high-reward peaks that are hard to reach. The researchers in this paper asked: Is there a way to teach the water system by simply rewarding the path it takes, rather than just checking the pressure at every stop? They propose a new way to train these networks that feels more like a video game character learning to run a maze by collecting points, rather than a mathematician solving an equation.

The New Way to Train the Flow

The authors, Puhua Niu and their team, have developed a fresh training method for GFlowNets that shifts the focus from "checking the math" to "following the reward." In the old school of thought, the network was trained to keep the flow of water balanced across the entire map, a method that resembles how traditional Reinforcement Learning (RL) used to work by estimating the value of every state. The new approach, however, treats the training process like a direct policy gradient problem.

To understand this, imagine you are teaching a dog to fetch a ball.

The Old Way (Value-Based): You stand at every possible spot in the yard and calculate exactly how much "fetching value" that spot has. You then adjust the dog's behavior to ensure the math adds up perfectly at every single location. It's precise, but it requires a lot of mental energy to calculate the value of every blade of grass.
The New Way (Policy-Based): You simply say, "Good dog!" when the dog runs toward the ball and "Bad dog" when it runs the wrong way. You don't need to know the value of every spot in the yard; you just adjust the dog's running style based on the rewards it gets along the path.

The paper introduces a special kind of "reward" that depends on the strategy (or policy) the network is currently using. By doing this, they bridge the gap between the complex flow-balance equations of GFlowNets and the simpler, more direct "reward-and-punish" learning style of modern AI. They found that this method allows the network to learn much faster and more robustly, especially when the "treasure" (the high-reward solutions) is hidden in isolated spots that are hard to find.

What They Found and What They Avoided

The researchers tested their new "reward-based" training on several different challenges, including simulating grids (like a giant chessboard), designing biological sequences (like DNA strings), and creating molecular structures (like new medicines).

In these simulations, their new method, which they call RL-G (when using a smart guide) and RL-T (using a "trust region" to keep changes safe), consistently outperformed the old methods.

Speed: The new methods converged (found the solution) much faster. In the 256x256 grid experiment, the new methods reached a low error rate in fewer steps than the old "Trajectory Balance" (TB) methods.
Accuracy: The final results were often more accurate. For example, in the 256x256 grid, their best method (RL-G) achieved a Total Variation error of roughly 0.439, while the next best traditional method (TB-U) was around 0.728. In the molecular design tests, their methods found more unique "modes" (different high-quality solutions) than the old ways.

Crucially, the paper argues against the idea that we must always rely on complex, off-policy samplers (like Thompson Sampling or random mixing) to explore the space. While those methods try to balance "exploration" (trying new things) and "exploitation" (using what works), the authors show that by using their policy-based approach with robust gradient estimation, the network can find the best paths naturally without needing those complicated external tricks. They didn't just suggest this; they measured it across multiple datasets and showed that the new strategies provide a more stable and efficient way to train these networks.

The "Trust Region" and the "Guide"

To make sure the network doesn't get confused or jump into a bad habit, the authors added two special ingredients:

The Trust Region (RL-T): Imagine you are teaching a dog to run. If you tell it to run too fast too soon, it might trip. The "Trust Region" is like a leash that limits how much the dog's running style can change in one step. This keeps the learning stable and prevents the network from making wild, bad guesses. The paper shows this makes the training smoother and more reliable.
The Guided Policy (RL-G): Sometimes, the dog needs a little hint. The authors introduced a "guided" policy that acts like a map, gently steering the network away from dead ends (low-reward areas) and toward the treasure. This helps the network avoid getting stuck in "reward deserts" where there are no good solutions nearby.

Why This Matters

The paper concludes that by re-framing the training of GFlowNets as a direct reward-optimization problem, we can build better, faster, and more reliable AI for generating complex objects. Whether it's designing a new drug, optimizing a supply chain, or understanding the structure of the universe, this method offers a more direct path to the solution. The authors are confident in their results because they backed up their claims with rigorous mathematical proofs and extensive experiments on real-world and simulated data. They didn't just guess that this would work; they showed that it does, offering a promising new direction for how we teach AI to create and discover.

Problem Statement
Generative Flow Networks (GFlowNets) are designed to sample combinatorial objects from a space $\mathcal{X}$ with a probability proportional to a non-negative reward function $R(x)$ . While effective, existing GFlowNet training methods are predominantly value-based, relying on maintaining flow balance equations (e.g., Trajectory Balance, Detailed Balance) across a Directed Acyclic Graph (DAG). These methods face challenges in designing powerful samplers that balance exploration and exploitation, particularly in high-dimensional spaces with isolated modes. Furthermore, standard value-based approaches can suffer from numerical instability due to the explosion of flow values and may get trapped in local optima when using typical annealing or random-mixing strategies. The paper posits that bridging GFlowNet training with policy-based Reinforcement Learning (RL) could offer more robust gradient estimation and improved training efficiency.

Methodology
The authors propose a new training framework that reformulates GFlowNet training as a policy-based RL problem over a special Markov Decision Process (MDP).

Policy-Dependent Rewards: The core innovation is the definition of policy-dependent reward functions. For a forward policy $\pi_F$ and a backward policy $\pi_B$ , the rewards are defined as the log-ratios of these policies (e.g., $R_F(s, a) = \log \frac{\pi_F(s, a)}{\pi_B(s', a)}$ ). This formulation allows the flow balance condition to be reinterpreted as minimizing the expected accumulated reward, effectively bridging the gap between flow balance and RL objectives.
RL Formulation: The training of the forward policy is cast as minimizing an expected value function $J_F$ , which corresponds to the Kullback-Leibler (KL) divergence between the forward trajectory distribution and a target distribution derived from the backward policy. Similarly, the backward policy design is formulated as an RL problem to minimize a separate objective $J_B$ .
Coupled Training Strategy: The paper introduces a coupled training strategy that jointly optimizes the forward policy and the backward policy design. Unlike previous two-phase approaches, this method allows for simultaneous updates, theoretically guaranteeing performance improvements.
Advanced Optimization Techniques: Drawing from established policy-based RL literature, the authors incorporate:
- Variance Reduction: Using functional baselines (critics) and generalized advantage estimation (controlled by a parameter $\lambda$ ) to manage the bias-variance trade-off in gradient estimation.
- Trust Region Constraints: Adapting Trust Region Policy Optimization (TRPO) and Proximal Policy Optimization (PPO) concepts to ensure stable, conservative policy updates, preventing the model from getting trapped in local optima.
Theoretical Guarantees: The paper provides theoretical proofs showing that the proposed policy-based objectives are equivalent to minimizing the KL divergence between distributions. It establishes convergence guarantees for the policy gradients under policy-dependent rewards and absorbing Markov chains (specific to the DAG structure of GFlowNets), which differ from standard ergodic MDP assumptions in traditional RL.

Key Contributions

Reformulation of GFlowNet Training: The authors successfully reformulate GFlowNet training as an RL problem with policy-dependent rewards, transforming the flow balance equation into an optimization of expected accumulated rewards.
Policy-Based Training Strategies: They derive new training strategies that directly optimize the forward policy via gradients, contrasting with existing value-based methods. This includes the adaptation of TRPO and PPO-style objectives for GFlowNets.
Coupled Backward Policy Design: The paper formulates the design of backward policies as an RL problem and proposes a joint training strategy that optimizes forward and backward policies simultaneously, improving training efficiency.
Theoretical Analysis: The work provides theoretical guarantees for the convergence of these policy-based methods, proving that they asymptotically converge to stationary points and establishing upper bounds on performance improvement.

Experimental Results
The authors evaluated their methods on three domains: hyper-grid modeling, biological/molecular sequence design (SIX6, QM9, PHO4, sEH datasets), and Bayesian Network structure learning.

Performance: Across simulated and real-world datasets, the proposed policy-based methods (specifically RL-G with guided policies and RL-T with TRPO-based updates) consistently demonstrated faster convergence rates and lower Total Variation (TV) distance compared to state-of-the-art value-based methods (such as TB-U, TB-B, TB-Sub, and DB-based methods).
Stability: The TRPO-based variant (RL-T) showed greater stability during training compared to vanilla policy-based methods, confirming the benefit of conservative updates.
Ablation: Experiments varying the bias-variance trade-off parameter $\lambda$ showed that the policy-based methods provided more robust gradient estimation than Trajectory Balance (TB) methods, particularly when $\lambda$ was tuned.

Significance and Claims
The paper claims that its primary significance lies in providing alternative, policy-based perspectives for GFlowNet training that offer robust gradient estimation. By reframing the problem through the lens of policy gradients, the authors argue that GFlowNets can leverage advanced RL techniques (like trust regions and variance reduction) to overcome the limitations of value-based samplers. The work suggests that this approach is particularly effective in environments with large, isolated modes where traditional exploration-exploitation balancing is difficult. The authors modestly note that while their methods do not require explicit sampler design (unlike value-based off-policy methods), they still face challenges with gradient variance in very large state spaces, leaving room for future work on further improving gradient estimation techniques. The code is made available to facilitate reproducibility.

The New Way to Train the Flow

What They Found and What They Avoided

The "Trust Region" and the "Guide"

Why This Matters

More like this