Joint Energy & GFlowNet Training

• The paper introduces a framework where joint training couples energy functions with amortized GFlowNet policies, mitigating sparse reward issues in long-horizon tasks.
• It employs trajectory-balance and local credit assignment objectives to ensure scalable, efficient sampling from multimodal and combinatorial energy landscapes.
• This approach enhances mode coverage, accelerates convergence, and demonstrates practical impact in applications like molecular design, peptide generation, and structured output modeling.

Joint training of energy functions and Generative Flow Networks (GFlowNets) integrates the specification or learning of an energy-based reward model $\mathcal{E}(x)$ directly with amortized generative sampling over discrete or structured objects via flow-consistent policies. This paradigm, motivated by the need for reward-sampler compatibility, dense credit assignment, and improved mode coverage, has rapidly advanced from foundational EBM-GFlowNet schemes to local-credit and learned-decomposition methods capable of leveraging partial or implicit reward information.

1. Theoretical Foundations and Motivations

Generative Flow Networks are designed to sample composite structured objects $x$ through sequences of actions in a construction DAG so that the marginal probability of generating $x$ is proportional to a user-specified nonnegative reward $R(x) = \exp(-\mathcal{E}(x))$, with $\mathcal{E}(x)$ interpreted as an energy function (Bengio et al., 2021, Pan et al., 2023). GFlowNets generalize MCMC sampling and reinforcement learning by amortizing exploration with a learned flow, providing efficient discovery and sampling from highly multimodal or combinatorial energy landscapes.

In canonical settings, the reward $R(x)$ is computed only at terminal states, causing delayed and sparse credit assignment—analogous to the sparse-reward RL regime—which retards convergence, especially for long-horizon domains. This also introduces reward-sampler incompatibility when rewards are learned independently and then fixed for sampling (Ekbote et al., 2022). Joint training, wherein both $\mathcal{E}$ and the GFlowNet policy are updated alternately or in tandem, resolves these issues by tightly coupling the induced distribution and the sampler. This approach maintains high entropy, prevents mode collapse, and adapts gracefully to evolving data or adversarial negatives (Zhang et al., 2022, Zhang et al., 2022).

2. Energy Functions as GFlowNet Rewards

Joint energy–GFlowNet training is predicated on re-expressing the unnormalized target distribution

$\pi(x) = \frac{\exp(-\mathcal{E}(x))}{Z}$

in terms of a reward $R(x)=\exp(-\mathcal{E}(x))$, where $Z$ is the (intractable) partition function (Bengio et al., 2021, Zhang et al., 2022). The GFlowNet induces a flow $x$0 over the construction DAG (states and transitions) via flow-matching or detailed-balance constraints

$x$1

ensuring that, at convergence, the marginal distribution over terminal states matches $x$2. The mapping from energy (unnormalized log-density) to flow-parameterized probabilities makes the GFlowNet an efficient and scalable sampler even for highly multimodal or discrete domains.

Energy functions may be fixed, externally supervised, or—crucially—jointly learned (EBM-style) from positive and negative examples, the latter commonly drawn from a negative sampler that is itself the current GFlowNet policy (Zhang et al., 2022, Ekbote et al., 2022). This tight co-adaptation ensures reward-sampler compatibility and minimizes pathologies such as over-dispersed, mode-seeking, or miscalibrated sampling when reward supervision is insufficient or unrepresentative.

3. Joint Training Mechanisms: Objectives and Algorithms

A variety of loss functions couple the GFlowNet and energy function via trajectory-wise or local constraints.

Trajectory-Balance Objective: For a full trajectory $x$3,

$x$4

where $x$5 is a learned scalar. Minimizing the TB loss enforces that the terminating distribution of $x$6 is $x$7 (Zhang et al., 2022, Zhang et al., 2022, Ekbote et al., 2022).

Energy Function Learning: The energy model $x$8 is updated by a (possibly approximate) maximum-likelihood or contrastive-divergence objective:

$x$9

where $x$0 are negative samples produced by GFlowNet-guided MCMC chains and $x$1 denotes their marginal distribution (Zhang et al., 2022, Ekbote et al., 2022).

Alternating Update Routine:

1. Roll out trajectories from $x$2; update GFlowNet parameters $x$3 using $x$4 (or local variants).
2. Draw negatives via GFlowNet-guided block Gibbs or backward/forward trajectories; update $x$5 using contrastive-divergence gradients.
3. Optionally, update backward policy $x$6 and/or partition function estimator $x$7.

Local Credit and Forward-Looking GFlowNets: When $x$8 can be evaluated at intermediate states, the reparameterization $x$9 allows per-transition credit assignment (Pan et al., 2023). The local loss

$R(x) = \exp(-\mathcal{E}(x))$0

with $R(x) = \exp(-\mathcal{E}(x))$1, can be applied even to incomplete trajectories, enabling dense supervision at every step.

Learned Energy Decompositions: When intermediate energies are unavailable or uninformative, learned transition potentials $R(x) = \exp(-\mathcal{E}(x))$2 decompose $R(x) = \exp(-\mathcal{E}(x))$3. The flow losses and a regularized least-squares decomposition constraint are optimized jointly (Jang et al., 2023).

4. Local Credit Assignment and Training with Incomplete Trajectories

The principal technical advance of the forward-looking GFlowNet (FL-GFN) is the ability to backpropagate local credit assignments using the incrementally accrued energy along each transition in a trajectory. When $R(x) = \exp(-\mathcal{E}(x))$4 is additive or evaluable at intermediate states, this property enables:

• Per-step computation of credit $R(x) = \exp(-\mathcal{E}(x))$5.
• Training GFlowNet parameters using observed transitions from both complete and incomplete (prefix) trajectories (Pan et al., 2023).
• Dense, local gradient information, resulting in accelerated convergence and improved credit propagation, especially in long-horizon DAGs.

Similarly, learned energy decompositions make it feasible to proxy local energy increments by a parameterized per-edge potential $R(x) = \exp(-\mathcal{E}(x))$6, trained to sum (in expectation) to terminal $R(x) = \exp(-\mathcal{E}(x))$7, thereby obviating the need for explicit intermediate energy evaluation (Jang et al., 2023). Variance reduction and smoothness regularization (e.g., via dropout masking) ensure robust and informative potential assignments.

5. Extensions: Joint Distributions and Structured Outputs

Joint energy–GFlowNet training naturally extends to distributions over tuples $R(x) = \exp(-\mathcal{E}(x))$8 or more general structured objects (Ekbote et al., 2022). In "Joint Energy-Based GFlowNets" (JEB-GFN), the framework models

$R(x) = \exp(-\mathcal{E}(x))$9

with the GFlowNet trained to sample $\mathcal{E}(x)$0 according to $\mathcal{E}(x)$1. The energy model is jointly updated via contrastive divergence using negative pairs drawn from the current GFlowNet policy.

This joint objective resolves mismatches in two-stage pipelines (e.g., first training $\mathcal{E}(x)$2, then fitting a sampler), preventing drift into spurious high-reward regions, mode collapse, or slow adaptation under distribution shift. Empirically, JEB-GFN improves recovery of conditional and joint densities and significantly boosts diversity and discovery rates in combinatorial design settings such as peptide sequence generation (Ekbote et al., 2022).

Joint training also supports specialized fusion with structured generative models such as diffusion-based sequence/structure co-design, in which each diffusion step is cast as a GFlowNet state, and terminal binding energy is backpropagated via trajectory-balance constraints (Abir et al., 18 May 2025).

6. Empirical Impact and Practical Considerations

Empirical studies across GFlowNet, EBM-GFN, FL-GFN, LED-GFN, and JEB-GFN consistently find that joint training:

• Accelerates convergence (environment interactions are halved in FL-GFN relative to TB/DB losses).
• Increases mode coverage, top-$\mathcal{E}(x)$3 reward, and sample diversity (notably in set/graph generation, molecular design, and sequence modeling) (Pan et al., 2023, Jang et al., 2023, Zhang et al., 2022, Ekbote et al., 2022).
• Is robust to long horizons and incomplete supervision.
• Enables active learning and efficient exploration, outperforming RL and amortized MCMC baselines under similar computational budgets.

Practical aspects include:

• The need for efficient negative sampling from the current GFlowNet policy or learned backward policies.
• Careful tuning of joint loss weights and regularization (e.g., spectral penalties on energy parameters, variance reduction for potentials).
• Batch-based or replay-buffer training to stabilize learned decompositions and facilitate compositional generalization.
• Monitoring both training loss (e.g., TB or local DB) and held-out EBM NLL as convergence/control metrics (Zhang et al., 2022, Jang et al., 2023).

7. Limitations and Future Directions

Challenges and open questions remain regarding:

• Scalability to continuous or mixed discrete/continuous domains (most current applications are discrete combinatorial).
• The computational overhead and stability of joint training, especially with deep or highly parameterized energy models.
• Situations where intermediate energies are expensive, unavailable, or highly noninformative (addressed in part by learned decompositions).
• Extensions to fully online or model-based active learning scenarios, multi-task and multi-modal settings.

Recent advances, including energy-decomposed GFlowNets, reward fusion with denoising diffusion, and joint models over auxiliary variables, are expanding the scope and robustness of joint training paradigms in amortized probabilistic modeling (Jang et al., 2023, Abir et al., 18 May 2025, Ekbote et al., 2022). Theoretical results guarantee that flow-matching constraints sufficed by TB/DB losses, combined with compatible energy parameterizations, yield exact sampling of the Boltzmann-Gibbs distribution at convergence (Bengio et al., 2021, Zhang et al., 2022).

---

References

• Pan et al. "Better Training of GFlowNets with Local Credit and Incomplete Trajectories" (Pan et al., 2023)
• Liu et al. "Consistent Training via Energy-Based GFlowNets for Modeling Discrete Joint Distributions" (Ekbote et al., 2022)
• Bengio et al. "Unifying Generative Models with GFlowNets and Beyond" (Zhang et al., 2022)
• Bengio et al. "GFlowNet Foundations" (Bengio et al., 2021)
• Zhang et al. "Energy-Based Generative Flow Networks" (Zhang et al., 2022)
• Zou et al. "Learning Energy Decompositions for Partial Inference of GFlowNets" (Jang et al., 2023)
• Wu et al. "AbFlowNet: Optimizing Antibody-Antigen Binding Energy via Diffusion-GFlowNet Fusion" (Abir et al., 18 May 2025)