Source-linked AI summary
Generative Flow Networks for Discrete Probabilistic Modeling
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, Yoshua Bengio
TL;DR
High-dimensional discrete distributions can have separated modes that are difficult for MCMC to mix between. EB-GFN jointly trains a GFlowNet sampler and an energy function, using compositional generation and GFlowNet-based proposals for mode exploration. The authors report competitive results across synthetic and real tasks.
Problem
High-dimensional discrete distributions may have rapidly growing, well-separated modes, making mixing between modes slow for MCMC methods.
Source: §1 · p. 1Method
EB-GFN jointly trains a GFlowNet sampler and an energy function from data, with the sampler learning to model the energy distribution.
Source: §1 · p. 1Results
EB-GFN achieved competitive results on a variety of synthetic and real tasks.
Source: §1 · p. 2Takeaways & Limitations
As GFlowNet training converges, its MCMC proposal approaches block Gibbs sampling and can make large jumps by exploiting learned compositional structure across modes.
Source: §6 · p. 8Contrastive-divergence-type algorithms use initialized previous states rather than true data at every step, producing biased gradient estimates that may not faithfully model the data distribution.
Source: §2.2 · p. 3
Abstract
from arXiv · showhide
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.
1. Introduction
EB-GFN jointly trains an energy function and a GFlowNet sampler for high-dimensional discrete data. Its compositional generation policy can support mode exploration beyond local MCMC search, and the paper reports competitive results on synthetic and real tasks.
1. Introduction
- High-dimensional discrete distributions can have rapidly growing, separated modes that slow MCMC mixing and contribute to spurious modes in learned energy distributions.Source: §1 · p. 1
- The method jointly trains a GFlowNet sampler and an energy function from data, adapting GFlowNets from querying a fixed energy function to learning from datasets.Source: §1 · p. 1
- When underlying data regularities are learnable, compositional generation may help the sampler guess unvisited modes without traversing long low-probability paths or requiring known a priori structure.Source: §1 · pp. 1–2
- A GFlowNet-based MCMC proposal uses the learned compositional policy to make efficient large jumps with low rejection probability.Source: §1 · p. 2
- The proposed joint-training procedure specifies conditions under which it estimates the true data log-likelihood gradient for the energy parameters.Source: §1 · p. 2
- EB-GFN achieved competitive results across a variety of synthetic and real tasks.Source: §1 · p. 2
2. Preliminaries
GFlowNets represent object generation as trajectories through a state DAG, with terminal-state probabilities matched to rewards. Energy-based models define distributions through energies, but exact normalization and negative sampling can be costly.
2.1. GFlowNets
- Objects are built from an initial state through discrete actions in a DAG, and multiple trajectories may reach the same terminal object.Source: §2.1 · p. 2
- Reward matching trains terminal probabilities to be proportional to nonnegative rewards; the forward policy samples trajectories and their terminal objects.Source: §2.1 · p. 2
- Flow learning can use statewise conservation errors, while alternative objectives parameterize log Z, a forward policy, and a backward policy.Source: §2.1 · p. 2
- For a Markovian flow, the route chosen after reaching a state is independent of the route used to reach it.Source: §2.1 · p. 2
- The approach uses the backward policy for local exploration in a contrastive-divergence-like algorithm.Source: §2.1 · p. 3
2.2. Energy-based models
- An EBM assigns probability proportional to exp(−Eφ(x)); finite-space flexibility leaves exact normalization potentially exponentially costly.Source: §2.2 · p. 3
- MLE's negative-sample term requires model samples; K-step MCMC contrastive divergence approximates updates and recovers the true gradient only as K→∞, assuming chain mixing.Source: §2.2 · p. 3
- CD-type methods can be computationally efficient but biased, while persistent CD reuses previous chain states; these methods may not model the true data faithfully.Source: §2.2 · p. 3
3. Methodology
The method represents discrete objects as partial constructions and trains GFlowNets to match a target distribution. It couples this sampler to an energy model, using GFlowNet-generated proposals and negative samples to train the energy model.
3.1. GFlowNet generative process
- The GFlowNet state is a partial binary vector with void entries; each action assigns one void entry, and every complete trajectory takes D actions.Different trajectories correspond to different orders of assigning the entries.Source: §3.1 · p. 3
3.2. GFlowNet training towards a target distribution
- Trajectory balance trains the GFlowNet so that, when its objective is globally minimized to zero, terminal-state probabilities are proportional to rewards.Training trajectories may be sampled forward from the current policy or backward from data examples; the backward samples visit states near data that may be poorly explored by forward sampling and could stabilize optimization.Source: §3.2 · p. 4
- The method uses a mixture of forward and data-conditioned backward trajectories for GFlowNet training, with a uniform backward policy as a canonical choice that enforces a maximum-entropy property.The paper reports this choice is beneficial in some domains.Source: §3.2 · p. 4
3.3. Interleaved updates of GFlowNet and energy
- A K-step proposal removes entries from x with the backward policy, then uses the forward policy to complete a new sample, changing at most K entries.With a perfectly reward-matching GFlowNet, the Metropolis–Hastings acceptance probability is one; vanilla block Gibbs instead requires exponential-in-K computation to evaluate and normalize all 2^K changes.Source: §3.3 · p. 5
- Alternating updates train the GFlowNet toward the energy-defined reward and train the energy model with approximate maximum-likelihood gradients using GFlowNet negatives.If the GFlowNet is perfectly trained, its terminal distribution matches the energy distribution, making the approximate MLE objective an unbiased estimate and amortizing MCMC sampling computation.Source: §3.3 · pp. 5–6
4. Experiments
Across Ising, synthetic, and discrete-image tasks, EB-GFN is evaluated against sampling baselines and shows strong qualitative and quantitative results. The experiments also test whether it can recover structure from data and model multimodal distributions.
4.1. Ising models
- For the Ising test, the method fits a symmetric matrix Jφ and a GFlowNet from samples alone, then evaluates matrix discrepancy by RMSE.The method has no access to the true data-generating matrix; the energy is parameterized by a single matrix rather than a deep model.Source: §4.1 · p. 6
- On Ising models, EB-GFN has an advantage over Gibbs and GWG, while Figure 4 shows faithful samples and near-complete recovery of the interaction matrix from 2000 samples.The learned matrix has 4950 or 4851 degrees of freedom, depending on the grid.Source: Figure 4 · p. 6, §4.1 · p. 7
4.2. Synthetic tasks
- On multimodal synthetic tasks such as checkerboard and 8gaussians, GFlowNets capture modes and their structure better than the baselines.The authors hypothesize that learned regularities let the GFlowNet infer unvisited modes from those already encountered.Source: §4.2 · p. 7
- EB-GFN outperforms baselines on all seven synthetic datasets and metrics except MMD on 2spirals and 8gaussians, where it exceeds comparable-size baselines.Table 2 reports NLL and MMD; ALOE+ has a parametrization thirty times larger than ALOE and EB-GFN.Source: §4.2 · p. 7, Table 2 · p. 7
4.3. Discrete image modeling
- The image experiments model four binary datasets with an MLP energy function and compare PCD with Gibbs sampling and GWG.Training uses PCD-100 with a replay buffer and selects the checkpoint with the best validation NLL.Source: §4.3 · p. 8
- EB-GFN surpasses GWG on three of four discrete image datasets in test-set NLL.The experiments use four binary image datasets and report NLL per sample.Source: §4.3 · p. 8, Table 3 · p. 8
5. Related Work
The related work situates EB-GFN among energy-based modeling and discrete inference methods, and distinguishes its treatment of generation order from prior approaches.
Energy-based models
- Prior work improved EBM objectives, training stability, and sampling, while learning a sampler or proposal distribution was also explored; this work extends that direction to discrete data.The cited related work includes contrastive-divergence methods and approaches to learn samplers or proposals.Source: §5 · p. 8
Related methods
- Unlike fixed-order autoregressive models, GFlowNets generate sequentially; the paper frames learning generation order as joint inference of forward construction and reverse erasure processes.The authors describe this interpretation as the first of its kind.Source: §5 · p. 8
Discrete inference
- Discrete inference methods include continuous relaxation for local energy landscapes and local-search strategies for initializing negative samples.The passage identifies Grathwohl et al. (2021b) and Dai et al. (2020), respectively.Source: §5 · p. 8
6. Conclusion
The conclusion extends GFlowNets to jointly learn a sampler and energy function from data, and introduces a proposal that approaches block Gibbs sampling as training converges.
6. Conclusion
- The GFlowNet-based MH proposal approaches block Gibbs sampling as training converges and can make large state-space jumps by using learned compositional structure.The authors contrast these jumps with simple Gibbs sampling and their potential to generalize across modes.Source: §6 · p. 8
- Future work could iterate proposals from a trained GFlowNet for efficient exploration instead of sampling complete trajectories with K = D.Source: §6 · p. 8
- The approach adds the cost of training a sampler, which can be amortized when it will be reused because GFlowNet sampling is often cheaper than MCMC for covering modes.The paper presents this as a cost and a conditional amortization benefit.Source: §6 · pp. 8–9
A. Summary of GFlowNet notation
This section defines notation for the discrete data space, the construction graph, and the policies and flows used to describe GFlowNet generation.
A. Summary of GFlowNet notation
- X denotes the data space {0, 1}^D, while S is the state set, A the action or transition space, and T the set of complete trajectories.Source: A. Summary of GFlowNet notation · p. 12
- The construction graph G is a directed acyclic graph; s0 is the initial state and x denotes a terminal state.Source: A. Summary of GFlowNet notation · p. 12
- PF is the forward policy over children and PB the backward policy over parents; F denotes trajectory, state, or edge flow depending on its domain.The passage identifies state and edge flow as used in the proof of Proposition 1.Source: A. Summary of GFlowNet notation · p. 12
B. Proofs of propositions
The proofs establish two properties of reward-matching GFlowNets: an entropy-maximizing backward policy and always-accepted Metropolis-Hastings proposals under perfect reward fitting.
B. Proofs of propositions
- Among Markovian flows satisfying reward matching, the flow induced by a uniform backward policy has maximal entropy.The proof connects maximizing forward-flow entropy to maximizing expected backward-policy entropy; the maximum occurs when the backward policy is uniform over each state's parents.Source: B. Proofs of propositions · pp. 12–13
- The proof relates forward-policy entropy to expected backward-policy entropy over trajectories sampled from the forward policy.The displayed derivation invokes linearity of expectation and the definition of state flow.Source: B. Proofs of propositions · p. 13
- With perfect reward matching, the Metropolis-Hastings rejection step always accepts proposals.The proposition states acceptance probability is 1 for moves along reverse and forward trajectories.Source: B. Proofs of propositions · p. 13
C.1. Ising models
The Ising experiments report that EB-GFN's benefit is clearest for settings with σ < 0, where the lattice model has more than two modes.
C.1. Ising models
- The GFlowNet uses a four-layer MLP with 256 units per hidden layer; EB-GFN and baselines are trained with Adam at batch size 256, and baselines use 100 MCMC steps (PCD-100).Training stops when RMSE between ground truth and learned J reaches its minimum.Source: §C.1 · pp. 13–14
- The appendix includes visualizations of ground-truth synthetic samples and learned energy functions from baseline methods.Figure C.1 shows ground-truth samples; Figure C.2 visualizes baseline energy functions.Source: Figure C.1 · p. 14, Figure C.2 · p. 14
- EB-GFN's benefit is clearest when σ < 0; for σ > 0, the lattice Ising model has only two modes.The authors relate this pattern to GFlowNet's inductive bias for multimode distributions.Source: §C.1 · p. 14
C.2. Synthetic tasks
On synthetic tasks, EB-GFN is reported to outperform PCD and basic ALOE on most datasets, while ablations identify the back-and-forth proposal as important for optimization.
C.2. Synthetic tasks
- The synthetic data begin as 2D floating-point samples and are transformed into 16-bit Gray code.The authors describe the task as challenging even before the highly nonlinear Gray-code transformation.Source: §C.2 · p. 14
- EB-GFN remains state of the art and surpasses PCD and basic ALOE on all reported datasets except swissroll.The comparison is reported for exponential Hamming-kernel MMD and also for linear-kernel MMD; Table C.1 reports the latter across seven 2D synthetic problems.Source: §C.2 · p. 15, Table C.1 · p. 15
- The learned EBM has slightly worse average NLL and MMD than the GFlowNet, but remains competitive with other baselines.The comparison covers both models in the EB-GFN framework on seven 2D synthetic problems.Source: §C.2 · p. 15, Table C.2 · p. 15
- Removing backward trajectory samples causes only little performance harm on checkerboard and moons.On moons, NLL changes from 19.732 to 19.746 and Hamming exponential MMD from 0.121 to 0.342; on checkerboard, NLL changes from 20.696 to 20.709 and MMD from 1.206 to 2.648, with MMD in units of 1 × 10^-4.Source: §C.2 · pp. 15–16
- Removing the back-and-forth proposal makes EB-GFN's training loss quickly diverge on checkerboard and moons tasks.The ablation always sets K = D; the authors say the proposal trick is crucial to a reasonable optimization landscape.Source: §C.2 · pp. 15–16
- The learned backward policy deletes the highest-magnitude Gray-code bits first, which correspond to bits generated last by the forward policy.The appendix visualizes damaged samples to illustrate the erasure policy's structure.Source: §C.2 · p. 16
C.3. Discrete image modeling
For discrete image modeling, the appendix describes EB-GFN's training setup, reports a benefit from LayerNorm, and notes that removing either of two training techniques causes static-MNIST divergence.
C.3. Discrete image modeling
- In static MNIST, EB-GFN diverges without either backward-sample training or the back-and-forth proposal.The authors state that these techniques are both important in this task, partly reflecting its difficulty.Source: §C.3 · p. 16
- The EBM uses an MLP with three hidden layers of 256 units, while the GFlowNet uses an MLP with three hidden layers of 512 units.EBM likelihood validation uses 300,000-step annealed importance sampling; training uses a canonical backward policy and a proposal K increasing from 1 to D.Source: §C.3 · p. 16
- Adding LayerNorm after each forward-policy linear layer except the last greatly benefits generative-modeling performance.The appendix presents a negative-likelihood ablation for this change.Source: §C.3 · p. 16
- The EBM experiments omit a ResNet-18 backbone because Gibbs-With-Gradients took two weeks to finish with the public code.The stated runtime refers to the original GWG public-repository code.Source: §C.3 · p. 16
- Figure C.3 shows pinwheel samples progressively damaged by voiding bits and replacing them with random Bernoulli values.The visualization includes up to 31 steps of this process.Source: Figure C.3 · p. 17