When do GFlowNets learn the right distribution?
Authors: Tiago Silva
Rodrigo Barreto Alves
Eliezer de Souza da Silva
Amauri H Souza
Vikas Garg
Samuel Kaski
Diego Mesquita
Venue: ICLR 2025 Spotlight (2025)
Keywords: GFlowNets
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We examine the limitations and the stability of GFlowNets in face of balance violations, introducing a novel metric for assessing their accuracy.

Abstract: Generative Flow Networks (GFlowNets) are an emerging class of sampling methods for distributions over discrete and compositional objects, e.g., graphs. In spite of their remarkable success in problems such as drug discovery and phylogenetic inference, the question of when and whether GFlowNets learn to sample from the target distribution remains underexplored. To tackle this issue, we first assess the extent to which a violation of the detailed balance of the underlying flow network might hamper the correctness of GFlowNet's sampling distribution. In particular, we demonstrate that the impact of an imbalanced edge on the model's accuracy is influenced by the total amount of flow passing through it and, as a consequence, is unevenly distributed across the network. We also argue that, depending on the parameterization, imbalance may be inevitable. In this regard, we consider the problem of sampling from distributions over graphs with GFlowNets parameterized by graph neural networks (GNNs) and show that the representation limits of GNNs delineate which distributions these GFlowNets can approximate. Lastly, we address these limitations by proposing a theoretically sound and computationally tractable metric for assessing GFlowNets, experimentally showing it is a better proxy for correctness than popular evaluation protocols.

Published as a conference paper at ICLR 2025
WHEN DO GFLOWNETS LEARN THE RIGHT DISTRIBUTION?
Tiago da Silva Getulio Vargas Foundation (tiago. henri rodri go alves, eliezer. si sivva] @fgvvbr.
Rodrigo Barreto Alves
Eliezer de Souza da Silva
Amauri Souza
Vikas Garg
Federal Institute of Ceará
Yai Ltd and Aalto University
amaurihelandad fce. edu.br
vgarg@csail mit edu
Samuel Kaski
Diego Mesquita
Aalto University, Manchester University samuel. www.kkkkaalttt..o. fi
Getulio Vargas Foundation diego mesquital fgv.bvgvbb...
ABSTRACT
Generative Flow Networks (GFlowNets) are an emerging class of sampling methods for distributions over discrete and compositional objects, e.g. graphs. spite of their remarkable success problems such as drug discovery and phylogenetic inference, the question of when and whether GFlowNets learn sample from the target distribution remains underexplored. To tackle this issue, we first assess the extent to which a violation of the detailed balance of the underlying flow network might hamper the correctness of GFlowNet's sampling distribution. In particular, demonstrate that the impact of an imbalanced edge the model's accuracy is influenced by the total amount of flow passing through it and, as consequence, is unevenly distributed across the network. We also argue that, depending on the parameterization imbalance may be inevitable. In this regard, we consider the problem of sampling from distributions over graphs with GelowNets parameterized by graph neural networks (GNNs) and show that the representation limits of GNNs delineate which distributions these GFlowNets can approximate. Lastly, we address these limitations by proposing theoretically sound and computationally tractable metric for assessing GFlowNets, experimentally showing it is better proxy for correctness than popular evaluation protocols.
INTRODUCTION
Generative flow networks (GFlowNets, Bengio al., 2021; 2023) reward-driven generative models for compositional objects (e.g., sequences or graphs) that have been successfully employed in several scientific domains (Deleu et al., 2022; 2023: da Silva al., 2023; Zhang al., 2023d; Jain et al., 2022; Bengio et al., 2021; Jain et al., 2023). In essence, GFlowNets cast sampling from an unnormalized distribution as solving network flow problem (Bazaraa et al., 2004). Starting from an initial state, GFlowNets create valid samples by drawing series of actions according to a (forward) policy network determined by the amount of flow between adjacent states. This process can be interpreted as spreading the total mass of the distribution (flow at the source) through trajectories that lead to elements in the target distribution 's support (sink nodes).
While most works on GFlowNets are primarily empirical, developing a deeper theoretical under- standing of GFlowNets is key to designing better models and assessment methodologies that are both theoretically sound and practically efficacious. this regard, Bengio et al. (2023; 2021) 1) laid out the technical foundations for GFlowNets, showing that model satisfying the so-called balance conditions samples from the target discrete distribution. Lahlou (2023) extended this theory to the context of probability measures supported on arbitrary topological spaces. Also recently, the relation- ship of GelowNets with variational inference (Malkin et al., 2023). reinforcement learning (Tiapkin al., 2024), and diffusion models (Garipov al., 2023; Lahlou, 2023) has been formally established.

---

show that Published as a conference paper at ICLR 2025
Table Main contributions of this work. Highlighted items represent methodooogical advancements.
Section 3
Sensitivity to local failures for general SGs and targets
Thm. Eq. 6, Fig. 4
Formulation of weighted DB loss and comparison against DB
Section 4
Universal approximation of distributions over trees Representational limits of 1-WL GFlowNets
Thm.
Thm.
Formulation of Look-Abead (LA) GFlowNets
Eq. 7
LA-GFlowNets Standard GFlowNets
Thm. 4, Fig. 6
Section 5
Definition of FCS tractable goodness-of- metric Relationship between FCS and TV
Def.
Thm. 5, Cor.
FCS is highly correlated to TV
Sec. 5.1
Inadequacy of commonly used evaluation protocols
Sec. 5.2
Despite these advances, a question of important practical implications remains: when do GFlowNet correctly learn its target distribution? More specifically, little is known regarding the sensitivity of GFlowNet's accuracy to balance violations. the possible causes of imbalance, how evaluate the distributional correctness of GFlowNets for large state spaces in a principled manner.
This paper establishes a series of results to address these fundamental questions. Firstly, we provide bounds on the total variation of GFlowNets as functions of balance fluctuations/vollatios. By considering tree-structured state graphs with identical rewards, we show that flow imbalances at different depths have non-uniform impact on the approximation eapabilities of GFlowNets more specifically, balance mismatches near the root state may have higher impact than those near terminal states. We also extend our analysis to show that similar results hold for general directed acyclic state graphs (DAGs) and multimodal target distributions. To illustrate the pragmatic benefits of these insights, we devise a novel family of learning objectives extending the traditional detailed balance loss. As demonstrate in Section 3, this approach often accelerates training convergence.
After delving into the consequences of an imbalanced flow network the GFlowNet's accuracy, take a closer look at its potential causes. Towards this objective, we study the distributional limits of GFlowNets when sampling -raph-structured objects. Notably, most applications of GFlowNets consist of sampling from distributions over graphs, which render graph neural networks (GNNs) (Gori et al., 2005; Gilmer al., 2017: Xu al., 2019) particularly convenient to parameterize policy networks. In fact, GNNs are often used to parameterize the policies in practice (Bengio et al., 2021; Roy et al., 2023; Zhu al., 2023; Zhang al., 2023d; Pandey et al., 2024). With this in mind, we provide constructions exposing their shortcomings. While GNN-based GFlow can express any distribution over trees under mild conditions, we show that there are simple state graphs and target distributions from which no GFlowNet correctly sample. We leverage analysis to introduce look-ahead GFlowNets (LA-GFlowNett)), a simple yet effective scheme to provably boost the expressiveness of GelowNets. In essence, LA-GFlowNet incorporates children-state embeddings inputs to the forward policy. This allows LA- LA-GFlowNer to distinguish actions that lead to distinguishable states but cannot be told apart by the Weisfeiler-Leman (WL) (Weisfeiler & Lehman, 1968) test.
Remarkably, these impossibility results underline the importance of having reliable metric for probing the accuracy of trained GFlowNet. In this sense, we also provide theoretically sound framework for the distributional assessment (i.e., goodness-of-fe state spaces, which we call flow consistency in sub-graphs (FCS) metric. Put simply, FCS consists of a Monte Carlo estimate the average LI error w.r.t. a distribution of "cuts" of the target's support. The FCS metric is proxy for the absolute error between GFlowNet's sampling and target distributions, and we empirically show that FCS highly correlates with the (often intractable) LI error while requiring up to three orders of magnitude less compute. In contrast, popular evaluation metrics do not accurately capture distributional correctness, e.g., the number of high-reward states visited during training and the average reward of the top-k scoring states
GFlowNets in high-dimensional
2

---

Published as a conference paper at ICLR 2025
(Jang al., 2024; Kim et al., 2024; Bengio et al., 2021). In our view, this contribution is extremely valuable for ensuring smooth progress in the GFlowNet literature.
In Table 1, we summarize the main contributions of this work. Section 3 analyzes the distributional correctness of GelowNets as function of balance violations and leverages these theoretical insights to propose a family of weighted detailed balance (WDB) losses. Section 4 discusses the representational limits of GelowNets for graph domains and proposes LA- LF-GelloNeets to boost the expressive power of GNN- -bsed GFlowNets. Finally, Section proposes FCS as a theoretically grounded metric to assess the accuracy of GFlowNets well-suited to high-ding dibensiinnii settings. Importantly, all sections of this work provide experiments to substantiate our theoretical analyses, illustrating the claims and demonstrating the practical relevance of the methodological contributions.
BACKGROUND
Notations. Let X be a finite set, which we call the set of terminal states, and R be an unnor- malized distribution over X, which is also called a reward function (Bengio et al., 2021). We define the set of states S as an extension of X' comprising two distinctive elements: an initial state, 8, S, and a final state, 5f E S. We hence define DAG G (S, E), termed state graph (SG), such that (i) there are no incoming edges to Soi (ii) for each x E X, there is a directed path from So to z; (iii) there is an edge from each x to $y, which is not directly con- nected to any other state in S; and (iv) there are no outgoing edges from sy. Figure 1 illustrates state graph in which the states represent multisets and the edges denote the addition of an el- ement. We say that a trajectory t = (3))}= on 9 is complete if it starts at 8 and ends at 8y. and write x to denote complete trajectory finished
81 pp(sala)
sy) by the transition (x, sy); for example, (s,, sy) in Figure is complete trajectory. A forward policy over G is a function pf: SxS + R4 for which (s, is prob- ability measure supported on s's children in g, denoted by child(s). We use pr(-1s) (8) and pr(s,) interchangeably. buckward policy pn a forward policy over G's trans- pose. For a complete trajectory т, we write (T) I PF (P, st). A flow is a function F: + I+ s.t. Fix = R. We denote the cardinality operator #. For a trajectory t, #7 denotes its number of transitions.
pr(silo.)
(0,13
po(a.la))
(0)
[11, 2)
88
8o
((2))
[ 2]
X
Figure 1: Illustration of a state graph.
GFlowNets. A GFlowNer learns forward policy pr and (sometimes) a backward policy pn and flow function F on a SG 9 such that the marginal distribution of pp over X. pr(x) = Σ Er (T()). matches given reward function R (up to a normalizing factor). We refer to pr as the GFlowNet's sampling distribution and to π x R as the reward-induces
probability measure over X. Also, we will often denote a GFlowNet by (g, pr, PB, F) or (when there no risk of ambiguity) just (pr, PB, F). GNN (Bengio et al., 2021; Zhang et al., 2023d)
In many applications, Pr and F are parameterized or as a transformer (Deleu et al., 2022; Kim et al., 2024) and pg is fixed as a uniform policy. Since the seminal work of Bengio et al. (2021), numerous learning objectives have been proposed to estimate the model's 's parameters (Malkin et al., 2022; 2023; Madan et al., 2022; Zhang et al., 2023b), most of which are based on the principle of network balance ensuring that the incoming and outgoing flows in state are equal. The detailed balance (DB) loss, for instance, enforces the detailed balance condition (())pr(s' (s) ((s'))((ssss) by minimizing the average log-squared difference between the incom- ing and outgoing flows of states within a trajectory (Bengio et al., 2023; Zhang et al., 2023d),
log ((s)pF(s((s) 2 ((sss))()ss')
LDB ,B, F) = E,
Σ (s.e))er
#T
with the hard-coded constraint F(s') R(s') when s' E X' is a terminal state; the expectation is computed with respect to any positive probability measure over trajectories. Other popular learning objectives, such as the trajectory balance (TB) (Malkin et al., 2022) and subtrajectory balance (SubTB) (Madan et al., 2022) losses, are reviewed in the supplement.
Assessment of GelowNets. For most problems, X' is intractably large and it is not possible to directly compare the GFlowNet's s sampling distribution to the target R. As a consequence, assessing the goodness-of-fi and convergence rate of GetswNets is challenging problem. To avoid this issue,
3

---

Published as a conference paper at ICLR 2025
common practice in the literature (Pan et al., 2023a;b; 2024; Zhang et al., 2023d; Jang et al., 2024) is to measure the count and average reward of the highest-scoring states found during training. The intuition is that well-fitted model will quickly locate high-probability regions of the target distribution. More precisely, let Xr C X be the terminal states found during training, R. E IR, be a hand-crafted thresh- old, and H(Kr, Ra) {x: x E Ar A R(x:) 2 R. be the samples in Xy with reward larger than JS ~ pr be S samples from the trained GelowNet. Then, we define
Ro. Also, let
R(x) and Acc(pr) = min $ Eisiss R(OA)) ... J
Avg(KI, = E
(2)
#H Ro)
[E<R(R(E)
This second metric, referred to as accuracy by Shen et al. (2023); Kim et al. (2024), measures the prox- imity of the GFlowNet the target based on the expected reward assuming that Ez<R [R(x)] can be computed. Importantly, we will show in Section that Avg(Ar, R.). #H( Ro), and Acc(pr) not necessarily connected the closeness of GFlowNet the global minimum of its learning objective and, therefore, may lead to misguided conclusions if not interpreted carefully.
The expressive power of GNNs and the 1-WL isomorphism test. Graph neural networks (GNNs) are the dominating paradigm for graph representation learning (Xu et al., 2019; Hamilton, 2020; Wang et al., 2024; Corso et al., 2024). Most GNNs employ a multi-layered interleaving neighborhood aggregation and update operations at each layer. Specifically, for each node at layer E, the aggregation a nonlinear function of the (8 - 1)-layer representations v's neighbors. The update step computes a new representation for based its representation at layer 8- and the aggregated messages (output of the aggregation step). Importantly, seessag-ppassing GNNs have well-known representational limits, which are upper-bounded by the first-order Weisfeiler-Lehne isomorphism test (1-WL) (Xu et al., 2019; Weisfeiler & Lehman, 1968). We provide more details regarding this relationship in Appendix B.
pressage-passingg scheme,
ON THE PROPAGATION OF ERRORS IN FLOW NETWORKS
Our investigation starts with the main source of distributional errors a GFlowNet, namely, the lack of balance in the underlying flow network. In this pursuit, the primary question we wish to address is: what is the impact of violations to the balance conditions on the goodness-of-fit of GFlowNets?
3.1 BOUNDS ON THE TOTAL VARIATION OF GFLOWNETS
To build intuition, our first result (Remark 1) quantifies the extent to which a violation to the detailed balance condition in single node might affect the TV distance between the GFlowNet's 's sampling distribution and a uniform target in the case of tree-structured SGs, which are often featured in appli- cations, e.g., (Jain et al., 2022; Jiralerspong et al., 2023; Liu & et al., 2023; Hu al., 2023).
Remark (TV for tree-structured SGs). Let (9, (G,pr, plPPF.P Pn, F) be a GFlowNet balanced with respect to a reward R, where G is a directed regular tree with branching factor g and depth h. and R is unnormalized uniform. Also, consider the GFlowNet (g, PF. PB, F such that i) Fs.)) = F(s.) + d and F(s*) F(s*) + ô for some s* child(s,) (s)) and 2 0: ii) F(s) F(s) for all s not reachable from s*; iii) F(s) = Σ (echild(s) F(s'); and iv) Pr (s, s') x F(s') V(s, s') E E()) (see Figure 2). Let fr be the marginal distribution induced by Pr. The TV between PT and π x R satisfies
e (b.g, F(8o)) < TV (ir, m) S <5. F(8o)), with x,, = (1 1/a) 6/cto.
(3)
Naturally, the upper and lower bounds are increasing functions of 6. Importantly, these bounds are tight, i... for any 8. there is a corresponding flow function for which the TV equals the stated bounds. Also, we note that the upper bound c(65, g², F(s.)) increases monotonically with the number of leaves g², i.e., the further the imbalanced edge (so, s*) is from the leaves, the greater the potential damage to accuracy. Notably, this demonstrates that the effect of balance violations the distributional approximation is het- erogeneously spread among the SG's edges. We ex-
F((o0) +
degroe g
F(ee) + 8
F(aa)
9
F(s.) gh
+ 81
F( g"
+ 82
F(sa)
F(s.)
g"
g"
Figure 2: Tree-structured SG w/ excess flow
8 from so to left child. We omit node labels. perimentally validate these findings for the benchmark tasks outlined in Section 2 in Figure 3.
4

---

Published as a conference paper at ICLR 2025
Sequeriess
Plyllogenetics
Sets
Hypergrid
10
Avg. Czale, +)
0
10
Depth (s
20 60
80
Epoch
Figure 3: Average LDB(S,8) 8') (log( (Fs)(Fs)) log(((s'))e (s)))s (s)) along randomly sampled trajectories during the early stages of training. As suggested by our analysis, the DB loss is unevenly distributed across trajectory, with different transitions influencing the loss in diverse ways.
As we shall show in Theorem 1, these intuitive results can extended the context of arbitrarily shaped SGs labeled with any target probability measure т. In this broader setting, of which Equation a particular case, the tree-inherited concepts of depth and branching factor of state are replaced by the total probability mass accumulated by the terminal descendants of s. The reader is invited observe that, under the assumptions of Remark 1, these properties are interchangeable. From practitioner's perspective, however, the exact computation of the quantities appearing in Theorem I unfeasible for most benchmark and realistic problems. Consequently, our empirical analysis in the following section leverages the insights from Theorem 1 and the computational tractability of Remark to derive a weighted detailed balance (WDB) loss which, by assigning different weights to different transitions based their distance the initial state of the SG, aims facilitating the search for a balanced flow assignment and speeding up the training convergence.
Theorem 1 (TV bounds for arbitrary distributions). Let (G, (G,PF, P.PPP,P PB, F) a GFlowNer with arbitrary state graph G satisfying the DB condition w.r.t. arbitrary reward R. Similarly to Remark 1, define (G, pr. Pg. F) by increasing the flow F(s) in some node s by and redirecting the extra flow flow to a direct child s* by properly adjusting py (s, Likewise, F is defined by propagating the extra flows to all states reachable from s* Also, let DD- C X be the set of terminal states reachable from 8*. Then, the TV between the distribution fr over X induced by pr and the normalized target xx R satisfies
8
1- (1-)
Σ T(क) RED,.
TV (pr. F(Fo) +8 (1
5
(1
(1- min a(x)) ZED,
(4)
F(6y) +
3.2 APPLICATION TO GFLOWNET TRAINING
Weighted DB. We note that, by default, the DB loss in Equation computes arithmetic average of the transition-level errors. Intrinsically, this design encodes that each transition has the same impact on our overall goal of approximating the target distribution. Nonetheless, as indicated by our theoretical and empirical analyses in Section 3.1, this is not the case. Therefore, we construct a family of weighted detailed balance (WDB) losses,
(log F(s)pp(('(s) F(s')pp(s(s")
FWDB = (.(..,)): (8, s(s
(5)
and train a GelowNet by choosing C.y E FWDB and minimizing the stochastic objective
Cwon(pF,pE.F) F): E,
(s, Eienrer (
Σ (x, (s,8")er
(6)
We are left with the task of choosing appropriate 7. Inspired by recent advances in diffusion probabilistic models (Kingma et al., 2021; Kingma & Gao, 2023), we might choose a ensuring that no term in Equation 6 dominates the loss. In light of Remark 1, any monotonically decreasing function on #D., (i.e., the number terminal descendants s') would be a principled choice for 7. Here, use 78, 8²) 1/#D,, but acknowledge that other y might be optimal for different tasks.
5

---

Published as a conference paper at ICLR 2025
Sequences
Phylogeneticss
Sets
Hypergrid
0.5
0.8
0.3
10
0.6
0.5
11.2
0.4
0.4
10000
200000
30000
20000
00000
60000
10000
20000
0.0
1.0
Number of observed trajectories
Subits
2b (can)
Figure 4: WDB performs competitively with or better than DB, SubTB, and TB. By weighting each (8, s') in inverse proportion to the number of terminal descendants of s' (i.e., 7(8,8") 1/#D,,) in the DB loss, a faster convergence terms of TV w.r.t the standard objective achieved.
Empirical illustration. We compare the performance of WDB against TB, SubTB, and the standard DB objective (with y = 1) in Figure 4 using four benchmark tasks for GFlowNets: autoregressive sequence design (Jain et al., 2022; Malkin et al., 2022: Jiralerspong
al., 2023), phylogenetic
inference (Zhou et al., 2024), set generation (Bengio et al., 2023; Pan et al., 2023b;a; Jang et al., 2024), and hypergrid navigation (Malkin al., 2022: 2023; Madan et al., 2022). We provide more details regarding the experimental setup in Appendix E. Remarkably, using WDB leads to faster convergence of the GFlowNet's sampling distribution to the target distribution for phylogenetic inference and set generation measured in terms of the TV distance. Note these two environments are exactly the ones for which early-stage transitions dominate the loss, as shown in Figure 3. For the two remaining cases, WDB performs approximately
par with standard DB. These results, which are discussed at length in Section E,4 (see Figure 11), suggest that one may drastically improve the DB loss by adequately weighting each transition within trajectory.
4 DISTRIBUTIONAL LIMITS OF GNN-BASED POLICY NETWORKS
Our analysis so far has focused on the impact of imbalanced nodes on the GFlowNet's sampling distribution. We now step back and analyze natural cause for this lack of balance: parametrization. Notably, some of the hottest applications GFlowNets lie in graph domains and leverage GNNs incorporate desirable inductive biases (e.g., Bengio al., 2021: Nica et al., 2022; Roy al., 2023; Zhang et al., 2023d; Zhu et al., 2023; Pandey et al., 2024). Thus, this section explores the representational limits of GNN-based GFlowNets. Towards this goal, we show their universal capacity of approximating distributions over trees. Then, we construct family of problems that GFlowNet based on 1-WL GNNs, termed as 1-WL GFlowNer, cannot solve, showing that balance violations may arise due to limited expressivity of the policy network.
Our first result (Theorem 2) demonstrates that, for any reward supported over trees, there is an 1-WL GFlowNet capable of sampling proportional to that reward. To achieve this result, construct simple generative process starting from a totally disconnected graph and adding one edge at a time, always yielding only one non-singieton component.
Theorem 2 (Universality of 1-WL GelowNets for trees.). S is collection of trees such that (8, 8') E & implies that GFlowNer equipped with 1-WL GNNs approximate any distribution over X C S.
(s is a proper subtree of s') with #E(s') = #E(s) + 1, then there is
Theorem 2 certifies that 1- GFlowNets can sample from arbitrary distributions over trees. However, the 1-WL test is not perfect oracle for isomorphism. A natural question is: are there limits to the representational power of 1-WL GFlowNets? Theorem shows a broad family of cases (i.e., combinations of SGs and re- ward functions) for which 1-WL GFlowNets fail. This result rests the fact that states must dis- tribute flow evenly to children if the actions lead- ing them are 1-WL. indistinguishable.
6
Rini) =
Ri
Figure 5: A combination of state graph and re- ward function that causes 1-WL GFlowNets to fail.

---

Published as a conference paper at ICLR 2025
Illustrating failure modes of 1-WL GFlowNets. Figure 5 provides a construction where 1-WL GFlowNets fails to achieve balance. Note that the actions leading to the children s' and s" of s (enclosed by a box) are 1-WL indistinguishable. Hence, 1-WL policies distribute the flow in equally among s' and 8", failing to match the distinct rewards R(s')= = 4 and R((*)) = 8. Theorem 3 extends this intuition
wider range of cases.
Theorem 3 (Limitations of -NN-based GFlowNets). Letg= (S, be a state graph and R. X S S Rx be reward function. Suppose G is a directed tree. Let T(s) S X for s E 8 denote the set of terminal states reachable by a directed path starting at 8. If there is a state & = EV, E S and two pairs of nodes (a. + (c, E V2 \ E that are not 1-WL distinguishable and ExET(e)) R(z) + Ener(e") R(x) with s' = (V, E' u ((), and s" = (V, E U {(c, d)}) (illustrated in Figure 5), then there is no 1-WL GFlowNet capable of approximating т & R with TV zero.
We now leverage these insights to propose more expressive GNN-based GFlowNet: Look-ahead GFlowNets (LA-GGlowNett). The rationale of Netss is to incorporate children's 's graph embeddings inputs to the forward policy. This allows LA- AetoNsteeets to disambiguate between children states obtained from I-WL equivalent actions, enabling assignment of uneven probabilities to non-distinguishable actions as long as the embeddings of corresponding children states differ.
More formally, let s' and s be two neighboring nodes in the SG, differing only by an edge (u, u) not in - recall that and s' are graphs themselves. Let also 6.IG be the 1-WL embedding of node within a graph G. Then, LA-GFlowNett forward policy can be described as
pf s') x exp {MLP (vi (fdu)s: )o(+)) || 4/2 (10u)e)
where thi and 1/2 are order-invarisn functions (Zaheer et al., 2017). Since child embeddings are added (via concatenation) to the original action embedding. there is no loss of expressiveness w.r.t. 1-WL GFlowNets. On the other hand, LA-GFlowNets can perfectly approximate cases like the one depicted in Figure 5. Theorem states the superior expressiveness of LA-GFlowNets.
Theorem 4 (LA-GFlowNe is more expressive than 1-WL. GFlowNet). If there is a I-WL.forwart policy inducing a sampling distribution proportional to a reward R. there a LA-GFlowNet forward policy over the same SG with sampling distribution proportional to R. The converse does not hold.
Empirical illustration. To demonstrate the limitations of 1-WL GFlowNets, we define next group G of SGs for which there are actions that, despite leading to non-isomorphic states, cannot be distinguished by a GNN-based policy. In this scenario, let RR,, be the set of regular graphs with n nodes of degree k. Then, let G be the set of SGs C1 + P - C2 such that P E Rr,, and Ci and C2 are non-isomorphic graphs differing from P by a single additional edge; see Figure 9 in the supplement for an illustration. Note that, due to the (graph- theoretic) regularity of P. (P. = Pr C2) for any GNN-based pF. Thus, the corresponding GelowNet is inher- ently unable to learn a non-uniform distribution on C2). Netss, in contrast,
LA-GFloda
GPhinN
0.4
0.2
0.0
-022
0.0
100
100
epochs
not constrained by such limited expressivity. As an example, we create four triples (C1, P,C2)) C2) with n 8, 3, R(C1)- = 0.1 and R(C₂) 0.9. Under these conditions, Figure 6 shows LA-GFlowNet can accurately sample from the target distribution. However, a standard GNN-based GFlowNet can only sample from a uniform, attaining (constant) L, error of 0.4 throughout training.
Figure 6: Illustrations in which LA- GFlowNets succeed but standard
GNN-based GFlowNet fail.
5 CONVERGENCE DIAGNOSTICS FOR GELOWNETS
Finally, with the understanding that there are distributions from which a GFlowNet cannot sample, we ask: how can we tractably assess the closeness of a GFlowNet's to its target? To answer this, we propose a provably correct and computationally amenable metric for probing the distributional incor- rectness of GFlowNets (Section 5.1), termed Flow Consistency in Subgraphs (FCS). Strikingly, we compare FCS against three popular techniques for assessing the convergence of GFlowNets. namely, the number of modes, average reward of top-scoring samples, and Shen's 's accuracy, and show that FCS is often the only metric accurately reflecting GFlowNet's goodness-of-fit (Section 5.2).
7

---

Published as a conference paper at ICLR 2025
5.1 PROBING GELOWNETS' DISTRIBUTIONAL INCORRECTNESS
Flow Consistency in Subgraphs (FCS). The basic principle of FCS is estimate the discrepancy between ratios of probabilities (Hyvärinen, 2007) instead of measuring the divergence between the intractable learned and target distributions. For this, recall that the marginal distribution pr of a GFlowNet (G.p,, Pl,PP.P Pn) over the terminal states X' can be computed as
(T(T)
pr(x)=
Σ pF(() E (r-pp((1)) Fo
(8)
(
for each a E X. For most benchmark tasks, c.g.. hypergrid environment (Malkin et al., 2023), set generation (Shen et al., 2023), and sequence design (Jain et al., 2022), we can exactly and tractably compute pr. In autoregressive problems (Jain et al., 2022). for instance, there is a single trajectory Tz leading to each graphs, the sum in Equation 8 can be explicitly calculated by enumerating the trajectories finishing x. When exact computation of pr is unfeasible, Monte Carlo estimator of the expectation in Equation 8 offers an accurate approximation. In this case, FCS consists of comparing restrictions of Pr (x) and R(x) to random subsets of X. We formalize this procedure the definition below.
X. Hence, Pr(x) (r( can be directly evaluated. Similarly, for small state
Definition 1 (Flow Consistency in Sub-graphs), Let P's be a positive probability distribution on 3-sized subsets of X. 2 2. For each S S X, define the restrictions of pr and R to the set S as
1 [(E8)Pr(z) (2)
and (x) = 1 (UES) R((x) for x E X.
( (x) Eyes PT)) PT
Eves R(y)
We define FCS as the expected TV between PT (9) and R(S);
FCS(pr, R) == Es-rs [V (5) RS)))
(10)
Clearly, FCS(pT, Rt) E [0, Moreover, Theorem shows that FCS(pr, R) only pr (x) R(x), asserting the conceptual correctness of our metric.
Theorem (Equivalence between TV & FCS). Let Ps be any sull-support distribution over (5 S X: #S = 3} for some 3 > 2. Also, let TV (pT, TT 1/2 ExEx (T( πz)] be the TV distance between pr and π = R/Z, Z E R(r). Then, TV(Pr, σ) = 0 and only if FCS(pr: R) 0.
metric (Hyvärinen, 2007) (8 2) and the TV distance (B = #x)). Here, we set / as the size of the batch of trajectories used in training. In this respect, Corollary clarifies the role of B3 in FCS in terms of its proximity to the TV distance.
Notably, 3 interpolates FCS between a sattio-matching--ae
Corollary 1 (Role of 3 in FCS). Let Ps(S; 1 ((88=8y (3=1)
Exes be distribution
over 3-sized subsets X' Also, let pT (S) Exes pT(x), and define (8)) similarly. Then,
TV (pT,T) R) - ES-PP(10) TV (T n(
#X
- 2 3
max (Pr(S) - T(S)-- (11)
SCX, $SX
An implementation of FCS. First, we emphasize that FCS can be easily extended to accommodate variably sized subsets of X. To see this, let Ps be any positive distribution {S C: a < 3) and
TS) note Es-P; (p))) R()))] = 0 only if FCS(PT. R) == (1155(()) [VV (p) R()) 0. In
this context, our implementation of FCS defines Ps as the distribution over at-most- al-most-3-sized subsets of X' corresponding to the terminal states of a batch of trajectories sampled from a fixed policy.
PAC statistical guarantees for FCS. From a statistical viewpoint, FCS approximates the distribu- tional accuracy of GFlowNet by probing the model on relatively small fraction of the state graph. One might wonder, under these conditions, how the empirical estimate compares to a deterministic goodness-of-fit measure. Corollary 2 addresses this issue from a probably approximately correct (PAC) perspective, showing that an estimate of FCS closely approximates TV when a sufficiently large number of subsets is sampled (large m) and the model is relatively accurate (small error).
Corollary 2 (PAC bound for FCS). Let Ps be as in Theorem 5. Then, for any ō E (0, 1), with probability least 1 over choosing i.i.d. B-sized subsets Si,..., Sm Ps of X::
TV(PT, x)
Σ
TV
(Sm) PT
R(Smm))
#X
max SCx#5=
(Pr(S) (T)) + v 220 t/m.
23
1545m
8

---

Published as a conference paper at ICLR 2025
Sets
Sequences
Sets
Sequences
Relative runtime
0.04
0.2
10
10a
10*
0.02
10
0.1
10P
0.025
0,050
FCS
TV
FCS
TV
FCS
FCS
Figure FCS a computationally feasible surrogate for the TV distance. (lett) FCS accurately rep- resents TV in the considered tasks (right) while being to three orders of magnitude faster to compute.
Sets
Bags
Sets
Bags
Top- 1000 sng.
Sees
Bags
g 0.25
modes
2500
50000
0.2
0.00
1000
2000
200000
b000000
25000
23000
Namber of samples
Objective
LED
FL
TB
Figure 8: FCS is the only metric correctly reflecting GFlowNet's distributional accuracy. Number of modes (columns 3-4) and average reward (columns 5-6) of the highest scoring samples found during training do not accurately reflect GFlowNet's goodness-of-fee -fftt, while FCS does (columns 1-2). Remark- ably, we consider the terminally unrestricted variants of LED- and FL-GFlowNets (see Proposition I).
Empirical illustration. Figure 7 (left) shows that FCS closely resembles the TV distance for the tasks of set and sequence generation. In fact, the Spearman correlation between FCS and TV for these tasks 0.99 and 0.90, respectively. Importantly, however, the estimation of FCS is up to three orders of magnitude faster than the computation of the TV distance (Figure 7 (right)). Remarkably, these results attest the usefulness FCS as a
goodness-of-fi metric for GFlowNets.
5.2 CASE STUDY: LED- AND FL-GFLOWNETS WITH UNRESTRICTED FLOWS
LED- and FL-GFlowNets. When training Gets the learning signal is sparse; only available at the end of each trajectory via the reward function. LED- (Jang et al., 2024) and FL- (Pan et al., 2023a) GFlowNets aim at reducing this issue by reparametrizing log F(s, s') as the residual of potential function (s, i.e log F(s,s' so 8¹) + log (s, and minimizing
(LED(S, (log (() (o s') log (ag 8) log F(5') oo s')
(12)
for every (s, 8'). For FL-GFlowNet, (o( -((s)))) - {(s) is fixed as the gap between hand-crafted energy functions satisfying {(x) - log R(x)) for E X and E(s.) 0 (Pan et al., 2023a, Eq. (5)). For GFlowNet, (s, s') is learned. Readers may consult Appendix for further details.
{(s')
LED- and FL-GFlowNets with unrestricted flows. Our findings below reveal that, even when the terminal flows F(x) for x E X' are not constrained to equal R(xx), both LED- and FL-GFlowNets greatly outperform standard GFlowNet according to conventional metrics. As we show both theoret- ically (Proposition 1) and empirically (Figure 8), however, constraining F(r) to R(x) is necessary to ensure GFlowNet's 's sampling correctness even under w.r.t.like parameterization. Importantly, we have significant reasons to believe that unrestricted F(x)) was part of some experiments in the original works Pan et al. (2023a) and Jang al. (2024) (see Section E.3), a fact that strengthens the need for standard, casy-to-compute, and sound metric for GelowNet assessment, such FCS.
Proposition (Unpredictabilityy of GFlowNets with unrestricted terminal flows). Consider an FL- or LED-GFlowNee achieving LLED s') for all transitions (s, and trajectories Then, the learned marginal distribution over X. satisfies pr (x) x R(x)) F F(x) for every & E X.
We refer to GFlowNets trained without enforcing F(x)=R(xr) as terminally unrestricted (TU).
Experimental setup. We empirically demonstrate that FCS is the only metric able to represent GFlowNet's s accuracy when compared to three popular alternatives for assessing these models: number of modes, average reward of top-scoring samples, and Shen's accuracy (see Equation 2; Shen et al., 2023). For this, we consider the standard tasks set and bag (multiset) generation, also featured in Jang al. (2024)'s experimental campaign, and design of DNA sequences bindable with

---

Published as a conference paper at ICLR 2025
yeast PHO4 transcription factor. For the latter, we omit results for the FL-GFlowNet due to the lack of a clear candidate for ç; results are presented in Figure 13. Appendix E contains further details.
Results. There are three main takeaways from Figure 8 and Table 2. First, our baseline model (TB-GFlowNet) accurately learns to sample from the target distribution (Figure 8, left), whereas the (terminally unrestricted) LED- and FL-GFlowNets variants do not. Second, both LED- and FL- GFlowNets find a significantly more valuable portion of the state space than their standard counterpart during training (Figure 8, middle, right), but fail to sample correctly. The large variance of the reported results is a consequence of the lack of a unique stationary solution to the models' respective learning objectives, as stated in Proposition 1, and was also observed by Pan et al. (2023a, Figure 2). Third, Shen's accuracy is not an appropriate proxy for goodness-of-fit All in all, our experiments show that usual metrics should be used carefully when comparing the convergence speed of GFlowNets. In contrast to conventional wisdom in the liter- ature, these quantities do not directly measure GFlowNet's closeness to a global minimizer of its learning objective. In view of this, our analyses highlight the importance of having a theoretically sound and computationally amenable metric for assessing GelowNets to drive progress in the field. Strikingly, FCS is as far as we know - the only alternative satisfying both of these constraints.
Table 2: Shen et al. (2023)'s accuracy metrie in- correctly assigns perfect score to the provably un- sound TU-FL and TU-LED GFlowNet's variants.
LED
FL
TB
Sets
100.00=0000
100.00=0.00
93. -0.08
Bags
100.00 10.00
100.00 10.00
81.38=66
PHO4
100.00=0.00
NA
96.98=0.77
CONCLUSIONS, LIMITATIONS, AND BROADER IMPACT
Conclusions. The learning objective of GFlowNets is to find a balanced flow assignment in a flow network. As such, the inaccuracies of a trained model are consequence of violations to the posited balance conditions. In this work, we first argued that the impact of an imbalanced node on the GFlowNet's distributional correctness heterogeneously distributed across the flow network. As consequence, we extended the DB loss by non- non-uniformly weighting the transition -iise terms to account for this hetcrogeneity, which was shown to be effective in practice. For graph-structured generative tasks, we proved that these violations to the balance might be associated to the limited expressiveness of the GNN that parameterizes the policy network, which limits the range of distributions that GFlowNet can sample from. To mitigate this issue, we introduced LA-GFlowNets to boost the expressive power of GNN-based GFlowNets by incorporating the embeddings of the children of a state into the policy network. Finally, we proposed FCS computationally amenable metric for probing the goodness-of-fit of GFlowNets to its target distribution when the learned flow assignment is potentially imbalanced. Notably, our experiments demonstrated that FCS is better proxy to the GFlow distributional accuracy than conventional diagnostic methods.
Limitations. Although our empirical analysis comprehends standard benchmark problems in the GFlowNet literature and is par with other works in terms of the variety of generative tasks, it does not consider specialized applications such as molecule generation (Pandey et al., 2024) and natural language processing (Hu et al., 2023). Additionally, our weighting scheme for the WDB learning objective was heuristically derived. Albeit effective, it could likely be improved by taking into account desirable properties of loss function, e.g., low gradient variance (Richter al., 2020; Malkin et al., 2023). Lastly, we established results concerning the capabilities and limitations of 1-WL baNdb GFlowNets due their widespread use in practice (Kipf & Welling, 2017; Veličković et al., 2018; Xu et al., 2019; Corso et al., 2024). However, extending this analysis to more expressive (e.g., ordher-order) GNNs is a promising direction (Morris et al., 2021).
Broader impact. In conclusion, we believe our work paves the road for many advancements in the GFlowNet literature. Firstly, the development of more effective weighting schemes for the DB objec- tive in Section 3 may speed up GFlowNet training. Similarly, the limited expressivity of GNN-based GFlowNets laid out in Section 4 is a cautionary tale for using equivariant neural networks when param- eterizing the model's policies and may be useful tool for explaining the difficulty in approximating certain distributions, as well as developing more expressive models. Last but not least, we expect FCS to greatly impact the assessment of Gets and the validation of novel learning objectives.
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 Es- tado do Rio de Janeiro FAPERJ (SEI-26000300077//0000),,, the São Paulo Research Foundation FAPESP (2023/00815-6), the Conselho Nacional de Desenvolvimento Científico e Tecnológico CNPq (404336/2023-0) and the Silicon Valley Community Foundation through the University Blockchain Research Initiative (Grant #2022-199610)
We acknowledge the Aalto Science-IT Project from Computer Science IT and FGV TIC for providing computational resources.
REFERENCES
Pierre Alquier. User-friendly introduction to pac-bayes bounds. Foundations and Trends@ in Machine Learning, 17(2):174-303, 2024.
Lazar Atanackovic and Emmanuel Bengio. Investigating generalization behaviours of generative flow networks. arXiv preprint arXiv:2402.05309 2024.
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-- 1454, 2016.
Mokhtar Bazaraa, John Jarvis, and Hanif Sherali. Linear Programming and Network Flows. Wiley-Interscience, 2004.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. In NeurlPS, 2021.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet Foundations. JMLR, 2023.
Lars Buesing, Nicolas Heess, and Theophane Weber. Approximate inference in discrete distributions with monte carlo tree search and value functions. In AISTA7S, 2020.
Gabriele Corso, Hannes Stark, Stefanie Jegelka, Tommi Jaakkola, and Regina Barzilay. Graph neural networks. Nature Reviews Methods Primers, 4(1):17, 2024.
Tiago da Silva, Eliezer Silva, Adèle Ribeiro, António Góis, Dominik Heider, Samuel Kaski, and Diego Mesquita. Human-in-thheloop causal discovery under latent confounding using ancestral 2023.
GFlowNets. arXiv preprint: /2032, 2309..223..
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 UAI, 2022.
Tristan Deleu, Mizu SokhikawaTToomer Jithendaraa Subramanian, Nikolay Malkin, Laurent Charlin, and Yoshua Bengio. Joint Bayesian inference of graphical structure and parameters with single generative flow network. In NewriPS, 2023.
Joseph Felsenstein. Evolutionary trees from DNA sequences: maximum likelihood approach.
Journal of Molecular Evolution, 1981.
Timur Garipov, Sebastiaan De Peuter, Ge Yang, Vikas Garg, Samuel Kaski, and Tommi S. Jaakkola. Compositional sculpting of iterative generative processes. NeurIPS, 2023.
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In ICML, 2017.
M. Gori, G. Monfardini, and F. Scarselli. A new model for learning in graph domains. In IJCNN, 2005.
Caterina Graziani, Tamara Drucks, Fabian Jogl, Monica Bianchini, Franco Scarselli, and Thomas Gärtner. The expressive power of path-based graph neural networks. In ICML. 2024.
11

---

Published as a conference paper at ICLR 2025
William L Hamilton. Graph representation learning. Morgan & Claypool Publishers, 2020.
William L. Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning large graphs. In NeurlPS, 2018.
Edward Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, and al. Amortizing intractable inference in large language models. In ICLR, 2023.
Aapo Hyvärinen. Some extensions of score matching. Computational statistics & data analysis, 51 (5):2499-2512, 2007.
Moksh Jain, Emmanuel Bengio, Alex Hernandez- Jarrid Rector-Brooks Bonaventure P. 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 ICML, 2022.
Moksh Jain, Sharath Chandra Raparthy, Alex Hernandez-Garria, Jarrid Rector-Brooks, Yoshua Bengio, Santiago Miret, and Emmanuel Bengio. Multi-objective GFlowNets. ICML, 2023.
Hyosoon Jang, Minsu Kim, and Sungsoo Ahn. Learning energy decompositions for partial inference in GFlownets. In ICLR, 2024.
Marco Jiralerspong, Bilun Sun, Danilo Vucetic, Tianyu Zhang, Yoshua Bengio, Gauthier Gidel, and Nikolay Malkin. Expected flow networks stochastic environments and two-player zero-sum games. In ICLR, 2023.
Thomas H Jukes and Charles R Cantor. Evolution of protein molecules. In Mammalian Protein Metabolism. Elsevier, 1969.
Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Yoshua Bengio Dinghuai Zhang, Sungsoo Ahn, and Jinkyoo Park. Local search GFlowNets. In ICLR, 2024.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2014.
Diederik P Kingma and Ruiqi Gao. Understanding diffusion objectives as the ELBO with simple data augmentation. NeurlPS, 2023.
Diederik P Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. density estimation with diffusion models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), NewrIPS, 2021.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In ICLR, 2017.
Salem al. Lahlou. theory of continuous generative flow networks. ICML, 2023.
Elaine Lau, Nikhil Vemgal, Doina Precup, and Emmanuel Bengio. DGFN: Double Generative Flow Networks. In GenBio workshop @ NeurIPS, 2023.
Elaine Lau, Stephen Zhewen Lu, Ling Pan, Doina Precup, and Emmanuel Bengio. QGFN: Control- lable Greediness with Action Values. In NeurIPS, 2024.
Dianbo Liu and al. GFlowOut: Dropout with Generative Flow Networks. In ICMI., 2023.
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 ICML, 2022.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment GFlownets. In NeurlPS, 2022.
Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. ICLR, 2023.
Christopher Morris. Martin Ritzert, Matthias Fey, William L. Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In AAAI, 2021.
12

---

Published as a conference paper at ICLR 2025
Andrei Cristian Nica, Moksh Jain, Emmanuel Bengio, Cheng-Hao Liu, Maksym Korablyov, Michael M Bronstein, and Yoshua Bengio. Evaluating generalization in gflownets for molecule design. In Machine Learning for Drug Discovery workshop @ ICLR, 2022.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GFlowNets with local credit and incomplete trajectories. In ICML, 2023a.
Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, and Yoshua Bengio. Generative
augmented flow networks. In ICLR, 2023b.
Ling Pan, Moksh Jain, Kanika Madan, and Yoshua Bengio. Pre-training and fine-tuning generative flow networks. In ICLR, 2024.
Mohit Pandey, Gopeshh Subbaraj, and Emmanuel Bengio. GFlowNet Pretraining with Inexpensive Rewards. In FM4Science workshop @ NeurlPS, 2024.
Pál András Papp and Roger Wattenhofer. A theoretical comparison of graph neural network extensions. In ICML, 2022.
Jarrid Brotor-Brooks,.. Kanika Madan, Moksh Jain, Maksym Korablyov, Cheng-Har Liu, Sarath Chandar, Nikolay Malkin, and Yoshua Bengio. Thompson sampling for improved exploration in gflownets. In SPIGM workshop @ ICML., 2023.
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco J. R. Ruiz, and Ömer Deniz Akyildiz. Vargrad: low-variance gradient estimator for variational inference. NeurIPS, 2020.
Julien Roy, Pierre-Luc Bacon, Christopher Pal, and Emmanuel Bengio. Goal-conditioned gflownets for controllable multi-objective molecular design. In Challenges in Deployable Generative Al @
ICML, 2023.
Max W. Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso Biancalani. Towards understanding and improving gflownet training. In /CML, 2023.
Amauri H. Souza, Diego Mesquita, Samuel Kaski, and Vikas Garg. Provably expressive temporal graph networks. In NeurlPS, 2022.
Balasubramaniam Srinivasan and Bruno Ribeiro. On the equivalence between positional node embeddings and structural graph representations. In ICLR, 2020.
Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry P Vetrov. Generative flow networks as entropy-regullrrized rl. In AISTATS, 2024.
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench: Benchmarks for data-driven offline model-based optimization. In ICML, 2022.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. In ICLR, 2018.
Nikhil Vemgal, Elaine Lau, and Doina Precup. An empirical study of the effectiveness of using replay buffer on mode discovery in gflownets. In SPIGM workshop @ ICML, 2023.
Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, Alexandre Adam, Jarrid Rector- coccorrBrooos, Yoshua Bengio, Glen Berseth, and Nikolay Malkin. Amortizing intractable inference in diffusion models for vision, language, and control. In NewrIPS, 2024.
Qing Wang, Dillon Ze Chen, Asiri Wijesinghe, Shouheng Li, and Muhammad Farhan. Simathscr[N]]: WL: A new hierarchy of expressivity for graph neural networks. In ICLR, 2023.
Zhe Wang, Petar Veličković, Daniel Hennes, Nenad Tomašev, Laurel Prince, Michael Kaisers, Yoram Bachrach, Romuald Elie, Li Kevin Wenliang, Federico Piccinini, et al. Tacticai: an ai assistant for football tactics. Nature communications, 15(1): 2024.
B. Weisfeiler and A. A. Lehman. A reduction of graph to a canonical form and an algebra arising during this reduction. Nauchno- Techhhne-Technniheskkayaaa Informatsia, 2291:12-16, 1968.
13

---

Published as a conference paper at ICLR 2025
K. Xu, W. Hu, J. Leskovee, and S. Jegelka. How powerful graph neural networks? In ICLR, 2019.
M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. Salakhutdinov, and A. Smola. Deep sets. In NeurIPS, 2017.
Bohang Zhang, Shengjie Luo, Liwei Wang, and Di He. Rethinking the expressive power of GNNs via graph biconnectivity. In ICLR, 2023a.
David Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan. Robust scheduling with gflownets. In ICLR, 2023b.
Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffu- sion generative flow samplers: Improving learning signals through partial trajectory optimization. In ICIR, 2023c.
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 NeurIPS,
2023d.
Ming Yang Zhou, Zichao Yan, Elliot Layne, Nikolay Malkin, Dinghuai Zhang, Moksh Jain, Mathieu Blanchette, and Yoshua Bengio. PhyloGFN: Phylogenetic inference with generative flow networks. In ICLR, 2024.
Yiheng Zhu, Jialu Wu, Chaowen Hu, Jiahuan Yan, Chang-Yt Yu Hsich, Tingjun Hou, and Jian Wu. Sample-efficient multi-objective molecular optimization with gflownets. In NeurIPS, 2023.
14

---

Published as a conference paper at ICLR 2025
A RELATED WORKS
GFlowNets. Generative Flow Networks (GFlowNets) (Bengio et al., 2021; 2023; Lahlou, 2023) were proposed as alternative Monte Carlo Tree Search (Buesing al., 2020) in DAG-structured envi- ronments. Presently, GFlowNets have been successfully applied with remarkable success in difficult problems such as causal discovery (Deleu al., 2022; 2023; da Silva et al., 2023), molecule discovery (Bengio et al., 2021; Nica et al., 2022; Atanackovic & Bengio, 2024; Pandey et al., 2024), text infilling (Hu et al.. 2023; Venkatraman et al., 2024), phylogenetic inference (Zhou et al., 2024), and combinato- rial optimization (Zhang et al., 2023b;c), Similarly. a large effort has been devoted to the development learning objectives with enhanced credit assignment (Bengio et al., 2023; Malkin et al., 2022; Madan et al., 2022; Pan et al., 2023b; Tiapkin et al., 2024) and better exploratory policies (Pan et al., 2023a; Vemgal et al., 2023; Rector-Brooks et al., 2023; Lau et al., 2023; 2024; Jang et al., 2024; Kim et al., 2024). In work, considered both the trajectory balance,
of pnller-to-minmmin
[[LS F(s.)Pr(T) (B(TT=)
CTB(DF.PB.F) F) = =E,
(13)
and the subtrajectory balance (Ti denotes the ith state of T and # the number of states in r),
Eisocistr+1 X log F(ri)p(TaglT)) 2 (T(TTsTTT))
X-i
(SubTB(PF, PB, F) Er, Σ
(14)
which are the most popular loss functions for training GFlowNets. Nonetheless, the problems of identi- fying the failure modes and soundly assessing the accuracy of GFlowNets have received far less atten- tion from the literature, with Shen et al. (2023)'s work being the most closely related to ours.
Expressiveness of GNNs. Graph neural networks (GNNs) are the leading approach for representation and predictive learning graph-structured data (Kipf & Welling. 2017: Veličković et al., 2018; Hamilton et al., 2018; Xu et al., 2019; Corso et al., 2024). Despite their outstanding performance, unveiling the limitations of GNNs remains active line of research. Notably, most works focus on the design of additional structural features of the graph to enhance the expressive power of the GNN (Srinivasan & Ribeiro, 2020; Souza al., 2022: Zhang et al., 2023a; Wang et al., 2023; Graziani et al., 2024). Usually, these features add a non-negligible cost to the evaluation of the GNN and, similarly to the proposed LA- L.ooGsloNeets, there a trade-off between computational complexity and expressivity (Moriis et al., 2021). We refer the reader to the work of Papp & Wattenhofer (2022) for a comparison between GNN extensions terms of their expressive power.
B MESSAGE- GNNS AND THE 1-WL TEST
Here, denote graph as a tuple G EV, where V = 21, n) is the set of nodes and E C V x V is the set of edges. More specifically, we consider attributed graphs - i.e., each node E V has associated features 2e E Rd. Also, we denote the set of neighbors of node the graph as N₁ = {u: : (u, v) E E}.
1-WL test. The Weisfeiler- all nodes in an attributed input graph G by applying the following iterative procedure:
isomorphism test (Weisfeiler & Lehman, 1968) assigns colors for
Initialization: The colors of all nodes in G are initialized using the initial node features: Yo V. Vvci(v) )) = T.. If node features are not available, all nodes receive identical colors; Refinement: At step f, the colors of all nodes
refined using hash (injective) function: for all v E V. we apply cf+ (v) = HASH(C (v), (c((u) (u. v) E EB);
Termination: The test is carried out for two graphs in parallel and stops when the multisets of corresponding colors diverge, returning non-isomorphic. If the algorithm runs until the number of different colors stops increasing, the test is deemed inconclusive.
Message-passing GNNs. Most popular GNNs can be described using the message-passing paradigm (Gilmer et al., 2017). Within this framework, GNN initializes each node's 's embedding
with its original features, i.e., (a) I, for all node E V. Then, each layer E, each node gathers messages from its neighbors u e NN., compiling them into an aggregated message m. (f).
m (
AGGREGATE, (H uu N..),
15

---

Published as a conference paper at ICLR 2025
where AGGREGATE, is an arbitrary function multisets, i..., it is order-invariant. Subsequently, each node uses the (so-called) UPDATE function (e.g., a feedforward neural network) to refresh its embedding in light of the aggregated message:
( = UPDATE (ml)) ( 1
When both the AGGREGATE and UPDATE functions incur no loss of information (i.e., they injective), the message-passing procedure coincides with the refinement step from the 1-WL test (Xu al., 2019). In this case, the GNN attains the same power of the 1-WI isomorphism test. Conversely, the 1-WL serves as upper-bound on the expressive power of message-passing GNNs.
C FORWARD-LOOKING AND
As mentioned earlier, both the forward-looken looking (FL-) and Learning Energy Decompositions (LED-) GFlowNets built upon the principle of enhancing credit assignment via reparamaterizing the flow function F as logarithmic residual of a (either hand-crafted or learned) basic potential function o, i... log F(s, s") log Φ(8, s') + log (sss, s'), and F(s, s') = (s'))p(s'(ss) in the usual policy-based parameterization. Under this novel perspective, the DB loss becomes
CDB(PF.PB. F) = E,
1
Σ (a,
((s,,')) + log F(s)p=(d'(s) F(s') F(s')p((s(s))
(15)
#T
with the constraint that E(s.''er o(s, s') = log R(x), in which x the unique terminal state in the trajectory т. Recall that, for FL-GFlowNets, ō(s, = E(s') E(s) for a hand-crafted energy function such that ((s.) 0 and {(x) = log R(x) (Pan al., 2023a, Assumption 1). For LED-GFlOwNess, (s,,')) is parameterized as an neural network taking as input a concatenation of the vectorial representations of and s'. The parameters of á are then learned by minimizing
HT {(0)) (1E50))
LIS(T) = Bem "inn Bere Bernoull((1-s))
Σ (s,5")ET
My (o(8,,'))
(16)
in which {m,,s' (e.er) is Dropout mask and C = E(ade)rr = Ma.s' is the number of unmaked transitions. During training, we interleave gradient- updates of the potential function & and (pF,PB, F)) until a chosen stopping criterion (e.g., maximum number of epochs) is satisfied.
D PROOFS
This section contains self-contained and rigorous proofs for our theoretical results.
D.1 PROOF OF REMARK
The terminal states of the modified flow network will have two types of nodes, with flow F go
and
F. + 6,, with 6; and Σ 5, 8. We normalize those probabilities to obtain the individual probabilities for each terminal state, which determines the density of each sample. From that, we can proceed to compute the total variation distance between pr and π.
gt
Pr T|lrv=
Σ ZE
(+((x) - (x))
F 1
F + ++ghh gh 1 1
(g) 1
9² F +
g"
gh F+6 8 gh
" h- 8 Σ9 1=
0g:
8
g' +
We can lower bound -h 0 8). by considering that Σ (g -76;-8) 8) g² n
18 taking
the absolute value of the result and each element of the to obtain n0 1グ 250-g~-1 h-18 18 Σi" 0g0
6, 5)
16

---

Published as a conference paper at ICLR 2025
Thus obtain the lower bound
g 8 g" 11-18
+ +g"0
8-
9" ~~1
2i=1 1g" 81 g¹
----18
18
S
g¹(F + 5)
gd [F + 8)
(1 -1- 9
1
8
S 11 x||TV.
F +
This lower bound is reached when all error terms in the terminal states have the same value 6; gh:
To upper bound Ig" 6l we apply the triangle inequality, obtaining Ig" - 61 s g" + 6 and
D Ig"d, 61 g"d - from which obtain the upper bound
g"t 8 g 18 ghō g"~18"
g
P a|llv S
(
8
F+0
To obtain a tighter bound we break the sum Ein1 1g" 61 by partitioning the sum into the first
terms SA = =g" E/= 1à; - f with o; < ÿ and subsequent g² terms SB = g" Σj h-1
jōj
9² with 6, 2 gt. By construction, we know that SA g Σl + Σ h1 Σj=141 6, SB = g g
88
simplifying to Sn SA 8(1² 1) We rewrite S + Sp = SB SA + 2SA 8(g"
-ー + A, and by triangle inequality on SA. we obtain the upper bound Ee= 1g"0; -
SA SB S g² p g +-16. Setting I = g² -1 (the biggest value it can have without breaking the constraints 6,), it simplifies SA SB S g" g" 0 28
1
g 8 n 8 +
h-1
|g* 0; 8)
PT क 7||TV
2
9" (F + +8)
1
g 8 9d hー 18 + n0 8 + 9" h-16 28
2
g¹ (F 6)
g
g (F +5)
1
th F+8
D.2 PROOF OF THBOREM 1
To demonstrate this result, we will need the following facts regarding the function f(x): E R + E [2, as} for positive constants a;,
Lemma (Convexity). Let Ant1 ExER": x, 2 0Λ
EA 2 = 1} and a E R". Then,
f: An41 defined by [(x) En [x; ail convex.
Proof. follows from flax+ (1-a)1) = = laz; 0a, + +1 - (1-a)y;-1e (1-a));-(1-a)o)) S E |x
a,l + (1-10) a) En 1v, a,l = (z) + 11-o0 af) for any E 10. and An+1.
Lemma 2 (Maximality edges). Let e R satisfy eij f j # and eū 1. Then, the function f from Lemma achieves its maximun at arg max fle)).
Proof. We will show that, for each x E An+1, there is a é for which f(et) 2 f(x). In particular, f is maximized at one of the ss. For this, note that
f(x) f Σ x;ci S
xxf(cs)
< max f(e)
(17)
1si5n
17

---

Published as a conference paper at ICLR 2025
due to the convexity of f. Thus, is upper bounded by maxisign f(4). Conversely, there is C for which this upper bound is attained. Hence, arg max, f(x) 2 arg max f(e)).
Lemma 3 (Minimality). Let f be the function of Lemma I and assume that a 2 0 and Eis aj S 1.
Then, f is minimized by Σi=1. ar.
Proof. Choose j E (I,...,n) arbitrarily. Since 2i = _i=1,
Tie
n |0
72
n
2 asl
2il + Og
1+
a
(18)
im1.igh
Correspondingly, the lower bound in Equation 18 is achieved when H; a; for i + j and x 1 1- 1i11 aj 2 0. This ensures that f is minimized by -- 0z.
In words, Lemma and Lemma 2 ensure that the TV distance between finitely supported distributions convex and attains its maximum at Dirac delta.
Proof of Theorem I. Initially, let & be the amount of extra flow reaching X and define 3z 6m/6. Then,
|P 1||rV
1 - Σ
(T(()) [x)]==
Σ (p)(x) [xx]++xx++++))+
Σ
[pT(x) [x])) (19)
2
xEX
sED,
DED:
Since (x) = #(2))+1/ for x E DT- and pr (x) F(z)/++8 F+8 for x E D...
Σ
pp((x) - [xx)
Σ (σx)
(20)
F + 8
MEDE
DEDE
On the other hand,
7(x) 6g 7(x)
7(x)
(pr(x) - x(x)|= = Σ
= Σ
32
(21)
F+8
F
F + DED,
F
DED
.ED,.
By Lemma 2, the function f: EXED,, 18p - -x(x)] maximized at
max f(ey) max ,ED,,
lezy (<:)|
EED,, + 2ED,*
max JED,,
Σ
π(Φ) + +11 - T()))
(22)
DED,,, .azy
=+
T(x) min T()).
vED,
RED,,
Similarly, Lemma 3 ensures that
min BEA#D,.11
(13)=1- Σ T(x).
(23)
ZED..
Thus, since EXED, T(x) = 1- EXEDS, T(x).
8
1- Σ T(00|
< 1|r =||rv - F+8
8 (1 min
Tl0)
(24)
F F+8 8
vED,,
DED,
18

---

Published as a conference paper at ICLR 2025
D.3 PROOF OF THEOREM 2
As stepping stones towards proving Theorem first lay down Lemma 4 and Lemma 5.
Lemma 4. Ler G (V, and G' = (V", (V'E) EE) be two non-isomorphic trees of size at most n. Let be the node embedding map of a 1-WL GNN with at least 2n layers. Then, de F Ou for all v E V and v' E EV'
Proof. Recall 1-WL GNNs can distinguish any pair of non-isomorphic trees. Let Tn and T, denote the of computation trees (CTs) for each node in G and G' after layers, respectively. Likewise, let -1 and - denote the sets of CTs after 2n + layers. Since both graphs are non-isomorphic, 1-WI. has already converged with n steps - the maximum diameter of a tree is n 1. Without loss of generality. Tn Tn * Ø, i.c., there is at least one CT in Tn that is not isomorphic to any tree in Tü The same holds for 2n 1 layers, i.c. T2n-1 T2m-1 # Ø. Note that a CT E Tx - Ty is also a subtree of any T T2 Since Tn Ty Tn is not subtree of any CT in T2n-1 otherwise would be in Tn too. In other words, T2n 0 T- 2, implying directly our claim.
Lemma' 5. Let G = (V, and G' = (V',E") be any two trees of size at most n, i.e. [V] and IV' S.n n. Also, let I I= (U, 8) and 1' = (U", .0) be graphs comprising isolated nodes, and & be the node embedding map of I-WL GNN with at least 2n - layers. If foo duf (dv &w for any (v, u) E VV x and (v',u') E V' x U", then the trees (V U {u), E [(e,u)) and (V' U (W',, (O(((','')) are isomorphic.
Proof. If (dv, 0n) (dor, pw). then we either have that i) é. Oi and o., du,' or ii) Φ D and o. 0n. In the first case, we can apply Lemma 4 to conclude that G = G" (with associated bijection 91). Since du du', we know that F. Zne' and the corresponding singleton graphs are trivially isomorphic as well (with bijection 92). Finally, we can build bijection g between the vertices of the merged graphs by making g(v) (v) if V and g(u) 92(u) = u'. For the second case, Lemma 4 implies G and G' are singletons with 2., Frp and J., T..- The result totally disconnected graph, except for an edge linking nodes with identical features in both graphs.
Armed with the previous lemmata, Theorem 2 is straightforward assuming GNN depth 2n 1. From Lemma 5, we know that the action embeddings for any two nodes have an empty intersection. Likewise, two actions have the same embedding only they leave from the same state and arrive at the same state. Therefore, all edges in the SG receive different embeddings. Recall that GNN embeddings are fed MLP layers, which are universal approximators given enough width. Therefore, a 1-WL GNN followed by MLP can approximate any policy forward pr. The same applies to the backward policy p-- We can use the same combination to get state embeddings, which allow approximating any node flow function F. Therefore, we can choose the triplet (pp, PB, F') respecting the DB conditions, for instance.
D.4 PROOF OF THROREM 3
Assume there is a 1-WL GFlowNet sampling from т. Since G is tree-structured, the mass arriving at T(s1) (82) must arrive through i.e., all paths from 80 to some x E T(s1) T traverse Furthermore, there is no directed path from s' to any terminal in T(s") vice-versa, otherwise the skeleton (i.e., undirected structure) would contain a cycle. Then, F(s, s') Exer(e) R(x)
and F(s, EmEr(>)) R(x), implying F(a, 8') * F"),
D.5 PROOF OF THEOREM 4
Since child embeddings are included as additional inputs to LA-GFlowNets, follows directly that LA-GFlowNets at least as expressive 1-WL GFlowNets. We are left with showing the converse does not hold. In Figure 5, we provide a construction for which I-WL GelowNets fail but LA-GFlowNets do not.
19

---

Published as a conference paper at ICLR 2025
D.6 PROOF OF THEOREM 5
Firstly, let S = f rBB C X' and
pT(z)
R(x)
e(S) - 2 PT(()
(25)
R(S)
in which (S) = EzESPP Pr (r) and R(S) =E L R(xx), as the TV distance between the restric- tions of Pr and R to S. For conciseness, we write Pr (x) pr())/p(5)) and (2) = R(a)/(()),
(3)
R(z)/((x) the normalized reward in X. Similarly, we define when TV H) O. For this, note that e(p) Es-ple(S)]. Then, we first show that c(p)
We also denote by π(z)
(V(rr, T)) = 0 implies pr ((r) for every x and hence pT (5) (S) AA C X. Thus.
e(p) Es-p
(8) PT (x) (5)) (z)1
0.
(26)
On the other hand, assume that e(p) = O. Recall that is distribution of full support over {S _ N: |S| = B} and that 2 2. particular, e(p) ensures that
(x
(x) = 0. (S)
Σ e(S, = 2 EE
(27)
(r(
Hence, (T())((2) p())pr(z) for each S and x E S. Write then S = (x) and conclude that (r('') = (S' (x) for every S' and $ S". Thus, by summing both members of this equality across x' & 5", we notice that
pr((')(1 - -()
(28)
ie.. PT() = T(S).. Thus, by iterating this procedure, we conclude that pr(x) = Txx) for all S and x & S'. Since S and x were chosen arbitrarily, pr(x) T(x) for every æ X. Consequently, TV(pr, πT = This ensures the equivalence between e(p) and TV(pr, π) in terms of characterizing the GFlowNet's distributional correctness.
D.7 PROOF OF COROLLARY
Recall the definition of e(S) in Equation 25. We start demonstrating that Ps(: is indeed probability distribution. Clearly, 3) 2 0 for every S C X. On the other hand,
Σ Ps(S;)) SCX
Σ
Σ (T) BES
B -
#SX
Pr(S)
Σ
n n-1)
(S)
(29)
B (--1)
SCX
Σ n1 n-1
8 8-
3 - pr(x)) = =1,
BEX
since each (T()) appears exactly (3-1) times on the sum above. Hence, Ps(: B) is probability
distribution. As for the rest, let ê = Es-ple (SS)], #X' = n, Pp = {S S X: #S = (8. and
n maxsep, (pr - T (S)). We will first show that 23
TV(pr, - A.A
(30)
Then, will verify that TV m) - é2 A. These inequalities will jointly imply Corollary In this scenario, note there are (1=1) subsets of X with 3 elements containing a X. Thus,
TV(rr.
1 ΣΣX
Σ
n b-1)
(pr(x) ([f())><())))) (x))--
(31)
SEP: ZES
20

---

Published as a conference paper at ICLR 2025
For conciseness, define drv TV(pr, T)) Hence,
B B-1)
-1 [pr(x)) - -x()) Ps(S) -P()))
Pr (x)
((x)
div - é= Σ
PT(S)
xTS)
SEPA FES
ΣΣ
n2 1-11 )
1)
((pr(e) PT(()
1 pr(S)
) ((S)
(2r)
+ x())
(S)
SEP BES
pr(()) (S)
T(S)
(T) x (x)
(r(
(32)
-1 ((x)
PT(
B
(SS)
SEP DES
n -
(S) *(S)
8 -
SEPa
n 1)
$
max (T( T(S) SEP
A
B
2
since Ps(S)/pr(S)
(B 1)
and there are (*) B-sized subsets of X. For the reverse inequality,
notice that
drv - ē=
Σ
3 11
n
(TT(T) x(x)] Ps(S)
(S(
(x T(x)
2
(r()) x ((S)
SEP= ZES
-1
Σ
3 1
11
[PT(T)) - x(x)]
SEP; ZES
pr(#))
(④))
Ps((S)
(33)
(S)
(S)
Pr( (5)
1
Σ
8 11)
(S()) Σ x(x)
2 SEPo
PT( (S)
xES
1 n - -1
)
Σ
(S( - -(S) 2
max (pr(5) =(S)
2 B-1
23 SEP
SEP₂
D.8 PROOF OF COROLLARY
Again, recall the definition of e(S) in Equation 25. We now provide a self-contained proof of Corollary 2, which follows from Corollary I and Hoeffding's 's inequality Alquier (2024). Firstly, let
è Es= (SS) and e; e(S,). Since é e, [1, Hoeffding's inequality yields
H
E exp
A
1
Σ
e
exp
X2 2m
(34)
m
Isism
Then, Chernoff's ' bound implies
X2
Pss....... S....
Σ
e
SE
exp
X
Σ e
ho S
exp
2s
m
m
2m
1345m
15:5m
due to Equation 34. This upper bound minimized when X = sm. this case, X²/2m 2s By letting log 6/mm we verify that
62122.
logg
IP PS.
Sm
e
+
< 6.
(35)
m
m
Issism
21

---

Published as a conference paper at ICLR 2025
Then, Corollary and the complementary of the preceding inequality imply
2 log à
S Sm
TV( S
Σ
SC2 (TT(S) (T()) +
21- (36)
1sism
D.9 PROOF OF PROPOSITION 1
As detailed Appendix C. the global minimizer of both FL- and LED.GGlloNeets learning objectives satisfy ((x)"))r Φ(8, s') = looo R(x) for every trajectory т. Since (LED(S, s') = 0,
F(s) exp(da(s 8') ())) 8') = ((())
for every trajectory finishing at æ. Therefore, for every trajectory ~4 x,
F(x) II exp(-o(s,s)
pr(r) pT(()
F(so) (8,8")
Fxx) exp
pn(((a)) F(So)
Σ ((000))
Fx) R(x).
<PB(T|z) Fs))
Hence,
pr(x)= = Pp Σ R(x) pB(T|>) x F(x)R(x) Σ F(s.)
= (((()))))(((>) (37)
Tume
THME
ensuring that the marginal distribution learned by terminally unrestricted FL- and LED-GFlowNets does not necessarily match GFlowNet's target distribution.
22

---

Published as a conference paper at ICLR 2025
E EXPERIMENTAL DETAILS
We provide further details regarding the experimental setup for each section below. Experiments were in cluster equipped with A100 GPUs, using single GPU per run. We also include more exten- sive discussion of existing empirical results and additional experiments in Section E. and Section E.2, respectively. Section E.3 describes our implementation of both FL- and minally unrestricted variants. All experiments relied on Adam (Kingma & Ba, 2014) with a learning rate of 10- for pp and 10-2 for log Z for stochastic optimization (Madan et al., 2022).
and their ter-
E. EXPERIMENTS FOR SECTION
Set generation. The support X' is defined as the collection of sets with 16 elements sampled from a deposit D = (1.....32) 32}. To define the reward function, we let f: - IR with f(d) (/0,, 1] and let log R(x:) = E dez f(d). We implemented an MLP with 2 256-dimensional hidden layers to parameterize both the forward policy and the flow function. For the weighting function y, we note
that #D.
1 (1) in which |s'| is the current state's 's size.
Sequence design. The support X' is defined as the collection of sequences of size up to 12 with elements extracted from a deposit D = fl....,4 4). We implemented an MLP with 2 256-dimensi hidden layers for both the forward policy and flow functions, both of which received as input
sequence of length 12 padded with 0s. Then, the reward function of a x E R is defined by
f: D + RR. and g: [[1, 12]] - TR,, with ((d), g(o) ~ U[] for d E 'D and i [[1, 12]], through log
(x))g( For the weighting function 7. we note that WD, 1+ 4 + 412-1''1 is the number of terminal descendants.
Phylogenetic inference. A phylogenetic tree is defined by a complete binary tree with labeled leaves corresponding to observed biological species and anonymous internal nodes corresponding to their evolutionary ancestors. Also, we consider a set Y E 1R32x7 of DNA sequences of size 32 associated to the 7 observed species; the likelihood of Y is defined by J&C69 (Jukes & Cantor, 1969y's mutation model and computed by Felsenstein's is algorithm (Felsenstein, 1981), and the reward function is the unnormalized posterior induced by an uniform prior distribution over trees. We adopt the iterative process proposed by Zhou et al. (2024) to sample phylogenetic trees with GelowNets, and use Graph Isomorphism Network (Xu et al., 2019) parameterize pr. For the weighting function y, we recall that #D. = (2 (7 |s')l - 1)!! is the number of terminal descendants of 8', with ss'l as the amount of connected components s' (?).
Hypergrid. The support X' is composed of the tuples (x, E [0, x [0,11] characterizing the 12 x 12 2-dimensional grid. The generative process is identical to that of Bengio al. (2021); Malkin et al. (2022; 2023). The reward function, in particular, is defined by
R(s) 10-3 + 0.5
II [[ +2 2 II 15i52
(11x/111E(03.0000)
(38)
15:52
for a state s = 82): 1a is the indicator function of the event A. Similarly to Madan et al. (2022), we use batch size of 16 trajectories and train the model for 62500 epochs (108 trajectories). We parameterize the forward policy with a MLP composed of 2 228-dimensional layers. Also, #Dx, = (11 (11 x2) + is the number of terminal states reachable from state s'.
Details on the experiments for Figure 3. To further understand Theorem 1 in the context of the training GFlowNets, we show in Figure 3 the average lqg-red balance violation along trajectories for the generative tasks considered in Section 3. As expected, the magnitude of the DB loss is mostly dominated by early-transitions of the generative process. Also, this dominance is more notorious for the problems of set generation and phylogenetic inference and less noticeable for the problems of se- quence design and hypergrid navigation, consistently with the results observed in Figure 4 concerning the improved performance of minimizing our weighted loss in Equation 6 with respect to the tradi- tional approach. In this regard, we emphasize that the design of sequences and hypergrid navigation are the only tasks in Figures 3 and 4 with variably sized trajectories. When training model for these tasks, we observe that most sampled trajectories are relatively short in the initial stages of training, which may hamper the improvements enacted by our weighted learning objective due to the simi- larities of the weights. Nonetheless, as we acknowledge in Section 6, a deeper understanding of the influence of each transition to the accuracy of the GFlowNet's sampling distribution is still required. In any case, is clear that the conventional uniform weighting in Equation 6 is sub-optimal. Importantly, we also believe that the best choice for y should be considered in haabllmm-b--rrbblee basis.
23

---

Published as a conference paper at ICLR 2025
Sequences
Phylogenetics
Sels
Hypergrid
0.3
0.77
Avg. Caugle
0.4
0.2
0.50
0.1
0.25
0.00
0.0
0.0
2
Depth (s.
20 40 60
80
Epoch
ran- dom trajectories during training. Results are averaged across runs. contrast to Figure 3, we ob- serve that y has deeeing--hen-smoollhe (where the training signal, R, is received), then equalizes the weights across transitions. This is notice- able for the Phylogenetics and Set tasks, which are the ones in which LDDB performs the best.
Figure 11: Average (D == (s,, (log(F())pe(s.s)) log((((s)po(sss))))
effect, i.e., it initially assigns large weights terminal states
E.2 EXPERIMENTS FOR SECTION
Setup for Figure 6. This experiment built upon simple 3- 3-state SGs with the form L +- P - R. in which P is a 3-regular graph of 8 nodes and L and P's non-isomorphic children obtained by the ad- dition of a single edge. In particular, we choose four different tuples (Li, P., R.) for the four plots of Fig- ure 6. See Figure 9 for an illustration of 2 of the im- plemented SGs. To parameterize the policies of both LA- and the standard GFlowNets, we use a 3-layer GIN (Xu et al., 2019) having 32-dimensional layer, followed by an MLP of 32-dimensional layers. For LA-GFlowNNt, the MLP's input size is twice as large the one for the standard model.
Figure 9: Examples of tuples (Li, Pi, R.). Added edges and their nodes highlighted.
Runtime analysis for LA-GFlowNets. From a theoretical standpoint, the time-complexity of LA-GFlowNet grows linearly with the maximum number C of children of state in the state graph. In
Runtime for 16 epochs (s)
LA-GFN
contrast, the cost of a conventional GFlowNet implementation constant with respect to C. Please Figure 10 for comparison between the training times of LA-GFlowNNt, a standard GelowNet
10-1
FM
TB
(trained by minimizing TB), and a GFlowNet trained via the flow matching objective (FM) - see (Bengio et al., 2021, Equation 11) for the set generation task with the same hyperparameters described in Section E. We report the average running time for 16 epochs averaged across 5 seeds. Investigating whether an equivalent boost in expressivity can be achieved at a lower complexity is promising future direction.
10a
10*
Batch size
Figure 10: Training times.
E.3 EXPERIMENTS FOR SECTION
FL- and LED-GFlowNets. We followed the experimental setup of Pan et al. (2023a) and Jang al. (2024) when implementing both FL- and ...D-GFl..Neess.... To avoid implementation bias, we reproduced our experiments using Pan et al. (2023a)'s publicly released code and obtained similar results. In particular, both pF and were parameterized with MLPs. For LED-GFlowNet. we carried out 8 stochastic gradient steps for learning á for each epoch during training. For the standard GFlowNet trained by minimizing the TB loss, we followed Malkin al. (2022)'s 's instructions.
We privately exchanged emails with the main author LED-GFlowNets (Jang et al., 2024) regarding whether their implementations enforced F(x) R(x) not. Importantly, he told us he verified their experiments and that their code for the set environment indeed did not enforce F(x:) R(x). thereby implementing the terminally unrestricted GelowNet variant described in Proposition but explained to us his remaining experiments were correct. He also told us that the base code for his experiments on the set enviromment was borrowed from Pan et al. (2023a)'s work.
Available online gituh..om//lnggnnn/llCCoo
24

---

Published as a conference paper at ICLR 2025
Objective with o0 ox
7(5, wh, 7a(s, f #D,
(5) /0D.
Sequences
Sets
0.8
0.4
0.6
0.2
0.4
10000
20000
30000
10000
20000 30000
Number of observed trajectories
Figure 12: Illustration of the effect of y on the learning convergence of GFlowNet. We introduce
two novel weighting functions: 7i and 72. On the one hand, Y1 weights each transition in inverse proportion to the square root of its number of descendants #D., and, expected, produces a behavior similar to that of the originally proposed ς. On the other hand, 7/2 - which is the inverse of ? significantly hampers the learning convergence of the trained GFlowNet.
Top 100 avg.
0.4
modes
FCS
500
2.15
0.2
2.10
0
200
20000
20000
Number of trajectories
Objective
LED
TB
Figure 13: Complement to Figure 8 for the PHO4 task (Shen et al., 2023, Section 7). Again, FCS is the only metric that properly reflects the GFlowNet's distributional accuracy. Indeed, despite quickly covering the high-probability regions of the target (mid and right panels), TU-LED-GFlow fails learn the right distribution (left panel). Results are averaged over three runs.
Set generation. The experimental setup is identical to the one described at Section 3. To compute Shen et al. (2023)'s accuracy, pre-computed the average of R(x) under the target distribution by extensively enumerating the SG's terminal states. For FCS, we randomly sampled 32 batches of terminal states of size up 128 for the Monte Carlo estimator.
Bag generation. The experimental setup is mostly the same used for set generation. However, due to the space of bags being significantly larger than the space of sets, we fix D = {l,..., generate 8-sized multisets with elements in D.
[1,...,16] 16) and
E.4 ADDITIONAL EXPERIMENTS
Transition-we wise losses for WWDE Figure 11 shows the averages of (s, as a function of the transition (s, s')'s "s depth along randomly sampled trajectories. We note that, while the standard DB loss is mostly dominated by earlier states on the initial training stages (see Figure 3), L is governed by near-terminal transitions. Hence, the primary training signal, which corresponds the reward of terminal states, receives a larger weight during the optimization process. On the other hand, as training progresses, the loss variability within trajectory decreases, guiding the GFlowNet to learn a balanced flow across the entire state graph essential for accurate distributional approximation (Theorem 1).
Additional weighting schemes for LYDB In the light of Figure 3, Theorem and of the above analysis, we deduce that an effective weighting function should prioritize near- termitamminl transitions. Figure 11 validates this intuition on the tasks of set generation and sequence design (see Figure 4) by
25

---

Published as a conference paper at ICLR 2025
showing that, when is a strictly increasing function of the transition's depth, the resulting LWDB performs competitively with or better than the standard DB loss. practice, these results can guide the design of appropriate weighting function. However, principled approach for optimally constructing 7 remains open, being an important venue for future investigations.
Effectiveness of FCS in real dataset. We reproduce the experiments in Figure for the task of sampling DNA sequences of length 10 in proportion to a reward function defined by wet-lab measurements of the sequence's binding affinity to yeast transcription factor (PHO4) (Shen et al., 2023; Jain et al., 2022; Barrera et al., 2016; Trabucco et al., 2022). For this, we follow the experimental setup described in Section E.3 for the sequence design task. Consistently with Figure 8, Figure 13 and Table 2 show that FCS is the only tractable metric able to correctly infer the provable incorrectness of TU-LED-GFlowNet. Indeed, in terms of mode discovery and Shen et al. (2023)'s 's accuracy, TU-LED- GelowNet outperforms the I'B-GFloNNet omit FL.-GFlowNG Net since there is natural candidate for the potential function.
despite being farther from the target distribution. We
26


Reviews:
[{"review_id": "hXw4QG3ozy", "reviewer": "Reviewer_ykTe", "summary": "The paper provides a detailed analysis of Generative Flow Networks (GFlowNets) to understand when and whether they accurately learn target distributions. The authors extend the DB loss by non-uniformly weighting the transition-wise terms to account for this heterogeneity. For graph-structured generation tasks, they introduce LA-GFlowNets to boost the expressive power of GNN-based GFlowNets by incorporating the embeddings of the children of a state into the policy network. The paper also introduces Flow Consistency in Subgraphs (FCS), a new metric for assessing GFlowNet performance, arguing that it provides a more reliable measure than popular protocols.", "questions": "1. Could the authors elaborate on the choice of the $\\gamma$ for the WDB loss? Are there alternative weighting functions that could improve performance in specific applications?\n2. As for other applications such as NLP and molecule generation, could the authors provide more insights on whether FCS metric can be generalizable?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The question is well-defined. The authors start with analysis and findings, and then propose their method and new evaluation metrics. \n2. As for the paper presentation, key contributions are well-articulated, with a clear breakdown of findings and their implications for GFlowNets.\n3.. This paper proposes a new metric FCS, which has the potential to standardize GFlowNet evaluation, addressing a critical gap in current methodologies. And the paper provides a detailed comparison of FCS with traditional evaluation metrics.\n4. The paper includes theoretical formulation and proofs as well as empirical results to validate their findings.", "weaknesses": "1. The legends, captions and ticks in figures are too small, which is not very readable for readers.\n2. Some expressions are not well-written or formal. e.g. 'Figure 7 and Table 2 teach us three facts.'\n3. Some figures lack legends (e.g. Figure 5). Though the authors use different colors in captions to distinguish, it's may still lead to confusion.\n4. The weighted detailed balance (WDB) loss design (e.g. the choice of $\\gamma$ ) appears heuristic without detailed discussion on optimizing these weights or examining the variance across tasks, which needs more detailed explanation.\n5. The experiments predominantly focus on controlled scenarios, limiting insight into real-world application suitability.", "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": "XmnHkXwyXW", "reviewer": "Reviewer_2KxY", "summary": "Generative Flow Networks (GFlowNets) are a recent class of deep generative models suitable for representing probability distributions over discrete and compositional data structures. The paper asks a temporary and practically important question, which, to the best of my knowledge, has not been answered until now: Do GFlowNets correctly learn the target distribution? The paper answers this question in three parts.\n\nThe first part investigates how a small perturbation of the detailed balance condition affects the total variation (TV) distance between the sampling and target distributions. The authors manifest their analysis by the average of the summands in the detailed balance loss for different depths of the state transitions. Nicely, this theoretical and empirical finding is then put to practical use by designing a weighted detailed balance loss, which assigns the state transitions with different weights. Consequently, the authors achieve better (or on par) performance than other state-of-the-art training objectives.\n\nThe second part focuses on the parameterization of GFlowNets as the cause of violating the detailed balance. The authors select the popular domain where the target distribution is a distribution over graphs, where the parameterization of GFlowNets is thus instantiated with graph neural networks, and are concerned with the standard questions about the one-hop Weisfeiler-Leman (WL) test and its impact on the representational power of GFlowNets. The authors show that a 1-WL GFlowNet can approximate any target distribution over trees but not graphs with two 1-WL indistinguishable nodes. This analysis leads to the design of the look-ahead GFlowNets, which offer more representational power than the 1-WL GFlowNet. This synthesis is again confirmed by experimental evidence.\n\nThe third part is concerned with assessing the goodness-of-fit of GFlowNets in cases where the learned distribution is intractable. The authors propose Flow Consistency in Subgraphs (FCS) as the expected total variation distance between $\\beta$-sized restrictions of the marginal distribution of the forward policy and the reward function. Moreover, the authors discuss the equivalence between FCS and TV distance and provide PAC statistical guarantees. The empirical results demonstrate that the FCS metric is a computationally efficient and close approximation of the TV distance.", "questions": "Line 370: ``*For most benchmark tasks, e.g., hypergrid environment (Malkin et al., 2023), set generation (Shen et al., 2023), and sequence design (Jain et al., 2022), we can exactly and efficiently compute $p_T$.*'' Please clarify if the statement about efficiently computing $p_T$ for benchmark tasks means you can compute arbitrary marginal distributions of $p_T$ over subsets of $x$. If so, could you provide more details on how this is done efficiently?\n\nFor Figure 2, please clarify if the results are averaged over multiple runs. If so, specify how many runs were performed and what aspects (e.g., initialization) varied between runs.\n\nCould it be possible to show an equivalent of Figure 2 for Avg. $\\mathcal{L}^{\\gamma}_{\\text{WDB}}$ to see a comparison to the standard DB loss? This could help illustrate the impact of the weighting scheme on different transitions.\n\nWhat is the shaded area for each curve in Figure 3? Is it the standard deviation or the interquartile range? How was it computed?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper is very dense and comprehensive.\n\nMost of the theoretical findings are nicely supported by experimental evidence.\n\nThe writing is excellent. Although some parts would deserve minor improvements (more on that below).", "weaknesses": "Minor:\n\nThe writing is excellent, as mentioned above. The whole paper works extensively with the sampling distribution and the target distribution. However, these need to be adequately defined. I would expect to see a more precise statement of these two distributions, e.g., in the first paragraph of Section 2. For instance, in line 143, the authors start using ``*GFlowNets sampling distribution*'' without showing the symbol for it. Right after that, the authors say that their target is $R$, and then, only in line 185, we see that $\\pi\\propto R$. The first mention that $\\pi$ is the target is in line 202. The first appearance of `$\\pi$ is defined to be equal to ...' can be seen in line 398.\n\nLine 120:  ``*Finally, a flow is a function ...*'' It should be a function $\\mathcal{T}\\rightarrow\\mathbb{R}_{+}$, where $\\mathcal{T}\\subset\\mathcal{S}$ is a set of complete trajectories (Definition 7 in Bengio et al. 2023).\n\nLine 123: *``a SG''* -> *``an SG''*\n\nLine 127: *``an uniform''* -> *``a uniform''*\n\nLine 200: Which tasks in Section 2 do the authors have in mind? Do the authors mean those mentioned in the last paragraph of Section 3? The authors first refer to Figure 2 in Section 3.1, but the four tasks are first introduced in Section 3.2. The sequence of ideas can be improved.\n\nFigure 4: *``state graph''* -> *``state graphs''* It would be more readable to separate the two state graphs. Their blending is confusing.\n\n*``i.d.d.''* is inconsistent with *``wrt''* throughout the text.\n\nLines 458 and 462: *``a FL-''* -> *``an FL-''*\n\nLine 459: *``part''* -> *``a part''*", "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": "xst4fhiCf4", "reviewer": "Reviewer_MUz7", "summary": "This paper investigates the theoretical foundations of Generative Flow Networks for distributions over discrete and compositional objects. The paper evaluates the impact of violations in the detailed balance of the underlying flow network on the correctness of GFlowNet's sampling distribution. They demonstrate that the effect of imbalanced edges is influenced by the total amount of flow passing through them. The paper also explores the representational limits of GNN-based GFlowNets, and shows that they cannot correctly sample from certain state graphs and target distributions. To address these limitations, the authors propose the new method, LA-GFlowNets, and a metric to evaluate GFlowNets.", "questions": "In Figure 7, why does FL-GFlowNets have such large variance compared to other methods?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper tries to understand theoretical foundations of the popular GFlowNets, which is important and insightful. It produces theoretical guarantees and evaluate it both theoretically and empirically.\n\n2. The paper clearly shows the contributions in Table 1, which makes the paper easier to follow.\n\n3. The experiments are organized and easy to follow. The evaluation process is logical and comprehensive.", "weaknesses": "Two tiny points:\n\n1. In Theorem 4, the paper says \"LA-GFlowNet is more powerful\". How to define \"powerful\" here?\n\n2. I think the notations in Background is a little complicated. It might be better to use a figure to illustrate the model.", "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 (Spotlight)