GFlowNet Training by Policy Gradients

TL;DR

This paper introduces a new policy-gradient training framework for Generative Flow Networks (GFlowNets), enhancing their ability to generate combinatorial objects by integrating reinforcement learning principles.

Contribution

It develops a novel policy-based training method for GFlowNets with theoretical guarantees, bridging flow balance with reward optimization, and jointly training forward and backward policies.

Findings

Policy-based GFlowNet training improves performance.

Joint training of forward and backward policies enhances efficiency.

Experimental results show robustness and better gradient estimation.

Abstract

Generative Flow Networks (GFlowNets) have been shown effective to generate combinatorial objects with desired properties. We here propose a new GFlowNet training framework, with policy-dependent rewards, that bridges keeping flow balance of GFlowNets to optimizing the expected accumulated reward in traditional Reinforcement-Learning (RL). This enables the derivation of new policy-based GFlowNet training methods, in contrast to existing ones resembling value-based RL. It is known that the design of backward policies in GFlowNet training affects efficiency. We further develop a coupled training strategy that jointly solves GFlowNet forward policy training and backward policy design. Performance analysis is provided with a theoretical guarantee of our policy-based GFlowNet training. Experiments on both simulated and real-world datasets verify that our policy-based strategies provide advanced RL perspectives for robust gradient estimation to improve GFlowNet performance.