📊 statistics

Proximal Policy Optimization for Amortized Discrete Sampling

This paper establishes theoretical links between GFlowNets and entropy-regularized reinforcement learning to derive and demonstrate the superior convergence and data efficiency of Proximal Policy Optimization (PPO) for training stochastic policies to sample from structured discrete distributions across various benchmarks.

Original authors: Anna Zykova-Myzina, Timofei Gritsaev, Daniil Tiapkin, Nikita Morozov

Published 2026-06-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Anna Zykova-Myzina, Timofei Gritsaev, Daniil Tiapkin, Nikita Morozov

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 you are a chef trying to recreate a specific, complex recipe for a dish that tastes exactly like a famous, award-winning meal. However, you don't have the recipe card. You only have a list of ingredients (the "reward") that tells you how good a dish tastes if you get it right, but you don't know the total number of possible dishes in the universe or the exact proportions needed to make the perfect one.

This is the problem the paper tackles: How do you teach a computer to randomly generate items (like molecules or DNA sequences) so that they follow a specific, desired pattern, rather than just finding the single "best" item?

Here is the breakdown of their solution using simple analogies.

1. The Problem: The "Best" Trap vs. The "Right Mix"

In many computer tasks, algorithms are trained to find the single best solution (like finding the highest peak on a mountain). But in fields like chemistry or biology, you often need a whole variety of solutions that follow a specific distribution. You don't just want the one molecule with the highest energy; you want a diverse set of molecules that matches a specific probability curve.

The authors use a framework called GFlowNets (Generative Flow Networks). Think of a GFlowNet as a factory assembly line. The machine builds an object step-by-step (adding one Lego brick at a time). The goal is to tune the machine so that if you run it a million times, the final pile of finished products looks exactly like the target distribution you want.

2. The Old Way: The "Mapmaker" Struggle

Previously, training these factories relied on "Value-Based" methods.

The Analogy: Imagine trying to navigate a maze by drawing a perfect map of every single path and calculating the exact "flow" of water through every pipe.
The Problem: To draw this map, you need to know the total size of the maze (a number called the "normalizing constant" or $Z$ ). In complex problems, calculating this number is incredibly hard, like trying to count every grain of sand on a beach to figure out the weight of the sand. If your map is slightly off, the whole training process gets stuck or becomes very inefficient.

3. The New Way: The "Trial-and-Error" Coach (PPO)

The authors decided to try a different approach used in Reinforcement Learning (RL), specifically an algorithm called PPO (Proximal Policy Optimization).

The Analogy: Instead of drawing a perfect map, imagine a coach standing next to the factory machine. The coach watches the machine build a few items, sees how good they are, and says, "Hey, when you added that blue brick, you were a bit too aggressive. Next time, be a little more gentle."
The Benefit: The coach doesn't need to know the total size of the beach or draw a perfect map. They just need to look at the immediate results and nudge the machine in the right direction. This is much more data-efficient.

4. The Twist: Why Standard PPO Failed

The authors tried using standard PPO, but it failed.

The Failure: Standard PPO is designed to find the single best outcome (the highest peak). If you just tell the factory "make the tastiest dish possible," it will stop making variety and just keep making the one dish that tastes the best. It collapses into a single mode.
The Missing Ingredients: The authors realized that to make PPO work for sampling (creating variety), two specific things were missing from the standard recipe:
1. The "Backward" Clue: You need to tell the machine not just about the reward at the end, but also about the "history" of how it got there. It's like telling the chef, "Not only is the cake good, but the way you mixed the eggs was also crucial."
2. The "Entropy" Bonus: You must explicitly reward the machine for being uncertain or exploratory. If the machine is too confident, you penalize it. This forces it to keep exploring different paths rather than settling on one.

5. The Solution: "Ent-PPO"

The paper introduces Ent-PPO (Entropic Proximal Policy Optimization). This is a custom-tuned version of the coach.

How it works: It combines the "clipping" mechanism of standard PPO (which prevents the coach from giving advice that is too wild and destabilizing the factory) with a new mathematical "trust region" derived from the entropy bonus.
The Result: This new coach successfully teaches the factory to produce a diverse, high-quality mix of items that perfectly matches the target distribution.

6. The Results: Faster and Better

The authors tested this on several "playgrounds":

Synthetic Grids: Simple digital mazes.
DNA Sequences: Creating strings of DNA that bind to specific proteins.
Molecules: Generating chemical structures.

The Findings:

Speed: Ent-PPO learned much faster than the old "Mapmaker" methods (like Trajectory Balance or Detailed Balance).
Efficiency: It needed far fewer attempts (samples) to get the job done.
Stability: The old methods often got stuck or produced bad results if the math wasn't perfect. Ent-PPO was robust and stable, even when the problems got very large and complex (like generating full molecular graphs).

Summary

The paper is essentially saying: "We took a powerful training tool used for large language models (PPO), fixed two specific bugs that made it fail at 'sampling' tasks, and proved that it is now the best way to teach computers to generate diverse, complex structures like molecules and DNA, beating the previous state-of-the-art methods."

They didn't just find a new way to do it; they found a way that is faster, uses less data, and is more stable, making it a significant upgrade for anyone trying to generate complex discrete data.

Technical Summary: Proximal Policy Optimization for Amortized Discrete Sampling

Problem Statement

The paper addresses the task of sampling from a probability distribution $p(x) = R(x)/Z$ over a finite discrete space $\mathcal{X}$ , where $R(x)$ is a positive reward function and $Z$ is an unknown normalizing constant. This problem arises in domains such as statistical physics, Bayesian statistics, and computational biology. The authors focus on Generative Flow Networks (GFlowNets), which model the generation of objects $x$ as a sequential process over a Directed Acyclic Graph (DAG). The goal is to train a stochastic forward policy $P_F$ to navigate this DAG such that the distribution of terminal states matches $p(x)$ .

While GFlowNet training has been dominated by value-based objectives (e.g., Trajectory Balance, Detailed Balance), the authors note that policy gradient methods, particularly Proximal Policy Optimization (PPO), remain underexplored in this context despite their dominance in large-scale Reinforcement Learning (RL) and diffusion model post-training.

Methodology

Theoretical Foundation: GFlowNets as Soft RL

The work builds upon the theoretical connection established in prior literature [42, 7] between GFlowNets and entropy-regularized RL. By setting the discount factor $\gamma=1$ and the entropy coefficient $\alpha=1$ , the authors demonstrate that the GFlowNet sampling problem is equivalent to maximizing the soft value function in an RL setting.

Reward Structure: Crucially, the per-step reward is defined as $r(s_t, s_{t+1}) = \log P_B(s_t | s_{t+1})$ for non-terminal transitions, where $P_B$ is a backward policy, and $r(s_T) = \log R(s_T)$ at terminal states.
Soft Policy Improvement: The optimal policy for this soft-RL objective coincides with the target sampling distribution $R(x)/Z$ .

Ent-PPO: A Novel Policy Gradient Algorithm

The paper derives Ent-PPO (Entropic Proximal Policy Optimization), a variant of PPO adapted to the soft-RL formulation of GFlowNets. The derivation highlights two critical components often missing in naive applications of PPO to GFlowNets:

Correct Reward Formulation: The intermediate rewards must be $\log P_B$ (as defined above) to align the soft-RL optimum with the sampling target. Without this, the objective reduces to standard RL with sparse rewards, optimizing for $\arg\max_x R(x)$ (mode collapse) rather than the full distribution.
Analytic KL Trust Region: Unlike standard PPO, which adds a heuristic entropy bonus $c_2 H(\pi)$ $c_{2} H (π)$ , Ent-PPO derives an analytic KL penalty directly from the soft policy improvement theorem.
- The surrogate objective is derived by combining the cross-entropy term (from the soft advantage definition) and the entropy term.
- This results in a per-step loss: $\ell_t(\theta) = \text{PPOClip}(\rho_t, \hat{A}_t) - \text{KL}(\pi_\theta(\cdot|s_t) \parallel \pi_{\text{old}}(\cdot|s_t))$ .
- The KL term acts as a trust region that is intrinsic to the soft-RL formulation, replacing the need for a separate, tunable entropy coefficient.

Variance Reduction

The authors systematically evaluate variance reduction techniques for the Vanilla Policy Gradient (VPG) baseline within this framework:

Estimators: They compare the simplest policy gradient, reward-to-go, reward-to-go with a learned value baseline, Generalized Advantage Estimation (GAE), and Subtrajectory Evaluation Balance (SubEB-GAE).
Findings: GAE with a learned value baseline is identified as the most effective strategy, offering a favorable bias-variance trade-off and significantly improving convergence speed and stability compared to simpler estimators.

Key Contributions

Derivation of Standard Policy Gradients for GFlowNets: The authors derive standard policy gradient objectives for discrete samplers under the entropy-regularized RL framework, experimentally validating the importance of variance reduction techniques (baselines and advantage estimation) for stable training.
Ent-PPO Algorithm: They propose Ent-PPO, the first successful application of PPO to GFlowNets. The derivation explicitly identifies the necessary intermediate rewards ( $\log P_B$ ) and the analytic KL trust region.
Empirical Validation: The paper demonstrates that Ent-PPO outperforms established off-policy GFlowNet objectives (DB, TB, SubTB) and naive PPO variants across a range of benchmarks, from synthetic grids to complex molecular graph generation.

Experimental Results

The authors evaluate their methods on several environments:

Synthetic/Small-Scale: Hypergrid, TFBind8 (DNA sequences), and String QM9.
Large-Scale: sEH (fragment-based molecular graphs) and QM9 (atom-by-atom molecular graphs).

Key Findings:

Superiority over Baselines: Ent-PPO consistently achieves lower Total Variation (TV) distance on small problems and better Evidence Lower Bound (ELBO) / Evidence Upper Bound (EUBO) on large molecular problems compared to DB, TB, and SubTB.
Sample Efficiency: Ent-PPO benefits significantly from multiple update epochs ( $K > 1$ ) per batch, showing improved sample efficiency and convergence speed, a trait common in large-scale RL but previously unexplored in GFlowNets.
Necessity of Corrections:
- KL Penalty: Removing the KL term (reverting to naive PPO) causes training to collapse or become unstable, particularly with multiple update epochs. The KL term is essential for preventing the policy from deviating too far from the behavior policy while maintaining the correct soft-RL objective.
- Clipping: Removing the PPO clipping mechanism while retaining the KL term leads to divergence, indicating that the analytic KL alone is insufficient to bound the importance ratio fluctuations.
Backward Policy Learning: While Ent-PPO does not natively support joint forward-backward learning like off-policy objectives, the authors show that using Trajectory Likelihood Maximization (TLM) to learn the backward policy separately significantly accelerates convergence.

Significance and Claims

The paper claims that Ent-PPO represents a significant advancement in GFlowNet training by bridging the gap between modern, scalable policy gradient methods and the specific requirements of discrete sampling.

Theoretical Rigor: The work corrects previous attempts to apply PPO to GFlowNets by rigorously deriving the necessary reward structure and trust region from first principles of entropy-regularized RL.
Practical Impact: The results suggest that policy gradient methods, when properly adapted, can outperform the dominant value-based GFlowNet objectives, particularly in terms of convergence speed and data efficiency.
Scalability: The ability of Ent-PPO to leverage multiple update epochs and scale to complex molecular generation tasks (QM9) suggests a path toward applying GFlowNets to larger, more complex problems where traditional value-based methods may struggle with credit assignment or partition function estimation.

The authors conclude that while their method currently focuses on discrete domains, the soft-RL framework offers a promising direction for adapting policy gradient methods to continuous sampling domains in the future.