Partial GFlowNet: Accelerating Convergence in Large State Spaces via Strategic Partitioning

Abstract: Generative Flow Networks (GFlowNets) have shown promising potential to generate high-scoring candidates with probability proportional to their rewards. As existing GFlowNets freely explore in state space, they encounter significant convergence challenges when scaling to large state spaces. Addressing this issue, this paper proposes to restrict the exploration of actor. A planner is introduced to partition the entire state space into overlapping partial state spaces. Given their limited size, these partial state spaces allow the actor to efficiently identify subregions with higher rewards. A heuristic strategy is introduced to switch partial regions thus preventing the actor from wasting time exploring fully explored or low-reward partial regions. By iteratively exploring these partial state spaces, the actor learns to converge towards the high-reward subregions within the entire state space. Experiments on several widely used datasets demonstrate that \modelname converges faster than existing works on large state spaces. Furthermore, \modelname not only generates candidates with higher rewards but also significantly improves their diversity.

• The paper introduces Partial GFlowNet, which uses overlapping action-masked regions, adaptive UCB-like action scores, and region-aware local search to make flow estimation more reliable in exponentially large spaces.
• Experiments show substantial gains: FM+PLS discovers 141,115 high-reward molecular modes versus 17,898 for vanilla FM, while TB+PLS finds 2,626 versus 8 for TB, with only about 1–10% training overhead.
• The method improves molecule, sequence, and RNA-binding generation, but its global convergence guarantees, sensitivity to masking and switching parameters, inter-action synergy modeling, and weak subTB performance remain open questions.

Motivation and problem statement

Generative Flow Networks (GFlowNets) learn a forward policy that samples terminal candidates $x$ with probability proportional to a reward $R(x)$, by estimating flows over a flow graph whose nodes are states and whose edges are actions. The size of this flow graph grows exponentially with trajectory length and polynomially with the branching factor, so in large state spaces the flow assignment becomes poorly conditioned and exploration stalls. The paper's motivating observation is stark: on the molecule design task of Bengio et al. with maximum trajectory length 8, GFlowNets identify more than 10,000 distinct modes with reward above 7.5, but when the trajectory length is extended to 12 — which strictly enlarges the state space and should only increase discoverable modes — existing methods find fewer than 10 such modes even after sampling $10^6$ candidates. This inversion is attributed to sparse high-reward regions: the actor becomes trapped in low-reward areas and cannot transition out.

Method

Partial search via action masking. The central idea is to restrict exploration to overlapping partial regions $\mathcal{R} \subseteq \mathcal{S}$. The paper decomposes each action into a state-irrelevant component $a^* \in \mathcal{A}^*$ (e.g., which building block or bit pattern to use) and a state-relevant component $a'$. A region is induced by sampling a valid subset $\mathcal{A}^*_v$, where each $a^*$ is valid with probability $p$; the region contains exactly the states reachable from $s_0$ using only valid actions. Since validity is applied at every step, the expected fraction of reachable states at depth $R(x)$0 shrinks as $R(x)$1, so modest values of $R(x)$2 yield regions small enough for reliable flow estimation. Setting $R(x)$3 recovers the full state space.

Planner. A planner with three modules governs region selection. The selecting module draws new regions proportionally to per-action scores, with an indicator $R(x)$4 trading off expected intersection against union size across consecutive regions — encouraging overlap that preserves useful flows while covering new territory. The updating module maintains UCB-like statistics, $R(x)$5, where HR accumulates rewards of trajectories containing $R(x)$6 and CNT counts occurrences. The decision module triggers a region switch when recent mode-discovery rates fall below historical averages, subject to a minimum dwell time; it is biased toward switching unless the current region is clearly productive.

Partial local search (PLS). Building on Local Search GFlowNets, the paper adapts trajectory reconstruction to constrained regions. After sampling a batch of complete trajectories under the masked forward policy, each trajectory is refined $R(x)$7 times: a $R(x)$8-step backward pass from the terminal state, a $R(x)$9-step forward reconstruction, then a full $10^6$0-step backtrack to $10^6$1 so the reconstructed trajectory is complete and consistent with the current region mask. Only reconstructions with higher reward than the original are retained for training. Training itself uses any standard objective — Flow Matching, Detailed Balance, Trajectory Balance, or Sub-Trajectory Balance — computed over the unrestricted action set, so the learned policy generalizes beyond the current mask.

Experimental results

The method is evaluated on three tasks with four configurations per objective: vanilla GFlowNet, +LS (local search), +P (partial search), and +PLS (partial local search), plus PPO and MARS baselines on molecule design.

On molecule design, where extending the building-block budget from 8 to 12 yields more than $10^6$2 terminal states, the results are pronounced:

Model	#Modes (R>7.5)	R_topk
DB	10	7.085
DB+PLS	78	7.606
FM	17,898	8.319
FM+PLS	141,115	8.508
TB	8	7.290
TB+PLS	2,626	8.029
subTB+PLS	14	7.402
PPO / MARS	0	7.360 / 7.322

FM+PLS discovers roughly eight times as many high-reward modes as vanilla FM, and TB+PLS improves over TB by more than two orders of magnitude. PPO and MARS find no modes above threshold at all. Notably, the relative benefit of partial search varies sharply by objective: it rescues TB and DB substantially but leaves subTB nearly non-functional (1–14 modes), suggesting the partitioning interacts unevenly with credit-assignment schemes — a point the paper does not analyze in depth.

On sequence generation ($10^6$3 states, 60 predefined modes, Levenshtein distance threshold 28), Partial GFlowNet variants improve mode discovery across all objectives, with TB-based configurations performing best. An ablation on synthetic target sequences built from patterns like 11111111 and 11110000 shows the planner's action scores converging, with the score of the 1111 block dominating as expected from its frequency in the targets — evidence that the updating module identifies genuinely relevant action components.

On RNA-binding (state space only $10^6$4, i.e., already tractable), partial search still helps: e.g., on RNA1, TB+PLS finds 36 modes versus 9 for TB+LS and 3 for vanilla TB, with top-$10^6$5 reward rising from 0.863 to 0.936. This indicates the gains are not solely an artifact of extreme state-space size.

Overhead. The planner adds roughly 1–10% training time per round depending on the objective (e.g., 1.27 → 1.39 units for FM), since its updates do not involve network evaluation.

Limitations and open questions

The paper concedes that its region construction treats state-irrelevant actions independently, ignoring synergistic effects between blocks — combinations of actions that jointly produce superior candidates are not captured by per-action scores. More broadly, several questions remain open: the sensitivity of performance to the validity probability $10^6$6 and the decision-module thresholds is not characterized; there is no theoretical guarantee that iterating over overlapping regions converges to sampling proportional to the global reward distribution rather than a region-biased mixture; and the failure of subTB under partitioning is unexplained. All experiments also lack error bars or significance tests, so run-to-run variability — particularly material given the authors' own report of seed-dependent failures among baselines — is not quantified.

Conclusion

This paper addresses the collapse of GFlowNet exploration in very large state spaces by constraining the actor to planner-selected partial regions defined through masks on state-irrelevant actions, combined with a region-aware variant of local search. The approach yields large empirical gains in both mode discovery and candidate quality on tasks with up to $10^6$7–$10^6$8 states, at negligible computational overhead, and remains beneficial even on smaller benchmarks. Its main open issues concern the absence of convergence guarantees for the resulting sampler, unexamined hyperparameter sensitivity, and the inability of the scoring mechanism to capture inter-action synergy.