Better Training of GFlowNets with Local Credit and Incomplete Trajectories

Abstract

Generative Flow Networks or GFlowNets are related to Monte-Carlo Markov chain methods (as they sample from a distribution specified by an energy function), reinforcement learning (as they learn a policy to sample composed objects through a sequence of steps), generative models (as they learn to represent and sample from a distribution) and amortized variational methods (as they can be used to learn to approximate and sample from an otherwise intractable posterior, given a prior and a likelihood). They are trained to generate an object x through a sequence of steps with probability proportional to some reward function R(x) (or exp(−E(x)) with E(x) denoting the energy function), given at the end of the generative trajectory. Like for other RL settings where the reward is only given at the end, the efficiency of training and credit assignment may suffer when those trajectories are longer. With previous GFlowNet work, no learning was possible from incomplete trajectories (lacking a terminal state and the computation of the associated reward). In this paper, we consider the case where the energy function can be applied not just to terminal states but also to intermediate states. This is for example achieved when the energy function is additive, with terms available along the trajectory. We show how to reparameterize the GFlowNet state flow function to take advantage of the partial reward already accrued at each state. This enables a training objective that can be applied to update parameters even with incomplete trajectories. Even when complete trajectories are available, being able to obtain more localized credit and gradients is found to speed up training convergence, as demonstrated across many simulations.

Emergence of Functional Circuits in the Early Visual Pathway

“…Partial inference is a promising paradigm to resolve this issue by incorporating local credits (Pan et al, 2023a). Specifically, the partial inference aims to evaluate individual transitions or sub-trajectories, i.e., local credits, and provide informative training signals for identifying the specific contributions of actions.…”

“…To sample from the Boltzmann distribution, GFlowNet trains the policy to assign action selection probability based on energy of terminal state (Bengio et al, 2021a;b;Malkin et al, 2022a), e.g., a high probability to the action responsible for the low terminal energy. However, such training has fundamental limitations in credit assignment, as it is hard to identify the action responsible for terminal energy (Pan et al, 2023a). This limitation stems from solely relying on the terminal energy associated with multiple actions, lacking the information to identify the contribution of individual actions, akin to challenges in RL with sparse reward (Arjona-Medina et al, 2019;Ren et al, 2022).…”

“…An attractive paradigm to tackle this issue is partial inference (Pan et al, 2023a) that trains flow functions with local credits, e.g., evaluation of the intermediate states or transitions. Such local credit identifies individual action contributions to the terminal energy before reaching the terminal state.…”

“…Such local credit identifies individual action contributions to the terminal energy before reaching the terminal state. To this end, Pan et al (2023a) proposed a forward-looking GFlowNet (FL-GFN), which assigns the local credit based on the energy of incomplete objects associated with intermediate states.…”

“…We extensively validate LED-GFN on various tasks: set generation (Pan et al, 2023a), bag generation (Shen et al, 2023), molecular discovery (Bengio et al, 2021b), RNA sequence generation (Jain et al, 2022), and the maximum independent set problem (Zhang et al, 2023). We observe that LED-GFN (1) outperforms FL-GFN when the assumption of intermediate energy does not hold, (2) excels in practical domains compared to GFlowNets and RL-based baselines, and (3) achieves similar performance to FL-GFN even when intermediate energy provides the "ideal" local credit.…”