Source-linked AI summary
A Variational Perspective on Generative Flow Networks
Heiko Zimmermann, Fredrik Lindsten, Jan-Willem van de Meent, Christian A. Naesseth
TL;DR
GFNs aim to sample composite objects according to reward-defined target distributions, while common training objectives match forward and backward trajectory models. This paper formulates GFN training with KL-based variational objectives, relates them to trajectory balance, and reports theoretical equivalences and numerical comparisons on synthetic tasks.
Problem
GFNs approximate reward-defined target distributions, and their usual flow-matching or trajectory-balance training objectives match forward and backward transition models over trajectories.
Source: §1 · pp. 1–2Method
The paper defines variational objectives using forward and reverse KL divergences, combines them convexly, and studies variance-reduction techniques for trajectory-balance gradients.
Source: §1 · p. 2, §3.1 · p. 5, §4.2 · pp. 7–8Results
In certain settings, optimizing trajectory balance is equivalent to optimizing forward or reverse KL, and the paper numerically compares the variational and trajectory-balance objectives on synthetic tasks.
Source: §1 · p. 2, §6 · p. 11Takeaways & Limitations
Variational-inference techniques, including control variates, can provide alternative objectives and variance-reduction approaches for training GFNs.
Source: §4.2 · pp. 7–8Score-function gradient estimators can exhibit high variance, making additional variance-reduction techniques often essential for stochastic-gradient learning.
Source: §3 · p. 5
Abstract
from arXiv · showhide
Generative flow networks (GFNs) are a class of models for sequential sampling of composite objects, which approximate a target distribution that is defined in terms of an energy function or a reward. GFNs are typically trained using a flow matching or trajectory balance objective, which matches forward and backward transition models over trajectories. In this work, we define variational objectives for GFNs in terms of the Kullback-Leibler (KL) divergences between the forward and backward distribution. We show that variational inference in GFNs is equivalent to minimizing the trajectory balance objective when sampling trajectories from the forward model. We generalize this approach by optimizing a convex combination of the reverse- and forward KL divergence. This insight suggests variational inference methods can serve as a means to define a more general family of objectives for training generative flow networks, for example by incorporating control variates, which are commonly used in variational inference, to reduce the variance of the gradients of the trajectory balance objective. We evaluate our findings and the performance of the proposed variational objective numerically by comparing it to the trajectory balance objective on two synthetic tasks.
1 Introduction
GFNs use reward-weighted trajectory sampling to approximate target distributions, connecting their training to probabilistic inference. This work relates trajectory balance to forward and reverse KL objectives and proposes a convex-combination objective, including study of variance reduction and synthetic-task performance.
1 Introduction
- Malkin et al. (2022b) proposed trajectory balance to address ineffective credit propagation in earlier GFN objectives; later work explored partial trajectories and optimal transport.Madan et al. (2022) proposed partial-trajectory optimization, while Do et al. (2022) proposed an optimal-transport-based objective.Source: §1 · pp. 1–2
- The proposed variational objectives connect GFN sampling to inference over a target distribution represented by a positive reward function.A positive reward function can be viewed as an unnormalized distribution to sample from.Source: §1 · p. 2
- The paper shows that, in certain settings, optimizing trajectory balance is equivalent to optimizing forward or reverse KL, and proposes a convex combination of the divergences.It also evaluates the theoretical findings and variational objective against trajectory balance on synthetic tasks.Source: §1 · p. 2
- Prior work connected hierarchical variational autoencoders to a special class of GFNs, with forward-KL training corresponding to ELBO maximization.The equivalence concerns the full backward and forward transition models.Source: §1 · p. 2
- The paper distinguishes its contribution from concurrent work by proposing a convex combination of forward and reverse KL and studying energy-based models and shared-parameter settings.Malkin et al. (2022b) independently derived equivalences between trajectory balance and forward and reverse KL.Source: §1 · p. 2
2 Generative Flow Networks
GFNs define trajectory sampling on a DAG so that terminal-state probabilities are proportional to reward. Their flow and transition models provide the framework for training with trajectory balance.
2 Generative Flow Networks
- GFNs sample trajectories through a DAG from a root to a terminal state, with the terminal-state distribution targeted to be proportional to its reward.The DAG's partial order means trajectories do not contain loops.Source: §2 · p. 2
2.1 Trajectory Flows
- A trajectory flow assigns non-negative mass to complete trajectories, inducing state and edge flows and forward and backward transition probabilities.For Markovian flows, transition probabilities depend on the current state rather than the preceding partial trajectory.Source: §2.1 · pp. 2–3
2.2 Training Generative Flow Networks
- GFNs parameterize Markovian flows with forward transitions and a normalizing constant; the backward model may be fixed or separately parameterized.For reward R, the terminal backward transition is targeted to R(sT)/Z; trajectory balance is optimized over sampled trajectories.Source: §2.2 · pp. 3–4
3 Variational Inference
The section frames matching GFN forward and backward trajectory distributions as variational inference, minimizing forward or reverse KL divergences. It then describes score-function variance reduction, including control variates that can connect variational estimators to training with trajectory balance.
3 Variational Inference
- Forward and reverse KL divergences provide alternative objectives for minimizing differences between GFN forward and backward transition probabilities.The divergences can be estimated with samples from the forward and backward models, respectively; backward-model samples are often unavailable and may require approximations.Source: §3 · p. 4
- Cross-parameter derivatives of the KL objectives require score-function gradients, which can have high variance during stochastic-gradient learning of variational approximations.Derivatives of reverse KL with respect to forward-model parameters and forward KL with respect to backward-model parameters involve score-function gradients.Source: §3 · pp. 4–5
3.1 Variance reduction techniques for score-function estimators
- Control variates can preserve a gradient estimator’s expectation while reducing its variance through an appropriate scaling.The scaling can be chosen to minimize the modified estimator’s variance.Source: §3.1 · p. 5
- The score function of the forward model is a useful, easy-to-compute control variate when optimizing reverse KL.With this control variate, the scaling can be incorporated into the negative log-importance weight.Source: §3.1 · p. 5
- Estimating optimal control-variate scaling with the same samples requires leave-one-out estimates to keep the gradient estimator unbiased.The optimal estimate requires per-sample gradients and can require S forward-backward passes; simpler non-optimal scaling choices avoid access to gradient information.Source: §3.1 · pp. 5–6
- Trajectory-balance training can leverage variational score-function variance-reduction techniques in settings where the objectives are equivalent.Source: §3.1 · p. 6
4 Variational Inference for Generative Flow Networks
The section compares trajectory-balance (TB) training with reverse- and forward-KL variational objectives under forward, backward, or mixed trajectory sampling. It identifies parameter settings where TB gradients correspond to KL gradients and describes a convex combination of the divergences.
4 Variational Inference for Generative Flow Networks
- A convex combination weights reverse- and forward-KL objectives by (1 − α) and α, respectively.The section compares this variational objective with TB objectives under different sampling and parameter-sharing settings.Source: §4 · pp. 6–7
4.1 Forward model and backward model with shared parameters
- The analysis distinguishes shared-parameter models from models with separate forward and backward parameters when comparing TB and KL gradients.The shared-parameter case uses η = φ = θ.Source: §4 · p. 7, §4.1 · p. 7
4.2 Forward model and backward model with distinct parameters
- With distinct forward and backward parameters, optimizing TB for the forward model is equivalent to reverse-KL optimization, while its backward-model gradient is equivalent to a forward-KL gradient.The score-function estimators use a learned scaling parameter.Source: §4.2 · pp. 7–8
5 Experiments
The experiments compare variational and trajectory-balance objectives on synthetic densities and an Ising model, including tests with control variates. They examine negative log-likelihood and expected log-weight, alongside qualitative samples.
5 Experiments
- The experiments represent terminating states as binary vectors and construct them sequentially from partially specified states.The state space includes an empty-bit value for positions not yet sampled.Source: §5 · p. 9
5.1 Synthetic densities
- At α = 1, αTB performs significantly worse than αKL on the fixed-energy synthetic task; for 0 < α < 1, their negative log-likelihoods are similar, with αTB slightly ahead.At α = 0, the objectives perform comparably within one standard deviation.Source: Figure 1 · p. 10, §5.1 · p. 10
- The synthetic-density experiments test both joint learning of the energy function and GFN and GFN learning with a previously learned fixed energy function.They evaluate αTB and αKL with learned and LOO control variates.Source: §5.1 · pp. 9–10
5.2 Ising model
- For the Ising model, lacking ground-truth samples, the experiments use α = 0, where αTB and αKL optimization are equivalent, to compare learned and LOO baselines.The study reports expected log-weights and shows approximate samples from the trained GFN and MH chains.Source: §5.2 · p. 11, Figure 2 · p. 11
- Across Ising models with different β, the learned baseline and leave-one-out (LOO) control variate show no significant performance difference.Expected log-weights are averaged over ten trained GFNs.Source: §5.2 · p. 11, Table 2 · p. 11
6 Conclusion
The conclusion connects GFN training to variational inference by expressing objectives through divergences between forward and backward trajectory distributions. It reports comparable expected log-weight performance in synthetic-density and Ising evaluations and identifies further variational methods to explore.
6 Conclusion
- GFNs can be trained by minimizing forward, reverse, or convex-combination KL divergences between forward and backward trajectory distributions.The reverse-KL objective is analogous to standard variational inference; forward KL yields an objective used in wake-sleep methods.Source: §6 · p. 11
- The reverse-KL gradient is proportional to the expected trajectory-balance gradient when trajectories are sampled from the forward distribution.Source: §6 · p. 11
- On synthetic densities and an Ising model, variational GFN objectives achieve comparable expected log-weight performance to trajectory-balance variants.Source: §6 · p. 11
- The authors identify credit-assignment and importance-sampling methods, including variational sequential Monte Carlo and nested variational inference, as opportunities to explore for GFNs.Source: §6 · p. 11