Generalization and Distributed Learning of GFlowNets
Authors: Tiago Silva
Amauri H Souza
Omar Rivasplata
Vikas Garg
Samuel Kaski
Diego Mesquita
Venue: ICLR 2025 Poster (2025)
Keywords: GFlowNets
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We derive non-vacuous statistical guarantees and introduce the first distributed learning method for GFlowNets with network-level parallelization.

Abstract: Conventional wisdom attributes the success of Generative Flow Networks (GFlowNets) to their ability to exploit the compositional structure of the sample space for learning generalizable flow functions (Bengio et al., 2021). Despite the abundance of empirical evidence, formalizing this belief with verifiable non-vacuous statistical guarantees has remained elusive. We address this issue with the first data-dependent generalization bounds for GFlowNets. We also elucidate the negative impact of the state space size on the generalization performance of these models via Azuma-Hoeffding-type oracle PAC-Bayesian inequalities. We leverage our theoretical insights to design a novel distributed learning algorithm for GFlowNets, which we call *Subgraph Asynchronous Learning* (SAL). In a nutshell, SAL utilizes a divide-and-conquer strategy: multiple GFlowNets are trained in parallel on smaller subnetworks of the flow network, and then aggregated with an additional GFlowNet that allocates appropriate flow to each subnetwork.  Our experiments with synthetic and real-world problems demonstrate the benefits of SAL over centralized training in terms of mode coverage and distribution matching.

Published as a conference paper at ICLR 2025
GENERALIZATION AND
DISTRIBUTED LEARNING OF GFLOWNETS
Tiago da Silva
Amauri Souza
Getulio Vargas Foundation
Federal Institute of Ceará
tiago. hwwri fg
omauriholandal fce. edu br
Omar Rivasplata
Vikas Garg
University of Manchester omar www.sppaaaaeenncchhstte....
Yai Ltd and Aalto University vgarg@esale mit edu
Samuel Kaski
Diego Mesquita
Aalto University, University of Manchester samuel www.kikaaalaat..
Getulio Vargas Foundation diego. messuuitafgg.....
ABSTRACT
Conventional wisdom attributes the success of Generative Flow Networks (GFlowNets) to their ability to exploit the compositional structure of the sam- ple space for learning generalizable flow functions (Bengio et al., 2021). Despite the abundance of empirical evidence, formalizing this belief with verifiable non- statistical guarantees has remained elusive. We address this issue with the first data-dependent generalization bounds for GFlowNets. We also elucidate the negative impact of the state space size the generalization performance of these models via Azuma-Hoese Hoeffding-tyre typpe oracle PAC-Bayesian inequalities. We lever- age our theoretical insights to design a novel distributed learning algorithm for GFlowNets, which we call Subgraph Asynchronous Learning (SAL). In nutshell, SAL. utilizes divide-and-conquee strategy: multiple GFlowNets are trained parallel on smaller subnetworks of the flow network, and then aggregated with an additional GFlowNet that allocates appropriate flow to each subnetwork. Our ex- periments with synthetic and real-world problems demonstrate the benefits of SAL over centralized training in terms of mode coverage and distribution matching.
INTRODUCTION
Generalization is long- problem in the machine learning literature, asking whether learning algorithm can reliably make predictions beyond the data it was trained on (Valiant, 1984; Vapnik, 2000; Catoni, 2007; Alquier & Guedj, 2017; Dziugaite al., 2021; Lotfi et al., 2024a). In an age of rapid deployment of AI models to end-users, there has been an emerging interest in the design of theoretically robust algorithms, with remarkable results for GANs (Mbacke et al., 2023), diffusion models (Li et al., 2024), transformers (Lotfi et al., 2024a;b), and graph neural networks (Ju al., 2023; Tang & Liu, 2023). In this pursuit for developing models with proven generalizability, rich set of tools has been created (Vapnik & Chervonenkis, 2015; Shalev-Shwartz & Ben- David, 2014), with McAllester (1998: 1999)'s PAC-Bayesian theorems often providing the tightest statistical guarantees (Pérez-Ortiz et al., 2021; Dziugaite & Roy, 2017; 2018; Lotfi et al., 2024b;a). Notably, however, there is no one-size-fits fits-all solution for understanding generalization: the diverse nature of data and learning algorithms demands distinctly unique approach to each problem.
In the realm of probabilistic methods, for example, it has been widely hypothesized that the outstand- ing performance of Generative Flow Networks (GFlowNets, Bengio et al., 2021; 2023; Lahlou at al., 2023), which have demonstrated exceptional results in problems such as design of biological sequences (Jain et al., 2022; Malkin al., 2022) and combinatorial optimization (Zhang et al., 2023a;b), to name a few, emerges from their potential to exploit the compositional structure of the underlying state space leamn generalizable flow assignment function in a flow network by only

---

Published as a conference paper at ICLR 2025
observing fraction of the network's 's nodes (Bengio et al., 2021; Nica et al., 2022; Shen et al., 2023; Atanackovic & Bengio, 2024; Krichel et al., 2024). Nonetheless, in spite of the wealth of empirical evidence indicating that generalization occurs in GFlowNet learning (Nica et al., 2022; Atanackovic & Bengio, 2024), no work so far has provided non-vacuous high-probabilityy empirical bounds on the population risk of GFlowNets, which might serve as statistical certificates for generalization.
Given the described scenario, in this paper we develop the first non-vacuous generalization bounds for GelowNets in the literature in Section 5. In doing so, the core questions we want address in this work are when GFlowNets (provably) generalize and which factors potentially contribute to diminish their generalization performance. To kickstart our analysis, we present in Section an example in which a GFlowNet catastrophically fails to generalize even after learning compatible flow assignment for over 90% of the flow network. This example demonstrates that, to properly understand the generalization of GFlowNets, we must consider not only the extension of the observed flow network but also the specific parts that have been encountered during training, a fact that is also implicit in popular techniques such as the replay buffer (Vemgal et al., 2023) and local search (Kim et al., 2024b). From a technical perspective, this implies that the development of meaningful statistical guarantees for GFlowNets must be based not on iata-agnostic theoretical results, but on data-dependent priors, in the fashion of Dziugaite & Roy (2018); Dziugaite et al. (2021); Pérez-Ortiz al. (2021). This observation guides the establishment of the non-vacuous high-probabiligg bounds for the population risk of GetowNets in Section 5.1 and distinguishes our approach from other investigations of GFlowNets generalization (Krichel et al., 2024).
The empirical results in Section 5.1, however, do not provide fine understanding of which charac- teristies of the flow network tend hinder the generalization of GelowNets. What effect do larger trajectory lengths, for instance, have on the provable learnability of generalizable flow assignments? Intuitively, generalization is harder in larger state spaces as, borrowing the terminology from the reinforcement learning (RL) literature (Bengio et al., 2021), the visited portion of an environment by an agent constrained by fixed time budget decreases with increasing environment's size and, therefore, the agent would have to rely on increasingly sparse information for learning from larger trajectories. In Section 5.2, we formalize this intuition through the lens of PAC- BaC-Bayesian bounds, revealing the increasing difficulty in obtaining tight statistical certificates for larger state spaces. To achieve this, our technical contributions are two-fold: First, an oracle concentration inequality for the forward Kullback-Leite (Theorem 5.2) inspired by Malkin et al. (2023)'s interpretation of GFlowNets as variational inference. Second, an Azuma-Hoeffding sampled martingales representing the extent to which the learned flow assignment violates the detailed balance condition (Bengio al., 2023) in the observed trajectories (Theorem 5.4).
(KL) divergence between the learned and targeted flow assignments
inequality (Seldin et al., 2012b) for independently
Motivated by these results, we pose the following question: what can we do to mitigate the issues raised by an inadequate coverage of the state space, as illustrated in Section and larger trajec- tory sizes, as analyzed in Section 5? As we show in Section 6, the answer lies in breaking up the flow network into mudti-source subnetworks grouped together by a small root network (Figure 3). In this new paradigm, finding compatible flow assignment becomes two-stage process. Firstly, a different GFlowNet is trained on each subnetwork in a distributed fashion. Secondly, an addi- tional GFlowNet trained on the root network learns to assign the correct amount of flow to each subnetwork, as estimated in the previous step. The resulting algorithm, which we call Subgraph Asynchronous Learning (SAL), has several advantages over the standard approach (Section 6.2). In particular, each GFlowNet within this framework needs to solve a problem that is relatively simpler than that of unique, centralized model. Similarly, the asynchronous nature of the algorithm im- plies that we are able to visit a considerably larger fraction of the original flow network within a fixed time window and, therefore, that we get significantly better coverage of the state space. As consequence of this, we are also able to drastically improve the discovery of value states within the flow network, which a metric of great interest in the GFlowNet literature (Bengio al., 2021; Shen et al.. 2023; Zhang et al.. 2023a; Pan et al., 2023a; Jang et al., 2024; Kim et al., 2024b).
In summary, our contributions are:
1. We construct a family of examples in which a GFlowNet does not generalize even after learn- ing a compatible flow assignment on arbitrarily large fractions of the flow network (Section 4);
2. We provide the first non-vacuous generalization bounds for GelowNets (Section 5.1);
2

---

Published as a conference paper at ICLR 2025
3. We derive oracle PAC-Bayese inequalities for the population risk of GFlowNets, emphasiz- ing the impact of the flow network's topology on generalization performance (Section 5.2);
4. We design the first distributed algorithm for learning GFlowNets with network- paral- lelization, and evaluate its performance in common benchmarks in the literature (Section 6);
The first part of the paper establishes the notation and terminology to be used throughout the paper and reviews relevant results in the GelowNet and PAC-Bayes literature, alongside an overview of our main contributions (Sections and 3). The second part provides a formal treatment of GelowNets from the viewpoint of the PAC-Bayesian theory (Section 5). The third and final part outlines the foundations of SAL and conducts an empirical evaluation of the algorithm in common benchmark problems (Section 6). We defer the proofs and details of the experiments to the supplement.
PRELIMINARIES
Notations and terminology. Let G = (V, (VED) E) by a directed acyclic graph (DAG). A forward policy over G is Markov transition kernel pp: V V + [0,1] supported on G's edges, i.c., such that (F(u.) -) is distribution over {u: (v, u) E}. for each vertex u. We interchangeably use pf(l and pr (c. for representing pr. A backward policy P over Gs forward policy on the transpose graph G (V, E with E = v(; (v., 2) E E}. The uniform policy assigns the same prob- ability mass to state's 's children, i.e., (s
L with Ch(s) = {s': (s,s') E). We say that G is pointed if there are nodes 5, and s,, respectively called initial (source) and final (sink) nodes, s.t. 8, (resp. 5y) the only node without incoming (resp. outgoing) edges and, for each E V. there is a trajectory (directed path) between s, and 8 containing In this case, trajectory T in is complete it starts at s, and finishes at s.. which we denote by T: so ~4 8y. Clearly, a for- ward policy induces distribution over trajectories starting at via (rapr((s)) = =II [s
ps(""ss);; will often omit s, from this notation. Lastly, for probability when T is unambiguously complete, measures P and Q on the same space, we let KL(PI (1)). (PIIQ) and TV (P, Q) respectively denote their Kullback-Leibler (KL) divergence, x² divergence, and total variation distance.
GelowNets. We represent a GFlowNet (Bengio et al., 2021; 2023; Lahlou at al., 2023) 9 tuple (S, X, G, A, T, PF. PB, R, F') consisting of set of states S, set of terminal states X C S, pointed DAG G = (8, E), which is called a state graph, an action mapping A: S 2- associating each state s with an abstract action space A(s) S AA that is isomorphic to the children of s in G, a transition function T: Uses ({s} x A(s) + S defining how a state s is affected by an action E As) pr and backward pn policies a positive value to each terminal state, and a flow function F: S + R__ such that Flx = R. Importantly, only the elements of X are connected to the sink node s of G. When there is no risk of ambiguity, we will simply write G = (pF,PB, F). The objective the marginal distribution Pr (x) == Pr pr(xxs.)) (x]s.) = Σr: Sy wax Pr (T) over X' matches R up to normal- izing factor (Bengio et al., 2021). In Appendix A, we illustrate how this abstract representation can be instantiated to accommodate three frequently considered use-cases (Malkin et al., 2022; 2023).
G, a reward function R: X - R2 attributing
GFlowNet to find a pF s.t.
Learning GelowNets. In this context, S, X, G, A. and R are problem-deeendent, while pr, PB. and F are unknowns that should be estimated. Remarkably, however, Pg is often fixed as uniform (Shen et al., 2023; Liu et al., 2023; Zhang et al., 2023a), assumption that we make throghout the paper, albeit most of our theoretical results and all our methods can be extended to the case of learnable pp. Under these circumstances, many learning objectives have been proposed for learning PF and F. Two popular choices, which we adopt here, are the trajectory balance (TB, LrB (Pp, F), Malkin et al., 2022) and detailed balance (DB, LDB (PF, F), Bengio et al., 2023) losses,
E Tropr
(108 ((S))DF(T) (())pn(T(x)
and
E TMPR
1
Σ (s.s")
log PF (5) 2
Pn(()s))((s))
in which |r| represents r's length, is r's (unique) terminal state, and pg is exploratory policy. Intuitively, Pg has the role of the data-generating distribution in a standard supervised learning context and is often defined PE (c)pu (1 e)pr, an e-greedy version of pr. with pu denoting an uniform policy; although more sophisticated techniques have been developed (Kim al., 2024b; Rector-Brooks al., 2023; Vemgal et al., 2023). In Appendix A we provide more
3

---

Published as a conference paper at ICLR 2025
thorough overview of GFlowNet learning, including the subtrajectory balance loss (SubTB, Madan et al., 2022) and divergence- based objectives (Malkin et al., 2023; Lahlou at al., 2023).
Generalization bounds for neural networks. The field of statistical learning theory ( (Vapnik, 1998; 2000) seeks to develop statistical certificates for the generalization of a learned model by providing high-probability upper bounds on the population error of an estimator as a function of the observed empirical risk. In the context of GelowNets, based on the observed trajectories, such as the losses in Equation or other locally computed metrics (see Sections and 5), are appropriate surrogates for the GFlowNet's overall distributional accuracy. In particular, we are interested in inductive statistical guarantees, namely, those based on the training set (as opposed the transductive setting, in which test set is used). To this end, the PAC-Bayes framework of McAllester (1998; 1999; 2013) often provides the tighest bounds Vérez-Ortiz et al., 2021; Lotfi et al., 2024b;a; Dziugaite & Roy, 2017). In a nutshell, consider data X [X: drawn from some data distribution, a significance level ò, an empirical loss 2(, X) and a population loss [()) = Ex [0, X)] associated to the model's 's parameters A. Given a 'prior' (independent of X) distribution Q over 0, PAC-Bayes bound typically assumes the form
ask whether an empirically measured imbalance
Ex {(()) - Eg- [((0, x + (6 m.Q.m).
(2)
The inequality holds with probability 1 - 5 over draws of X, simultaneously for all 'posterior' distributions P over 0: and is term penalizing the model's complexity (McAllester, 1999). We direct the reader to Alquier (2024) for a comprehensive introduction to PAC-Bayesian analysis. For bounded C, the right-hand side of Equation is termed vacuous if is larger than an upper bound of C. Although McAllester's s original works posited that the data were independent and identically distributed (i.idd)) and that the risk function was uniformly bounded (McAllester, 1998; 1999), recent advances relaxed these assumptions by deriving generalization bounds for non-i.i.d. data (Seldin et al., 2012b; Barnes et al., 2022), with applications to multi-armed bandits and RL (Fard & Pineau, 2010; Beygelzimer et al., 1: Tasdighi et al., 2024), and for unbounded losses limited by high-probability bounds (Alquier & Guedj, 2017; Haddouche & Guedj, 2023; Casado et al., 2024; Mbacke et al., 2023). To the best of our knowledge. however, this is the first work promoting the development of PAC-Bayesian bounds for understanding the generalization of GFlowNets.
3 OVERVIEW OF OUR RESULTS
Before delving into the details of our work in Sections 4, 5, 6 (and further details in the appendices in the supplement), we provide below brief discussion around our technical results under the light of the formalism presented in Section 2, alongside the main ideas they were built upon.
Non-vacuout generalization bounds for GFlowNets. The learning objectives in Equation 1, due to the unboundedness the logarithm, cannot be directly incorporated into standard PAC-Bayesian theorems (McAllester, 2013), which assume that the risk function has at least bounded exponential et al., 2024a;b). To circumvent this issue, our empirical analysis in Section 5.1 adopts the recently proposed FCS metric (Silva et al., 2024) as the risk functional measuring the accuracy of a trained GFlowNet, which may be written as
moments (Casado et al., 2024; Rodríguez-Gáäve
LFCs = E(T..r)).... [TV (p RR )] E [0, 1],
(3)
where pf and R are the respective restrictions of pT and R to the B-sized multiset (fa1...., C X' of terminal states, and TV is the total variation distance. However, in spite of easily computable, Lecs is not an appropriate learning objective for GelowNets due to the potential numerical instability of the non-log-domain. Instead, we minimize LTB as a surrogate objective for LFCS during training and evaluate the generalization bound on Lecs in the inductive fashion mentioned in Section 2. Importantly, Figure shows that the resulting bounds are remarkably tight.
Oracle generalization bounds for GFlowNets. As a complement, we also establish non-empirical high-probability upper bounds on the population risk of GFlowNets by assuming that a potentially intractable quantity bounds the corresponding loss function. In Section 5.2, we follow this rationale and demonstrate that there always is an a 0 for which the set of policy networks of the form apu + (1 a)pr contains the solution to the flow assignment problem. Armed with such a fam- ily of models, which guarantee log probabilities uniformly bounded away from zero, consider
4

---

Published as a conference paper at ICLR 2025
the reverse KL divergence risk (Malkin et al., 2023) avoid explicitly bounding the flow func- tion F. Although informative, the resulting Theorem 5.2 only considers the trajectories- and not transitions-as data points. As the number of observed transitions is significantly larger than that of trajectories, we enrich our results by constructing a martingale difference sequence based on the DB loss and adapting Azuma's inequality (Azuma, 1967; Seldin al., 2012b) to the context of indepen- dent martingales to derive a transition-level generalization bound for GFlowNets. Both approaches, which are respectively encapsulated in Theorems 5.2 and 5.4, show that the population risk can be bounded with nigh-probabiligg as, apart from technical nuances,
(Loge fm n1
Eg (C(0) S Ee-r|((0)]
in which & is an empirical measure of risk, fm the maximum trajectory length of the state graph, and is the number of observed data points: either trajectories (Theorem 5.2, 0.5) or transitions (Theorem 5.4, a = 1). From an analytical perspective, these results suggest that learning provable generalizable flow assignments is increasingly harder for state spaces having longer trajectories.
4 WHEN DO GFLOW NOT GENERALIZE?
To start our discussion on the generalization of GFlowNets, introduce simple, but non-trivial, examples in which a GFlowNet does not learn a generalizable policy network even after minimizing the loss on an arbitrarily large portion of the state space, raising the questions of when do GFlowNets generalize and how to measure such generalization, which we investigate in Sections 5.1 and 5.2.
A non-generalizable data distribution. To concretize our arguments. we recall the task of set generation for GFlowNets (Pan al., 2023a;b; Bengio al., 2023; Jang et al., 2024). Each state corresponds to subset of a set W = {1... W for given W; the generative process starts at an empty set so = 0 and iterativey adds elements from W to s, until a prescribed size T is achieved. For our purposes, we fix a function u: W + [0, 1). representing the log-utilin of each w E A(s,) == W and define the reward R associated to S as R(S) = 1 (##=2) exp u(w)). Also, let pg be forward policy s.t. (-(ss) is supported on A(s) \ {1} = W. U{I)Us) Us) for every i.e., the support of the marginal PF,r of pr on X' is the set X" of subsets of (2. (2...., W). We next
show that X' covers an arbitrarily large portion of X' for specific choices of and W.
Lemma 4.1. For each & E (0, 1), there exist T and W such that | 2 E|X
The (straightforward) proof of Lemma 4.1 can be found in Appendix D. Obviously, we cannot hope that GFlowNet trained by minimizing an empirical risk defined on trajec- tories sampled from Pr would generalize unseen states, as no information regarding u(1) would be available during training. To empirically validate our reasoning, we show in Figure that a GFlowNet trained on samples from pæ fails to learn the right distribution, whereas standard e-greedy strat- egy succeeds. It is remarkable, however, that GFlowNet unable to successfuly sample from the target distribution even after minimizing the empirical risk samples covering over 90% of the state space. From a statistical viewpoint, this behavior can be explained via a change of measure inequality: preference over states is not properly captured by the sampling distribution (pg). We formalize this intuition in the proposition below.
visited
60.62 vicited
0.25
0.00
200000
200000
Truising spoch
Figure 1: Convergence speed when
actions are masked (bluc) or not (or- ange) for different state space sizes.
Proposition 4.2 (Generalization depends on the sampling distribution). Let (pr: PB. R) be GFlowNet and p be (any) distribution over X. Also, recall π(x) represents the normalized target and pr the learned marginal. Define T as an uniform PMF on X. i.e. qE,((z) 1/1x, Then,
[ (1og
(F(T) T()))p(T(a)
27
TV (PPIT) (1 ((EErlIEE
Errpa(rle)
(5)
in which (P)()) represents the x² divergence between P and Q.
5

---

Published as a conference paper at ICLR 2025
We interpret Equation 5 in the following way: If the sampling policy (pg,T) greatly deviates from the uniform (qE.r), then a small empirical risk does not necessarily ensure an accurate distributional approximation. In contrast, Equation 5 does not entail that the uniform distribution is the optimal choice for sampling trajectories, as it does not address the algorithmic difficulty of minimizing the empirical risk via SGD. Illustratively, we show in Table in Appendix B the values of X"E, (PE.T) when PE,T is far away from QE.T and of (qE,,Tllpe,"" for the e-greedy policy considered in Figure 1. On a fundamental level, these examples underline the importance of taking into account the data distribution for understanding generalization performance. Section 5 elaborates on this problem through the lens of PAC- PAC-Bayes bounds, albeit with data-dependent priors (Dziugaite & Roy, 2017; 2018; Dziugaite et al., 2021; Pérez-Ortiz et al., 2021).
5 PAC-BAYESIA GENERALIZATION BOUNDS FOR GFLOWNETS
Towards the objective of understanding GFlowNet generalization. we construct high-probability upper bounds on different risk functions. In Section 5.1, we build upon McAllester's empirical bound and Dziugaite's s data-dependent priors to derive the first non-vacuous generalization bounds for GFlowNets. Then, to gain a clearer understanding of the factors hindering the generalizability of these models, we provide both trajectory- and transition-levw oracle bounds in Section 5.2 by draw- ing upon the martingale-basee PAC-Bayesian theory for non-i.i.d. data (Beygelzimer et al., 2011).
5.1 NON-VACUOUS EMPIRICAL GENERALIZATION BOUNDS
GFlowNet learning as supervised learning. To rigorously address the generalization of GFlowNets, we firstly frame the training of these models as supervised learning problem (Shalev-Shwartz Ben-David, 2014; Atanackovic & Bengio, 2024). For this, we assume that a set of independently sampled complete trajectories, Tn [TI....,TTI I is drawn from a fixed distribution and that each trajectory т; is annotated with a noise-free target, y/4 = pn (T (r, R(a:), with F2 representing 's's unique terminal state. Importantly, the only supervision during training comes from the reward function; we do not make assumptions on the distribution over T In this context, minimizing LTB corresponds to finding the least-squares solution to the equation log Z + log (T) = log (T(())) (T)x)))()))) R(x) in pr. Importantly, this setting differs from conventional GFlowNet training algorithms, for which the sampling policy depends on the trajectories observed so far, that is, the trajectories are not independently sampled. Nonetheless, the question of whether GFlowNets generalize remains relevant even under our relatively simplified conditions, which may be seen as a single-iteration of an e-greedy strategy (Krichel et al., 2024).
A bounded risk functional for GelowNets. As mentioned, PAC-Bayesian theory originally relied on the assumption of bounded risk functions (McAllester, 1998). Despite recent advances in extending the theory to unbounded losses (Casado et al., 2024; Rodríguez-Gáil Gálvezz et al., 2024a;b; Haddouche & Guedj, 2023), most generalization bounds still depend on technical and hard-to-vert verrfy conditions, e.g., on exponential moments. For this reason, we use the FCS metric as measure of risk (Silva et al., 2024); see Appendix B for an unbiased estimator (PF, Tn) of (PF)
Data-dependent priors for PAC-Bayes. For this, we first recall the techniques originally developed by Dziugaite& & Roy (2017; 2018); Dziugaite et al. (2021) a striking series of papers for probing the generalization of overparameterized neural networks in the supervised learning context. To start with, we state below Dziugaite et al. (2021)'s empirical PAC-Bayes bound, which combines results from McAllester (2013), Rivasplata et al. (2019), and Boucheron et al. (2013). For completeness, we also provide a self-contained proof of Proposition 5.1 in Appendix D in the supplement.
Proposition 5.1 (Empirical PAC- BAC-Bayesian bounds). For any distribution on parameters 0 of pp, let LVcs(()= = Eg-,; Lrcs RR, pr)) and define Lrcs(ç, Tw.) similarly. Also, let (0,1) and let P be a distribution on 0 learned on an uniformly random [[1 - a:mm] -sized subset Ti-a of Tn. Then,
nnn+ + 2L (P, (PTiia)))
LFCs(P) Lecs((P TT-a) min
n
(6)
with probability at least 1 - 8 over Ti-a in which = KLIP 2010 awnl/s and Q is distribution that does not depend Ti oa a but may depend on To Tn Ta-
[I=a)n]

---

Published as a conference paper at ICLR 2025
When the prior distribution Q is naively chosen (e.g., standard Gaussian dis- tribution), the KL divergence in Equation 6 often dominates the right-hand side of the equation and results in vacuous bounds, i.., (P) S a 0.8
for some a > 1. To address this issue, the influential work of Dziugaite & Roy (2017) proposed the use of data-dependent learned by minimizing the empirical risk functional on a fraction of the data and, after learning P by minimizing Equation evaluating the generalization bound on the remaining (1 a) portion of the data, as presented in Proposition 5.1.
Risk (Lrcs)
0.6
Bound
0.4
0.2
Empirical results. We follow a similar approach to derive the first non-vacuous generalization bounds for GelowNets in the literature. For this, we disjointly partition the dataset Tn with = 3 10 into sets To and Ti with o = 0.6. We learn an isotropic Gaussian prior Q on To and then a diagonal Gaussian posterior P on Ta U UI-a by minimizing the bound in Equation 6. Finally, the bound is evaluated on T to obtain the statistical certificate (Pérez-Ortiz et al., 2021). Results in Figure 2 for the tasks of set generation (Pan et al., 2023a; Bengio et al., 2023), bag generation (Shen et al., 2023; Jang et al., 2024), sequence design (Malkin et al., 2022; 2023; Madan al., 2022) with additive rewards, and SIX6 (Jain et al., 2022; Malkin al., 2022; Shen et al., 2023) highlight the non-vacuousness of Equation 6 and the generalizability of the trained models. Please refer to Section 6 and to Appendix B for a detailed description of the experimental setup. Appendix A describes the design of the GFlowNet for each of these generative tasks.
0.0
Sets
Bags
Seq. SIX6
Figure 2: Non-vacuous gener- alization bounds for the FCS risk functional in Eq. 10.
5.2 ORACLE GENERALIZATION BOUNDS
Although the previous section's 's empirical results certify the generalization of the learned policy network to novel trajectories, they do not necessarily shed light on which characteristics of the generative task are hindering the model's generalization capability. In the remaining of this section, thus derive generalization bounds that, despite not being directly computable, provide finer understanding of which factors play role when the goal is to learn a generalizable policy. In particular, we observe that larger trajectories and peakier target distributions tend to make general- ization harder when a fixed sampling budget is available. In Section 6, we will see how a distributed algorithm may alleviate these issues (Yagli et al., 2020; Barnes et al., 2022; Sefidgaran al., 2022).
Trajectory-level bounds. We start by deriving generalization bounds for GFlowNets when the trajectories are independently sampled (see Section 5.1) and the risk functional is the KL divergence between the forward and backward policies, i.e., (T.(pollrr). in which PB(T) x (T( R(x)). This choice is motivated by Malkin et al. (2023)'s interpretation of GFlowNets as hierarchical variational inference algorithms and by the ability of (L(pBlpp)) to focus the model on high- probability regions of the target, which is desirable trait of GFlowNets. Remarkably, we show in Lemma B.1 that KL(pg (pF) can be bounded by sensibly reparameterizing pr as a mixture policy. Then, as shown in Theorem 5.2 below, this reparametrization enables developing oracle generalization bounds in the fashion of the tight results we derived in Section 5.1.
Theorem 5.2. Let G (DF,PR, F) be a GFlowNet with policy network pr parameterized as in Lemma B. Also, let be probability distribution over the parameters 0 of pg. Denote H Hppl] -Eropa [log (p(p) for pa 's entropy and Mr = max, (|T log(a maxxer (Ch(s))). Then,
1
PB(TT) log
Ee-r [[[(T])PT] S Ea-p
+ +-HH[pa) + Mr) (7,
(7)
m
(r(
15
in which we recall that TP,,Q
+lg 2VT// and (x) x R(x) is the target.
A few remarks on the excess risk upper bound of Theorem 5.2. Firstly, the assumption that trajectories are sampled according to (T(T) x ps(T(z)R(a) is consistent with popular strategies for learning GFlowNets that focus sampling trajectories leading to high-reward states more often than those leading to low-reward states, e.g., using a replay buffer (Deleu et al., 2022). Secondly, in alignment with well-established practical knowledge, the result in Equation 7 shows it is harder to
7

---

Published as a conference paper at ICLR 2025
achieve tighter generalization bounds when the target distribution is spiky with small entropy term H [pe) and when the generative task is composed of longer trajectories or larger action spaces.
Transition-level bounds. For many applications, the number of observed complete trajectories be orders of magnitude smaller than the number of collected
when training GFlowNets state transitions. In this context, one may obtain significantly tighter generalization bounds by interpreting the transitions, and not the complete trajectories, as data samples (Lotfi al., 2024a;b). Indeed, it is assumed that GelowNets' outstanding potential emerges from its capacity to exploit the compositional structure of the space characterized by the state graph (Bengio et at., 2021; Nica al., 2022: Shen et al., 2023; Atanackovic & Bengio, 2024). To incorporate this structure into our theo- retical bounds, we shift our focus to the design of cuuuaa-Hoffdiig-tyyee concentration inequalities (Azuma, 1967; McDiarmid, 1998; Boucheron et al., 2013) applied to the stochastic process induced by the Markov Decision Process (MDP) governing the data-generating process. For this, we start defining martingale difference sequence based on the DB loss (Bengio et al., 2023, Example 5).
Definition 5.3 (A martingale difference sequence for the DB loss). Recall the detailed balance loss (DB(S, s = (log ((())p((s(()) MSS e.) CD U, E Bsp-pe(-1 [CDn (s,, S,-1))S,-il.
log F( (p)))))(ssss))))))))))))))))))) For a fixed sampling policy PE, we let
(8)
where Szi {S 1/-1}. Also, we define the natural filtration F; OS St) generated by the first states of the Markov process {S,).2 Clearly, each M (S,,, Sex)) is F-mmasurable and Es, [M Feed = i... (M, Sei) is martingale difference sequence.
From this definition, it is immediate that M, = Eigigt (S,,, Sz)) is a martingale w.r.t the filtra- tion {Fi Fe21211. We defer to Appendix B the discussion regarding its properties and the assumptions imposed on it for proving Theorem 5.4 below. Additionally, we define
-
C()) = T>PE
1
LE [CDDB(S,, (--1)S,-111 and (() = Σ IsiSirI
Σ
Σ
CD ((2))
1SJSn ISiSt,
as the population and empirical bB-based risk functionals for GFlowNets. Under these conditions, Theorem 5.4 complements Theorem 5.2 with a generalization bound based on the DB loss.
Theorem 5,4 (Transition-level generalization bounds for GFlowNets). Let (f) be the martingale arising from Definition 5.3, with 9 representing the parameters of the forward policy. Also, let Q be
distribution on 0. Assume that (SB,S S,-1) < U uniformly on (S;, S;-1) and that Mt,, (0)2 < K, which im is the maximum trajectory length. Similarly, define X S 1/20 and B (0, and let P be a data-dependers posterior distribution 0. Then, with probability at least 1 - 8 over the set of independent martingales (So. (S..)()) ................................ St, Jisisn such that S. So almost surely,
Ea-p [C(0)] < _Bemp [2(0] + aTT (KL(PIIQ) + log +
1 3
2
log tm
AK
BTX
13T
in which T is the number of observed transitions, ar,,
U 2)( + BTX
and 7 = - 2.
Similarly to Theorem 5.2, Theorem 5.4 implies that obtaining tighter generalization guarantees is harder for larger state spaces with longer trajectories when the sampling process is limited by a maximum number of observable transitions (or states) T. which is an often imposed constraint for comparing the sample-efficiency of different learning objectives for GFlowNets in the literature (Pan et al., 2023b;a; Madan et al., 2022; Malkin et al., 2022: 2023). In the next section, we show how these issues can be addressed via distributed learning scheme with network-leve parallelization.
6 DIVIDE AND CONQUER: DISTRIBUTED LEARNING OF GFLOWNETS
In light of the above analysis, the diverse exploration of state graphs (Section 4) with smaller trajectory sizes (Section 5) is beneficial for the successful training of GFlowNets. In what follows, we show how these features can be efficiently implemented by recasting the GFlowNet training as an embarrassingly parallel divide-and- learning (SAL). This is, to the best of our knowledge, the first method enabling the distributed learning of GFlowNets with network-level parallelization. We remark that previous work on the topic (da Silva et al., 2024) promoted only the partitioning of the reward function for parallel Bayesian inference and that, in stark contrast to SAL., each client learned from the same state graph.
algorithm, which we call subgraph asynchronous
8

---

Published as a conference paper at ICLR 2025
U.S
Algorithm Subgraph Asynchronous Learning
1: S S, , Sj b Fixed-horizon partition
2: Z, = S, n ns, for (!....,m)
3: > Local training
Sampling (mgg
parfor E [I.....m) do
5: Minimize LTT in S with SGD
a ampp(lla)
6: (p.. arg min L (P
7: end parfor
Figure 3: A fixed-horizon DAG partition with three leaves (S1, S2, Sa) and one root (S.). For inference, the sampling policy is chosen based on the current state.
t 8: R x I R(z) 1=ز (uEZ,) P((x)
9: (pG. arg aigmin (PB F,R")
10: return {(p.. Fo)) U (1pp Fj)}
6.1 SUBGRAPH ASYNCHRONOUS LEARNING
Overview. There are two ingredients making up SAL: fixed- horizon partition (FHP) and an assign- ment function (AF). In short, a FHP defines the state graph split explicitly, while an AF indirectly encodes it by assigning states to partitions. Here, we formally define the former concept. We intro- duce the idea of an AF and provide a comprehensive theoretical analysis of SAL in Appendix B.4.
Convergence guarantees. We introduce the notion of FHP of a pointed DAG below. In the flow network perspective, such a partition can be viewed source subnetworks, termed leaves, which are grouped together by a single-source network, referred to as root. We use the term fixed-horizon due to the fixed distance of the subnetworks' sources to So. Also, note that a FHP is only a partition in the set-theoretical sense when the state graph is a tree.
collection of possibly overlapping multi-
Definition 6.1
DAG partition). We say that S S, u (Uisism S,) is a FHP of the state space S, with leaves {5, mm and root S,, when it satisfies the conditions below:
1. (Disjointnesss of sources) 8, E S, and the sets (I, == S. ))) are pairwise disjoint.
2. (Completeness) If s Sj for a 2 1, then all descendants of are Sy.
3. (Regularity) If d denotes the shortest-path distance, d(s.. d(s..I,) for all i.j.
Under Definition 6.1, we let X, S, n X be the set of terminal states reachable from I,. For conciseness, we denote {S, im {S,} = FHPP m) when (S, { a FHP S with m components. and the Iy. represented by the doubly
We then illustrate FHP of tree in which circles, are singletons for the blue and green leaves in Figure 3. We are now ready to define SAL.
Definition 6.2 (SAL). Let {S = FHP m)) For each 1 S S m, let Gj = (wr, P's Fj) be GelowNet and F'g be any forward policy over the Sy-induced subgraph of the state graph. Finally, let qj be any distribution with full support on Zy. Then, define
(1oS B(O))((()))) 27 R()))
"1
Cxuulrr P2) Bengy B (oogg()))
in which r represents r's terminal state, as the amortized trajectory balance (ATB) objective. For the root, let pf be a policy in S, and Go (pf. p.. R²). For each E (X - UU; I,) let (x) = R2 (x) if x e I₁ for some j and (z)x) = R(x) otherwise. In this context, SAL follows two-step procedure: first, m models are trained in parallel by minimizing Equation 9; then, a global model is estimated by optimizing the TB loss with reward R which we denote by LTa(PF, F, R?) for a GFlowNet (pF.p8, F). We summarize this approach in Algorithm 1.
Clearly, any flow- learning objectives (e.g., SubTB (Madan et al., 2022), Munchausen DQN (Tiapkin et al., 2024), GAFlowNets (Pan et al., 2023b)), parametrizations (e.g.. forward-looking (Pan et al., 2023a), LED (Jang et al., 2024), temperature-sealed (Kim al., 2024a)), and off- policy sampling strategies (e.g., replay buffer (Vemgal et al., 2023) and local search (Kim et al., 2024b)) could be employed for estimating both the root and leaf GelowNets in Definition 6.2. In

---

Published as a conference paper at ICLR 2025
SAL
Tapa
Figure 4: SAL leads to faster mode discovery and more accurate approximations given a fixed training time budget (right-most plot). Results for centralized GFlowNet, for algorithm (SAL), and the target reproduced from Malkin et al. (2022, Section 5.1) are shown from left to right. Run- ning time for SAL equals the running time of the longest client plus that of the aggregation phase.
Appendix C, we demonstrate the soundness of SAL and conduct an extensive theoretical analysis on the character of the local distributions and error propagation within this framework. Similarly, Appendix E. discusses these issues from an empirical viewpoint.
Recursive SAL. As defined Definition 6.2 and shown in Figure 3, SAL has a single layer: each leaf is directly connected to the root in the underlying FHP. Nonetheless, there is obstacle preventing us from building SAL upon a multi-layered partition of the state graph, as illustrated in Figure 11 in the supplement. For this, recursively learn a flow assignment for each partition by starting at the lowest levels of this hierarchy and moving upwards, in the fashion of backward-induction algorithms. Each learning step is based on minimizing the amortized trajectory balance loss in Equation 9 via SGD. We demonstrate the correctness of the resulting algorithm, termed Recursive SAL, in Proposition C.6, which follows from Theorem C.I and an inductive argument. Although we do not provide an empirical assessment of Recursive SAL this work, Appendix C.3 considers potential implications.
must first define a hierarchy of partitions. Then,
6.2 EMPIRICAL ILLUSTRATION
Experimental setup. We evaluate the performance of SAL in six different generative tasks encom- passing both synthetic and real-world problems (Appendix C.I). In Appendix C. we extensively dis- cuss how to implicitly define a FHP via an assignment function, which allows for an efficient imple- mentation of SAL. Also, please refer to Appendix for detailed account of the experimental setup.
Results. As expected, Figure above, and Figures 7, 9, 10, and 13 in the supplement, show that SAL drastically speeds up the discovery of high-value states for all considered generative problems under varying computational constraints. Complementarily, Figure 6 and Table 2 in Appendix C.2 underline that our distributed algorithm achieves more accurate distributional approximations than its centralized counterpart. We discuss these promising results more detail in Appendix C.2.
7 CONCLUSIONS
Discussion. We developed the first PAC-Bayesian bounds and non-vacuous statistical guarantees for the generalization of GFlowNets in the literature. Additionally, our theoretical results provided deeper insights into the negative effect of the trajectory length on the proven learnability of generalizable policy. Inspired by these conclusions, our distributed algorithm SAL, which is also the first its kind, exhibited promising performance both synthetic and real-world problems.
Future works and limitations. We discuss the limitations our work large in Appendix In particular, we acknowledge that deeper theoretical understanding of advanced sampling techniques still required. From a practitioner's perspective, we believe that SAL can greatly improve the performance of GelowNets in specialized domains, e.g., NLP (Hu et al., 2024) and drug discovery (Bengio et al., 2021), which are beyond the scope of our work. Finally, it has not escaped our atten- tion that SAL is related to Mankowitz et al. (2016)'s s Adaptive Skills, Adaptive Partitions (ASAP) framework for learning temporally extended actions in MDPs and may find fruitful applications in multi-task RI. by interpreting each leaf (resp. root) GFlowNet an intra- (resp. inter-) skill policy.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work was supported by the Fundação Carlos Chagas Filho de Amparo à Pesquisa do Estado do Rio de Janeiro FAPERJ (SEI-260003//000099223333 the São Paulo Research Foundation FAPESP (2023/00815-4) and the Conselho Nacional de Desenvolvimento Científico e Tecnológico CNPq (404336/2023-0), We acknowledge the UKRI EPSRC grant EP/Y028783/1.
We acknowledge the Aalto Science-IT Project from Computer Science IT and the FGV TIC for providing computational resources.
REFERENCES
Pierre Alquier. User-friendly introduction to PAC-Bayes bounds. Foundations and Trends@ in
Machine Learning. 17(2):144-333, 2024. URL https: / axxiv.orgi........ abs/ wwwp.v..or//bbb//11....oo/ 11216.
Pierre Alquier and Benjamin Guedj. Simpler PAC-Bayesian bounds for hostile data. Machine
Learning, 187-5) December 2017. ISSN 1573-0565, doi: 10. 1007/510994-17-55699900
URL hwtp::tpp//.. /x doi org/ / 1. 1007 1 001--400077-6000000
Pierre Alquier. Xiaoyin Li, and Olivier Wintenberger. Prediction of time series by statistical learn- ing: general losses and fast rates, 2012. URL /ttp:p::///xxx..rrr/e// htgg/bb/////a 1211.1847.
Lazar Atanackovic and Emmanuel Bengio. Investigating Generalization Behaviours of Generative Flow Networks. ICML 2024 Workshop on Structured Probabilistic Inference Generative Mod- eling, 2024. URL https: / //aaiiiivvco.. //taaiiiwgggabbbaaa/aaaoooa.a/aa abbs /200///24422 2202 05309.
Kazuoki Azuma. Weighted sums of certain dependent random variables. Journal, Second Series, 19(3):357-367, 1967.
Tohoku Mathematical
Akshay Balsubramani. PAC-Bayes Iterated Logarithm Bounds for Martingale Mixtures. arXiv preprint arXiv:1506.06577 2015. URL https: //a /arxivv.v..// .ww.eooor.ooooooooooooe.oooe.ee. /abs///////// 115.6.6066666666
Leighton Pate Barnes, Alex Dytso, and Harold Vincent Poor. Improved Information-Theorettee Gen- cralization Bounds for Distributed, Federated, and Iterative Learning. Entropy, 24(9):1178, Au- gust 2022. ISSN 1099-4300. doi: 10.3390/e24091177 178. URL http://dx: doi org/10. 3390/e2409117..
Luis A Barrera, Anastasia Vedenko, Jesse V Kurland, Julia M Rogers, Stephen S Gisselbrecht, Elizabeth Rossin, Jaie Woodard, Luca Mariani, Kian Hong Kock, Sachi Inukai, al. Survey of variation in human transcription factors reveals prevalent DNA binding changes. Science, 351 (6280):1450-14) 2016.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow Network based Generative Models for Non-Iterative Diverse Candidate Generation. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Yoshua Bengio and Nikolay Malkin. Machine learning and information theory concepts towards an AI Mathematician. Bulletin the American Mathematical Society, 61(3):457 2024.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet Foundations. Journal of Machine Learning Research (JMLR), 2023.
Alina Beygelzimer, John Langford, Lihong Li, Lev Reyzin, and Robert E. Schapire. Contextual Bandit Algorithms with Supervised Learning Guarantees. In Artificial Intelligence and Statisties, pp. 19-26. PMLR, 2011. 1. URL https: / /ttp/aaavvvv/..////v///.//...avs org/
http://wwwww..../////////oo/a//t 101.0.004440440 4058.
Felix Biggs and Benjamin Guedj. On Margins and Derandomisation in PAC-Bayes. In International Conference on Artificial Intelligence and Stotistics, pp. 3709-3731. PMLR, 2022. URL https:
/ //proceedings mlr. prwps/eeeeeee.pee. /vllll1ll//// htt./////lg/////eeeeee http//iggggaaaaaa html.
Felix Biggs and Benjamin Guedj. Tighter PAC-Bayes generalisation bounds by leveraging example difficulty. In International Conference on Artificial Intelligence and Statistics, pp. 8165-8182. PMLR, 2023. URL https:ttpp... http::::////oppediipoaooo ngs ml wtt.:sssrr..prree/ee0m /iiees//////0//gggg23333 html.

---

Published as a conference paper at ICLR 2025
Stéphane Boucheron, Gábor Lugosi, and Pascal Massart. Concentration inequalities: A nonasymp- totic theory of independence. Oxford University Press, 2013.
loar Casado, Luis A Ortega, Aritz Pérez, and Andres R Masegosa. MAC-Bayes-Chernen bounds for unbounded losses. In Advances in Neural Information Processing Systems, 2024.
Olivier Catoni. PAC-Bayesian Supervised Classification: The Thermodynamies Statistical Learn-
ing. Institute of Mathematical Statisties, LECTURE POTESS-ONOGRAPH SERIES, 56, 2007. URL https://////.cc arriiv..v... org/ abs/ 0712. 0248.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang. Sampling is as easy as learning the score: theory for diffuusin models with minimal data assumptions. In International Conference on Learning Representations, 2023.
Yihang Chen and Lukas Mauch. Order-Preserving GFlowNets. In International Conference on /2ds/2211 00386.
Learming Representations, 2024. URL /
David Cohn, Les Atlas, and Richard Ladner. Improving generalization with active learning. Ma- chine Learning, 15:201-221, 1994. URL htttt:://iiim springer wttp:/:...o/aaattteeeeeeeeeeeeee
10C7/bE00993277.
Tiago da Silva, Eliezer Silva, Adèle Ribeiro, António Góis, Dominik Heider, Samuel Kaski, and Diego Mesquita. Human-in-the-Loww Causal Discovery under Latent Confounding using Ances- tral GFlowNets. arXiv preprint, 2023. URL httttt////aaat
org/ aaaaab//200 2309. 12032.
Tiago da Silva, Luiz Max Carvalho, Amauri Souza, Samuel Kaski, and Diego Mesquita. Embar- rassingly Parallel GFlowNets, 2024. URL https: org/ /tsg/bb//////a/// / 2406 03288.
Tristan Deleu and Yoshua Bengio. Generative Flow Networks: Markov Chain Perspective. arXiv preprint arXiv:2307 012307.01420 0122011222222 2023, URL https: www.oogg//////oooo//////eeeeeaae hatgg/abs/23000 02307.004222
Tristan Deleu, António Góis, Chris Chinenye Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian Structure Learning with Generative Flow Networks. In Uncertainty in Artificial Intelligence (UAI), 2022.
Tristan Deleu, Mizu Nishikawa-Toome Toomey, Jithendaraa Subramanian, Nikolay Malkin, Laurent Char- lin, and Yoshua Bengio. Joint Bayesian Inference of Graphical Structure and Parameters with a Single Generative Flow Network. In Advances in Neural Processing Systems (NeurIPS), 2023.
Tristan Deleu, Padideh Nouri, Nikolay Malkin, Doina Precup, and Yoshua Bengio. Discrete Proba- bilistic Inference as Control in Multi-path Environments, 2024. URL https://////ee. http://aaxvvvorooe
abs/2402. 10309.
Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA:
Efficient Finetun- ing of Quantized LLMs. In Advances in Neural Information Processing Systems, volume 36, pp- 10088-10115, 2023. URL hww.t:tpsttto....m /ttptpp::///xii..//// http//aawwv.rrr..a..////// /abs//////// 2305 14314.
Yilun Du and Leslie Kaelbling. Compositional Generative Modeling: A Single Model Not All You Need, 2024. URL https::t///.o. //apiiaaaaivv.//aaa//aa/aaaaaiia org/ htgg/bs//////o abs//402/ 01103.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Uncertainty in Artificial Intelligence (UAI), 2017.
Gintare Karolina Dziugaite and Daniel M Roy. Data-dependent PAC-Bayes priors via differential privacy. In Advances in Neural Information Processing Systems, volume 31, 2018.
Gintare Karolina Dziugaite, Kyle Hsu, Waseem Gharbieh, Gabriel Arpino, and Daniel M. Roy. On the role of data in PAC-Bayes bounds. Artificial Intelligence and Statisties, pp. 604-612. PMLR, 2021. URL https:/ peatts:://rxtte /waxiviivvv. org/ wtgg/bs/////oooooooo//o///aaa
10929.
M. Fard and Joelle Pineau. PAC-Bayesian Model Selection for Reinforcement Learning. In J. Laf- ferty, C. Williams, Shawe-Taylor R. Zemel, and A. Culotta (eds.), Advances in Neural Infor- mation Processing Systems, volume 23. Curran Associates, Inc., 2010.
12

---

Published as a conference paper at ICLR 2025
Ronald Aylmer Fisher and Frank Yates. Statistical tables for biological, agricultural and medical research. Edinburgh: Oliver and Boyd, 1963.
Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep Bayesian Active Learning with Image Data. In International Conference on Machine Learning, pp. 1183-1192 PMLR, 2017.
Timur Garipov, Sebastiaan De Peuter, Ge Yang. Vikas Garg, Samuel Kaski, and Tommi S. Jaakkola. Compositional Sculpting of Iterative Generative Processes. In Advances in Neural Information Processing Systems, 2023.
Pascal Germain, Francis Bach, Alexandre Lacoste, and Simon Lacoste-Julics PAC-Bayesian The- ory Meets Bayesian Inference. In Advances in Neural Information Processing Systems, 2016.
URL https://///. http:///wwwwwoo///..at/ooooooo/t
wrgg/bssgsaaa... abs/ 605. 08636.
Benjamin Guedj. A Primer PAC-Bayesian Learning. In Proceedings of the French Mathematical Society, volume 33, pp. 391-414. Société Mathématique de France, 2019. URL https:// arxiv.org/abs 901 05353.
Maxime Haddouche and Benjamin Guedj. Online PAC-Bayes Learning. In Advances in Neural Information Processing Systems, 2022. URL wwtph::tpppaaataeee
abtg//bs/22066 200240
Maxime Haddouche and Benjamin Guedj. PAC-Bayes generalisation bounds for heavy-tailed losses through supermartingales. Transactions on Machine Learning Research, 2023.
Maxime Haddouche, Paul Viallard, Umut Simsekli, and Benjamin Guedj. A PAC-Bayesian Link Between Generalisation and Flat Minima. In Algorithmiic Learning Theory, 2024. URL https:
/ www.iiv.rggrooooooooooooo /2gg////2////22222/////// 2402. 08508.
Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural computation, 9(1):1-42, 1997.
Matthew Holland. PAC-Bayes under potentially heavy tails. In Advances in Neural Information Processing Systems, volume 32, 2019. URL /estts://arxtee org/ 0ttps/1955.09999
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax Flows: Learning Categorical Distributions with Normalizing Flows. Third Symposium on Advances in Approximate Bayesian Inference, 2021.
Edward J Hu, Nikolay Malkin, Moksh Jain, Katie Everett, Alexandros Graikos, and Yoshua Ben- gio. GFlowNet-EM for learning compositional latent variable models. In International Confer- ence on Machine Learning (ICML.,, 2023.
Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, and et al. Amortizing intractable inference in large language models. In International Conference on Learning Representations, 2024. URL https: arrxiv org/ abs 2210/ 01.300000 04363.
Moksh Jain, Emmanuel Bengio, Alex Hernandez-Garcha Jarrid Dector-Brooke Bonaventure F. Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological Sequence Design with GFlowNets. In Inter- national Conference Machine Learning (ICML), 2022.
Moksh Jain, Sharath Chandra Raparthy, Alex Hernandez-Garee Garcia, Jarrid Rector-Brooks Yoshua Ben- gio, Santiago Miret, and Emmanuel Bengio. Multi-Objective GFlowNets. In International Con- ference on Machine Leurning (ICML), 2023.
Hyosoon Jang, Minsu Kim, and Sungsoo Ahn. Learning Energy Decompositions for Partial In- ference GFlowNets. In International Conference on Learning Representations, 2024. URL
https: http:s///peerrvvie. net/ http:eew.nnntoootutnteeeeeeeeeee wtt.:/id-eeimmiimaaa..
Marco Jiralerspong, Bilun Sun, Danilo Vucetic, Tianyu Zhang, Yoshua Bengio, Gauthier Gidel, and Nikolay Malkin. Expected flow networks in stochastic environments and two-player zero- sum games. In International Conference on Learning Representations, 2024. URL https: / //arasiv.org/aba//3300000
13

---

Published as a conference paper at ICLR 2025
Haotian Ju, Dongyue Li, Aneesh Sharma, and Hongyang R Zhang. Generalization in graph neural networks: Improved PAC-Bayeste bounds on graph diffusion. In International Conference on Artificial Intelligence and Statistics, pp. 6314-6341, PMLR, 2023.
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On Large- Training for Deep Learning: Generalization Gap and Sharp Minima, 2017. URL https::////. wawiiriiiivv../.ia .rgg abs/ 609. 04836.
Minsu Kim, Joohwan Ko, Taeyoung Yun, Dinghuai Zhang, Ling Pan, Woochang Kim, Jinkyoo Park, Emmanuel Bengio, and Yoshua Bengio. Learning to Scale Logits for Snaperature-Condiiionn GFlowNets, 2024a. URL https: /aaxiviiivv.////////aaaaa//aaava wrxiv.orgg//ooo......o...ao.oo.r abs/ 2210..110000 02823.
Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local Search GFlowNets, 2024b. URL htt.ttttt//aaaee wwwpaawwwwwwooaaaaaaaarorrreea org/ wbwg/ba/////// 02710.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representatioss 2015.
6980.
Anas Krichel, Nikolay Malkin, Salem Lahlou, and Yoshua Bengio. On Generalization for Generative Flow Networks, 2024. URL https:tt//// http://axxvvvrroo http/bbs/220 /2007 03105.
Salem Lahlou at al. A Theory of Continuous Generative Flow Networks. In International Con- ference on Machine Learning (ICML), 2023. URL ht https mlr, p.rssp/ w202/lahl oo23a html.
Elaine Lau, Nikhil Murali Vemgal, Doina Precup, and Emmanuel Bengio. DGFN: Double Genera- tive Flow Networks. In NeurlPS 2023 Generative Al and Biology (GenBio) Workshop, 2023.
Puheng Li, Zhong Li, Huishuai Zhang, and Jiang Bian. On the generalization properties of diffusion models. In Advances in Neural Information Processing Systems, volume 36, 2024.
Jens Liebehenschel. Ranking and unranking of lexicographically ordered words: An average-case analysis. Journal of Automata, Languages and Combinatorics, 2(4):227-268, 1997.
Dianbo Liu et al. GFlowOut: Dropout with Generative Flow Networks. In International Conference on Machine Learning. PMLR, 2023, URL
mlr wwwp://ss.........eeeeeeeee
liu23r. htmll.
Ben London, Bert Huang, Ben Taskar, and Lise Getoor. PAC-Bayesian Collective Stability, In Samuel Kaski and Jukka Corander (eds.). Proceedings of the Seventeenth International Confer- ence on Artificial Intelligence and Statistics, volume 33 of Proceedings of Machine Learning Research, pp. 585-594, Reykjavik, Iceland, 22 25 Apr 2014. PMLR.
Sanae Lotfi, Mare Finzi, Yilun Kuang. Tim G. J. Rudner, Micah Goldblum, and Andrew Gordon Wilson. Non- Vacuous Generalization Bounds for Large Language Models. In International Conference on Machine Learning, pp. 32801-328 PMLR, 2024a. URL org/abs/ 2312. 17173.
Sanae Lotfi, Yilun Kuang, Brandon Amos, Micah Goldblum, Marc Finzi, and Andrew Gordon Wilson. Unlocking Tokens as Data Points for Generalization Bounds on Larger Language Models. In Advances Neural Information Processing Systems, volume 37, pp. 9229-9256, 2024b. URL. https: : aarxiv... org/ alss / 2407 18158.
Jianzhu Ma, Jian Peng, Sheng Wang. and Jinbo Xu. Estimating the Partition Function of Graphi- cal Models Using Langevin Importance Sampling. In Carlos M. Carvalho and Pradeep Raviku- mar (eds.), Artificial Intelligence and Statistics, pp. 433-441. PMLR, 2013. URL http: / //proceedings mlr. press 11311mmalaa.pd...
Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, An- drei Cristian Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. In International Conference on Machine html.
Learning, 2022. URL https: httpps://pooceeeeeeeeeee http///oooediioeocoo ngs. ml
14

---

Published as a conference paper at ICLR 2025
Eran Malach. Auto- Auto-Regressive Next-Token Predictors are Universal Learners. In International Conference on Machine Learning, pp. 34417-34431. PMLR, 2024.
Shreshth A. Malik, Salem Lahlou, Andrew Jesson, Moksh Jain, Nikolay Malkin, Tristan Deleu, Yoshua Bengio, and Yarin Gal. BatchGFN: Generative Flow Networks for Batch Active Learning, 2023. URL https::///.. : / arxiv. org/ atsgg//b////e wtw.ab////.... 15058.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory bal- ance: Improved credit assignment in GFlowNets. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. In International Conference Learn- ing Representations (ICLR), 2023.
Daniel J. Mankowitz, Timothy A. Mann, and Shie Mannor. Adaptive Skills, Adaptive Parti- tions (ASAP). In Advances in Neural Information Processing Systems, 2016. URL https: //arxiv, orgg abs 0tt.//16033333 03351.
Andreas Maurer. A Note on the PAC Bayesian Theorem. arXiv preprint cs/0411099, 2004. URL
https://arxia: org/ abs cs//cc//// 0411099.
Sokhna Diarra Mbacke, Florence Clerc, and Pascal Germain. PAC-Bayesian Generalization Bounds for Adversarial Generative Models. In International Conference on Machine Learning, pp-
24271-24290. PMLR, 2023. URL /estps://arxi...e. org/ abs abs//3/2/222222 08942.
David McAllester. A PAC-Bayesian Tutorial with A Dropout Bound. arXiv preprint arXiv:1307.2118, 2013. URL
/ttgg//b//i/ 117111370000 2118.
David A McAllester. Some PAC-Bayesian theorems. In Proceedings of the eleventh annual confer- ence on Computational Learning Theory, pp. 230-234, 1998.
David A McAllester. PAC-Bayesian model averaging. In Proceedings of the twelfth annual confer- ence on Computational Learning Theory, pp. 164-170, 1999.
Colin McDiarmid. Concentration. In Probabilistic Methods for Algorithmic Discrete Mathematics, pp. 195-248, Springer, 1998.
Kohei Miyaguchi. PAC- Bayesian Transportation Bound. arXiv preprint arXiv:1905.13453 2019.
URL https: /ttp://aaavvvv... org/ abs 1905.134335
Wendy Myrvold and Frank Ruskey, Ranking and unranking permutations in linear time. Information Processing Letters, 79(6)221-228) 2001.
Andrei Cristian Nica. Moksh Jain, Emmanuel Bengio. Cheng-Hac Liu, Maksym Korablyov, Michael M Bronstein, and Yoshua Bengio. Evaluating generalization GFlowNets for molecule design. ICLR2022 Machine Learning for Drug Discovery workshop, 2022.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better Training of GFlowNets with Local Credit and Incomplete Trajectories. In International Conference on Machine Learning (ICML), 2023a.
Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, and Yoshua Bengio. Generative Aug- mented Flow Networks. In International Conference on Learning Representations (ICLR), 2023b.
Mohit Pandey, Gopeshh Subbaraj, and Emmanuel Bengio. GFlowNet Pretraining with Inexpensive Rewards. NeurlPS 2024 Workshop AIDrugX, 2024.
María Pérez-Ortiz, Omar Rivasplata, John Shawe-Taylor, and Csaba Szepesvári. Tighter risk cer- tificates for neural networks. Journal of Machine Learning Research, 22(227):1-40, 2021. URL
http:// /mlr org/ papers v 20 879 html..
org/papers
Jarrid Rector-Brooks, Kanika Madan, Moksh Jain, Maksym Korablyov, Cheng-Hao Liu, Sarath Chandar, Nikolay Malkin, and Yoshua Bengio. Thompson sampling for improved exploration in GFlowNets, 2023. URL https: //arxiv org/ abbs /230 17693.
15

---

Published as a conference paper at ICLR 2025
Omar Rivasplata, Vikram M Tankasali, and Csaba Szepesvari. PAC-Bayes with backprop. arXiv preprint urXiv:1998.0775 073300 2019. URL https: https://arxx...orr//aa/a/////
07380.
Omar Rivasplata, IIja Kuzborskij, Csaba Szepesvari, and John Shawe-Taylor. PAC-Bayes Analysis Beyond the Usual Bounds. In Advances in Neural Information Processing Systems, 2020. URL
attpps:/aaree org/ abs 2006. 13057.
Borja Rädve, Omar Rivasplata, Ragnar Thobaben, and Mikael Skoglund. A note on generalization bounds for losses with finite moments. In 2024 IEEE International Symposium on Information Theory (ISIT), pp. 2676-2681. IEEE, 2024a.
Borja Rodríguez-Gälvez, Ragnar Thobaben, and Mikael Skoglund. More PAC-Bayes bounds: From bounded losses, losses with general tail behaviors, to anytime validity. Journal of Machine Learning Research, 25(110):1-43 2024b. URL https: //axivi.vv./ org/ hbtp/abs/220 2306 12214.
Julien Roy, Pierre-Luc Bacon, Christopher Pal, and Emmanuel Bengio. Goal-conditioned GFlowNets for controllable multi-objective molecular design. ICML 2023 Workshop Deploy- ableGenerativeAl, 2023. URL
wwtg/bsg//////o//e/..ooeee 04
Otmane Sakhi, Pierre Alquier, and Nicolas Chopin. PAC-Bayesian offline contextual bandits with guarantees. In International Conference on Machine Learning, pp. 29777-29799, PMLR, 2023.
Milad Sefidgaran, Romain Chor, and Abdellatif Zaidi. Rate-Distortion Theoretic Bounds on Gener- alization Error for Distributed Learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Yevgeny Seldin, Nicolò Cesa-Bianchi, Peter Auer, François Laviolette, and John Shawe-Taylor PAC-Bayes-Bernsee inequality for martingales and its application to multiarmed bandits. In Proceedings of the Workshop on On-line Trading of Exploration and Exploitation 2, pp. 98-111. PMLR Workshop and Conference Proceedings, 2012a.
Yevgeny Seldin, François Laviolette, Nicolò Cesa-Bianchi, John Shawe-Taylon and Peter Auer. PAC-Bayesian Inequalities for Martingales. IEEE Transactions on Information Theory, 58( (1):;) 7086-7093, 2012b. URL htttps://arii.. http:///wwxx......aa./oo http:o////abbaaaaooe/ 6886.
Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector-Brooks Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. Improved off-policy training of diffusion samplers, 2024. URL / wttpaaawww...rrraaaaaaa/aaaaaaae 2402 050988
Shai Shalev-Shwartz and Shai Ben-David. Understanding Machine Learning: From Theory to Algorithms. Cambridge University Press, 2014.
John Shawe-Taylor and Robert C Williamson. A PAC analysis of Bayesian estimator. In Proceed- ings of the tenth annual conference on Computational Learning Theory, pp. 2-9, 1997.
Max W. Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso Biancalani. Towarrss Understanding and Improving GFlowNet Training. In Interna- tional Conference on Machine Learning, 2023.
Tiago Silva, Eliezer de Souza da Silva, Rodrigo Barreto Alves, Luiz Max Carvalho, Amauri H Souza, Samuel Kaski, Vikas Garg, and Diego Mesquita. Analyzing GFlowNets: Stability, Expres- siveness, and Assessment. In ICML 2024 Workshop on Structured Probabilistic Inference & Gen- erative Modeling, 2024, URL httpt:tt////.. /ttp:////eeeeeiieeeeemme .et Leeeet/orrmmmmd-BXXXmma B8==8KX iF
Huayi Tang and Yong Liu. Towards understanding generalization of graph neural networks. In International Conference Machine Learning, pp. 33674-33719, PMLR, 2023.
Bahareh Tasdighi, Abdullah Akgül, Manuel Haussmann, Kenny Kazimirzak Brink, and Melih Kan- demir. PaC-Bayesian Soft Actor-Critic Learning. In Symposium Advances in Approximate Bayesian Inference, pp. 127-145. PMLR, 2024. URL https:/ /tttp:://xxivv orgg abs/2301.. 12776.
16

---

Published as a conference paper at ICLR 2025
Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry Vetrov. Generative Flow Networks as
RL. In International Conference on Artificial Intelligence and Statistics, pp. 4213 4221. PMLR, 2024. URL https://aaxxy.. .rg/ abs 22310.12934...
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-Bench: Bench- marks for Data-Driven Offline Model-Based Optimization. In International Conference on Ma- chine Learning, pp. 1665-221676. PMLR, 2022. URL https: / wwwaa.aivooo.............. abs/2202. 08450.
Dustin Tran, Keyon Vafa, Kumar Krishna Agrawal, Laurent Dinh, and Ben Poole. Discrete Flows: Invertible Generative Models of Discrete Data. Advances Neural Information Processing Systems, volume 32, 2019. URL https:://///. //appss///xxiiii//// org abs
Leslie G Valiant. A theory of the learnable. Communications of the ACM, 27(11): 1142, 1984.
Vladimir Vapnik. Statistical Learning Theory. John Wiley & Sons, 1998.
Vladimir Vapnik. The Nature of Statistical Learning Theory. Springer New York, 2000. URL
http:/ /dxxxxx///////////// doi. wttp///.///////////.oo
Vladimir N Vapnik and A Ya Chervonenkis. On the uniform convergence of relative frequencies of events to their probabilities. In Measures of Complexity: Festschrift for Alexey Chervonenkis, pp- 11-30. Springer, 2015.
Nikhil Vemgal, Elaine Lau, and Doina Precup. An Empirical Study of the Effectiveness of Using Replay Buffer on Mode Discovery in GFlowNets. ICML 2023 Workshop on Structured Prob- abilistic Inference & Generative Modeling, 2023. URL https: arxiv. org/ abs/2307. 07674.
Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, Alexandre Adam, Jarrid Rector- Brooks, Yoshua Bengio, Glen Berseth, and Nikolay Malkin. Amortizing intractable inference in diffusion models for vision, language, and control. Advances in Neural Information Processing Systems, 2024. URL ptttts://arxiv.oot/ orgg/ abs 420 20971.
Yi-Shan Wu, Andres Masegosa, Stephan Lorenzen, Christian Igel, and Yevgeny Seldin. Chebyshev- Cantelli PAC-Bayes-Bennet Inequality for the Weighted Majority Vote. In Advances in Neural Information Processing Systems, 2021.
Semih Yagli, Alex Dytso, and H. Vincent Poor. Information-Theoreee Bounds on the Generalization Error and Privacy Leakage in Federated Learning. In 2020 IEEE 21 International Workshop Signal Processing Advances Wireless Communications (SPAWC), pp- 1-5. IEEE, 2020. URL https: arxiv org/ abss/ 20005 02503.
David W Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan. Robust scheduling with GFlowNets. In International Conference on Learning Representations (ICLR), 2023a.
Dinghuai Zhang, Hanjun Dai, Nikolay Malkin, Aaron Courville, Yoshua Bengio, and Ling Pan. Let the Flows Tell: Solving Graph Combinatorial Optimization Problems with GFlowNets. In Advances in Neural Information Processing Systems (NeurIPS), 2023b.
Ming Yang Zhou, Zichao Yan, Elliot Layne, Nikolay Malkin, Dinghuai Zhang, Moksh Jain, Math- ieu Blanchette, and Yoshua Bengio. PhyloGFN: Phylogenetic inference with generative flow networks. In International Conference on Learning Representations, 2024.
Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Hoi, and Weinan E. Towards Theoretically Understanding Why SGD Generalizes Better Than ADAM in Deep Learning. In Advances in Neural Information Processing Systems, volume 33, pp. 21285-21296, 2020. URL https: // //aaxivv wttpabs///aa/aaaoo/o 05627.
17

---

Published as a conference paper at ICLR 2025
SUPPLEMENTARY MATERIAL FOR "GENERALIZATION AND DISTRIBUTED LEARNING OF GFLOWNETS"
Background and related works
19
A.1 Directed Acyclic Graphs
19
A.2 Generative Flow Networks
19
A.3 Learning GelowNets
20
A.4 Related works
20
A.5 Additional review of PAC-Bayes bounds
B Experimental details and additional discussions
23
B.1 A non-generalizable distribution
23
B.2 Non-vacuous generalization bounds
23
B.3 Oracle generalization bounds: Lemmata
24
B.4 Subgraph Asynchronous Learning
25
C SAL: Implementation and Theoretical Analysis
26
C.1 Summary of the experimental campaing
26
C.2 An efficient implementation of SAL
26
C33 Theoretical analysis and extensions..
30
C.4 Conditional SAL
32
C.5 SAL and EP-GFlowNets
32
Proofs
34
D.I Proof of Lemma 4.1
34
D.2 Proof of Proposition 4.2
D.3 Proof of Proposition 5.1
35
D.4 Proof of Lemma B.1
36
D.5 Proof of Theorem 5.2
36
D.6 Proof of Lemma B.2
36
D.7 Proof of Theorem 5.4
36
D.8 Proof of Theorem C.I
38
D.9 Proof of Lemma C.3
39
D.10 Proof of Proposition C.5
D.11 Proof of Proposition C.6
40
Limitations and future works
41
E.1 Additional experiments
41
18

---

Published as a conference paper at ICLR 2025
A BACKGROUND AND RELATED WORKS
For probability measures P and Q on the same space X, we recall for convenience that the Kullback-Leibler (KL) divergence is KL(PIIQ)
En> [log (4P/aQ) (z)], the chi-squared divergence is (²())P))= = B.~P ((0Q/aP) (x)) -1] - and the total variation distance is given by (P,Q) SupACa IP(4) Q(A)- When X' is finite, QP) = 1/2 Enex IP(x) - (x))- There are other notions of "divergence" for probability measures; we have mentioned here the ones used in our paper. Readers are refereed to Boucheron et al. (2013) for further details on the topic.
A.1 DIRECTED ACYCLIC GRAPHS
We briefly recall the definition of a pointed directed acycliy graph. For this, we firstly define the concept of a finitely abxorbing Markov transition kernel (Lahlou at al., 2023) in a topological space. Henceforth, we let ({s,) US {sy}, V) be topological space endowed with topology V and two special elements, So and sy- We denote S {s. (sy()y)) } the state space; 8o and 8 the initial final states, respectively. We also assume that both {5o} and {8y} are open sets with respect to V. Definition A.1 (Finitely absorbing Markov transition kernel (MTK)). Consider the measure space (S, M, (1) with measure il and a Borel o-algerba Σ. Let k: Σ + R. be a reference kernel, i.c., n(s..) E R+ is measure absolutely continuous with respect μ for all s, and we recursively define x **' As A)= J (,8t-((s, ds')dss dss', A) for measurable A E Σ. We say that pr: S x Σ + R+ is a finitely absorbing Markov transition kernel if the following conditions are satisfied.
1. pF(8, :): + + is an absolutely continuous probability measure with respect to k((s:);
2. there is tm < oc such that Ztmm (s, {s = 1 for every E US and pr {8y}) = 1;
3. + p= B) is continuous for every measurable B E Σ;
4. if pF(ss { > 0, then pf(s, {sy}) =
5. for every A Σ, there is at < fm. such that pf (so, A) > 0.
In this work, S is always finite, V is the discrete topology, and is the counting measure. The state graph G is induced by pF. i.e., (u, v) is an edge in G if and only if pr {u}) > 0. Acyclicity is ensured by the finitely absorbing property of Pr (item 2 of Definition A.I). Notably, the finite S assumption covers the vast majority of use-cases for GelowNets. Under these conditions, we say k is a backward reference kernel in S with respect to k if n(u, {e}) = k (v, {u}) for all (u, v) E S S. We refer the reader to (Lahlou at al., 2023) for an overview of GFlowNets in infinite spaces.
A.2 GENERATIVE FLOW NETWORKS
A GFlowNet can be seen as tuple (fso} UUSJUSUUSSS (sy), Pp, PB, PF, PB,PP,, K, RT Hr, R) for which
1. x is a forward reference kernel on S:
2. kT is backward reference kernel in S with respect to K;
3. PF (resp. PB) is a finitely aborsbing MTK with respect to n (resp. x ))
4. R: Σ 3-42 R4, is measure such that R < μr;
5. PF: S x Σ + 1R, is MTK, called the forward policy, such that (s < pF(s,-);
6. PB: S Σ RR is a MTK, called the backward policy, such that Ps(s. < pB(S.-).
We denote by PF and pB the densities of Pp and PB with respect to their respective reference kernels. For simplicity, we interchangeably let R(x) be the density of R with respect to μ. In this scenario, the set of terminal states X' is defined by X = {x E S: Pr(s, {sy}) > 0}. In practice, Pr is parameterized by neural network and its parameters are estimated to ensure that the marginal of (r.s..)) .) over X' matches R up to a normalizing constant. In the terminology of Section 2, the abstract actions space A would correspond A U.eS (s, u): Pr(s, {w}) > 0} and A(s) {(s, u): (s. {u)) 0}. For most problems. we identify the edge (s. u) with an entity representing the difference between u and s, e.g., a nucleotide base when S is the space of nucleotide strings. We complement the discussion in Section 2 on how to learn GFlow Net in the next section.
19

---

Published as a conference paper at ICLR 2025
A.3 LEARNING GFLOWNETS
Below, we illustrate our definition of Gets for three common generative tasks. These tasks encompass a large number of applications, e.g., Jain al. (2022); Shen et al. (2023); Hu et al. (2024; 2023); Liu et al. (2023); Malkin et al. (2022); Pan et al. (2023b); Madan et al. (2022).
1. Autoregressive generation. Each object in S is string of length up to a I., and G is a tree rooted at So- Also, action sets A(s) represent an alphabet and transition T(s, a) appends the character a to the string s. Here, (s a)) = 1 for every s E S and a E A(s).
2. Set generation. Each s E S is a subset of W = = ......, .,,,,,,,,, W with X' containing those s of size T. Action sets are A(s) W and transitions T(s, add the element to 5; see Figure 5,
3. Hypergrid environment. Each s E S is point within (O..... H - 11ah x {0, 1 for given H (size) and d (dimension); 8, = O2+1 and the last coordinate indicates whether A(s) {e;: s² < H H-1 ( with ej denoting the i-th canonical vector in IR and T stop action. Transitions (ss, a) either add 8, if e, for some é; or set sd+1 1, a = T.
X. Also,
Notably, -(x,-) E {sy} for every terminal state x E X. We provided examples of R. F. and pr throughout the main text; in particular, see Sections 4 and 6.2 and Appendix B. Figure illustrates the state graph for the set generation task (omitting 8y). To learn a forward policy pF, we minimize a stochastic objective based on the observed trajectories. Besides the ones shown in Equation 1, many loss functions has been recently proposed. The SubTB loss (Madan et al., 2022), for instance, is defined by
n.23
0
{2}
(1. s1
(2.
Am-n
log F(Tn )PF (Tuu)))m) PB (Trmm) (Tm)
Figure 5: State graph for the set generation task (W = 3, T 2).
Ey-pe
Σ
E Eignemsih| An-
ISn<mElT
with the constraint that F(x)= = R(x)) for x E X' and Tn representing the nith element within the trajectory т. Corresponningly, the VarGrad (Zhang et al., 2023a) and contrastive balance (da Silva et al., 2024) objectives avoid the estimation of F by minimizing
Ey [(oggp(()) loggp(r') logp logpa(7'lz") - log R(x) log R(2'))1
which led to faster training convergence in some cases. On the same page, Malkin et al. (2023) considered series of divergence-based loss functions for training GFlowNets, showing that the on-policy version of the TB loss (Equation 1) corresponds to the reverse KL between the forward and backward policies in terms of the gradients. In particular, we note that these learning objectives can only be used for estimating the parameters of the root network in SAL. For the leaf networks, which must provide an estimate of the flow function these loss functions cannot be used. Nonetheless, flow-based learning objetctives such as TB and SubTB often exhibit a convergence speed comparable to that of variational alternatives and are frequently implemented for large-scale applications (Nica et al., 2022; Jain et al., 2022; Hu et al., 2024; Zhou et al., 2024). Learning objectives aside, there a growing interest in the literature in the development of more effective parametrizations for GFlowNets, with remarkable results for the forward-looking GFlowNets (Pan et al., 2023a) and LED-GFlowNess (Jang et al., 2024), which residually reparameterize F as log F(s) = log (b() + log ((s) for a (given or learnable) é, Sewaatture-salle---lllooooo
for the aggregation step,
(Kim et al., 2024a), in which (F(''(s) [s) x exp(o(3)) </sss(s)) [s)) for neural networks & and & and an inverse-temperature parameter 3 > 0, and QGFN (Lau et al.,
2023), which learns Q-function concomitantly to F' and pp and prune the values of pv based on Q during inference time for controlable greediness.
A.4 RELATED WORKS
GFlowNets (Bengio et al., 2021; 2023; Lahlou at al., 2023) were canonically proposed as a rein- forcement learning algorithm for sampling compositional objects (e.g., graphs) proportionally to prespecified reward function. From a theoretical perspective, the relationship between GFlowNets and variational inference (Malkin et al., 2023), entropy- regula--e Q-learning (Tiapkin al., 2024; Deleu et al., 2024), and diffusion models (Lahlou at al., 2023; Sendera et al., 2024; Venkatraman et al., 2024) has been thoroughly established. From practitioner's viewpoint, GetowNets have
20

---

Published as a conference paper at ICLR 2025
been successfuly applied many problems including, but not restricted to, causal discovery (Deleu et al., 2022; 2023; da Silva et al., 2023), Bayesian phylogenetic inference (Zhou et al., 2024; da Silva et al., 2024), language and image modelling (Hu et al., 2024; Liu et al., 2023; Hu et al., 2023; Venkatraman et al., 2024), combinatorial optimization (Zhang et al., 2023a;b), and drug discovery (Bengio et al., 2021; Nica et al., 2022; Vemgal et al., 2023; Pan et al., 2023a). Indeed, we are confident that problems such as language modelling and drug discovery could greatly benefit from SAL if appropriate policy networks and fixed-horizon partitionings are designed. Nonetheless, given the open-endedness and specialized nature of these applications, we believe that they would be more suited for future, dedicated works and are, hence, not addressed in this text. Correspondingly, recent work by Jiralerspong et al. (2024) highlighted the competitive performance of stochastic GFlowNets in two-player zero-sum games, specifically, Tic-Tac-Toe and Connect-4, and we are optimistic that an extension of SAL to stochastic environments would exhibit promising results for games having larger trajectories, c.g., Chess and Go. Orthogonal to these advances, the issue of generalization in GFlowNets has also received significant attention in the literature (Atanackovic & Bengio, 2024; Krichel et al., 2024). In sharp contrast to previous works, ours is the first one that derives PAC-Bayesian bounds and provides non-vacuous statistical guarantees for GFlowNets, along with a theoretical analysis that highlights which factors are potentially harmful to the model's generalization performance. Notably, a recent discussion by Bengio & Malkin (2024) provides an interesting perspective on generalization, active learning, and GFlowNets in the context of abstract reasoning for muuchine-learing learning-based based theorem proving and conjecture formation. Concomitantly, we note that there is a well-established interest in the community towards the development of more sample-efficient learning objectives for speeding up training convergence (Malkin et al., 2022; Deleu et al., 2022; Madan et al., Zhang et al., 2023a; da Silva et al., 2024; Tiapkin et al., 2024).
A.5 ADDITIONAL REVIEW OF PAC-BAYES BOUNDS
Historically, McAllester (1998; 1999)'s s PAC-Bayesian theorems, which were inspired by the work of Shawe- Taylor& & Williamson (1997), were developed towards the objective of providing Probably Approximately Correct (PAC) guarantees to Bayesian algorithms with potentially misspecified prior distributions. Recently, the relationship between PAC-Bayesian theory and (approximate) Bayesian algorithms has been made explicit by Germain et al. (2016). From this perspective, Alquier (2024) provides an informative and comprehensive account of the literature on PAC-Bayes bounds. both theory and applications. In what are now well-established references, Catoni (2007) gives a rigorous foundation of PAC-Bayes bounds in supervised classification, including results on the form of the distributions that optimise the bounds; and Guedj (2019) provides a nice concise exposition of the essential form of PAC-Bayesian inequalities. In the context of contemporary machine learning, PAC-Bayesiaa theory has found enormous success in the development of numerical generalization bounds for overparameterized neural network classifiers, achieving non-vacuous results (Dziugaite & Roy, 2017; 2018) and tight certificates (Pérez-Ortiz et al., 2021), which subsequent works have applied even for large language models (Lotfi et al., 2024a;b) with billions of parameters through appropriate compression techniques (Dettmers et al., 2023). In a recent work, Malach (2024) intro- duced the notion of length complexity for next-token autoregressive learning on Chain-of-Thougeee data, referring to the minimum number of iterations required by an AR learner to compute target function, which is (vaguely) connected to our results regarding the harmful effects of the maximum trajectory size on GFlowNet learning. Importantly, the advantegeousness of distributed approaches for the generalization performance of learning algorithms was already pointed out by Yagli et al. (2020); Barnes et al. (2022); Sefidgaran al. (2022); similarly SAL, these authors consider the problem of training a set of models in parallel and subsequently aggregating them with a (possibly randomized) estimator in a central server. In spite of these advances, the development of tighter PAC-Bayes bounds with weaker assumptions on the risk functional, e.g., heavy tailedness instead of boundedness, is still an active research field (Holland, 2019; Wu et al., 2021; Balsubramani, 2015; London et al., 2014; Biggs & Guedj, 2023; Rivasplata et al., 2020; Rodríguez-Gáález et al., 2024a;b). Also, the development of PAC-Bayesian theory in the setting of non-i.i.d. data is still relatively underdeveloped when compared against other branches of machine learning, albeit there are interesting results in online learning (Haddouche & Guedj, 2022), reinforcement learning (Fard & Pineau, 2010; Beygelzimer et al., 201 I; Sakhi et al., 2023), and time series (Alquier et al., 2012). Finally, TAC-Bayesian theorems provide statistical guarantees for stochastic predictors, which are arguably not frequently used in practice, and the problem of derandomizing the resulting bounds
21

---

Published as a conference paper at ICLR 2025
is still mostly open. Notably, the derandomization of PAC-Bayes bounds has non-negligible cost, and we refer the reader to Miyaguchi (2019); Biggs & Guedj (2022) for further details on this topic.
22

---

Published as a conference paper at ICLR 2025
B EXPERIMENTAL DETAILS AND ADDITIONAL DISCUSSIONS
All experiments were conducted on a single Linux machine with 128 GB of RAM and featuring a NVIDIA RTX 3090 GPU and 12th Gen Intel(R) Core(TM) 19-12900K CPU. Unless specified otherwise, the code for reproducing the experiments below was executed on this GPU.
B.1 A NON-GENERALIZZABE DISTRIBUTION
(W.
(9E. (IpT,) x x2(1 (4..rlppp
For the experiments in Figure 1, we considered the set generation task (see Appendix A) with W E {32, 64] elements to choose from and set size S = 6, and the forward policy was parameterized by an MLP with 2 64-dimensional layers. The elements' log-utilities u were sampled from [-1. 1] prior to training and the resulting values were normalized that the largest reward of set was 5. For both settings in Figure I, the models were trained for 1500 epochs with a batch size of 128. To compute the quantities in Table 1, we compared the uniform policy of an untrained GelowNet against a policy pe such that the (unnormalized) logit corresponding to the addition the element 1 is set to log (l|s)) -1.55 = log 10 10-5, Table shows the large discrepancy between the resulting Pe and an uniform policy, providing a taste for the upper bound in Proposition 4.2.
(32,6) 6)
1.20 10
1.32
(64, 6)
4.56 10
1.24
Table 1: x² divergence between the ex- ploratory (pruned. PE,T. and e-greedy, Pe,r) and uniform (qE,T) distributions.
B.2 NON-VACUOUS GENER GENERALIZATION BOUNDS
A bounded risk functional for GFlowNets. We start recalling the definition of the flo-cconsisteng in subgraphs (FCS) metric (Silva et al., 2024). Given a policy pE, the FCS is defined as
pr(x;)
R(x;)
Lecs R, (R,r)) ET THOPE
2
Σ 1S
EisiSBPr(Z,)
EiSISB R(I,)
(10)
in which B 2 2 is a (typically small) given integer. Equivalently, FCS may be seen the expected total variation distance between the learned pr and target R distributions over random subsets of X. It was shown by Silva et al. (2024) that LFcs(R, Pr) O if and only if (x) x R(x), i.e., the model samples correctly from the distribution proportional to R in X. Then, equipped with the dataset Tn described in Section 4, an unbiased estimate of LFcs is
LPcs
1
Σ
PT (T(
R(ar.) ΣS R(zx,)
(11)
2N
PT(,))
kL..... KB (1,.. (I...... ISISB
in which the outer summation covers N uniformly random B-sized subsets of (1,...,,1] represents the k;th observed terminal state in Tn- Importantly, Lrcs E [0, 1] for any R and pr, which enables the implementation of well-known algorithms for tightening PAC- PAC-Bayesian general- ization bounds through the adoption of data-dependent priors (Dziugaite al., 2021; Maurer, 2004).
_n and IE
Experimental details for computing non-vacuous bounds. To achieve the results illustrated in Figure 2, we use Ta to learn an isotropic Gaussian prior Q with variance 10 over the parameters of an MLP with 3 x 28-dimensional layers defining the forward policy by minimizing the expected TB loss on Ta under Q. For each problem, we used the same architecture of the neural network, changing only the input and output dimensions, and the resulting models were trained for 64 epochs on their respective datasets. Then, we freeze 0 and learn both the mean and the diagonal covariance of Gaussian posterior P over the parameters of policy network by minimizing the upper bound in Equation 6 with LPcs substituted by an unbiased estimate of the TB loss on Ta U Ti-a. Finally, we evaluate the upper bound in Equation 6 on Ti to certify its tightness. We closely followed the experimental setup of Dziugaite et al. (2021); Pérez-Ortiz et al. (2021) for conducting these ex- periments. In particular, the data-splitting protocol for learning the prior, learning the posterior, and evaluating the bound is analogous to the one used by Pérez-Ortiz et al. (2021), Similarly, in contrast to the other experiments, which rely the Adam optimizer (Kingma & Ba, 2015), we use SGD with a fixed learning rate of 10-3 that presumably achieves a flat minimum (Keskar et al., 2017) with potentially better generalization properties (Hochreiter & Schmidhuber, 1997; Zhou et al., 2020; Haddouche et al., 2024). Finally, we acknowledge Dziugaite et al. (2021) for making their code publicly available and adhering to the best current practices of scientific reproducibility.
23

---

Published as a conference paper at ICLR 2025
B.3 ORACLE GENERALIZATION BOUNDS: LEMMATA
Trajectory-level bounds. The technical lemma below ensures that KL(pnl(pr) can be directly bounded by adopting a mixture transition policy, sometimes called an (r-uniform policy (Hu et al., 2023), that keeps the trajectory-level probabilities away from zero and ensures the boundedness of the log-probe rog-prbbabilities (Dziugaite et al., 2021; Lotfi et al., 2024a) without limiting the GFlowNet's ability to learn the correct solution that samples from X' proportionally the reward.
Lemma B.1 (Realizabilitt of mixture policies), Let (-(ss)) denote the uniform policy on the state space S with reward R, i.e., po (s'ls) [Ch(s)]' Then, there is a a E (0, 1] s.t. the family {pF: pr(-)s) apu + (1 a)pp(-1s)) [-}] contains a policy sampling from X' in proportion to R.
In the classical statistical learning terminology, the result above states that the family of a-uniform policy networks realizable, meaning that member of this family satisfies the desired balance conditions. However, as we note in the proof of Lemma B. finding such depends on the knowl- edge of the minimum value of R(x) on X, which may be an NP-hard problem for some generative instances (Zhang et al., 2023b; Ma et al., 2013) that cannot be swifitly solved. Since the resulting generalization bound depends on hardly computable quantities, we call it orucle bound, similarly to the iestribution-dependde PAC- -ayessia inequalities in, e.g., (Alquier al., 2012; Alquier, 2024).
Transition-level bounds. From Definition 5.3, readily conclude that the stochastic process
M = Σ ((,, = Σ (D, - -B, -Bs.pp((() (Cn S(-1)l 1885
(12)
ISigt
is martingale with respect the filtration {Fi)e21 In Theorem 5.4, we developed concentration inequalities for M, to derive transition-levee generalization bounds for GFlowNets. Complementar- ily, the lemma bellow shows how the martingale Mt is connected to the traditionally implemented trajectory-wide DB loss (see Equation 1). There, we assume that trajectories have fixed length, an assumption that was also considered by Malkin al. (2023) when showing that GFlowNet can be seen as an instantiation of a hierarchical variational inference model.
Lemma B.2. Let pg be the sampling distribution and pg,r be the corresponding marginal over terminal states. Then, by denoting (S1,..., St) with fixed
EyspE
Σ ISiSt
LDB(S Sei)
[CDD ESEE (o 1815
(13)
In other words, the trajectory-wise objective in the left-hand side of Equation 13, which is often used as a learning objective for GFlowNets (Pan et al., 2023a;b; Madan et al., 2022; Bengio et al., 2023; Jang et al., 2024; Silva al., 2024), corresponds the transition-wise objective in Equation 8 when the trajectories are sampled in Markovian fashion. Under these circumstances, we defined the risk functional associated specific parameterization @ of the policy network as
((0)= = E, Brapp
Σ ISPSIT
(s,-PE((1))-)) (S.
(14)
which LDB implicitly depends on 0 via the forward policy pp. Importantly, we take the trajec- tory's ss length |r| into account when defining (()), which is often done in practice (Zhang et al., of independently sampled trajectories, we define 2023b). Then, given a set (so, s(5) ,...,, s(4) )
(
Σ Σ 1515 t
(DB(S .SS) IsiSt,
(15)
as the empirical estimate of L(0). Under these conditions, Theorem 5.4 established high- probability upper bound of C(0) as a function of L(() and of some characteristics of the generative process. We recall, however, that two assumptions were required to achieve this: that the DB loss and thus the martingale difference sequence M are almost surely bounded and that the training is constrained by pre-specified transition budget. In practice, the boundedness be achieved by
24

---

Published as a conference paper at ICLR 2025
either clipping the loss function (McAllester, 1999; 2013) or, when more detailed information about R is available, constraining the output the neural networks in the fashion of Lemma B.1 with the knowledge that the optimal flow F* (s) satisfying the detailed balance is bounded by (s) [mingex R(x), E R(x)] for each state (Bengio et al., 2023), or more refined version of this constraint (with upper and lower limits possibly depending on s). On the other hand, due to the CPU- bounded nature of GelowNet transition sampling (which cannot be easily parallellized in GPU), we assume that training is computationally limited by a fixed number of observed transitions. Hence, to promote equitable assessment of different generative tasks, assume in Theorem 5,4 that the number v. of sampled trajectories for training depends on a fixed budget of sampleable transitions T.
B.4 SUBGRAPII ASYNCIRONNOUS LEARNING
Please refer to Section C.2 for a detailed experimental evaluation of SAL. We would like to emphasize that all experiments below are based on standard practices for GFloNet training, with tra- jectories sampled from an e-greedy sampling policy, and not on the simplified setting of Section B.2.
25

---

Published as a conference paper at ICLR 2025
C SAL: IMPLEMENTATION AND THEORETICAL ANALYSIS
Sampling correctness. We first recall how to sample a X' in the context of SAL. For given collection {p%-} [ of forward policies trained the style of Definition 6.2, we do so by starting at So and following the root policy until we reach either a terminal state E X' or leaf partition Sy. In the former case, we interrupt the generation and return x as sample. In the latter, we proceed to X by following the leaf's policy p'y. as shown in the highlighted trajectory in Figure 3. In Theorem C.I, demonstrate that this approach samples R(x)) when both the leaf and root policies globally minimize their respective learning objectives.
X proportionally to
Theorem C.1 (Sampling correctness of SAL). Let {5) = FHP(S, m) and {9, Yo=o be the corresponding GFlowNets. Let FF and {Pr }==1 be global minimizers of their respective learning objectives. Then, the marginal distribution over X' induced by the learned policies {p' *. 1m
(+(
ΣE Σ DF (T(8))
Σ
DF (7'0)))))
(16)
ISJSm
8ome
matches the target distribution TTx) R(x)/2, 2, with Z ΣzEN R(z).
Remarkably, Theorem C.1 establishes SAL as the first asymptotically correct general-purposs distributed learning algorithm for GFlowNets. On the other hand, successful implementation of SAL requires having an efficient mechanism concomitantly enabling to sample states from given partition (to minimize CATB) and to recover the partition of a state (for inference). This is the reasoning behind what we name, and have long named, an assignment function. In Section C.2, we develop such mechanisms for some commonly considered generative tasks in the GFlowNet literature and provide empirical analysis asserting the effectiveness of the resulting algorithm.
C.I SUMMARY OF THE EXPERIMENTAL CAMPAING
We provide summary of the considered generative tasks for assessing the performance of SAL. For a comprehensive discussion on the implementation SAL, please consult the sections below.
1. Hypergrid (Bengio et al., 2021; Malkin et al., 2022; 2023; Pan et al., 2023b; Krichel et al., 2024). We consider both 8 and 64 64 hypergrid environment (Section with (Malkin et al., 2022, Section 5.1)'s reward function, which is illustrated in Figure for 8.
2. SIX6 (Jain et al., 2022; Malkin et al., 2022; Shen et al., 2023; Chen & Mauch, 2024; Kim et al., 2024a). We generate 8-sized nucleotide strings. The reward represents wet-lab DNA binding measurements to a human transcription factor (Barrera et al., 2016; Trabucco et al., 2022).
3. PHO4 (Jain et al., 2022; Malkin et al., 2022; Shen et al., 2023; Chen et al., 2023). Similarly, construct 10-sized nucleotide strings; the reward reflects wet-lab measurements of DNA binding activities to a yeast transcription factor (Barrera et al., 2016; Trabucco et al., 2022).
4. Bit sequences (Malkin et al., 2022; Madan et al., 2022; Rector-Brooks Brooks et al., 2023; Tiapkin et al., 2024). We produce 60-sized binary sequences. Given subset M of such sequences, we define R(x) = exp( - min dy, (z. m)), in which d,, is the edit distance.
5. Sequence design (Jain et al., 2022; da Silva et al., 2024). We build 8-sized sequences of Also, R(x)= = Σx= g (x:), with f and g being (-1, 1]-valued functions.
{1, {1,...,6), 6).
6. Set generation (Bengio et al., 2023; Pan et al., 2023a). We assemble 16-sized subsets of a fixed 32-sized set. We employ the same additive reward function described in Section 4.
C.2 AN EFFICIENT IMPLEMENTATION OF SAL
We start defining assignment function.
Definition C.2 (Assignment function). Let f: S - {0, 1,...,m, mm] [m], in which m is the number of available computational units. Assume that f satisfies the following conditions.
1. (Completeness). fr)( # for each j, i.e., assigns least state each available unit;
26

---

Published as a conference paper at ICLR 2025
2. (Consistency). (S, = ff is
partition of S.
Then, f is called assignment functtiont and {S,)3=o is the fixed-horizon partition associated f.
Condition (I) above, which we call completeness, ensures that no computational unit is wasted, whereas condition (2) - consistency guarantees that the partition of induced by f is a FHP. In this context, we denote by (V(S) the space of S-valued random variables (measurable functions). Then, we say that function g: [ma] -> (V) is a stochastic inverse of f if g()) E (J) with probability one for each E [ma]: similar concept exists in the literature of discrete normalizing flows (Hoogeboom et al., 2021 I; Tran et al., 2019). Notably, the distribution qi over the subnetworks' sources in the definition of SAL (see Eq. 9) corresponds to the PMF of the random variable 9()).
Importantly, to efficiently implement SAL, one only needs to develop an f that is both fast to compute and casily stochastically invertible. In this section, show how design such an assign- ment function for the problems of autoregressive design and set generation and for the hypergrid environment. The reader is reminded to recall Appendix A for an overview of each generative task.
SAL for autoregressive models. We first illustrate the concept of an assignment function for autore- gressively generated objects, which are very com- mon in applications (Jain et al., 2022; Malkin et al., 2022: Jiralerspong et al., 2024; Hu et al., 2024). For this problem, each state is represented as an ele- ment of the set [[0, k 1]] for fixed & (the vocab- ulary size, e.g., k = 4 for nucleotide strings) and L (the sequence's 's length). Then, to construct fixed-horizon partition, we choose a distance from the initial state and define
Sets
Sequences
Centralized SAL 0.072=0.00 00.00
0.092
0.126 0.09440.000
Table 2:
Total variation distance between target and learned measures for the central- ized model (top row) and SAL (bottom row).
h(s)=
Σ ISiSD
84 k--
(17)
as the k-ary representation of Then, f(s) ( (h(s) (mod m) + 1) is our assignment
function. Importantly, both (s) and (z) add negligible computational overhead the train- ing procedure. Indeed, to sample & from f (i) for 2 1, we first define l(s)) = m- (j - with & randomly sampled from [[0, [4--1/m]]] Thus, to recover from h(s), we only need to solve (triangular) linear system; the details are provided next. Let (s) = m + + (j-1) (mod k"). Also,
(s)
Σ sik' (mod kn) = Σ Isisa
s,k- 1 S k.
(18)
ISiSD
when S D. Clearly, (s) = and, recursively, (s) = kn-1s, 8n + hn-1(s) Therefore, s (si..... .3......88 (SD) jointly satisfy the triangular system Ts in which TG, = 1 (12)) b and h: h(s) (mod k+) for i,j E (1..... D). This system can be efficiently solved via backward substitution in parallel for batch (h.s").....,l
... h)) (b(s)).....h((')) of B sequences.
SAL for set generation. As in Section 4, we also consider the problem generating S-sized sets with elements extracted from source W {1...., W} of fixed size W. In this setting, each C W can be uniquely represented
Sets
Sequences
Contratred
binary vector E {1, 0) W
with s; = indicating that i E W is a member of Notably, the elements s at distance D < S the initial state can also be completely described by ΣW 8 = D. For these elements, the prefix $1.D (1.0) has at least max(0, 2D W} components equal to 1. Then, similarly to Equation 17, we de- fine the assignment function for each E 03.
422
Tagger
3.30
SAL
with DM, s; D as the binary representation of modulo the number of computational units,
Leamed
f(s) Σ 2 s; (mod m) 1.
(19)
Figure 6: Target vs. learned distributions.
1Si5D
27

---

Published as a conference paper at ICLR 2025
105
2010
1
1080
3.00
Tita wadden Apacls
Figure 7: SAL enacts faster mode discovery under varying time budgets. The horizontal axis represents the total number of epochs used by the centralized model, which is set to the sum of the number of epochs for each leaf GFlowNet and the root GFlowNet to ensure the approaches are fairly compared. We provide additional evidence for the enhance performance SAL in Figure 13.
On the other hand, let Vmin 2 30 1m/m] and Vmax 122-1/mm] To stochastically reverse f. define h()) m (j 1) with & P(umin Vmmax 2) sampled from Poisson truncated at [umin, [maz] and parameterized by X. Then, we let [h(i)] E {1,0) P be the corresponding bit-wise representation of h and yj (1,03W-B be random binary vector with exactly D - Ei=(hh()ls components equal to 1, obtained via Fisher- Yates' shuffling algorithm (Fisher Yates, 1963). Finally, we construct a sample ([h(i), yj) E {1,03 by concatenating [k(j)] and yý- Notably, an assignment function is closely related to the concept of ranking and unraking functions computational combinatorics (Myrvold & Ruskey, 2001). ). To see this, we recall that a ranking (resp. unranking u) function of a set S (resp. {O..... |8| 1}) injectively maps each member of S to an element of {O...... S] 91-11 1] (resp. S). A natural choice for is based on the lexicographic order on S (Liebehenschel, 1997); the corresponding u, however, may not be efficiently computable. Given a ranking function and a number m of computing nodes for training GFlowNet, we may defined an assignment function as f(s r(s) (mod m) + We also provide an implementation of lexicographic-basee ranking and unranking functions for the set generation task to support future research on the development of more effective partioning schemes for SAL..
SAL for the hypergrid environment. In conclusion, we also consider the difficult-to-explree hy- pergrid environment, which is defined by distribution supported on [[0, H 1]]4 for fixed H (the grid's size) and (the grid's dimension) (Bengio et al., 2021). For this problem, we note the states at distance D to the initial state can be fully described by the equation Eisisa I, D. Equivalently, each x satisfying the above equation can be injectively mapped to a point within the (A: simmix..ex. Hence, we define the assignment function f over states at distance D from 8, as
f(x) min <m
1/-1
Σ1
[++1] m
j
1/8-1
+ m 4 1.
(20)
The exponent 1/(d-1) is meant to ensure the workload is approximately homogeneously distributed among the computational units. To sample from f(((j) for j 2 1, we let
Umin
D (i-1/m)
and Vmax
D (5/m) 1/2 +-==m m
(21)
and pick xy uniformly random from [[umin [maz]] Then, (322 F&) is drawn from Dirichlet- multinomial with number of trials D - 21 and concentration parameter a set (arbitrarily) to 1.
Remarkably, the hypergrid environment illustrates an approach differing from the strategy of encoding- encoding-as-integer and computing- -eaputinn--th---rindeeee that was implemented for the other tasks. Figure shows the partition for d = 2, H- 8, and m 2, which was the setup for Figure 4. We represent red and teal the sources of the subnetworks assigned to the and j = 2, respectively, and in blue the remaining states. Recall that, by definition, all descendants partition. In this scenario, we hope that the development of sophisticated and expert-driven partitioning techniques will greatly benefit the use of GFlowNets in specialized domains, e.g., drug discovery.
leaves j
state s are members s's
Figure 8: FHP for the hypergrid.
SAL results in better distributional approximations, Figure 4 shows that, when compared against a centralized approach, SAL achieves a better distributional approximation for the hypergrid envi- ronment under fixed time-budget (for SAL, the training time is the longest client's training time
28

---

Published as a conference paper at ICLR 2025
Figure 9: SAL improves the discovery of high-value states for all considered tasks. For a fair comparison, the centralized model is allowed to explore for twice the number of epochs permitted each client, ensuring the training times are roughly the same for SAL and the standard GelowNet.
24.50
SAT.
Model
Figure 10: SAL results in a more accurate approximation than centralized approach for a similar time budget on the x2x 12 grid. Complementarily to Fig. 4, all models are trained by minimizing the SubTB (instead TB) objective with À = 0.9. The running time for SAL is determined by the train- ing time of the longest client (blue columns) plus the much faster aggregation step (green column).
plus the time for aggregation). Figure and Table 2 corroborate this claim for the set generation and sequence design tasks, showcasing that SAL learns distribution that matches the target more closely than standardly trained GelowNet. For the other tasks, learning an accurate distributional approximation not as important as finding high-value objects, and that the reason we not consider them here. Notably, these results are consistent with Theorems 5.2 and 5.4: by reducing the size the state graph that each model needs to focus on, SAL potentially faciliates the learning of generalizable policy network and leads to more accurate approximation to the target.
SAL greatly improves mode-discovery. In the GFlowNet literature, mode is often defined as a state æ whose associated reward R(xx) is larger than a predefined threshold t; see, e.g., (Bengio et al., 2021; Pan et al., 2023a;b; Madan et al., 2022; Jang et al., 2024; Malkin et al., 2022). For our experiments, we fix f = 0.1 for the hypergrid environment and, for the other generative tasks, sample initial batch of 2 104 from the uniform policy and set t the 0.99 quantile of the observed rewards. Importantly, the same threshold is used for both the centralized, leaf, and root GFlowNets. Under these conditions, Figures 4, 7, and 13 show that SAL enacts drastic improvement of the mode-discovery rate over centralized approach for varying computational budgets in all considered generative problems, leading the discovery of up to 8x more modes. There are two reasons for this. Firstly, the distributed nature of SAL ensures that a much larger portion of the state space is explored in a significantly shorter amount of time. Secondly, each client model focus on a subset of the state graph and may be regarded as specialist in the corresponding subtask. By collecting the samples fostered by these local specialists, we end up with a significantly more diverse and valuable collection than the one that would be obtained by, e.g., independently training multiple GFlowNets parallel. Remarkably, this interpretation highlights the relevance of appropriately defining
partition of the state graph, issue that defines key future direction for our work, as we discussed in Section 7 of the main text.
Alternative learning objectives for SAL. For simplicity of exposition, we outlined SAL in Def- inition 6.2 collection of TB-minimizing GFlowNets. However, previously discussed, one can straightforwardly adapt alternative learning objectives (e.g.. SubTB (Madan et al., 2022)) and sampling techniques (e.g., replay buffer (Deleu al., 2022)) the context of SAL. We illustrate
29

---

Published as a conference paper at ICLR 2025
these extensions here. Firstly, the SubTB objective for the jth partition would take the form
EiS N- logg F( (T Fr)) F(Ta) (Trmm (T)
Anmm m
2
(PFT F) = E 80 T (11s)
Σ [Sm<nnSr]]
in which Oj and F'e are a distribution over initial states of and a sampling policy for the jth partition, respectively. Secondly, the replay buffer would store the trajectories T leading high-value states within the jth partition, as measured by either I (for leaf partitions) or It (for the root partition; see Algorithm 1). Figure 10 compares the accuracy of SAL against a centralized GFlowNet, both of which trained by SubTB minimization (A = 0.9), for the 12 x 12 hypergrid environment. Similarly to Figure 4, SÁL achieves better distributional approximation in this case. Additionally, we found that SubTB leads to faster convergence with respect to (A)TB (not reported) in this particular problem, consistently with the evidence at Madan et al. (2022, Figure 1). On the other hand, our experiments did not provide evidence in favor of using the replay buffer. However, we acknowledge that a deeper empirical investigation, in the fashion of Vemgal et al. (2023)'s 's work, is required.
C.3 THEORETICAL ANALYSIS AND EXTENSIONS
This section aims answer core questions regarding the nature of SAL. From sampling perspective, we ask which distribution each client learns and suggest potential diagnostic techniques to evaluate their distributional accuracy. From a distributed learning standpoint, we assess the extent to which local errors are propagated to the global model. Additionally, we formally extend SAL to accommodate the learning over multi-
partitions of the state graph.
Local sampling distributions. Each leaf GFlowNet in SAL learns a distribution over a subset X; of the set of terminal states X. The character of such distribution, however, was not considered in the foregoing discussion, and one may wonder whether it just corresponds to the restriction of the target to Xy. As we show below, this is not generally the case: the optimal leaf distribution, which we de- note by p depends both Xj and on the specific structure of the state graph induced by the leaf Sj- Proposition C.3 (Local sampling distributions). Let { = FHP(S) and {gyby-n be the corresponding GFlowNets. Also, denote by TDc(s) the set of terminal descendants of s on the original state graph, i.e., x E TDc(s) if E X and there a directed path from to x. Then, for fixed backward policy Pn, the solution that globally minimizes Equation 9 satisfies
R(x)
F5(s)=
Σ #ETDe(s)
R(I)
Σ PB(T|z) 8owe
and pp(a)))
Σ PB(T(x)
(22)
F2(s)
for each E (1,...,m), m). S Ij, and x E X3 == U.ez, TDc(s).
When the sum over backward trajectories in Equation 22 does not depend on J, e.g., for autore- gressively generated object (in which (T .)) for the unique trajectory 7 connecting s to a)) and sets (in which the depends only on the depth of when PB is fixed to an uniform policy), Proposition C.3 says that p can be nicely interpreted as the restriction of the original target R to the induced terminal set Xy. We emphasize this fact in the corollary below.
Corollary C.4 (Local sampling distributions for autoregressive models). In the context of Proposi- tion C.3, assume that the state graph is represented as a tree. Then, (n(T(z) = 1 for every T and
R(x)
F5(3)
Σ R(x) and pf(s)s) #ETDo(s)
(23)
EmETDO(s) R(x)
For the set generation task, it also holds that p x R(x) for each x E TDc(s) and E Iz.
Interestingly, the result above suggests straightforward procedure for assessing the goodness-of fit of the locally trained GFlowNets. When Xr is considerably smaller than X, we can compute the normalized target in Equation 23 and directly compare it against the learned distribution in X. Otherwise, any technique for diagnosing GFlowNets can be readily applied to probe the accuracy of pr. e.g., measuring the Spearman correlation between log (-ss) and log R(x) for x E X, (Malkin et al., 2022; Madan et al., 2022; Shen al., 2023; Tiapkin al., 2024; Chen & Mauch, 2024).
30

---

Published as a conference paper at ICLR 2025
S St USS U &u Uوي Sp
Sampling:
Leool
Level
Figure 11: Illustration of Recursive SAL. We show a two-level partition with 2 models within the first level and m2 models within the second one. For training, we first train models at the bottommost layer (represented in blue, red, and green) and recursively proceed upwards towards the middle (magenta and yellow) and top (root partition, shown in pink) layers. For the non-root lay- ers, learning is based minimizing LXTB with the reward defined as in Equation 25; for the root, minimize LTB instead. For inference, we start at 8, and iteratively select the policy based on the cur- rent state, as illustrated in the highlighted trajectory and the annotated text the top-right corner.
Sensibility to error propagation. Proposition C.3 raises an important question: how do the errors of the leaf models (- affect the global goodness-of-fit To address this issue, the next proposition shows that the contribution of Gj to the overall distributional error is increasing function of the probability mass associated to the jth leaf by the root model, 9o. and of how inaccurate 9j 's itself is.
Proposition C.5 (Sensibility to error propagation). Lel (S, ) = FHP(S, m) with GFlowNets { Assume that Go satisfies its balance condition. Also, define
ZR = Σ R(x), ZF
Σ
F2(#), and Z = Σ ZR,
(24)
ZEANS.
xEU1SIS
ZEX
and note that ZR + ZF is the partition function associated to G.. Then, the TV distance between the learned distribution pr in Equation 16 and the target π (xx) x R(x)) for x satisfies
TV (pr, R) S
ZR 1
Σ DEX'
E
E.
T(x)
7 + 7 PT (())
ZF
(2/6))
7
Z + 7F
Frmr extmaring 7F
Ernor of the local approximations
which x is the restriction of per to Uisism In and f is the assignment function.
Importantly, the bound above is tight in the sense that, when the root and leaf models satisfy their balance conditions, Z ZR + ZF and Txx) E, Ep f(s) (x]ss)] as show in the proof of Proposition C.5 in Appendix D. There, we also provide an alternative, trajectory-based upper bound on the TV distance that similarly highlights the relatively large impact the distributional errors associated with wange-probability leaves to the overall accuracy. Heuristically, this suggests that SAL may benefit from FHP that approximately homogeneously distribute the probability mass among the leaf partitions, ensuring that no client has a disproportionate role on shaping the accuracy of the aggregated model. To achieve this, however, one needs prior knowledge the reward function; the definition of good fixed-horizon partition should be done in boollembyy-pooleee basis. We also believe that a human expert would have a remarkable impact on the effectiveness of SAL for the highly-specializee, molecular- molecular-biology -aaeed tasks in which GFlowNet are often implemented.
Recursive SAL. In Section 6, introduced extension of SAL to recursively defined FHPs. Proposition C.6 formalizes this procedure and demonstrates through an inductive argument that the resulting model samples correctly from the target distribution.
Proposition C.6 (Recursive SAL). Let S be the vertices of a state graph with diameter D. Then, for sequences = do < di < d2 ... < dy < D and {m, L.m..... my we define Uisism, Inj as disjoint ma-partition of the states distanced de from sa. Also, let X {x X: d(x, 8.) 2 dx) and, for < k, let X; {s:8 8 Ietis (d(s, S d, A X}}. Finally,
31

---

Published as a conference paper at ICLR 2025
p. .B F.j): 1 S m.) as set of GFlowNets trained state graph with initial states U, Ing. terminal states X₁ and reward function R, such that
define 9 Up' 3.
(s) (1))(8), if E Iit1 and i < k, R((s) ifs X.
(25)
Then, when the GFlowNets - satisfy their respective balance conditions, the generative process starting al So and recursively following Pr until either reaching Ii+1,a at which point the guiding forward policy is changed to p for 0 S i S k, reaching X, signaling to stop the generation and return the sampled object, samples each x E X proportionally to R(z).
Remarks on Recursive SAL. In plain English, the above proposition says that we can use the learned flow function Fi+1 at the (i 11th layer as the reward function of the GFlowNets within the ith layer to obtain a correct sampler when training the GFlowNets in a hierarchical fashion. In computational terms, the number of trained models grows linearly with the depth & and width max; m, of the multi-layered partition. In the light of Theorems 5.2 and 5.4, however, each model significantly
would have to solve considerably simpler problem and we may be able to smaller neural network to parameterize the corresponding forward policies, with advantageous con- sequences for both generalization--via the KL term in Theorems 5.2 and 5.4, which increases with the number of estimable parameters-and storage. Albeit we do not provide an empirical evaluation of Recursive SAL in this work, we believe its implementation could be beneficial for problems with very large trajectory sizes and are optimistic about its potential applications in future endeavors.
C.4 CONDITIONAL SAL
SAL and state-conditional flows. Bengio et al. (2023, Section 4.3) introduce conditiondi flows as a family {Fs)ses of flow functions defined on the subgraphs G, induced by {s' E S: 8 2 s'}. in which 2 s'' means that there is a path from to s' on the original state graph. Remarkably, a FHP m) can be interpreted as subset {F'] Iz of state-conditional flow. In spite of these similarities, which serve only to strengthen the foundations of our work, we emphasize the novelty and demonstrated effectiveness of our distributed strategy for learning state-conditional flows.
Learning reward-vonniiioneed GFlowNets with SAL. Recently, there has been growing interest in reward-conditioned flows, in which we learn family {FeJeEc of flow functions conditioned on some information e E C given as an additional input to the neural network. In most applications, corresponds to either temperature parameter (Zhang et al., 2023a; Kim et al., 2024a) defining the peakiness of the target distribution or pharmaceutical properties (Roy et al., 2023; Pandey al., 2024) guiding the drug discovery process. In view of this, we extend SAL accommodate the distributed learning of reward-conditioned GFlowNets on a conditioned FHP, i.e., a FHP that depends on the conditioning information. This may be formally expressed as follows.
Remark C.7 (Reward-conditionee SAL). Let C be a set of conditioning information and {Re X + R+: c E C} be the corresponding family of conditioned rewards. Also, let F: C x S + R and pF: x S x S > [0. 1] be conditional flow function and conditional forward policy, that is, (c is forward policy for each c. Finally, let {S} o FHP(S, m, c) be conditioned FHP and {G, 1- be a family of neward-conditionss
GelowNets. Following the arguments for demon- strating Theorem C. it is easy to see that SAL samples correctly in proportion to R, when each Gj satisfies its respective balance condition with respect to R. for every e E and S S m.
Conditional GFlowNets are commonly implemented for controllable generation by setting the conditioning information at inference time (Roy et al., 2023); see also Lau et al. (2023)'s QGFN. In fact, Pandey et al. (2024) recently explored this principle for the effective exploration of chemical space at an atomic-level given some desirable pharmacological properties, e.g., synthetizability. In this regard, Remark C.7 ensures that most of these approaches can ba adapted to the distributed set- ting via SAL, and we believe that assessing the resulting methods is an important research direction.
C.5 SAL AND EP-GFLOWNETS
The reward function of GFlowNets can often be decomposed as R(z) = II R₂(x) (Jain et al., 2023; Deleu et al., 2023; Zhou al., 2024; Pandey al., 2024), e.g., in multi-objective problems in
32

---

Published as a conference paper at ICLR 2025
which each R₂ is an objective and our goal is finding samples that concomintantly high valued for every R2 In these cases, da Silva et al. (2024) proposed a divide-and srride-and-comee and--un algorithm for learning K GFlowNets in parallel, each targeting reward R2, and then aggregating them with an extra GFlowNet. The resulting model, termed EP-GFlow trains (K + 1) GFlowNets on the same state graph, in sharp contrast SAL which trains GFlowNet for each component of FHP. To further highlight the distinction between SAL and EP-GFlowNets, we show below that these approaches can be implemented in complementar manner.
Proposition C.8 (EP-SAL), Let R(x) = I (a( be a multiplicative decomposition of R. For each S i S K, let (S Also, let 9 = {PP PB F m be the root- and leaf- GFlowNets corresponding to the i-th FHP and denote by p'g and P'g the induced distributions over trajectories. Assume that each p'y samples terminal objects X proportionally RG. Finally, let pg be any positive probability measure over trajectories. Then, if GFlowNet (pF, (P.,B 9) globally minimizes
1= (5
FHP,(S, m,) be fixed-horizon partition of the state space S.
log IEEPPTT + log
PF(T) PB(T)
rape
(26)
the marginal pr of Pr over X matches R K R₃ up to a normalizing constant.
Proof. The result follows directly from Theorem C.I and (da Silva al., 2024, Theorem 3.1).
For the sake of completeness, let Ti:d denote the first d transitions of r and Ta: be its complementar.
Then, let d, be the distance from the initial state 8, to the sets Zı, the underlying state graph for j .1...., mz (recall Definition 6.1, and note that d, does not depend on j). Under these conditions, the sampling distribution pip in Proposition C.8 can be formally written as
PT((T) [T)(T) if if d4
(27)
(T En E,ET,
E.er, L(EE))PH (Ta,ls), otherwise.
Importantly, Equation 27 can be efficiently computed by keeping track of the partition associated to each sampled state. In practice, the computational overhead wrt directly evaluating py is negligible.
Empirical illustration. From an empirical standpoint, Proposition C.8 says that SAL can be used to learn a set G, of GFlowNets jointly sampling x E X' in proportion to R2 (z). Then, given G. for 1 < < K, a GFlowNet sampling K in proportion I R3 can be obtained by minimizing EP-GFlowNet's learning objective. Figure 12 empirically validates this result for the task of set generation with the same hyperparameters considered in Figure 6 with K = and m, = 2 for each i E {1, 2). Each GelowNet was trained for 512 epochs and the log-utilities defining R; were independently sampled for each client. Nonetheless. in spite of its soundness, the effectiveness of this mixed approach in realistic problems remains to be assessed. Looking at the bigger picture, these observations emphasize the composability of GFlowNets (Garipov et al., 2023), which might be relevant for the design of more data-efficient algorithms (Du & Kaelbling, 2024).
x10-
Target
2
Learned K10-
Figure 12: SAL of EP-GFlowNets,
33

---

Published as a conference paper at ICLR 2025
D PROOFS
D.1 PROOF OF LEMMA 4.1
We simply note that the space of T-sized subsets of {L...., W} has size (*) and the space of T-sized subsets of (2....,Wi W has size 11 Since
W
- W
-1
(28)
W
when W + co, we can always find for any & (0, 1) a W and a T, both of which potentially depending for which |x4| 2 f/,2). For the cases considered in Figure 1, in particular, we compute the following proportions: 61/32 = 81. 81.25% and (64 (64-6)664 R 90.63%.
D.2 PROOF OF PROPOSITION 4.2
Our proof has three steps. Firstly, use Hölder's inequality bound the expectation of [F(x) pr (z)]. Secondly, we rely on Jensen's s inequality to bound [xx) pr (z)] with an expectation of log-scale with simple
[pr (TF(T)/BB(TT)) π (x)] over т. Thirdly, convert the probabilities technical argument based on the Taylor expansion of log. For this, let 6(x) = [T()) (p)( Then,
Ex [o(z)]
Σ
p(z) (EZ)q,,
xEX
(x)
Σ
((x)
ZEx
(29)
Σ
ò(x) ( p())
Σ (
Pe:
(EE((x))
xex
[ [ (2)
= (EEpPE, (0(x) t/ Euwper
1/2
( (2)
for any p, q > 1 such that 1/2 + 1/g 4 = 1. For = == 2, this bound becomes
Ex-qE. [o(x)] S (Ex~pE, (0(x) (x² (gg,T + 1))
(30)
For GFlowNets, may write (x E Eropa (px(r)/pn(rlu)) Hence, by Jensen's inequality,
PF(T)
Ez-pe,r [p(x)*] Eusperr
Eropn
((x)))
PE(((x)
(31)
PF((T)
-T(x) '1
Excper
Eynoe
(PB(T(Z)
In conclusion, we show that
PF(T) (B(T]z) ((x)) < log
PF(T) log क (())
(32)
pB(T(z)
In fact, let M maXr. instance, M < for autoregressive generative tasks (i.e., when pe (T(x)) 1). Thus,
pe(r) p which always exists due to the finiteness of the state space. For
T(x)) I M² (PP(T) PF(T) T(x) p
PF(T)
MMp(T(>) M
The lemma below, which is direct consequence of the mean value theorem, ensures that the quantity above is bounded above by the log-squared difference between (=(7)/pa(7(a) and ((x). Lemma D.I (Lipschitzness of x e²), For every x. y (0,1) 1), log log 2 |r yl-
34

---

Published as a conference paper at ICLR 2025
Proof. Consider f: (-20, + R, f: c', and notice that [t) |e'| I 1. Consequently, by the mean value theorem, f is 1 1-Lipstchitz and le" - -c*l - [t - s for every t, E (-o0, 0j. By letting log z = t and log y
conclude that |x yl log logyl for x, E (0, 1].
n
In summary, we have shown that
1/2
(F(() log (()) (x (x(gerlleE)) + +1) PB(T(z)
[o(x)] -
Exxpp, Erope
loog
(33)
The statement thereby follows by considering uniform reference distribution, (F)((x)
[
|
TV (pr, T)
Ex E
2
1/2
(10g PF(T) - log n| -log ((x)) (x2 (PE,)) + +1) pB(
2 E p
D.3 PROOF OF PROPOSITION 5.1
For completeness, we provide proof of Proposition 5.1. Clearly, enough to show that
LPcs(P) S LFcs(P)+ + Vz and LFcs(P) S LFcs(P) + V nnn + 2Lecs(P)), (34)
in which we omit the dependence of Leos on the dataset Tn for conciseness. We recall that rj (L(P))Q))oog 2ura/s with na = [(]1 - a)n]. is the complexity term that depends on the prior Q, posterior P, confidence 8, and the number of data points na- Notably, both inequalities directly follow from Maurer (2004, Theorem 5) bound: with probability 1 - 8 over T..
"g
kl(LFes (FCs))((((P)) S n.
(35)
in which kl represents the binary KL divergence, i.c., kl(p lq) = p plog & + (1 - p) log Le Below, we show that kl(pllq) is greater than or equal (1-0)²/2g when p q.
T-q
Lemma D.2. (Boucheron et al., 2013, Exercise 2.8). Let h(t) (1 (1- log(1 (1-(11))+0) +t and p: {1,0} +
[0, (resp. q) represent the PMF of Bernoulli with parameter p E (0, 1]. Then,
BewBe(g) (1-
p(xx) = kl(pllq) ((x)
(36)
and M(1) 2 2
fort for [0, 1]. In particular, kl(pllq) 2 (p-9)² //2 when p S q.
Proof. Equation 36 follows from direct algebraic manipulation of the left-hand side. On the other
hand, define
g(() h(()
(37)
2
for & E [0. Then, is continuous, g(0) = 0. and g(t) + 1/2 when + 1. Also, g'(t)
log(1 - - 0 for E [0, 1] since log(1 t) |log(1 - f)] In conclusion,
E Be(g) (1 p(xx)
E BynBa(a) (1 p(x)
(y p)²
(38)
((x)) 2
q(x)
2q
when p S g.
By the symmetry of Equation 35 with respect to LECs and LEcs, conclude that
Lrcs(P) - - I 22Lrcs (P)
(39)
Under these circumstances, the inequality LFcs(P) S Lrcs(P) +n+ vn(++ 2Lrcs(())) is obtained by solving the above quadratic inequality on VL Through similar reasoning. kl(pllq) 2 2(p q)² by Pinsker's inequality and, consequently, LFcs(P) & (Pcs(P) + vn/2. These results jointly entail Proposition 6.
35

---

Published as a conference paper at ICLR 2025
D.4 PROOF OF LEMMA B.1
Recall that (r() = Er- PF(T) for P..P ap-u+(1-a)))) in which we make the dependence of pr on and the (unconstrained) policy pr explicit. Let F(a, pr) be the family of such policies and F(a) be the set of a--reedy policies. It is straightforward to see that F(a)) is a convex set. Clearly, is enough to ensure that mina>0,p ming PT p.." (x) S ming (2)), namely, that the rarest object can be sampled correctly by properly adjusting pF and a (non-zero) a. Indeed, Bengio [ref, Theorem 8] showed that, for each given backward policy pg and positive reward R, there is unique forward policy Pr for which the marginal (r R(x) for each z E X. Hence, since Pr = aPu.T + (1 P.P,, with pU,T being the marginal of pu over X, the realizability of F(a) ensured when satisfies min opt,T(x) < minaex (2x), i..., min, 7(x)/min, pu,r(w), in which case may set pP such that (a)) = I-a (I optrr(a)) As example, consider the set generation task, the details of which are provided in Section 4. There, Pu induces uniform distribution over X' and we may set = 8/2min, π (az) < min, =(.)/min, (x), Importantly, our analysis is not considering the (limited) expressivity of the chosen parametric model for the policy network, which touches on mostly open problem the deep learning literature. Rather, we are concerned with the feusibility of finding transition policy PF convisteni and compatible with the given target distribution R, in the sense of Bengio et al. (2023, Definition 4, Definition 20).
D.5 PROOF OF THEOREM 5.2
We first show that the risk function is bounded. Then, Equation 7 follows directly from Maurer (2004, Theorem 5) and Jensen's 's inequality. Under these conditions, notice that (L(pllerF) = Er~re (loggp(()))
(ogpr(r)) = H [[n) - Er-rp (oggr(() (40)
Also, by definition,
pr(T)= = II (s((sss)
(s,0")ET
II (spo(s'(s) a)pr(811s))
(41)
(s.5))E
|rl
2 arl
Ch(())
max.er [Ch(s)
(s,s))er
and, consequently,
- BE-pplloggrle -min |r log
a
max |o| log
max.er [Ch(s)
max Ch((s)
ie., (L(PelIPF) < -H + My. In conclusion, the convexity of the KL divergence along with the the fact that pr and π are respectively convex functions of pF and pg imply that (L(Tl(Pr) s KT(pal (pallp))) The rest follows from Maurer (2004, Theorem applied to (T(psller).
D.6 PROOF OF LEMMA B.2
The result follows directly from the Markov property of the MDP. Equivalently, we note that
Es Es....... S.
Σ CDB(S,,
Σ
Es s, (CDB(S,, S.-1)
1sigs
15ist
(42)
Σ E E.
[BS,Np)()) (LD S,1))
(5)
15i51
D.7 PROOF OF THEOREM 5.4
Our proof has three main ingredients. Firstly, we build upon Azuma- Azuma-Hoeffilae Hoeffding-type inequality to bound the expected transition-level error with the observed empirical error. Secondly, derive
36

---

Published as a conference paper at ICLR 2025
trajectory-level bound of the crnsitiion-eevd results by relying McAllester's linear PAC-Bayes inequality. Thirdly, we combine these results with a standard union bound argument. To start with, Beygelzimer et al. (2011, Theorem 1) shows that the martingale M, Eisigt M (S,, S(-1) defined
above, with A, I M (S,, S,-1) s B, and B₁ A. S C, satisfies
E exp (AM, (e 2)A 1
(43)
in which V22 = Eisise Si-) and X [0; 1/c]- In context, , 8;-1)] S 20 by the triangle inequality, and we can take C' 2U. By assumption, V, K for all S tm
Then, for the martingale M₁(0) and corresponding (O),), with representing the parameters of the forward policy, by Donsker- Varadhan's 's variational formula, we notice that
Ba-P (M,( (e 2)) KL(PIIQ) log Eenq /exp { (MM,(0) ((--2))2V((0)]
Similarly (Seldin et al., 2012b, Theorem 1), let 6 8/2tm, so that Eigist. 6e = 6/2. Then, by Markov's inequality, with probability least 1 82.
E [exp (0) - (e- 2)A² s - 8,
[Ee-@ [exp (AM((0) - (e- 2) 2 vo(0)]]];
where the outer expectation is with respect to the joint distribution of (Si..... St). By Tonelli's theorem and Equation 43, the right-hand side of the equation above satisfies
1 B (AM((0) (e 2))}
E [o [exp (fM, (c 2)A
8,
1
2
8.
Consequently, by Donsker- Varadhan's formula applied to XM,(0) - (e - ()),, bounding its exponential moment as above, and union bound over I, yields with probability at least 5.
Eg [M)( S (e - 2) [V[(0)] KL(PIIQ) + log!n, + log 3/2
(44)
Hence, by the definition of M,) and the bounded-variance assumption,
1 Σ LDD (S; S:
Ea-p
Σ (
Isi5t
isist
(45)
K KL(PIIQ) + log tm + log 2/6
+ (e - 2)
Nextly, let 5(5) be independent samples from forward policy pr((1s)) for j s and (5 s(,) be the correspondingly observed trajectories. Also, we recall that
((0) = Es, St
Σ 18i5e
[SCDB(S, St
(46)
and define
i((0)
Σ 1
1 Σ n 15J5
Te ( 15ists
Σ
(47)
the inner expectations are computed with respect the Markovian data-generating process (recall that the conditional expectation E[LDB (S, 1 is random variable). By assumption, L(0) < U. Hence, McAllester's linear PAC-Bayes inequality (McAllester, 2013. Theorem 2) entails, with probability at least & over draws of (S...... S.).
Ea~ [C())] < EomP [200] +
KL Q) + log 3/6
(48)
2 3)
37

---

Published as a conference paper at ICLR 2025
Under these conditions, equations 45 and 48 jointly imply that, by standard bouuddbound argument,
1 Σ
- + (e - -2-2KK
BaNp [c(0)] - BamP
Σ
CDB(S
t
15i5n
15ist
QT(PIIQ) log Im + log 2/6
]
U KI((PII)
log 2/8
tzA
2 3)
n
with probability 1-6 over draws of transitions (n is between [T/lmin] and [T/bm]]
St). Since nt, 2 T for all Ł, as we observe T 1. with bmin being the minimum length of a complete trajectory). and 1, S lm. as mm is the trajectory's 's maximum lenght, the result above equivalent to
Eg [g-.P([(0]]
BB
Σ
Σ 1sist
(SB ...2))
(c-2)AK T
3
1815
=(
KL((PIQ) + log tm llog 2/8
U
KL + log 2/8
+
(49)
TX
2 - 3)
n
By aggregating the terms corresponding KT((PIQQ) and log 2/8, we derive the desired upper bound on the expected risk of the DB loss.
D.8 PROOF OF THEOREM C.I
Intuitively, when each balance condition is satisfied, each state on Ij is sampled in proportion to (s) and, conditioned on s, each terminal state will be sampled in proportion to (()) make
implying that, marginally, each is sampled proportionally this argument rigorous. We first consider the case which & E X' S.- As we are assuming that ()) R((r())R(a)) for each trajectory starting at Ij and finishing at æ, must conclude that
R(x). In the following,
R(x) Ff(s)
(+(s))
Σ PF((|s) suar
Σ PB(T(x).
(50)
sMa
On the other hand, since (o(s))pf(T(s)= = (T)) for E Ij.
Fy(s) Σ PB(T|s)= Fy(s) Fo(so)
(s(so)
E pF(T(s0)
(51)
Fo(s,))
the probability of reaching 8o by starting from and following Pg is equal one since 8, is the only sink state of the transposed state graph. In this context,
pr(x|so)
Σ Σ
(clss)))
SEI,
(s)
R(x)
Σ Σ
Σ Pb(T(T))
ISISmSE,, Fo(so)
F5(s)
T: smer
(52)
R(x)
Σ
Σ Σ
(T(x)
Fo(So)
15i5 SET, T: s
R(x)
R(x) Fo(so)
Σ
(T())
Fo(5o)
&EU I,
ie., (r(ss)) samples x proportionally to R(x). For the forth line above, we relied on the fact that the probability of eaching UI, is equal to when starting at E X' ) S. and following Pg- Correspondingly, when x E X, it follows from the satisfiability of the trajectory balance condition that pr(()).) x R(x). This ensures SAL is sound distributed learning algorithm for GetswNets.
38

---

Published as a conference paper at ICLR 2025
D.9 PROOF OF LEMMA C.3
The global minimizer of Equation 9 satisfies, for every j. (s())(((T) = R(x) for every
trajectory T: I starting E In and finishing Xz. Consequently,
(x1s) = E (b)()))== = E
(T((()))(())
R(x)
Σ (T())))
(53)
()()) Fg(8)
T. aw0z
Similarly,
Σ (T()))((T(s)- E (T(())))(())
(54)
T:
implies that
Py(s) E (T())
(55)
since Er: swa (T(s) for every s. These equations jointly entail the proposition.
D.10 PROOF OF PROPOSITION C.5
As in the demonstrations above, we consider two cases in separate. First, when æ E X' n&, then Pr (z) ml)/Z++Zp Zn due to the satisfiability of the balance condition by the model. Hence,
R(x)
R(x)
Σ (pT(x) - x(x)==
Σ
Zp +
Zn
Z
zEX'ns,
zEXNS.
(56)
Σ
R(x)
1
1
nn-
ZF ZR
ZF + Zn
Z
=EXNS.
Second, when E X S.. we note that
pr(x) Σ Σ
Σ
Pp(T(so)
15i5m SES, 7: 8omgg
Σ
Σ
Σ
pb((110))
Σ ptc(t((s)
(57)
ISSSmEES
Ej(s)
= Σ Σ
pf (cl))
Σ Σ
ZF + 7n (x)s).
SES,
ISismeES,
Similarly, for any va-ull function f.
F5(()
f(x) Σ
f(x);
(58)
ZF
15J5m SES;
hence,
Ff(s)
F2(s)
π(x) - Pp))) = Σ
Txx)
pl (xls)
ZF
ZF Zn
Isism SES,
(s(s)
Zr
Σ
[(x))
1e (2(())
(59)
Zr
ZF + Zn
15 13 BES,
Eaapg
[((x)-
ZF
pr (a(s)
(1+)
ZF + Z
By recalling that (V(T,p2) = (EEex (xx) (x))) this result, along with Equation 56 and Jensen's inequality applied the function x - |ar|, implies the proposition. To further strengthen intuition, we also consider directly bounding the accuracy of Go function of the trajectory- level inaccuracies each Gy. For this, we re-write Txx) as
(σ) = Σ Σ Txx) Σ (T())
(60)
ISSSmAES,
39

---

Published as a conference paper at ICLR 2025
Conrespondingly, by recalling the property pr(z)
Eisism Exes, (9)s) (xlss, conclude
((x) pr((===== =
Σ Σ Σ
I8ism sESj T: 8veD
(61)
ΣΣ "SJSm&ESS 8-a2
Enror associated to the jth client
Hence, the total variation distance between and pr is bounded above by
(V(,,PT)
ZR Z
Σ E -oo)
Zn
ZF
SEXVS, ISJSm SE
associated jth
For tree-shaped state graphs, the second term of the equation above can be significantly simplified by noticing that (i) each is uniquely associated to a relationship which we denote by g(x) j. and (ii) that p/ (T(z) and (T(s) = (x|s). Under these conditions,
Zn R + ZF + 3
TV(T,Pr) S
Σ
(()) - 1 (cls)) -7
(62)
2
TEA'S xESgte)
Error associated the )= model
D.11 PROOF OF PROPOSITION C.6
We proceed by strong induction on the number of Exed-horizon partitions. For 1, the result above is equivalent to Equation C.I. Assume, then, that the statement holds for j fixed-horizon partitions of the state graph for all j < k. Let Gi 0 S S k, be sequence of GFlowNets satisfying the amortized trajectory balance condition. By induction, each E Uisist-1 Xj is sampled proportionally to Eisjsk=1 1[2 (s) In particular, if E X UUUSS Xz, then sampled proportionally R(x). For what remains, let x E X' U Nj. Hence, for each state Uisism Ixas S. Xx-1 and each trajectory T: ~m W.
(s(5)p((T(s) = ((())))((z),
(63)
ie., pF(((s) = po(rlo)R(e)//A())) Thus, by marginalizing out the non-terminal components of T, pr(x(s) = R(x)
Σ PB(T(z)
(64)
(R(s)
T: 8:wwx
and, since each is sampled proportionally to 14)1(s) Fx)s),
pr(x) x
Σ pr(x)s)FF(s)
sEU sEUscimm Inis
R(x)
= Σ
Fa(s) (s(s)
Σ (B(T(z)
sEUisismi Ih.j
T: g-00
(65)
R(x)
Σ
Σ
pp(1|x) R(x).
SEUisSmm Iu,
This ensures that each e X Uisisk-1 X, is sampled proportionally to R(z). By induction, each X' is sampled proportionally R(x). Hence, the recursive instance of SAL is sound approach for sampling objects proportionally to reward function.
40

---

Published as a conference paper at ICLR 2025
Nommprdad 1050
300
000
300
624
000
Figure 13: Complementary results for Figure with different random seeds. Notice that, except for the hypergrid task, the threshold defining a mode is a random variable, which explains the variability (albeit consistency) of the number modes found within the same column.
E LIMITATIONS AND FUTURE WORKS
We presently discuss the limitations of our work and the consequent opportunities for future re- search. Although informative, our theoretical analysis was limited the case of i.i.d. sampled trajectories which, as highlighted in Section does not necessarily reflect the nature of usual strate- gies for learning GFlowNets, e.g., e-greedy sampling. Indeed, the typical training of GFlowNets is closer in spirit to active learning (AL) (Cohn et al., 1994; Gal et al., 2017) in the sense that batch of trajectories is sampled from policy that is dynamically updated more data points observed, and it would be interesting and important to pursue an investigation in this direction (Jain et al., 2022; Malik al., 2023). In contrast to AL., however, there is explicit acquisition function guid- ing GFlowNet training. On the other hand, Deleu & Bengio (2023)'s interpretation of GFlowNets as Markov chains the trajectory-space could be used together with Azuma's inequality (Azuma, 1967), in the fashion of Theorem 5.4, as useful starting point for this. More specifically, one could consider that a sequence
trajectories (Te (Tehe21 is observed during training and use the same tech- niques enabling the proof of Theorem 5.4, namely, constructing martingale difference sequence M, (T.) E,, [CTB(TV) [T.c)) and applying Theorem of Seldin al. (2012a) both within and between trajectories, the results of which would then be unified via union bound argument.
Additionally, the promising results SAL pave the road to a range of interesting investigations. Most promitently, we believe the development of principled partitioning methods can lead to substantial improvements scaling GFlowNet training. Although our discussion is constrained to fxed-horizon partitions for easeness of exposition and implementation, the algorithm could in prin- ciple be extended to more general settings. Theorem 5.2 suggests that a good partition would ensure close uniform. Intuitively,
that the within- and prrtenn-patitionn distributions that the target distribution of both the leaf and root GFlowNets are relatively simple to approximate when compared against the original target and that the most important regions of the state space, as measured by the reward function, are appropriately covered. The best way to ensure these properties, however, remains an open problem and we think a promising venue for future endeavors.
would like
E.1 ADDITIONAL EXPERIMENTS
Robustness of SAL with respect to the FHP's size. It is intuitively clear that an increase in the number of partitions in a FHP would improve the coverage of the state graph and accelerate mode discovery. In doing so. however, we also enlarge the memory cost the algorithm due to the necessity of aggregating a larger number of leaf GFlowNets in the server. To shed light on the effect of m on SAL's performance, Figure presents the number of modes found during training for the tasks of SIX6 and PHO4, As anticipated, SAL drastically improves upon centralized GelowNet,
41

---

Published as a conference paper at ICLR 2025
SIX6
PHO4
Number of mondes
400
4000
SAL
Centralized
2000
0
FHP' size
Figure 14: SAL consistently outperforms a centralized GFlowNet irrespective of the number of components defining the underlying FHP. As expected, the coverage of the state graph an increas- ing function of the FHP's size. Results for the centralized model are included solely for comparison.
SAL
Intentionaaly mode--ollaapedd SAL
Target distribution
Figure 17: Extremely inaccurate estimation of the flow function might lead mode collapse. (left) Standard SAL-trained GFlowNet for 8 x grid with 2-sized FHP. (midkle)) SAL-trained GFlowNet when the flow function Fj is severely inaccurate. (right) Target distribution.
notably regardless of the size of the underlying FHP. Overall, our experiments throughout this work indicated that SAL is notably robust to the choice of hyperparameters defining its implementation.
50-sized subsets of (1. (I..... 100)
Intentionally bad purtition sampling
6000
55S3
20000
15000
SAL 0.045 10.001
Centralized
0000
0.061 10.009
10000
6170
5000
Table 3: SAL improves upon a centralized GFlowNet for large- scale set generation in terms of FCS (Silva et al., 2024).
352
SAL
Cent.
Figure 15: SAL leads signif- icantly faster mode discovery for large-scale generation. Results averaged across runs.
Figure 16: SAL underperforms when the within-subgraph sampling distribution is inade- quately designed.
Experiments on large-scale set generation tasks ((x') 1030). To underscore the scalability of SAL with respect to the underlying domain's size, Figure 15 and Table 3 respectively show that our distributed algorithm entails faster mode discovery and lead to a more accurate distributional approximation for the problem of generating 50-sized subsets of [1..... 100); see Section A.3 for definition of this task. For fair comparison, the centralized GFlowNet is trained for 60 seconds twice the time allocated each leaf GelowNet and the aggregation step. Besides, adopt the experimental setup described Section B. Importantly, these results suggest that the applicability of SAL is not limited by the domain sizes presented in the main text.
In the following paragraphs, we examine two potential causes for catastrophic failures of the aggregated model, mentioned in Proposition C.5 on Section C.3.
42

---

Published as a conference paper at ICLR 2025
Insufficient training of leaf GFlowNet. As the aggregation phase in SAL relies on the locally estimated flows as a surrogate reward for the root GFlowNet (recall Algorithm 1), a natural driver of catastrophic failures is inaccurately approximated flow function by leaf GFlowNet. In particular, a local flow Fj is significantly larger than its correct value, the resulting model might allocate substantial probability mass to the subgraph associated to the jth leaf GelowNet. In this case, high-probability regions of the remaining subgraphs might be completely missed by the global GelowNet. We illustrate this effect Figure 17, which shows the learned distribution over a 8 x 8 grid for (i) SAL-trained GFlowNet and (ii) SAL-trained GFlowNet with substantially over-estimated F1- To emulate (ii), we train each leaf GFlowNet in standard fashion and, during the aggregation phase, multiply the learned F1 by 100. In both cases, we considered 2-sized FHPs and followed the experimental setup of Section 6.2. Notably, Figure confirms that severely inaccurate flow function F3 may lead to a misrepresentation of the target distribution's 's high-pre prob-pbbbbyility regions by the global GFlowNet. In spite of these results, we stress that we did not observed this pathological behavior throughout our experiments in the main text and in Section C.
Insufficiently diverse suthin- sampling distributions. Recall in Definition 6.2 that SAL depends on a distribution over the initial states Z, of the jth subgraph defining the FHP; see Figure 3. Clearly, when q, probability mass is overly concentrated in relatively small subset of Ij. the corresponding leaf GFlowNet might fail to accurately learn from the target distribution due restricted exploration of the state graph this was observed in Section 4. Also, the efficiency of mode discovery during training would be significantly hindered. In this context, Figure 16 illustrates the mode discovery rate for the task of generating 16-sized subsets of {l..... 32} with the reward function described in Section A when each 9; is a truncated Poisson distribution with mean equal to 1/s000-th of the size of the corresponding Ij. Importantly, similarly Figure 17, this is an extreme corner case that did not pose an issue in our experiments. In practice, we suggest setting qj as an uniform distribution to maximize the diversity of the explored partitions.
43


Reviews:
[{"review_id": "gQPS3moevz", "reviewer": "Reviewer_7dqe", "summary": "The authors develop the generalization bounds for GFlowNets, providing empirical guarantees that the models can generalize beyond the training data, using the PAC-Bayes technique. The authors explore the way how the structure of GFlow's network and trajectory lengths affect its generalization performance. Then the authors suggest a novel distributed algorithm called Subgraph Asynchronous Learning (SAL), which is shown to improve over its centralized counterparts in some experimental settings, such as subset generation, sequence design, and Hypergrid environments (Section 6).", "questions": "1. Is it possible to provide larger-scale experiments? For example, consider the applications of Subgraph Asynchronous Learning to molecular generation problems, such as QM9 or sEH, see [Mohammadpour et al,2024].\n\n2. Can the authors provide more details on the way the Fixed-horizon DAG partition (definition 6.1) should be constructed? May be there are some heuristics or theoretically justified approached to it? It seems that the choice of partition should be of crucial importance for algorithm's performance, otherwise this partition itself can yield mode collapse and other issues with diversity of the samples.\n\n***References***\n\n[Mohammadpour et al, 2024] Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, and Pierre-Luc Bacon. Maximum\nentropy gflownets with soft Q-learning. In International Conference on Artificial Intelligence and Statistics, pp. 2593–2601. PMLR, 2024.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "To the best of my knowledge, neither the PAC-Bayes generalization bounds, nor the distributed training of Gflownets were previously considered in the literature, so both main contributions of the paper appears to be new. The literature review performed by the authors, is exhaustive and well-written, as well as the main text of the paper.", "weaknesses": "From the theoretical point of view, theoretical contribution of the paper is rather incremental, as the main theoretical findings of the paper are mostly corollaries from known results in PAC-Bayes literature. Moreover, the supervised setting for training GFlowNets is questionable and is of limited interest, since its implementation requires a sampler from the target distribution.\n\nFrom the experimental point of view, I would also say that the proposed experimental investigation of the Subgraph Asynchronous Learning algorithm is insufficient. All considered environments can be considered as rather toyish, with actual number of terminal states less than $10^{10}$. This is quite far from the most challenging environments, for example, see for example the molecular generation problems from the recent papers [Jang et al, 2024], [Mohammadpour et al, 2024].\n\nAlthough minor, I would also highlight some stylistic nuances in the exposition. For instance, the authors describe their own results in Theorem 5.2 as 'insightful' (lines 204 and 2130). It is uncommon in the literature to use such language when referring to one's own findings. Similarly, the statistical guarantees presented in the current submission, are often referred to as \"non-vacuous\" bounds. At the same time, it is not clear to me what exactly \"non-vacuous\" means in the context of statistical guarantees, which seems to be an overstatement.\n\n***Minor typos*** \n1. I guess, that in line 148 $p_E$ should be defined as $p_E = (1-\\epsilon)p_F + \\epsilon p_{U}$;\n2. page 3, line 149 - polciy -> policy;\n3. Proposition 5.1 - $\\tau_{\\alpha}$ is undefined, only $\\tau_{1 - \\alpha}$ previously appeared;\n4. page 9, line 223: assigment function -> assignment function;\n5. line 979-980 - [ref] undefined;\n6. General remark - the authors should be more accurate when using the $\\lesssim$ symbol. At least it should be indicated, which problem-specific constants are omitted when it is used. \n\n***References*** \n[Jang et al, 2024] Hyosoon Jang, Yunhui Jang, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. Pessimistic backward\npolicy for Gflownets, 2024. \n\n[Mohammadpour et al, 2024] Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, and Pierre-Luc Bacon. Maximum\nentropy gflownets with soft Q-learning. In International Conference on Artificial Intelligence and Statistics, pp. 2593–2601. PMLR, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "18rGivZTIZ", "reviewer": "Reviewer_ZwMs", "summary": "The paper contributes to the theory of Generative Flow Networks (GFN) on finite directed acyclic graphs.  \nThe key contribution of the work consists in the formulation and the proof of several new so-called generalization bounds controlling the accuracy of the sampling distribution compared to its target.  The author uses these bounds as a motivation to introduce a divide-n-conquer distributed GFN training algorithms: SAL and recursive SAL. These algorithms pack several GFNs together by assuming that the initial flow of the one is the reward of another.", "questions": "1) Could you provide insight and/or examples illustrating the meaning of the bounds you introduce?\n2) I am confused about the link between the FCS and TB losses. Could the author clarify this link and how they differ? Is FCS loss used for training? Is it computable as a metric during training? \n3) Are you able to provide deeper insight into the structural properties of the graphs or into the reward leading to lower generalization, deeper than the dimensionality curse?\n4) How could one make your bounds actionable or act as a guardrail ?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper provides new generalization bounds for Generative Flow networks together with rigorous proofs leveraging sophisticated pre-existing bounds. \nThe paper also introduces an interesting recursive GFN training: recursive SAL algorithm\nThe appendix is rich with further theoretical analysis.", "weaknesses": "We feel the authors do not properly exploit the bounds they produce:\n- A full section is devoted to describing when GFN do not generalize. Still, the conclusion that the number of states or the length of trajectories is detrimental to training is disappointing.\n- motivating corner cases are not  discussed\n- The SAL algorithm does not need the bounds to be motivated, as divide and conquer is a very common strategy in CS in general.\n\nTo elaborate on this, let's consider equation 5. The hypotheses seem to allow any $p_{E,T}$, but the context provided above Lemma 4.1 suggests it is intended as the forward policy of the GFN at hand. Now, in practical cases, the forward policy will be very different from the uniform policy to the point that $\\chi^2(q_{E,T}||p_{E,T})$ may diverge. I tend to agree that target distribution is key, but it is not discussed how one may leverage this. For instance, the bound suggests one may use the product $\\mathcal L_{TB} \\times \\chi^2(q_{E,T}||p_{E,T})$ as a regularization during training, or even that modifying the reward as $R_{\\delta}(x):= R(x)+\\delta \\varphi(x) $ for some summable function $\\varphi$ and let $\\delta$ go to zero controlled by $ \\chi^2(q_{E,T}||p_{E,T})$ would allow controlling the generalization better. Coming back to the declared intent (describing when GFN does not generalize well), I am still looking for specific examples (even in the appendix) of GFN together with couples (graph, reward) that would typically make the bound high yet have a low loss. This would be very useful to showcase when one should be careful when using GFNs. \n\nWe feel the first four pages are too wordy. They announce the results several times, leaving too little space for the interesting matter: an insight discussion of the generalization bounds. \n\nSection 4 is unclear; the set generation problem is barely understandable.\n\nSome mathematical results are strangely placed. For instance, the purpose of Lemma 4.1 is unclear in the section context, while Lemma D.1 is both trivial and stated after it is used.\n \nThe recursive SAL algorithm (which we find of interest) is only introduced in the appendix.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bcNurrPpoF", "reviewer": "Reviewer_NTon", "summary": "The authors of this paper investigate the generalization properties of Generative Flow Networks (GFlowNets), focusing on providing theoretical guarantees for their performance. They introduce the first data-dependent generalization bounds for GFlowNets, leveraging PAC-Bayesian inequalities to analyze how the state space size affects generalization. Building on these theoretical insights, they propose a distributed learning algorithm, which employs a divide-and-conquer strategy to train multiple GFlowNets on subgraphs and then aggregate them to improve mode coverage and distribution matching.", "questions": "1. While the paper addresses generalization and distributed learning, it would be useful if it discussed the potential limitations of SAL more explicitly. For instance, are there scenarios where SAL might fail to converge, or does it require a careful balance of parameters to function well? \n\n2. Additionally, are there inherent limitations to the generalization bounds derived, particularly when applied to different classes of GFlowNets or domains with very large or complex state spaces?\n\n3. The paper notes that traditional training of GFlowNets resembles active learning. Could you expand on how you envision incorporating active learning strategies into your framework? How might this impact the performance of GFlowNets compared to the current approach?\n\n4. In the experiments, were there any signs of mode collapse when aggregating the outputs of subgraph-trained GFlowNets? What strategies would you recommend or envision to mitigate this issue during the aggregation phase?\n\n5. Are there empirical limits which have been identified in the experiments regarding the effectiveness of SAL as the complexity of the underlying GFlowNet increases?\n\n6. SAL relies on fixed partitions to divide the state space into subgraphs, which may limit the flexibility and adaptability of the approach. In real-world applications, the optimal partitioning of the state space may not be obvious and may vary over time as the policy evolves. A rigid partitioning could result in subgraphs that are either too sparse (leading to inefficient training) or too dense (leading to high computational cost). How do the authors envision adapting the partitioning scheme to be more flexible or dynamic as the policy evolves during training?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The authors address the gap in the theoretical understanding of GFlowNets, namely, their generalization properties, by deriving the first data-dependent generalization bounds, which is a novel and impactful contribution. Traditional analyses of generalization have primarily focused on models like neural networks or reinforcement learning agents, so extending this to GFlowNets represents a fresh and innovative application of statistical learning theory.\n\n2. By proposing a divide-and-conquer approach, where multiple GFlowNets are trained on different subgraphs of the flow network and then aggregated through a coordinating GFlowNet, the authors introduce a novel parallel training strategy that could address scalability challenges.\n\n3. The quality of the research presented is high, with clear efforts to rigorously derive and justify the theoretical claims. By framing the generalization bounds in terms of the state space size, the authors have identified a key factor that impacts GFlowNets' generalization capabilities and provided a principled way to analyze it.\n\n4. The structure of the paper is logical, and the methodology and experiments are presented in a well-organized manner.\n\n5. The significance of this work is substantial. GFlowNets have shown empirical promise for a variety of applications, including distribution matching and probabilistic modeling, yet a rigorous understanding of their generalization capabilities has been lacking. By providing data-dependent generalization bounds, this paper addresses a core theoretical question that could impact how GFlowNets are used and optimized in the future.\n\n6. The divide-and-conquer approach of SAL not only enhances scalability but also has potential benefits for robustness and mode coverage in the generated distributions. These improvements could make GFlowNets more viable for applications in domains like natural language processing, biology, and engineering, where large datasets and complex distributions are the norm.", "weaknesses": "1. While the paper presents experimental results for both synthetic and real-world tasks, the scope of these experiments might be limited in terms of diversity of datasets, complexity of environments, and breadth of tasks. Here are some suggestions: \n\n(i) Combinatorial Optimization over Graphs\nDataset: TSP, MAX-CUT, or Graph Coloring datasets commonly used in combinatorial optimization.\nTask: Solve classic graph-based optimization problems, where the state space represents possible solutions (e.g., routes, cuts, or colorings) and the reward function evaluates the quality of each solution. For instance, in a MAX-CUT problem, the reward would be based on the cut size.\nImportance for Distribution Learning: For combinatorial optimization tasks, covering a variety of valid solutions (e.g., different feasible paths with minimal distance in a graph problem) is crucial. Testing SAL here would show whether it can capture a distribution over multiple near-optimal solutions.\n\n(ii) Financial Portfolio Optimization\nDataset: Historical stock or cryptocurrency price data.\nTask: Generate diverse portfolio allocations that balance risk and reward, aiming to maximize expected returns with minimal risk. Rewards could be based on risk-adjusted performance metrics, like the Sharpe ratio.\nImportance for Distribution Learning: Portfolio generation requires balancing exploration of different investment strategies with optimization for high returns. Distribution learning helps to capture diverse investment strategies and avoid over-optimizing for a narrow range of allocations, making it a good test for SAL.\n\n(iii) Neural Architecture Search (NAS)\nDataset: CIFAR-10 or ImageNet for evaluating architectures.\nTask: Generate neural network architectures that maximize performance on a validation set. Rewards could be based on model accuracy, computational efficiency, or other architecture characteristics.\nImportance for Distribution Learning: NAS requires exploration across a diverse set of neural architectures rather than converging on a single high-performing design. This task would test whether SAL can generalize across different architectural choices and find multiple good architectures, covering a broad distribution in the search space.\n\n2. The SAL approach, while promising for distributed training, may have significant computational and memory overhead due to the need to train multiple GFlowNets in parallel and then aggregate their results. This could be a limiting factor in resource-constrained environments or for applications where computational efficiency is critical.\n\n3.  The generalization bounds do not take into account potential dependencies in the data that could arise in scenarios where GFlowNets are used for tasks like sequential decision-making. These dependencies could weaken the validity of the i.i.d. assumption and potentially affect the reliability of the bounds. It will be useful to discuss how the authors might extend their analysis to account for dependencies in sequential decision-making scenarios.\n\n4. Although SAL is designed to improve mode coverage, aggregating multiple GFlowNets could still risk mode collapse if subgraph-trained models fail to adequately capture all modes of the target distribution. If some subgraphs do not represent certain regions of the distribution well, the aggregation process could overlook these regions, resulting in poor overall coverage.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "w3qY8jci5n", "reviewer": "Reviewer_76jg", "summary": "This paper investigates the Generative Flow Networks (GFlowNets), and presents several novel contributions to the understanding and development of GFlowNets: (i) Data-Dependent Generalization Bounds, which is the first data-dependent generalization bounds for GFlowNets, providing statistical guarantees for their learning capabilities; (ii) Impact of State Space Size,  which explores the influence of state space size on generalization performance, utilizing PAC-Bayesian inequalities to formalize these effects. (iii) Subgraph Asynchronous Learning (SAL) Algorithm, which is proposed to improve training efficiency and generalization in GFlowNets by partitioning the network into smaller subnetworks.", "questions": "See it in the Weakness above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper offers a unique perspective on GFlowNet generalization through PAC-Bayesian bounds.\n- This paper is theoretically solid and well-orginazied.\n- The authors validate their findings with experiments on both synthetic and real-world tasks, strengthening the reliability of the conclusions.", "weaknesses": "- The main idea of SAL (i.e., divide and conquer) is novel for GFlowNet. However, I wonder whether the aggregation accuracy will be compromised for large networks. In other words, is there a trade-off in the theoretical results about $m$?\n- The time cost of the Fixed-horizon partition does not appear to be addressed. Does this operation impact model efficiency? More precisely, how can a Fixed-horizon partition be found efficiently?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)