PaperlayerView exact revision on arXiv

Source-linked AI summary

A theory of continuous generative flow networks

Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-García, Léna Néhale Ezzine, Yoshua Bengio, Nikolay Malkin

arXiv:2301.12594v2cs.LGstat.ML

TL;DR

GFlowNets had been restricted to discrete spaces despite applications involving continuous variables. This paper develops a generalized theory for arbitrary, including hybrid, spaces and finds that key discrete-GFlowNet advantages transfer to continuous settings.

  • Problem

    GFlowNets were restricted to discrete spaces, although probabilistic inference and modeling problems such as molecule design and Bayesian structure learning involve continuous variables.

  • Method

    The paper extends GFlowNet training objectives to arbitrary spaces using measurable pointed graphs, Markov kernels, and generalized flow-matching, detailed-balance, and trajectory-balance conditions.

  • Results

    Experiments confirm that advantages established for discrete GFlowNets transfer to general state spaces, including greater stability and better target fitting for trajectory balance when off-policy exploration matters.

  • Takeaways & Limitations

    The theory widens GFlowNets toward probabilistic inference and modeling problems involving continuous or hybrid choices, including Bayesian posteriors, molecular conformers, and inverse problems.

  • Takeaways & Limitations

    The practical theory requires care with change-of-variables arguments because omitting the Jacobian term can make the transformation invalid, and prior continuous-GFlowNet attempts violated critical assumptions.

Abstract

from arXiv · show

Generative flow networks (GFlowNets) are amortized variational inference algorithms that are trained to sample from unnormalized target distributions over compositional objects. A key limitation of GFlowNets until this time has been that they are restricted to discrete spaces. We present a theory for generalized GFlowNets, which encompasses both existing discrete GFlowNets and ones with continuous or hybrid state spaces, and perform experiments with two goals in mind. First, we illustrate critical points of the theory and the importance of various assumptions. Second, we empirically demonstrate how observations about discrete GFlowNets transfer to the continuous case and show strong results compared to non-GFlowNet baselines on several previously studied tasks. This work greatly widens the perspectives for the application of GFlowNets in probabilistic inference and various modeling settings.

1. Introduction

GFlowNets offer off-policy learning and multimodal sampling for compositional distributions, but existing theory has been restricted to discrete spaces. This paper extends their theory to arbitrary spaces and reports experiments showing that discrete-case advantages transfer to continuous and hybrid settings.

  • 1. Introduction: Continuous extensions matter because applications such as molecule design and Bayesian structure learning require modeling continuous variables alongside compositional structure.Examples include molecular torsion angles and parameters of graphical models.
  • 1. Introduction: Prior work showed continuous GFlowNet losses could provide off-policy benefits, but a sound theory for continuous actions remained lacking and a later integral extension relied on invalid assumptions.The paper identifies these gaps as motivation for its generalized formulation.
  • 1. Introduction: The paper extends GFlowNet training objectives from discrete spaces to arbitrary spaces using measurable pointed graphs based on Markov kernels.It generalizes flow matching, detailed balance, and trajectory balance, while recovering discrete losses as special cases.
  • 1. Introduction: Experiments validate the theory across domains with discrete, continuous, and hybrid action spaces and examine caveats specific to continuous settings.The comparative experiments are intended to test both theoretical claims and transfer from discrete GFlowNets.

2. Background and related work

GFlowNets build on sequential sampling and flow-based objectives, contrasting with conventional sampling methods through learned policies and off-policy exploration. Their discrete formulation learns transition probabilities whose induced terminal distribution matches the target when training reaches a global optimum.

  • 2. Background and related work: Related samplers include MCMC, sequential Monte Carlo, annealed importance sampling, learnable-kernel methods, and stochastic normalizing flows, but they differ in learned-policy structure or objective.Stochastic normalizing flows share the use of Markov kernels and Radon-Nikodym derivatives with the generalized theory.
  • 2. Background and related work: Discrete GFlowNets train a sequential sampler for an unnormalized target reward using local flow-matching consistency rather than ordinary reward maximization.They were initially framed as reinforcement learning methods with discrete states and actions.
  • 2. Background and related work: Applications of discrete GFlowNets include biological design, Bayesian structure learning, and robust combinatorial optimization.These methods target compositional spaces such as causal graphs and combinatorial objects.
  • 2. Background and related work: The forward policy induces a terminal-state distribution, and stochastic updates may use on-policy trajectories or off-policy modifications that encourage exploration.At a global minimum, the learned forward policy can represent the target terminal distribution.
  • 2. Background and related work: Detailed-balance and trajectory-balance objectives introduce a backward policy to avoid summing likelihoods over potentially exponentially many trajectories.This auxiliary distribution specifies parents of each state in the DAG.

3. A theory for generalized GFlowNets

The paper generalizes GFlowNet theory from discrete DAGs to measurable pointed graphs, extending its training conditions and sampling guarantees to continuous and hybrid spaces. The framework depends on explicit structural and density assumptions, and its zero-loss objectives recover reward-proportional terminating-state samples.

  • Training objectives: The generalized theory extends detailed-balance and trajectory-balance losses by replacing discrete probability mass functions with probability density functions.The framework also provides a dictionary connecting generalized and discrete GFlowNets.
  • Assumptions: Soundness requires every state to be reachable from the source, bounded trajectory lengths, and learned measures expressible as densities over states rather than actions.These conditions are automatic in suitable discrete pointed DAGs but must be verified in infinite spaces.
  • Previous continuous formulations: The previous continuous flow-matching extension can fail because its action-to-state change of variables omits a Jacobian term that need not equal 1.The omission is invalid in general, although it can hold in special cases such as additive transitions in Euclidean spaces.
  • Structured state space: Measurable pointed graphs generalize discrete DAGs by representing accessible states with transition kernels and defining flows and rewards relative to reference measures.The reference kernel encodes the state-space structure, while the reference measure supports densities for flows, rewards, and policies.
  • Theoretical guarantees: Detailed balance implies flow matching, while trajectory balance implies both flow matching and reward matching in the generalized framework.These propositions connect the local and trajectory-level objectives to the conditions needed for correct terminating-state sampling.
  • Theoretical guarantees: Zero FM, DB, reward-matching, or trajectory-balance losses imply their corresponding generalized consistency conditions, and flow plus reward matching yields sampling proportional to the reward measure.Theorem 1 states that recursively sampling the learned forward kernel from the source produces the normalized reward over terminating states.

4. Experiments

The experiments test generalized GFlowNets across synthetic continuous, stochastic-control, Bayesian structure-learning, and diffusion-model settings. Results extend several discrete-GFlowNet observations to continuous or hybrid spaces and show competitive performance against relevant baselines.

  • 4.2. Low-dimensional stochastic control: Generalized GFlowNets learn discretizations of stochastic differential equations for sampling black-box target densities.The forward policy is a conditional density over successive Euclidean states and becomes a T-step Euler–Maruyama discretization when Gaussian.
  • 4.2. Low-dimensional stochastic control: TB is more stable and fits targets better than objectives requiring importance weighting when off-policy exploration matters, while on-policy TB and reverse KL perform similarly when stable.TB remains competitive with PIS despite not using reward-density gradients.
  • 4.4. Posterior over continuous parameters in Bayesian structure learning: The Bayesian structure-learning experiment evaluates both graph-edge marginals and continuous parameter samples against exact posterior quantities using variational-inference baselines.The GFlowNet is reported as significantly more accurate on the discrete component.
  • 4.5. Connections with diffusion models: MLE-GFN is applied to DDPM training on ImageNet-32, where Table 4 compares the method with DDPM using FID and NLL.The generalized framework is therefore evaluated beyond fitting a sampler to a reward function.

5. Conclusion

The paper develops a measurable-space framework that contains finite-state GFlowNets as a special case. Its transition-kernel formulation supplies the measure-theoretic machinery needed to define trajectories, terminating-state measures, and reverse dynamics.

  • 5. Conclusion: The framework defines n-step transition and terminating-state measures induced by a forward kernel, enabling continuous analogues of trajectory and terminal distributions.Absolute continuity of one-step forward kernels transfers to n-step measures.
  • 5. Conclusion: Reverse kernels provide a construction route for backward reference kernels, with existence guaranteed under stated Polish-space and topological conditions.The paper notes that reverse kernels need not generally be unique, while Lemma 3 supplies a guaranteed construction under its assumptions.
  • 5. Conclusion: Finite pointed DAGs are special cases of finitely absorbing measurable pointed graphs.A finite graph becomes a measurable pointed graph under the discrete topology, with transition kernels counting trajectories.

C. Experimental details

The experimental evaluation approximates the Jensen–Shannon divergence between a learned terminating-state sampler and a normalized target reward distribution. It uses KDE-based discretizations of samples from both distributions.

  • C. Experimental details: The target distribution is sampled by rejection sampling with a uniform proposal, then a KDE is fit to those samples.A second KDE is fit to samples from the learned terminating-state measure.
  • C. Experimental details: Both KDEs are evaluated on a fixed discretization grid, normalized into probability mass functions, and compared using JSD.This procedure approximates the divergence between the learned sampler and normalized reward distribution.

C.2. A synthetic continuous environment

The synthetic continuous environment uses forward and backward policies defined as densities relative to reference kernels on a measurable pointed graph. Its implementation combines Beta-mixture parameterizations with JSD evaluation.

  • C.2. A synthetic continuous environment: The forward and backward kernels are represented by densities p_F and p_B with respect to reference kernels κ and κ_b.This supplies the density parameterization used by the continuous environment.
  • C.2. A synthetic continuous environment: The topology is the disjoint union of the source and sink states with the continuous state space, and all states are reachable within a bounded number of steps.These are among the structural conditions used for the measurable pointed graph.
  • C.2. A synthetic continuous environment: The initial forward density uses four-component Beta mixtures for radius and angle, while transition densities use two-component Beta mixtures over state-dependent angular intervals.The forward network also outputs a termination probability, with Jacobian corrections accounting for transformed Beta supports.
  • C.2. A synthetic continuous environment: The reward is a density with respect to Lebesgue measure on X=(0,1)^2, and training uses 128 sampled trajectories per iteration for 20,000 iterations.Both TB and DB use learning-rate annealing during optimization.
  • C.2. A synthetic continuous environment: Performance is evaluated by approximating JSD between the learned terminating-state distribution and the target distribution.The approximation follows the KDE-and-grid procedure described in the experimental details.

C.3. Low-dimensional stochastic control

The experiments compare continuous GFlowNet estimators of log-partition functions on Gaussian-mixture and funnel densities, including exploration and off-policy variants. Results indicate that off-policy forward-KL and trajectory-balance methods continue improving with training, whereas on-policy methods can collapse modes.

  • Experimental setup: The Gaussian-mixture and funnel experiments use exploration noise annealed from ϵ=0.1 to 0 and distinct density scales σ=5 and σ=1.The policy and log-partition parameters use separate learning rates of 10^-2 and 10^-1, respectively.
  • Results: Models trained for 1500 batches learn samples from a nine-Gaussian target density, while longer training improves the precision of exploratory TB policies.Figure C.1 compares the target density with samples from several algorithms.
  • Results: Off-policy forward-KL and TB algorithms continue improving with longer training, unlike on-policy algorithms that collapse modes and stop discovering new density regions.All models were trained for 1500 batches, while higher learning rates caused unstable optimization.
  • Methods: The study estimates log-partition functions using simple and importance-weighted bounds, with extended results reported in Table C.1.The estimates are defined from trajectories, terminating states, and products of forward or backward Gaussian policy densities; both are exact for a perfect sampler.

C.4. Stochastic control on a torus environment

The torus experiments test generalized GFlowNets on continuous multimodal and molecular-conformation reward densities. The learned terminating distributions closely match both targets, with JSD 0.063 for the synthetic task and 0.009 for alanine dipeptide.

  • Method: The GFlowNet parameterizes forward and backward policies with mixtures of independent von Mises distributions to represent angular variables on the torus.Inputs use trigonometric features of both angles to account for the torus topology.
  • Molecule conformation task: Alanine dipeptide’s torsion angles 𝜑 and 𝜓 vary continuously over [0, 2π], whereas 𝜃1 and 𝜃2 are constrained near 0 or π by energy barriers.The experiments generate 𝜑 and 𝜓 while keeping the other coordinates fixed.
  • Results: JSD is 0.063 for the synthetic multimodal task and 0.009 for the molecule conformation task, showing that generalized GFlowNets model densities on non-Euclidean spaces.The divergence compares the learned terminating-state distribution with the normalized reward distribution.

C.5. Posterior over continuous parameters in Bayesian structure learning

This section uses one GFlowNet over a mixed state space to approximate Bayesian-network graphs together with their continuous parameters. The construction alternates discrete graph growth with continuous parameter sampling and evaluates both components against exact posteriors.

  • GFlowNet over a mixed state space: A single GFlowNet approximates the joint posterior P(G, 𝜃|D) over DAGs and their associated continuous parameters, unlike approaches that update parameters with Variational Bayes.Terminating states have the form (G, 𝜃), with 𝜃 belonging to the parameter space determined by G.
  • GFlowNet over a mixed state space: The generation process first adds DAG edges from the empty graph, then samples parameters after selecting a special stop action.Intermediate states are (G, ♯), while terminating states are (G, 𝜃).
  • Transition model: Forward transitions either add an edge to the graph or sample parameters from a Normal distribution whose mean and diagonal covariance come from a neural network.The reference kernel combines discrete graph transitions with a continuous measure over parameters, while the backward kernel is discrete.
  • Experimental setup: Datasets contain 100 observations generated from Erdős-Rényi DAGs with normally sampled edge parameters, while the GFlowNet observes only the resulting data.Additional graph metrics include directed-path and Markov-blanket marginals, RMSE, and Pearson correlation.
  • Evaluation: The evaluation compares graph marginals and parameter samples with exact posteriors in a linear-Gaussian model, using edge-marginal RMSE and parameter negative log-probability.Lower parameter negative log-probability means samples are more likely under the exact posterior.

D. Additional lemmas and propositions

The additional results establish technical properties needed for generalized GFlowNet theory on measurable pointed graphs. They connect trajectory measures, flow matching, reward matching, and source–sink flow to the total reward.

  • Proof strategy: The proofs use finite absorption, absolute continuity, continuity of transition probabilities, and strong induction over measurable state layers.These assumptions and decompositions ensure finiteness and establish the stated flow identities.
  • Support properties: The support of the backward kernel from the sink is the closure of the terminating-state space X.The proof uses positive reference measure and the positive probability of transitioning from terminating states to the sink.
  • Trajectory-measure lemmas: The lemmas characterize finite-horizon trajectory measures through induction and establish properties of forward and backward kernels on measurable pointed graphs.These results support later propositions by controlling n-step measures and the accessibility structure of states.
  • Flow and reward matching: Under reward matching, the source and sink flows equal the total reward R(X), so the sink mass satisfies 𝜇({⊥})=R(X).This links the generalized flow construction to the partition-function interpretation used in discrete GFlowNets.
  • Relation to prior theory: A proposition generalizes a prior GFlowNet lemma from Bengio et al. (2021b) to measurable pointed graphs and is used to prove a main proposition.The extension is stated as crucial for the generalized theory’s proof sequence.

E. Proofs of results in the main text

The proofs establish that zero training loss for FM, DB, RM, or TB yields the corresponding generalized GFlowNet consistency conditions. They also show that detailed balance implies flow matching, trajectory balance implies flow and reward matching, and reward matching produces a valid terminating-state distribution.

  • Theorem 1: Reward matching makes the terminating-state measure a probability measure satisfying the target-reward relation.The proof verifies flow matching and reward matching separately, then applies the resulting identity to f(x)=1 to establish normalization.
  • Proposition 1: Detailed balance implies flow matching, because applying the detailed-balance identity to suitable bounded measurable functions recovers the flow-matching conditions.The proof constructs g(s,s′)=f(s′) with f(s0)=0 and then identifies the resulting integral equality with flow matching.
  • Proposition 2: Trajectory balance implies both flow matching and reward matching for a flow whose state measure is defined from the trajectory-balance normalizer.The resulting flow uses μ({⊥})=μ({s0})=Z and satisfies both condition families with respect to R.
  • Theorem 2: Zero FM, DB, RM, or TB loss almost surely implies the corresponding flow, detailed-balance, reward-matching, or trajectory-balance conditions.The TB implication holds for every trajectory length n≥0, with the initial normalizer Zν({s0}).
  • Theorem 2: The proofs repeatedly convert density equalities under reference measures into measure-level GFlowNet conditions using absolute continuity and kernel-measure identities.For DB and TB, the arguments use the equivalence of reference-measure products and replace densities with their corresponding measures.

F. Proofs of lemmas and propositions in §A

The appendix proofs establish positivity and support properties of the reference measure and transition kernels through induction and the reverse-kernel definition. These properties ensure that reachable sets receive positive measure and that later measure-theoretic constructions are well defined.

  • Reverse-kernel arguments: Sets on which the backward kernel assigns positive termination probability have zero νκ-measure under the stated assumptions, completing the appendix support argument.The proof decomposes the relevant integral and shows both terms vanish, yielding νκ(B)=0.
  • Support and positivity: Every nonempty measurable target set is reachable from the initial state within at most N steps, so some κ^n(s0,B) is positive.This is the stated equivalent form of the support condition used in the appendix.
  • Termination: Induction shows that transition kernels after sufficiently many steps terminate at ⊥, since κ^n(s0,−)=δ⊥ for n≥N.This reduces later arguments to trajectories that terminate within the bounded horizon.
  • Support and positivity: The reference measure ν is strictly positive as a consequence of the finite-step reachability property.The appendix derives this directly after establishing positivity of an appropriate transition iterate.
  • Reverse-kernel arguments: The reverse-kernel construction transfers positivity between forward and backward transitions while preserving the relevant measure-zero sets.The proof applies the reverse-kernel definition to open sets and bounded measurable functions, then uses induction hypotheses on κ^n.
Loading 2301.12594v2…