Controlling Exploration-Exploitation in GFlowNets via Markov Chain Perspectives

TL;DR

This work addresses the rigidity of equal forward/backward policy mixing in Generative Flow Networks (GFlowNets) by revealing a Markov chain reversibility grounding for their objectives. It introduces α-GFN, a generalized framework that blends forward and backward transitions with a tunable parameter $\alpha$, and proves convergence to unique flows via MC reversibility, complemented by a scheduling strategy to combine multiple $\alpha$ regimes. Theoretical results unify various GFN objectives under a single MC perspective and derive gradient dynamics that explain the empirical gains. Empirically, α-GFN consistently improves mode discovery and average reward across Set Generation, Bit Sequence Generation, and Molecule Generation, while maintaining diversity and compatibility with existing GFlowNet enhancements and large-scale exploration tasks like FlowRL.

Abstract

Generative Flow Network (GFlowNet) objectives implicitly fix an equal mixing of forward and backward policies, potentially constraining the exploration-exploitation trade-off during training. By further exploring the link between GFlowNets and Markov chains, we establish an equivalence between GFlowNet objectives and Markov chain reversibility, thereby revealing the origin of such constraints, and provide a framework for adapting Markov chain properties to GFlowNets. Building on these theoretical findings, we propose $α$-GFNs, which generalize the mixing via a tunable parameter $α$. This generalization enables direct control over exploration-exploitation dynamics to enhance mode discovery capabilities, while ensuring convergence to unique flows. Across various benchmarks, including Set, Bit Sequence, and Molecule Generation, $α$-GFN objectives consistently outperform previous GFlowNet objectives, achieving up to a $10 \times$ increase in the number of discovered modes.

Figures (19)

• Figure 1: (Left) Illustration of $\alpha$-GFNs. Vanilla GFlowNets implicitly assigns equal weights (0.5/0.5) to the forward policy $P_F$ and the backward policy $P_B$, while $\alpha$-GFNs assign $\alpha$ and $1-\alpha$ to $P_{F}$ and $P_{B}$ respectively. (Right) The performance gain with $\alpha$-GFN objectives in Set Generation pan2023better. With flexible exploration-exploitation trade-offs enabled by $\alpha$, GFlowNet training achieves significantly higher average reward of all generated samples.
• Figure 2: Spearman correlations ($P_F^{\top}$ vs. $R$) for $\alpha$-GFNs under (Left) DB, (Center) FL-DB, and (Right) TB objectives in large sets.
• Figure 3: Entropy of the forward policy $P_F$ for $\alpha$-GFNs under (Left) DB, (Center) FL-DB, and (Right) TB objectives in large sets. Increasing $\alpha$ reduces entropy, indicating a shift toward stronger reward-exploitation, whereas decreasing $\alpha$ promotes exploration.
• Figure 4: A case study on (a) the mode curves of $\alpha$-FL-SubTB, (b) the loss curves of $\alpha$-FL-SubTB($\lambda$) and (c) the Spearman correlation curves of $\alpha$-DB in Molecule Generation during training. The stage 1 and stage 2 in Alg. \ref{['alg:two-staged-training']} are marked in the figures.
• Figure 5: Average sample length and forward policy entropy with $\alpha$-FL-SubTB in Molecule Generation. We observe that both the length and the entropy increases with $\alpha$ in this case.

Theorems & Definitions (32)

• Proposition 3.1
• Theorem 3.2
• Definition 3.3: $\alpha$-GFN objectives
• Proposition 3.4
• Proposition 3.5: Gradient of $\alpha$-GFN objectives, proof in App. \ref{['app:proof for gradient of alpha-gfn objectives']}
• Definition 1.1
• Definition 1.2
• Definition 1.3
• Definition 1.4
• Definition 1.5