Learning Energy Decompositions for Partial Inference of GFlowNets

TL;DR

This paper introduces LED-GFN, a novel method that learns energy decompositions to improve partial inference in GFlowNets, enabling more efficient and stable sampling from complex distributions.

Contribution

We propose learning energy decompositions with potential functions to enhance GFlowNet training, addressing evaluation costs and fluctuations, while preserving optimal policies.

Findings

LED-GFN outperforms existing methods in multiple tasks

It effectively handles large energy fluctuations

It maintains optimal policy during training

Abstract

This paper studies generative flow networks (GFlowNets) to sample objects from the Boltzmann energy distribution via a sequence of actions. In particular, we focus on improving GFlowNet with partial inference: training flow functions with the evaluation of the intermediate states or transitions. To this end, the recently developed forward-looking GFlowNet reparameterizes the flow functions based on evaluating the energy of intermediate states. However, such an evaluation of intermediate energies may (i) be too expensive or impossible to evaluate and (ii) even provide misleading training signals under large energy fluctuations along the sequence of actions. To resolve this issue, we propose learning energy decompositions for GFlowNets (LED-GFN). Our main idea is to (i) decompose the energy of an object into learnable potential functions defined on state transitions and (ii) reparameterize the flow functions using the potential functions. In particular, to produce informative local credits, we propose to regularize the potential to change smoothly over the sequence of actions. It is also noteworthy that training GFlowNet with our learned potential can preserve the optimal policy. We empirically verify the superiority of LED-GFN in five problems including the generation of unstructured and maximum independent sets, molecular graphs, and RNA sequences.