Avoid What You Know: Divergent Trajectory Balance for GFlowNets

Generative Flow Networks (GFlowNets) are a flexible family of amortized samplers trained to generate discrete and compositional objects with probability proportional to a reward function. To this end, they learn a policy function over an intractably large state graph by minimizing a stochastic objective over sampled trajectories. However, learning efficiency is constrained by the model’s ability to rapidly explore diverse high-probability regions during training. To mitigate this issue, recent works have focused on incentivizing the exploration of unvisited and valuable states via curiosity-driven search and self-supervised random network distillation, which tend to waste samples on already well-approximated regions of the state space. In this context, we propose Adaptive Complementary Exploration (ACE), a principled algorithm for the effective exploration of novel and high-probability regions when learning GFlowNets. To achieve this, ACE introduces an exploration GFlowNet explicitly trained to search for high-reward states in regions underexplored by the canonical GFlowNet, which learns to sample from the target distribution. Through extensive experiments, we show that ACE consistently and significantly improves upon prior work in terms of approximation accuracy to the target distribution and discovery rate of diverse high-reward states.

Many important problems in science and engineering involve searching through huge spaces of possibilities. For example, researchers may want to find useful biological sequences, promising molecules, or strong solutions to a complex design problem. Since testing every possibility is impossible, machine learning models are used to suggest candidates worth considering. A common problem is that these models can become too focused on the first good solutions they discover. They may keep producing similar candidates instead of continuing to explore other promising parts of the search space. We propose a new training method that helps the model look for valuable candidates it is currently missing. The main idea is simple: exploration should not just mean looking randomly, but looking where the model is still weak. In our experiments, this approach finds better and more diverse candidates faster than previous methods, including in tasks related to biological sequence design and optimization.