PaperlayerView exact revision on arXiv

Source-linked AI summary

Particle GFlowNets: Rethinking Generative Marginalization Models

Tiago da Silva, Diego Mesquita, Salem Lahlou

arXiv:2609.11538v1cs.LG

TL;DR

MaMs lack a unified interpretation within GFlowNets and are limited by factorized, autoregressive domains. The paper identifies them as permutation-conditioned GFlowNets and introduces Particle GFlowNets with persistent Gibbs sampling and R-hat-triggered rejuvenation. P-GFlowNets significantly accelerate wall-clock learning convergence, especially in large, long-horizon spaces, while open questions remain about optimal rejuvenation and diagnosing distributional accuracy.

  • Problem

    MaMs were considered distinct from GFlowNets and rely on factorized domains, leaving their applicability to non-factorized generative processes unresolved.

  • Method

    The paper interprets MaMs as permutation-conditioned GFlowNets and introduces Particle GFlowNets with stochastic reverse transitions, persistent Gibbs sampling, and R-hat-triggered chain rejuvenation.

  • Results

    P-GFlowNets significantly accelerate learning convergence in wall-clock time compared with conventional GFlowNet training algorithms.

  • Takeaways & Limitations

    The MaMs–GFlowNet equivalence extends the computational benefits of persistent Gibbs-based training beyond autoregressive modelling.

  • Takeaways & Limitations

    The paper leaves optimal rejuvenation criteria and reliable assessment of GFlowNet distributional accuracy as open questions.

Abstract

from arXiv · show

Generative Marginalization Models (MaMs) have been recently introduced as efficient neural sampling models for any-order autoregressive modelling of discrete distributions. By learning both the marginal and conditional probabilities of a persistent-block Gibbs sampler, MaMs enable fast posterior evaluation with a single neural network forward pass. While prior work has considered MaMs to be distinct from Generative Flow Networks (GFlowNets), a well-established paradigm for inference in discrete stochastic models, we show that they are equivalent. Then, we also extend MaMs' sampling strategy to non-autoregressive generative processes. In particular, we describe an automatic criterion for full-state rejuvenation of the Gibbs sampler, derived from the Gelman-Rubin statistic, which plays a key role in speeding up learning convergence. Our experiments show that our method, called Particle GFlowNets, markedly accelerates training in large combinatorial spaces.

1 INTRODUCTION

The paper reframes MaMs as permutation-conditioned GFlowNets and extends their persistent Gibbs-sampling strategy to non-autoregressive domains through Particle GFlowNets. It uses rejuvenation to address correlated chains and reports faster wall-clock convergence, especially when long-horizon MDP simulation is costly.

  • MaMs and GFlowNets: MaMs are equivalent to permutation-conditioned GFlowNets, with MaMs’ learning objective matching the expected detailed balance loss.This connection identifies MaMs as a computationally efficient GFlowNet training strategy rather than a separate paradigm.
  • Particle GFlowNets: Particle GFlowNets extend persistent Gibbs sampling beyond autoregressive models by using a stochastic reverse transition for non-injective forward mappings.The method retains detailed-balance optimization while accommodating non-deterministic backward moves.
  • Chain rejuvenation: Rejuvenating the persistent Gibbs process when Gelman-Rubin statistics exceed a threshold addresses chain correlation that reduces gradient-estimate accuracy and slows convergence.The reported rejuvenation criterion uses the R-hat statistic and is intended to refresh inefficient exploration during training.
  • Computational efficiency: P-GFlowNets asymptotically reduce the average forward passes per gradient step as the MDP horizon increases.This targets the repeated neural evaluations required by conventional trajectory-based GFlowNet training.
  • Empirical results: P-GFlowNets significantly reduce wall-clock time to learning convergence, particularly when MDP simulation costs more than target-distribution evaluation.The experiments cover established benchmarks and novel tasks, with the stated regime typical of large, long-horizon state spaces.

2 PRELIMINARIES

The preliminaries define compositional discrete spaces, reward-normalized target distributions, GFlowNet state graphs and policies, MaMs, and persistent-block Gibbs sampling. They distinguish tree-structured autoregressive processes from DAGs with multiple trajectories to the same state.

  • Notation: A compositional object is represented as a collection of components, with [K]^d providing a sequence-based example whose dimension is d.The same object can be viewed either as a set of indexed components or as a sequence.
  • State and reward definitions: The target distribution is π(x) ∝ R(x), where R is a positive reward function, and S contains partial component sets that belong to at least one complete object.The state space combines partial states S with terminal objects X, starting from the empty state.
  • State graphs: A generative process is autoregressive when its state graph is a tree rooted at the empty initial state; multiple trajectories to one state instead produce a non-tree DAG.The figure illustrates the distinction between unique and merged generation paths.
  • GFlowNets: A GFlowNet learns a forward policy over state-graph children so that the induced terminal-state marginal matches the reward-proportional target.For non-autoregressive graphs, a backward policy and importance sampling handle potentially many trajectories to the same terminal state.
  • MaMs: MaMs jointly learn marginal and conditional distributions for persistent-block Gibbs sampling, enabling marginal evaluation after training for arbitrary subsets of coordinates.Their consistency condition makes the Gibbs transition kernel stationary with respect to the learned marginal when satisfied.
  • Gibbs sampling: Persistent-block Gibbs sampling updates groups of variables in one step while retaining the chain state across neural-model gradient updates.Because the model changes during training, the persistent chain is not stationary until training converges.

3 RETHINKING MAMS AS GFLOWNETS

The paper establishes a bijective correspondence between MaMs and permutation-conditioned GFlowNets, showing that their marginal and conditional representations satisfy the same detailed-balance structure. This connection yields a constant-forward-pass consistency objective while preserving direct evaluation of any-subset marginals.

  • 3.1 Permutation-Conditioned GFlowNets: A PC GFlowNet fills sequence coordinates in a prescribed permutation order, with each nonterminal state having a unique predecessor under its deterministic reverse transition.The permutation controls navigation through partially filled sequences, while the underlying object space and reward remain unchanged.
  • 3.1 Permutation-Conditioned GFlowNets: The PC-GFlowNet flow at a partial state equals the corresponding marginal probability under the reward distribution, establishing the bridge between detailed balance and MaM marginalization.The result applies to states formed by assigning the first coordinates according to the permutation.
  • 3.1 Permutation-Conditioned GFlowNets: MaMs have a unique functionally equivalent PC GFlowNet, and each detailed-balance PC GFlowNet induces a unique consistent MaM.The correspondence identifies MaM marginals with PC-GFlowNet flows and MaM conditionals with PC-GFlowNet forward policies.
  • 3.1 Permutation-Conditioned GFlowNets: Both MaMs and PC GFlowNets evaluate any-subset marginals with one neural-network forward pass after training.Their remaining distinction is the learning objective rather than operational marginal evaluation.
  • 3.2 Revisiting MaMs’ Objective: Because approximate consistency makes persistent-Gibbs samples yield biased KL estimates, the paper derives an unbiased estimator that preserves MaMs’ constant evaluation cost.The objective is motivated by the connection between MaMs and PC GFlowNets.
  • 3.2 Revisiting MaMs’ Objective: The proposed consistency-only objective enforces terminal flows equal to R, corresponds to a weighted detailed-balance loss, and needs a constant number of forward passes as dimension grows.This replaces direct KL minimization with terminal reward matching while retaining computational efficiency.

4 PARTICLE GFLOWNETS

Particle GFlowNets extend persistent-block Gibbs sampling to non-autoregressive generative processes, while supporting ergodic sampling and automatic chain rejuvenation. The resulting procedure reduces per-step computation and addresses exploration problems caused by correlated persistent chains.

  • Particle GFlowNets: O(1) forward passes per gradient step replace O(d) computation in persistent-block Gibbs sampling, regardless of the number of components d.The method applies to non-autoregressive processes such as set generation.
  • Particle GFlowNets: Sampling k components from a truncated Poisson distribution with average log d emphasizes near-terminal states while covering the entire state graph.This weighting is used when constructing the persistent Gibbs updates.
  • Particle GFlowNets: P-GFlowNets maintain a persistent Gibbs sampler over terminal states and use a detailed-balance objective to train the underlying Markov chain.The construction allows multiple trajectories to lead to the same object.
  • Particle GFlowNets: The resulting Markov chain is ergodic with respect to the target distribution when the P-GFlowNet satisfies detailed balance.This supports both correlated Gibbs samples and independent samples generated through MDP simulation.
  • Chain rejuvenation: The method refreshes the persistent chain when the split R-hat statistic exceeds 1.1 to counter inadequate exploration and accelerate convergence.For discrete states, R-hat is computed from last-layer embeddings of the forward policy network.

5 EXPERIMENTS

The experiments evaluate when P-GFlowNets accelerate convergence, how rejuvenation affects training, and whether exploration improves across combinatorial tasks. Results show the largest benefits for long horizons, expensive sampling, and relatively cheap reward evaluation.

  • Overall findings: The experiments confirm that P-GFlowNets reduce training time when policy evaluation is the bottleneck, particularly for long trajectories and cheaper reward queries.This conclusion is supported by the runtime and exploration experiments in Figures 4 and 6.
  • Set generation with log-additive rewards: P-GFlowNets widen the runtime gap over standard GFlowNets as the MDP horizon increases across tested dimensions.The evaluated settings are (d, K) ∈ {(32, 64), (64, 128), (128, 256)}.
  • Bayesian variable selection: P-GFlowNets consistently outperform standard GFlowNets when likelihood evaluation uses few samples, but their computational benefits dwindle as n grows.The benefit decreases when posterior evaluation becomes more expensive than trajectory sampling.
  • Exploration: P-GFlowNets improve exploration on Ising and bit-sequence tasks, especially when long trajectories make sampling expensive.The evaluation tracks average log reward among the 10 most rewarding states found during training.

6 DISCUSSION

The discussion identifies Particle GFlowNets as a non-autoregressive extension that reframes MaMs as conditional GFlowNets with persistent Gibbs exploration. It reports faster wall-clock convergence while highlighting open questions about rejuvenation, diagnostics, and latent dynamics.

  • Discussion: MaMs are characterized as conditional GFlowNets using a persistent Gibbs sampler, and this strategy extends beyond autoregressive modelling.The paper presents Particle GFlowNets as the resulting generalized method.
  • Discussion: Particle GFlowNets significantly accelerate learning convergence in wall-clock time compared with conventional GFlowNet training algorithms.This is the paper’s principal experimental conclusion.
  • Future directions: The work leaves open how to choose rejuvenation times, diagnose distributional accuracy, and determine when latent dynamics accelerate GFlowNet training.These questions connect the method to broader Markov-chain methodology.

(Supplementary Material)

The supplementary material contains the training-loop pseudocode and proof fragments supporting the method’s normalization and marginal-distribution properties. It also specifies persistent-chain updates, rejuvenation, and non-deterministic transition handling.

  • Training loop: Each persistent Gibbs step removes one component through the backward policy and adds one through the forward policy.The pseudocode places these operations inside the chain transition used to produce samples for the complete loss.
  • Training loop: Algorithm 2 initializes particles by MDP simulation, alternates complete and intermediate losses, advances a persistent Gibbs chain, and updates parameters.The loop accepts a reward function, batch size, Gelman-Rubin threshold, and truncated-Poisson weight vector.
  • Training loop: The rejuvenation step computes split ˆR from recent last-layer embeddings and resamples all particles when ˆR exceeds the threshold.The pseudocode uses the 512 most recent states for the diagnostic.
  • Proof material: The proof establishes that the initial-state flow equals the target partition function and is independent of the permutation, while induction verifies the intended flow values over partial states.The argument uses disjoint terminal-state sets and the detailed-balance condition.

B.2 PROOF OF COROLLARY 3.1

The proof fragment uses equality of terminal-state sets under different permutations to transfer the relevant flow expression between representations.

  • Proof of Corollary 3.1: Equal terminal-state sets under two permutations imply the corresponding partial-state flow arguments can be identified.The proof invokes T(σ1, s1,(i)) = T(σ2, s2,(i)).

B.3 PROOF OF COROLLARY 3.2

The proof establishes the stated transition relation by applying the shared-prefix condition to the relevant flow values.

  • B.3 PROOF OF COROLLARY 3.2: The shared-prefix equality σ1(≤i+1) = σ2(≤i+1) yields matching flow values at both consecutive states.The proof invokes Corollary 3.1 for F σ1(s1,(i)) = F σ2(s2,(i)) and for the corresponding next states.

B.4 PROOF OF PROPOSITION 3.2

This section proves that MaMs and permutation-conditioned GFlowNets are equivalent, then verifies that the induced Markov chain is ergodic with stationary distribution proportional to the reward.

  • B.4 PROOF OF PROPOSITION 3.2: The argument is presented as a comprehensive demonstration of the equivalence proposition and its associated chain properties.The surrounding discussion explicitly identifies the result as establishing equivalence.
  • B.4 PROOF OF PROPOSITION 3.2: The construction maps a consistent MaM’s marginal and conditional probabilities to PC-GFlowNet flows and transitions.For each subset J and added variable i, the proof defines F σJ (xJ) from pθ and the forward transition from pϕ.
  • B.4 PROOF OF PROPOSITION 3.2: The construction is permutation-independent, so the resulting PC-GFlowNet is equivalent to the MaM.Corollaries 3.1 and 3.2 establish independence from the permutation choice.
  • B.4 PROOF OF PROPOSITION 3.2: The induced chain is irreducible and aperiodic because valid transitions have nonzero probability and returns are possible after any number of steps.These properties follow from the non-degeneracy of the forward and backward transition distributions.
  • B.4 PROOF OF PROPOSITION 3.2: Detailed balance implies that the chain is stationary with respect to π(x) proportional to R(x), making it ergodic with respect to π.The proof combines stationarity from detailed balance with irreducibility and aperiodicity.

D EXPERIMENTAL DETAILS

The experiments use controlled implementations and repeated runs, while measuring runtime and convergence for P-GFlowNets against several GFlowNet baselines.

  • D EXPERIMENTAL DETAILS: P-GFlowNets reduce per-training-step computation cost by several orders of magnitude by avoiding complete trajectory sampling.Figure 8 compares per-step runtime for ACE, AT, TB, and P-GFlowNets.
  • D EXPERIMENTAL DETAILS: P-GFlowNets converge faster than AT, ACE, and TB GFlowNets while achieving more accurate marginals and larger log-predictive density.Figure 9 reports this pattern for set generation and variable selection tasks within a shorter time span.
  • D EXPERIMENTAL DETAILS: All plots average three independent runs, with error bars representing one standard deviation from the average.The Gelman-Rubin statistic is evaluated asynchronously using the 512 most recent states, and rejuvenation occurs when R-hat exceeds 1.1.

E ADDITIONAL EXPERIMENTS

Additional experiments compare P-GFlowNets with exploration-based GFlowNet methods on set, variable-selection, and phylogenetic tasks, while identifying a reward-evaluation regime where those alternatives may be preferable.

  • E ADDITIONAL EXPERIMENTS: P-GFlowNets converge significantly faster than AT and ACE in set generation and variable selection with respect to wall-clock time.AT and ACE increase training cost by evaluating exploratory and target GFlowNets on forward and backward trajectories.
  • E ADDITIONAL EXPERIMENTS: Artificial-curiosity exploration strategies are most appropriate when reward evaluation is extremely expensive relative to policy-network evaluation and cannot exploit caching.This condition defines a practical boundary for the comparison with AT and ACE.
  • E ADDITIONAL EXPERIMENTS: The parsimony reward is defined as R(T) = exp{−ParScore(r(T), T)}, with low parsimony score corresponding to more parsimonious trees.The score is recursively computed over child subtrees and aggregated across sequence elements.
  • E ADDITIONAL EXPERIMENTS: P-GFlowNets improve upon AT and ACE at similar wall-clock budgets in phylogenetic inference.The comparison uses parsimony-based rewards and examines state-space exploration during training.
  • E ADDITIONAL EXPERIMENTS: Modular parsimony scoring reuses most computation between neighboring trees, reducing the cost of reward queries.This reuse supports the phylogenetic comparison under P-GFlowNets’ backward-forward kernels.
Loading 2609.11538v1…