Local Search GFlowNets

TL;DR

This paper proposes to train GFlowNets with local search with backtracking and reconstruction guided by backward and forward policies, which allows biasing the samples toward high-reward solutions, which is not possible for a typical GFlowNet solution generation scheme.

Abstract

Generative Flow Networks (GFlowNets) are amortized sampling methods that learn a distribution over discrete objects proportional to their rewards. GFlowNets exhibit a remarkable ability to generate diverse samples, yet occasionally struggle to consistently produce samples with high rewards due to over-exploration on wide sample space. This paper proposes to train GFlowNets with local search, which focuses on exploiting high-rewarded sample space to resolve this issue. Our main idea is to explore the local neighborhood via backtracking and reconstruction guided by backward and forward policies, respectively. This allows biasing the samples toward high-reward solutions, which is not possible for a typical GFlowNet solution generation scheme, which uses the forward policy to generate the solution from scratch. Extensive experiments demonstrate a remarkable performance improvement in several biochemical tasks. Source code is available: \url{https://github.com/dbsxodud-11/ls_gfn}.

Figures (18)

• Figure 1: Strategy of LS-GFN.
• Figure 2: Illustration of a GFlowNet in a toy environment with two objects $x_1, x_2 \in \mathcal{X}$. $Z = R(x_1) + R(x_2)$ is the total amount of flow in the source, and the sink is the storage for the flow of the terminal state $x_1$ and $x_2$. The generative probability is: $p(x_1) = R(x_1)/\left(R(x_1)+R(x_2)\right)$, and $p(x_2) = R(x_2)/\left(R(x_1)+R(x_2)\right)$.
• Figure 3: Illustration of Local Search GFlowNet (LS-GFN) algorithm.
• Figure 4: Illustration of the $3$-step refinement process of LS-GFN.
• Figure 5: Accuracy of GFlowNet on various tasks. Ours stands for TB + LS-GFN.