Avoid What You Know: Divergent Trajectory Balance for GFlowNets

TL;DR

This paper introduces ACE, a new exploration algorithm for GFlowNets that enhances their ability to discover diverse high-reward states by explicitly training an exploration GFlowNet, leading to better approximation and diversity.

Contribution

The paper proposes Adaptive Complementary Exploration (ACE), a novel method that trains an exploration GFlowNet to improve the exploration of high-reward regions in GFlowNets.

Findings

ACE outperforms prior methods in distribution approximation accuracy.

ACE increases the discovery rate of diverse high-reward states.

Experiments demonstrate significant improvements in exploration efficiency.

Abstract

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. 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 significantly improves upon prior work in terms of approximation accuracy to the target distribution and discovery rate of diverse high-reward states.