Energy-Based Generative Flow Networks

• EB-GFN is a probabilistic model combining energy-based models with trajectory-based flow samplers to approximate Boltzmann distributions efficiently.
• It leverages trajectory balance loss and negative-sample amortization to ensure diverse, multimodal sampling from high-dimensional, discrete spaces.
• Innovations like LED-GFN enable local energy decomposition and smoothness regularization for improved training stability and sample quality.

Energy-Based Generative Flow Networks (EB-GFN) are a probabilistic modeling paradigm that combines energy-based models (EBMs) with the sequential, trajectory-based sampling framework of generative flow networks (GFlowNets). The central objective is to train models capable of sampling high-dimensional, often discrete objects $x$ with marginal distribution proportional to the Boltzmann weight $\propto \exp(-E(x))$, where $E(x)$ is a parameterized energy function or neural potential. Unlike conventional MCMC-based approaches, EB-GFN amortizes the typically expensive negative-sampling and mixing behavior into a learned stochastic policy over a finite sequence of actions, enabling scalable and diverse sampling from complex, multimodal distributions (Zhang et al., 2022).

1. Mathematical Foundations

Energy-based generative flow networks formalize sampling from a discrete input space $\mathcal{X}$ using both an energy function $E_\theta(x)$ (EBM) and a flow-based sampler (GFlowNet). The energy model defines the unnormalized distribution

$p_\theta(x) \propto \exp(-E_\theta(x)),$

with intractable normalizer $Z_\theta = \sum_x \exp(-E_\theta(x))$.

A GFlowNet models the construction of $x$ as a trajectory $\tau = (s_0 \to s_1 \to \cdots \to s_D = x)$ over a directed acyclic graph, with $s_0$ the root and $\propto \exp(-E(x))$0 the terminal state. The GFlowNet defines a flow $\propto \exp(-E(x))$1 over trajectories, parameterized as forward and backward policies $\propto \exp(-E(x))$2, $\propto \exp(-E(x))$3, and scalar $\propto \exp(-E(x))$4, such that

$\propto \exp(-E(x))$5

The induced marginal distribution $\propto \exp(-E(x))$6 is trained to approximate $\propto \exp(-E(x))$7 (Zhang et al., 2022).

2. Training Objectives and Joint Algorithms

Training of EB-GFNs typically employs a joint optimization of both energy and flow parameters:

• Trajectory balance (TB) objective for the GFlowNet, requiring that for each $\propto \exp(-E(x))$8 terminating at $\propto \exp(-E(x))$9,

$E(x)$0

with $E(x)$1.

• Approximate MLE for the EBM, replacing the intractable expectation over $E(x)$2 with negative samples from $E(x)$3:

$E(x)$4

Algorithmic training alternates between GFlowNet and EBM steps, frequently using GFlowNet-guided large-block Gibbs moves to sample negatives efficiently—even allowing amortization of negative-sample generation into the learned forward policy $E(x)$5 (Zhang et al., 2022, Ekbote et al., 2022).

3. Energy Decomposition and Partial Inference

A principal innovation in EB-GFN methodology is the transition-based energy decomposition, notably the Learned Energy Decomposition GFlowNet (LED-GFN) variant (Jang et al., 2023). Instead of assuming terminal energy evaluation is always accessible and informative, LED-GFN posits a learnable decomposition: $E(x)$6 with $E(x)$7 parameterized as a neural network over state transitions. This reparameterization yields a set of flow-consistency constraints: $E(x)$8 where $E(x)$9 absorbs accumulated local potentials.

To enforce a stable, dense local credit, LED-GFN applies a smoothness regularization: $\mathcal{X}$0 and fits $\mathcal{X}$1 via a least-squares objective with Bernoulli-masked terms to avoid degenerate decompositions (Jang et al., 2023).

4. Variational and Flow-Contrastive Perspectives

The relation between GFlowNets and variational inference is captured through trajectory-wise KL divergence objectives. The forward KL ($\mathcal{X}$2), reverse KL ($\mathcal{X}$3), and mixed convex combinations are all interpretable as generative flow network losses, specifically trajectory balance (Zimmermann et al., 2022). Control variates and learned scalar baselines are deployable to reduce the variance of TB/variational-gradient estimators.

Alternatively, flow-contrastive estimation (FCE) establishes an adversarial joint training between energy-based and flow-based models, coupling their updates such that the EBM performs noise-contrastive estimation with the flow as a noise source, while the flow model approximately minimizes Jensen-Shannon divergence to the data (Gao et al., 2019).

5. Joint Modeling of Structured Outputs and Conditional Distributions

Energy-based GFlowNets have been effectively extended to the modeling of joint distributions over complex objects and attributes. For example, Joint Energy-Based GFlowNets (JEB-GFN) optimize a joint energy $\mathcal{X}$4 over pairs $\mathcal{X}$5, treating $\mathcal{X}$6 as GFlowNet reward, and train both sampler and energy synchronously. This joint modeling enables both unconditional and conditional sampling, avoids reward–sampler mismatch, and yields consistent marginals (Ekbote et al., 2022).

Empirical studies in applications such as antimicrobial peptide design demonstrate that JEB-GFN approaches produce higher-scoring, more diverse, and more novel candidates than either baseline GFlowNets with fixed predictors or semi-supervised variational autoencoders, especially in active learning settings (Ekbote et al., 2022).

6. Practical Considerations and Empirical Performance

EB-GFNs have been validated on a range of high-dimensional discrete tasks, including Ising models, multimodal synthetic densities, discrete image data (Omniglot, MNIST variants), molecular docking, RNA sequence design, and combinatorial optimization problems (Zhang et al., 2022, Jang et al., 2023). Key findings include:

• Significant gains in mode discovery, sample diversity, and coverage over standard GFlowNet, forward-looking GFlowNet, MCMC-based contrastive divergence, and other learned local samplers.
• Stability in training due to mixing forward and backward sampled trajectories and smoothness-regularized potentials.
• Efficient amortization of negative-sample generation, eliminating the necessity for lengthy MCMC.
• Robust matching or improvement over oracle-aided baselines even when intermediate energy estimates are available (“ideal local credit”) (Jang et al., 2023).

7. Limitations and Future Directions

EB-GFN methodologies are currently most mature for purely discrete domains. The challenges in scaling include computational costs of MCMC-based negative sampling for long trajectories, limitations of discrete transition modeling, and integration with continuous-variable spaces. Active research directions involve hybrid discrete–continuous proposal mechanisms, improving negative sampling efficiency, and theoretical analysis of sample complexity and convergence rates in joint energy/sampler optimization (Ekbote et al., 2022).

A plausible implication is that as these architectures advance, EB-GFNs may present a general-purpose, amortized alternative to conventional MCMC and EBMs, particularly in structured and combinatorial domains that resist both tractable likelihood modeling and efficient direct sampling.