Dynamic Backtracking in GFlowNets: Enhancing Decision Steps with Reward-Dependent Adjustment Mechanisms

TL;DR

This work tackles the forward-only exploration bias in Generative Flow Networks (GFNs) by introducing Dynamic Backtracking GFN (DB-GFN), a reward-guided mechanism that backtracks during Markov-flow construction to correct suboptimal decisions. The method defines a dynamic backtracking process with adaptive step counts, and three backtracking selectors—Reward-Choose, Pearson-Choose, and MH-Choose—to decide whether to accept alternative trajectories. Across six biochemical and four genetic design tasks, DB-GFN achieves higher sample quality, more high-reward samples, faster convergence, and substantially stronger alignment between forward sampling probabilities and rewards (notably a near fourfold gain in Pearson correlation with LS-GFN). The approach is orthogonal to existing GFNs and RL methods, indicating strong potential for integration with other strategies to further boost search performance in high-dimensional design spaces.

Abstract

Generative Flow Networks (GFlowNets or GFNs) are probabilistic models predicated on Markov flows, and they employ specific amortization algorithms to learn stochastic policies that generate compositional substances including biomolecules, chemical materials, etc. With a strong ability to generate high-performance biochemical molecules, GFNs accelerate the discovery of scientific substances, effectively overcoming the time-consuming, labor-intensive, and costly shortcomings of conventional material discovery methods. However, previous studies rarely focus on accumulating exploratory experience by adjusting generative structures, which leads to disorientation in complex sampling spaces. Efforts to address this issue, such as LS-GFN, are limited to local greedy searches and lack broader global adjustments. This paper introduces a novel variant of GFNs, the Dynamic Backtracking GFN (DB-GFN), which improves the adaptability of decision-making steps through a reward-based dynamic backtracking mechanism. DB-GFN allows backtracking during the network construction process according to the current state's reward value, thereby correcting disadvantageous decisions and exploring alternative pathways during the exploration process. When applied to generative tasks involving biochemical molecules and genetic material sequences, DB-GFN outperforms GFN models such as LS-GFN and GTB, as well as traditional reinforcement learning methods, in sample quality, sample exploration quantity, and training convergence speed. Additionally, owing to its orthogonal nature, DB-GFN shows great potential in future improvements of GFNs, and it can be integrated with other strategies to achieve higher search performance.

Figures (7)

• Figure 1: A motivating experiment on TFbind8, a common nucleotide generation task for the performance assessment of GFNs. Here, the Pearson correlations between reward values and sampling probabilities are displayed, where $\log{p(x)}$ is the log-likelihood for forward sampling $P_F$. The results show a significant advancement in fitting the reward function with DB-GFN over LS-GFN, with all selection strategies of DB-GFN demonstrating about a fourfold increase in correlation compared to LS-GFN.
• Figure 2: The structure of DB-GFN is illustrated with terminal states $s_{f\{1..4\}}$, where $s_{f1}$ and $s_{f4}$ represent the initially generated terminal states. The size and shading of states indicate the magnitude of rewards, and the transparency and thickness of flows represent the probability magnitude. In this illustration, $s_{f4}$ with a smaller reward undergoes more backtracking, in this case, to $s_2$, and then transitions to new terminal states. After applying the Reward-Choose selection algorithm, one of three such algorithms considered, the newly generated terminal state $s_{f3}$ is less favorable than the original $s_{f4}$, so the original trajectory is kept. Similarly, $s_{f1}$ has a larger reward than $s_{f4}$ and undergoes fewer backtracking steps. The new outcome $s_{f2}$, produced after backtracking, is preferred in the selection algorithms, leading to an update and the choice of a new trajectory ($s_0\rightarrow...\rightarrow s_2\rightarrow...\rightarrow s_{f2}$) over the original one ($s_0\rightarrow...\rightarrow s_2\rightarrow...\rightarrow s_{f1}$).
• Figure 3: Accuracy of GFNs on QM9 and SEH tasks, with our DB-GFN model significantly outperforming other models.
• Figure 4: Accuracy comparison of GFNs on the TFbind8 and RNA1-3 tasks, where our DB-GFN model significantly outperforms the other models
• Figure 5: Accuracy for the TFbind8 and RNA tasks with DB-GFN represented by solid lines. Notably, integrating DB-GFN into the same models (indicated by matching colors) consistently yields better performance than both the original baselines and those augmented with LS-GFN, as demonstrated by the remarkable results indicated by the solid lines .