Towards Understanding and Improving GFlowNet Training

TL;DR

This paper investigates the training of Generative Flow Networks (GFlowNets), introduces evaluation strategies, and proposes new methods to improve their sample efficiency and practical performance in complex tasks.

Contribution

It introduces an efficient evaluation strategy, clarifies the importance of learned flows, and proposes novel training techniques including prioritized replay, relative edge flow parametrization, and a guided trajectory balance objective.

Findings

Improved sample efficiency on biochemical design tasks.

Proposed methods enhance practical GFlowNet training.

Clarified the role of learned flows in generalization.

Abstract

Generative flow networks (GFlowNets) are a family of algorithms that learn a generative policy to sample discrete objects $x$ with non-negative reward $R(x)$. Learning objectives guarantee the GFlowNet samples $x$ from the target distribution $p^*(x) \propto R(x)$ when loss is globally minimized over all states or trajectories, but it is unclear how well they perform with practical limits on training resources. We introduce an efficient evaluation strategy to compare the learned sampling distribution to the target reward distribution. As flows can be underdetermined given training data, we clarify the importance of learned flows to generalization and matching $p^*(x)$ in practice. We investigate how to learn better flows, and propose (i) prioritized replay training of high-reward $x$, (ii) relative edge flow policy parametrization, and (iii) a novel guided trajectory balance objective, and show how it can solve a substructure credit assignment problem. We substantially improve sample efficiency on biochemical design tasks.