GFlowNets and variational inference

1 INTRODUCTION

• Hierarchical VI and GFlowNets are closely connected despite differing foundations and typical applications to continuous-valued distributions and discrete structures, respectively.Both model richer distributions by chaining simple conditional choices.Source: §1 · p. 1
• Special cases of VI and GFlowNet objectives coincide in expected gradients, including hierarchical and nested VI with trajectory and detailed balance objectives.The paper also identifies gradient-variance reduction and off-policy learning without importance reweighting as differences between the families.Source: §1 · p. 1
• Experiments provide a hierarchical VI baseline for discrete problems previously studied with GFlowNets and use the comparison to examine similarities and differences.The experiments also test off-policy learning and its relevance to multimodal targets.Source: §1 · pp. 1–2
• Exploratory off-policy sampling benefits GFlowNets on discrete probabilistic modeling tasks, especially when the target distribution has many modes, without introducing high gradient variance.The paper contrasts this with VI's importance-sampling-based off-policy learning.Source: §1 · p. 2

2.1 GFLOWNETS: NOTATION AND BACKGROUND

• A GFlowNet samples trajectories through a forward policy over DAG states, with a learned backward policy and partition-function estimate used in trajectory balance.On-policy training samples trajectories from the forward policy; off-policy training samples them from another behavior distribution, such as a tempered forward policy or the backward policy.Source: §2.1 · p. 2

2.2 HIERARCHICAL VARIATIONAL MODELS AND GFLOWNETS

• Hierarchical models and GFlowNets correspond by mapping sequence variables to DAG states and conditional distributions to forward-policy transitions.The correspondence extends to non-graded DAGs by inserting dummy states to make the graph graded.Source: §2.2 · p. 3
• VI objectives compare learned and target trajectory distributions, but tractable unbiased gradient estimators may be unavailable when the backward-policy normalization is unknown.Some objectives, including forward KL, require importance-weighted estimators; reweighting can introduce variance that grows with the discrepancy between the backward policy and sampling distribution.Source: §2.2 · p. 4

2.3 ANALYSIS OF GRADIENTS

• Special cases of hierarchical and nested VI correspond in expected gradients to trajectory-balance and detailed-balance objectives, with an extension to subtrajectory balance.The subtrajectory-balance equivalence extends to a nested variational objective over subtrajectories.Source: §2.3 · p. 4
• For on-policy trajectories, the forward-policy update under reverse-KL VI with a global baseline matches the trajectory-balance update.The global baseline is updated like log Z in trajectory balance, which provides a control variate for the VI score-function gradient.Source: §2.3 · p. 5
• With a fixed backward policy and on-policy sampling, trajectory balance and reverse-KL VI differ in how they reduce gradient variance; learning the backward policy or sampling off-policy introduces further differences.The paper notes that REINFORCE gradients for reverse KL can have high variance and that a baseline can reduce it without biasing the estimator.Source: §2.3 · p. 5

3 RELATED WORK

• Black-box VI uses stochastic gradient estimators such as REINFORCE, while path-wise estimators can have lower variance but apply to a more limited range of distributions.Prior work also develops estimators to reduce REINFORCE variance and combines score-function and path-wise approaches for particular distribution families.Source: §3 · p. 5
• Beyond ELBO and KL objectives, VI studies tighter evidence bounds, other divergences, and Stein discrepancy; the paper views trajectory balance as a novel pseudo-divergence criterion.The related work notes mean-seeking behavior and posterior variance underestimation as issues in KL optimization.Source: §3 · p. 6
• Wake-sleep methods alternate optimization of generative and inference models, while GFlowNets have been applied to molecule discovery, biological design, and Bayesian structure learning.Prior work had mentioned and expanded the theoretical connection between GFlowNets and variational methods; a concurrent paper explored interpolations between forward and reverse KL.Source: §3 · p. 6

4 EXPERIMENTS

• Across the experiments, on-policy VI and TB can behave similarly, but HVI becomes more stable near convergence and is sensitive to learning rates.When exploration matters, off-policy TB outperforms on-policy TB and VI without the possible high gradient variance from importance sampling.Source: §4 · p. 6
• Off-policy TB is the best method among those tested for accurately fitting targets and finding modes, especially in molecule generation and causal graph discovery.The paper reports these as challenging domains where mode finding is particularly important.Source: §4 · p. 6

4.1 HYPERGRID: EXPLORATION OF LEARNING OBJECTIVES

• On the hypergrid, TB converges to lower JSD, finds all four modes, and models them precisely; forward-KL methods find modes faster but reach a local minimum without modeling them precisely.Reverse-KL methods can achieve low loss without finding every mode.Source: §4.1 · p. 7, Figure 1 · p. 7

4.2 MOLECULE SYNTHESIS

• On molecule synthesis, off-policy TB gives the best, lowest-variance fit, while off-policy REVERSE KL performs poorly against its on-policy counterpart, especially for smoother, more diverse targets.On-policy TB and REVERSE KL differ little in performance, though REVERSE KL is more learning-rate sensitive.Source: §4.2 · p. 8, Figure 2 · p. 8

4.3 GENERATION OF DAGS IN BAYESIAN STRUCTURE LEARNING

• As Bayesian-network graphs grow, off-policy (Modified) DB and TB remain accurate, while off-policy REVERSE KL degrades significantly in JSD to the target posterior.On-policy TB and REVERSE KL also degrade, but less; their performance is similar across settings.Source: Table 2 · p. 9, §4.3 · p. 9

5 DISCUSSION AND CONCLUSIONS

• Off-policy GFlowNet objectives may be advantageous for highly multimodal targets, balancing mode-seeking REVERSE KL and mean-seeking FORWARD KL.Source: §5 · p. 9
• Choosing the behavior policy for off-policy training remains a research direction for efficiently exploring and discovering modes.Source: §5 · p. 9
• The experiments focus on discrete variables; continuous-action GFlowNets as alternatives to VI in continuous-variable domains remain for future investigation.Source: §5 · p. 9

A CANONICAL CONSTRUCTION OF A GRADED DAG

• The canonical conversion inserts dummy states to grade a DAG, and applying it to a graded DAG leaves it unchanged.Source: Figure A.1 · p. 14, A CANONICAL CONSTRUCTION OF A GRADED DAG · p. 14

B PROOFS

• The proof defines a trajectory-level log ratio and uses it to express the forward- and backward-policy gradients of the hierarchical VI objectives.The proof identifies the two divergences as forward KL and reverse KL, respectively.Source: B PROOFS · p. 14
• The expected on-policy Trajectory Balance gradient does not depend on the estimated partition function Z.The proof states this as an immediate corollary for the Trajectory Balance parameter Z.Source: B PROOFS · p. 15

C.1 A VARIATIONAL OBJECTIVE FOR SUBTRAJECTORIES

• SubNVI defines objectives over partial trajectories between selected junction layers, using positive state-flow functions to define paired distributions over those trajectories.The SubTB objective reduces to Detailed Balance for length-one subtrajectories and to Trajectory Balance for complete trajectories.Source: C A VARIATIONAL OBJECTIVE FOR… · p. 15, §C.1 · p. 16
• When the hub-state flow distributions agree at every junction, the forward policy's terminating-state distribution matches the target reward distribution.The proof obtains the Trajectory Balance constraint for every complete trajectory as an intermediate condition.Source: §C.1 · pp. 16–17
• SubNVI matches NVI when every layer is a junction layer, and matches HVI when only the source and terminating layers are junction layers.These choices correspond respectively to K = L and K = 1.Source: §C.1 · p. 17

C.2 AN EQUIVALENCE BETWEEN THE SUBNVI AND THE SUBTB OBJECTIVES

• For specified f-divergences, the subtrajectory variational objectives are equivalent to the Subtrajectory Balance objective on a graded DAG with junction layers.The proposition states the equivalence for any forward and backward policies and positive hub-state flow functions.Source: §C.2 · p. 17

D.1 HYPERGRID EXPERIMENTS

• The hypergrid has non-terminating states, increment actions, and an exit action that leads to a terminating state.R0 controls exploration difficulty: lower values indicate harder exploration.Source: §D.1 · p. 18
• Figure D.1 compares local and global baselines for the three HVI algorithms using score-function gradient estimators.The comparison includes both on-policy and off-policy runs.Source: Figure D.1 · p. 19, §D.1 · p. 19
• HVI objectives were brittle to hyperparameter choices, so the authors tuned settings separately for each objective and sampling mode.The tuning procedure used Bayesian search and JSD-based selection.Source: §D.1 · p. 19
• On 64×64 and 8×8×8×8 grids with R0 = 0.1, REVERSE KL and REVERSE WS reach lower JSD than FORWARD KL and WS, and are on par with TB.This contrasts with the harder 128 × 128 setting, where R0 = 10−3.Source: §D.1 · p. 19

D.2 MOLECULE EXPERIMENTS

• The molecule experiments used the stated reward model, held-out molecule set, and GFlowNet model, and trained for up to 50,000 batches.Metrics were computed every 5,000 batches; some runs stopped early because of numerical errors.Source: §D.2 · pp. 19–20
• The figure plots the evolution of JSD between the learned sampler and a target distribution.The supplied caption does not provide further details of the comparison.Source: Figure D.2 · p. 20
• The figure shows average final learned distributions alongside the target distribution, with shaded standard errors across five runs.The plotted distributions include the 8 × 8 × 8 × 8 and 64 × 64 grids.Source: §D.2 · p. 20

D.3 BAYESIAN STRUCTURE LEARNING EXPERIMENTS

• The state space consists of DAGs built by adding acyclic edges, with a stop action that returns the current graph as a sample.The experiments use a fixed backward policy that is uniform over the graph's parents.Source: §D.3 · pp. 20–21
• The target posterior is approximated using a reward equal to the joint distribution over graphs and data, with a uniform graph prior.The experiments use linear-Gaussian models, for which the marginal likelihood is available in closed form.Source: §D.3 · p. 20
• Off-policy Modified DB and TB perform well across settings on edge-marginal RMSE; off-policy REVERSE KL degrades most as complexity increases.On-policy TB and REVERSE KL perform similarly and also degrade as the number of variables increases.Source: Figure D.3 · p. 21, §D.3 · p. 22
• Edge-marginal accuracy is measured by RMSE between target-posterior and approximate-posterior edge probabilities, with lower values indicating better performance.An edge marginal quantifies the probability that an edge is present in the Bayesian-network structure.Source: Figure D.3 · p. 21, §D.3 · p. 22

• When the state space is small enough, a recursively defined flow function lets the method propagate flows to compute the terminating-state distribution.Source: E METRICS · p. 22
• Dynamic programming in topological state order computes the full terminating-state distribution in time linear in |S| + |A|.Source: E METRICS · p. 22
• When the target distribution can be evaluated exactly, the Jensen-Shannon divergence is directly evaluable and is chosen as a symmetric divergence.Source: E METRICS · p. 22

• The experiment tests whether observations from discrete-space GFlowNets also hold in continuous settings using stochastic control.Source: F EXTENSION TO CONTINUOUS DOMAINS · p. 23
• The agent starts at (0, 0) and takes 10 Gaussian transitions over [0, 1], with learned means and fixed isotropic variance, to produce a final point.Source: F EXTENSION TO CONTINUOUS DOMAINS · p. 23
• The learned drift determines the final-point density, which is trained toward a reward proportional to an eight-Gaussian target with modes 2 units from the origin and standard deviations 0.25.Source: F EXTENSION TO CONTINUOUS DOMAINS · p. 23
• REVERSE WS and FORWARD KL are omitted because their gradients reach NaN values early in training, even with a 10× lower learning rate.Source: F EXTENSION TO CONTINUOUS DOMAINS · p. 23
• Off-policy exploration is essential for finding modes in the continuous task, and trajectory balance fits the target better.WAKE-SLEEP finds all modes but models them less precisely, with off-centre, oblong shapes and slightly higher MMD.Source: F EXTENSION TO CONTINUOUS DOMAINS · p. 23, Figure F.2 · p. 24
• The authors conjecture, but do not prove, that the necessary GFlowNet theory extends to probability densities; the experiments provide supporting evidence.Source: F EXTENSION TO CONTINUOUS DOMAINS · p. 23
• Figure F.1 shows intermediate-position KDEs matching the target closely and includes sampled points and trajectories from the trained model.Source: Figure F.1 · p. 24
• Figure F.2 compares learned marginal densities across algorithms and exploration policies and reports Gaussian-kernel MMD estimates from 2560 samples.Source: Figure F.2 · p. 24