Generative Flow Networks for Discrete Probabilistic Modeling (2202.01361v2)

Abstract: We present energy-based generative flow networks (EB-GFN), a novel probabilistic modeling algorithm for high-dimensional discrete data. Building upon the theory of generative flow networks (GFlowNets), we model the generation process by a stochastic data construction policy and thus amortize expensive MCMC exploration into a fixed number of actions sampled from a GFlowNet. We show how GFlowNets can approximately perform large-block Gibbs sampling to mix between modes. We propose a framework to jointly train a GFlowNet with an energy function, so that the GFlowNet learns to sample from the energy distribution, while the energy learns with an approximate MLE objective with negative samples from the GFlowNet. We demonstrate EB-GFN's effectiveness on various probabilistic modeling tasks. Code is publicly available at https://github.com/zdhNarsil/EB_GFN.

• The paper introduces a novel non-autoregressive model that frames discrete data generation as a sequential process managed by a GFlowNet.
• It presents a new MCMC proposal mechanism that enables large transitions with minimal rejection by leveraging learned compositional structures.
• Joint training of the energy function and GFlowNet yields competitive likelihood improvements and efficient mixing across benchmark tasks.

Generative Flow Networks for Discrete Probabilistic Modeling

The paper introduces a novel approach to probabilistic modeling in high-dimensional discrete spaces using Energy-Based Generative Flow Networks (EB-GFN). Building upon the foundations of generative flow networks (GFlowNets), the authors propose a method that models the generation process via a stochastic policy, effectively amortizing the costly exploration traditionally associated with Markov Chain Monte Carlo (MCMC) methods. This approach allows for approximate large-block Gibbs sampling and facilitates efficient mixing between distribution modes, addressing a key limitation of existing MCMC techniques.

Key Contributions

The paper makes several important contributions:

1. Non-Autoregressive Sequential Generation Model: It frames the generation of high-dimensional discrete data as a process managed by a GFlowNet. This model leverages the compositional nature of the data space, allowing for efficient exploration and generation of discrete structures.
2. Novel MCMC Proposal: A new method is introduced for generating MCMC proposals using GFlowNets, enabling large transitions with minimal rejection probability. This mechanism exploits the learned compositional structure to enhance sampling efficiency.
3. Joint Energy Function and GFlowNet Training: The authors propose a framework for concurrently training the GFlowNet with an energy function. This integrated training allows the GFlowNet to draw samples from the energy distribution and inform the energy-based model's learning process.
4. Empirical Validation: The effectiveness of EB-GFN is demonstrated across a range of probabilistic modeling tasks, indicating competitive performance against existing methods.

Methodology

The EB-GFN framework combines components from GFlowNets and energy-based models to achieve efficient probabilistic modeling. In this context, GFlowNets learn to generate data by incrementally constructing it, represented by a sequence of actions within a directed acyclic graph (DAG). The transition probabilities between states in this DAG are governed by a forward policy $P_F$, which is optimized to ensure distributional consistency with the target reward. Complementarily, a backward policy $P_B$ is employed to aid in effective exploration during training.

The energy-based component is parameterized as a function that assigns likelihoods to data structures, encouraging the GFlowNet to sample states reflecting this distribution. Training consists of iterative updates where each component—GFlowNet and energy model—refines its parameters iteratively. Training trajectories are sampled using a balanced mix of forward and backward exploration, stabilized by the trajectory balance objective.

Results and Implications

The empirical evaluations conducted in the paper underscore the potential of EB-GFN in handling complex, multimodal distributions. By providing mechanisms for efficient discovery and traversal of the solution space, EB-GFN demonstrates improvements in model likelihood and sample quality across several benchmark datasets. This method not only presents practical advantages in sampling efficiency but also opens theoretical avenues for expanding flow-based models into challenging high-dimensional discrete domains.

Future Work

The integration of EB-GFNs into broader applications remains an open frontier. Potential future work could explore the scalability of these methods to even larger state spaces or investigate adaptive strategies to dynamically adjust sampling and energy evaluation mechanisms. Additionally, the interplay between learned structure and sampling efficacy suggested by the results could inform novel architectures for GFlowNets, potentially benefiting other areas of machine learning reliant on efficient sampling.

Overall, this paper provides a substantial contribution to the field of probabilistic modeling, delivering innovative pathways for dealing with the combinatorial complexity inherent in discrete data spaces. The development of EB-GFN sets a foundation for future explorations into the utility of GFlowNets within discrete generative modeling contexts.