Proximal Policy Optimization for Amortized Discrete Sampling

2026-06-14

AI Summary

This work addresses the slow convergence and low sample efficiency of Generative Flow Networks (GFlowNets) in structured discrete sampling tasks by leveraging their theoretical connection to entropy-regularized reinforcement learning. For the first time, Proximal Policy Optimization (PPO) is integrated into the GFlowNet training framework. Through a systematic design of policy gradient updates, advantage estimation, and baseline mechanisms, the proposed approach substantially enhances training stability and data efficiency. Experimental results on benchmark tasks—including synthetic energy functions and molecular graph generation—demonstrate that the method achieves faster convergence and higher sampling efficiency compared to standard GFlowNet objectives.

Abstract

This paper explores policy gradient algorithms for training stochastic policies to sample from structured discrete probability distributions under the Generative Flow Network (GFlowNet) framework. Building on extensive theoretical connections between GFlowNets and entropy-regularized reinforcement learning, we derive equivalents of standard policy gradient algorithms for training GFlowNets, as well as experimentally explore their various methodological aspects, including baseline training and advantage estimation. Most importantly, our work is the first to derive and successfully apply proximal policy optimization to GFlowNets, showing its improved convergence speed and data efficiency compared to standard GFlowNet training objectives on benchmarks ranging from synthetic energies to molecular graph generation.

Problem

• GFlowNet
• discrete sampling
• policy gradient
• stochastic policies
• structured distributions

Innovation

• Proximal Policy Optimization
• GFlowNet
• policy gradient
• discrete sampling
• entropy-regularized reinforcement learning