Improving Generative Flow Networks with Path Regularization
Authors: Anh Do
Duy Dinh
Tan Minh Nguyen
Nguyen Duy Khuong
Stanley Osher
Nhat Ho
Venue: Submitted to ICLR 2023 (2023)
Keywords: generative flow networks
path regularization
optimal transport

Abstract: Generative Flow Networks (GFlowNets) are recently proposed models for learning stochastic policies that generate compositional objects by sequences of actions with the probability proportional to a given reward function. The central problem of GFlowNets is to improve their exploration and generalization. In this work, we propose a novel path regularization method based on optimal transport theory that places prior constraints on the underlying structure of the GFlowNets. The prior is designed to help the GFlowNets better discover the latent structure of the target distribution or enhance its ability to explore the environment in the context of active learning. The path regularization controls the flow in GFlowNets to generate more diverse and novel candidates via maximizing the optimal transport distances between two forward policies or to improve the generalization via minimizing the optimal transport distances. In addition, we derive an efficient implementation of the regularization by finding its closed form solutions in specific cases and a meaningful upper bound that can be used as an approximation to minimize the regularization term. We empirically demonstrate the advantage of our path regularization on a wide range of tasks, including synthetic hypergrid environment modeling, discrete probabilistic modeling, and biological sequence design.

Under review as a conference paper at ICLR 2023
IMPROVING GENERATIVE FLOW NETWORKS WITH PATH REGULARIZATION
Anonymous authors Paper under double-blind review
ABSTRACT
Generative Flow Networks (GFlowNets) are recently proposed models for learn- ing stochastic policies that generate compositional objects by sequences of actions with the probability proportional to a given reward function. The central problem of GFlowNets is to improve their exploration and generalization.
this work, we propose novel path regularization method based on optimal transport theory that places prior constraints on the underlying structure of the GFlowNet. The prior is designed to help the GelowNet better discover the latent structure of the target dis- tribution or enhance its ability to explore the environment the context of active learning. The path regularization controls the flow in the GFlowNet to generate more diverse and novel candidates via maximizing the optimal transport distances between two forward policies or improve the generalization via minimizing the optimal transport distances. In addition, we derive an efficient implementation of the regularization by finding its closed-form solutions in specific cases and a meaningful upper bound that can be used as an approximation to minimize the regularization term. We empirically demonstrate the advantages of our path reg- ularization a wide range of tasks, including synthetic hypergrid environment modeling, discrete probabilistic modeling, and biological sequence design.
INTRODUCTION
Recently proposed by Beeniio et al. (2021a), Generative Flow Networks (GFlowNets) are genera- tive models for compositional objects, which learn stochastic policy that sequentially modifies temporarily constructed object through a sequence of actions to make the generating likelihood pro- portional to given reward function. Specifically, GFlowNets aim to solve the problem of generating diverse set of good candidates. In biological sequence design, diversity is a crucial consideration because of improving the chance of discovering candidates that can satisfy many evaluation criteria later in downstream phases (Jain et al. (2022)). Especially in the multi-round active learning setting, where the generator was iteratively improved by receiving feedback from an oracle on their pro- posed candidates. the effect of diverse generation becomes apparent because more diversity means more exploration and knowledge gained. Besides, the generalization ability of GFlowNets (Zhang al. 2 Malkin al. (2022) over structured data makes them good framework for discrete probabilistic modeling.
The central problems of GFlowNets are improving exploration and generalization. In this work, propose to train the GFlowNet with an additional path regularization via optimal transport (Villani 2003)) which as prior constraint on its underlying structure. The prior is designed help the GFlowNet better discover the latent structure of the target distribution or enhance its ability to explore the environment in the context of active learning. Precisely, the path regularization via OT can help the GFlowNet generate more diverse and novel candidates via maximizing OT distances between two forward policies or improving generalization via minimizing the OT distances.
To improve GFlowNet's generalization, we propose the following prior con- straints: (i) The forward policies of two neighbor states are expected to be similar in the way that they both have the focused tendency of choosing the next action, which implicitly forces the GelowNet to find states with high rewards rather than exploring, especially in sparse environments. (ii) Tra- jectories related to positive objects (both have high rewards) must share their paths. As result, the similarity of states along trajectories with high flow is higher than in other places. From a proba- bilistic perspective, propose measure the similarity of states and s' in the Gel by the
For generalization:

---

Under review as a conference paper at ICLR 2023
transition probability from to s'; (iii) When the GFlowNet learns something, the sparse flow is expected to generalize better. Thus, although many solutions exist for learning a GFlowNet, our proposal priors promote refining the GFlowNet's flow, ie., enhance flow on high flow trajectories and vice versa.
For diversity and exploration: To encourage the GFlowNet's policy to generate more diverse candidates, such as the multi-round active learning settings, we propose to put prior constraint on the forward policies of two neighbor states. Specifically, this prior constraint intentionally promotes the "dissimilarity" between the forward policies of two neighbor states. In other words, it forces the children states of two considered neighbor states far from each other in terms of probabilistic transition, which will help the GFlowNet generate more diverse and novel candidates.
Why OT is a good solution?'ndeed, we need measure of "distance" between pairs of probability distributions. The optimal transport (OT) theory (Villani (2003) studies how probabiliitic mass can be optimally transported from the supports of one probabilistic distribution the supports of another distribution given a cost function. The minimum transportation cost, called distance, can be used as metric that quantifies the distance between two probability distributions. the context of GFlowNets, we want to affect nearby states, which can be done by regularizing the OT distance between the forward policies Pp (-|s) and P=(-(s') of two neighbor states s and s'. To compute the OT distance, we solve an OT problem between two discrete probability measures, whose support points are the child states of & and s correspondingly, given the transportation cost c(u,, ny) from each child u, of to each child fy of s'. While the weakness of KI divergence is that it requires two interested distributions to share the same set of supports, OT can deal with this problem efficiently. Another reason is that the cost used in our OT distance can capture the given DAG's structure and the GFlowNet's s flow, while directly using KL divergence cannot.
Contributions. In this work, we develop novel path regularization based on OT theory for either helping the GFlowNet better discover the latent structure of the target distribution or enhancing its ability to explore the environment in the context of active learning. Our contributions can be summarized as follows:
1. We propose to train the GFlowNet with an additional path regularization via OT, which acts prior the underlying structure of the GFlowNets for either improving the general- ization capability or enhancing the exploration ability of the Gel
2. We define a new directed distance between two arbitrary states in the GelowNet, which can be naturally chosen as the transportation cost for computing the OT distance and link the proposed regularization to entropy terms.
3. We also derive an efficient implementation of the proposed regularization by finding its closed-form solutions in specific cases and meaningful upper bound that can be used as an approximation when we want to minimize the regularization term.
Organization. The paper is organized follows. In Section we provide the background of GFlowNets and OT. In Section 3 we propose new directed distance between two arbitrary states in the GFlowNet and then derive the formulation of path regularization via OT. We also explain why it is the natural and optimal choice for constructing the transportation cost between states. Then, we derive the upper bound and efficient implementation of the proposed path regularization. We provide extensive experiment results of our path regularization via OT in Section (aatoon/aandd conclude the paper with a few discussions in Section [ Theoretical proofs, as well as experimental settings and additional results, are provided in the Appendix.
2
BACKGROUND
2.1
GFLOWWERS
Given a compositional space X,, where each object J E X' can be constructed by taking a sequence of discrete actions from the action space A. Specifically, the construct of each object begins from the source state so and ends in the final state 8,. Incrementally, the generation process modifies a tem- porarily constructed object, which is called a state E S. In addition, a specific action determines that the object is completely constructed and represents a terminal state, such that states and actions correspond to the vertices and edges directed acyclic graph G = (S,A). The construction of object x E X' defines complete trajectory, which is sequence of transitions
E X. These
2

---

Under review as a conference paper at ICLR 2023
Sn = = T is defined as the set of all complete trajectories. Follow- we may assume that each terminal state 8 E X' has only one outgoing edge, Flows Following Bengio et al. 2021b), a trajectory flow is a nonnegative function F; 7 R+ which represents the probability mass of each complete trajectory т. Consequently, the flow through each state can be defined as F(s) = -ET,,rrr (T)), as well as the flow through each edge F(s + s') = -ESTT F()). We can associate probability measure P with the trajectory flow F. In which, there are two important conditional probabilities, the forward transition proba- bilities (forward policy) Pp(s'(s) (s) F(s + F(s)((s) is related adding an element to build the objects, and the backward transition probabilities (backward policy) PB (s(s)s) = F(s + sFF(((sss) related to removing an element.
(5g $1 ing Malkin et al. (2022).
which 18 -+ Sy.
Learning Objective Theoretically, if the training objective such as flow matching objective (Bengio et al. 20 detail balance objective (Bengio et al., 202Tb), or trajectory balance objective (Malkin et al. 2022) is achieved on all states and possible trajectories respectively, then a GFlowNet can be trained completion, i.e., perfectly generating objects proportional to their rewards. In this paper, we use trajectory balance objective because it brings more efficient credit assignment and faster convergence (Makiin et al. (2022)) ). We provide more background of GFlowNets in Appendix B
2.2 OPTIMAL TRANSPORT DISTANCE
Transportation plans and joint probabilities For two discrete probability measures a and /3 over some space X, the admissible couplings set, which can be interpreted as the set of transportation plans or joint probability distributions, is defined as:
IT (a,,3 = { E ke wl, = ar π T 1 -33-
Optimal transportation The Kantorovich optimal transport (Peyré & Cuturi (0119) between & and 3 is defined as follows:
OTC (a, =
min (C, xEll(a.3)
where C is the cost matrix and Cty describes the cost of transport mass from the support 1th of toward the support jihh of B. Whenever the matrix C is itself metrie matrix, the optimum of this problem, OTc (ox, (3), can be proved to be also a distance. Assuming that h = = d, the worst-case complexity of computing that optimum with any of the algorithms known so far scales in O (dd log d) and turns out to be super-cubic in practice (Pele & Werman (2009), $2.1).
PATH REGULARIZATION VIA OPTIMAL TRANSPORT
3.1 OPTIMAL TRANSPORT FORMULATION OF THE PATH REGULARIZATION
Turn transition probability into directed distance We define a new directed distance between two arbitrary states in the GFlowNet, which is used as transportation cost to compute OT distance. Specifically, the directed distance from state 8 to another state s' is designed to be inversely propor- tional to the probability of going from to s'. And, because we consider two arbitrary states 8 and 8', is trivial that there does not always exist a sequence of forward transitions T from s to such as T (s So 51 - ... Sn s') where st + $t41 A. This will result in many entries of the cost matrix may have infinite values, which can also make the OT cost infinite. Therefore, we consider generalized notion of 7 as a sequence of transitions
(s 8o -> 81 -> ... 8,, 8') where 8 + St-1 can be a forward or backward transition, i.e, r can be back-and-forth trajectory Zhang et al. (2222). ). In fact, this sequence of transitions always exists because when we do not regard the direction of each edge, the given DAG can be considered an undirected connected graph.
Definition 3.1 (Directed distance in the GFlowNet) Let T = (s 50 + 31 + - S., = =') be the sequence of transitions from to s' where + $t41 can be forward or backward transition. The length of the trajectory is defined as follows:
Len(T) == -log {s)),
and the directed distance from s to s' is also defined as follows:
d(s, 8')
min
log (or(P(r s))-
7= -=(s
where (s so 81 ... Sn s') is sequence of transitions from s' where 8, - S141 can be a forward or backward transition and d(s, s') = 0 when = s'. Intuitively, d(s, 8') is the shortest path length from to s'.
3

---

Under review as a conference paper at ICLR 2023
Although, d(ss, s') does not satisfy conditions of being "distance", it is still a pseudoquasimetric, which motivates us to use the term "directed distance" as an analogy with "directed distance in Higraphs""Chaarrandd "Chartrand & Tian (1997). Indeed, the proposed directed distance is the natural choice for distance in the GelowNet. Firstly, let's consider two states 8 and s'. They can be considered the equivalent of one another, and have zero distance, if P(s'(s) [s), the probability of transitioning from state s to state s', is equal to 1. In contrast, if the transition probability is equal to O, we cannot reach s' from s by following the GFlowNet policy, then the distance must be infinite. Another reason is that distances are additive, whereas probabilities are multiplicative. Consequently, if we want the length of trajectory to be related to its likelihood, the transition probabilities between states should be changed to their distance by a "negative" logarithmic scale.
digraphs"
Why do we use the proposed directed distance as transportation cost to compute OT distance? Several reasons exist for using the proposed directed distance as transportation cost to compute the OT distance. First, as discussed above, is the natural choice for turning transition probability into distance. Second, with the transportation cost defined our directed distance, we can decompose the OT distance into entropy and other terms. Consequently, this property not only gives us an upper bound for the OT distance but also imposes sparsity prior the GFlowNet's structure when minimizing the OT distance or improves exploration when maximizing the OT distance. Besides, the construction of the transportation cost makes minimizing OT distance correspond to maximizing the transition probability between children of neighbor states. Moreover, it also allows us to derive closed-form solutions for the OT distance under some conditions.
Optimal transport formulation of the path regularization Once we have defined the directed distance between Iwo states in the GFlowNet, we can define the OT distance between two forward policies of neighbor states. Consider two neighbor states and s' in trajectory т, such that 8 7 A. The forward policy P-((ss) is a discrete probability measure supported by Child(s) {uI,.... ..., Ug] and Pr(+ss') ) is discrete probability measure supported by Child(s') = (V....., )
The OT distance between (-|(ss and P=((ss') can be defined as: OTc (Pr(-1s), Pr(((s') ==
(E]((Pt((b))P<(11))) (C,T),
min
(5)
where the set of admissible couplings is defined as: [[Pr(-Is), P=(+ls') PF((ss) == {- Raxi HI, PF(-1s), n T1, P=((1s)) and C is a cost matrix whose each entry is the length of shortest path from un to ny:
(6)
min Cij c(uu, u3) = d(u,, v3) --(--....-0))
log( ((r u:));
However, the underlying DAG is unavailable during training progress because of the enormous number of states and edges connecting them. Then C.3 in Eqn. can can only be approximately computed by using trajectories in the sub-gran nub-graph containing 8, s', their child states, and the edges that connecting them. Specifically, rather than going directly from u, to vj if this edge exists, we can always move from u, v, along back-and-forth trajectory, i.e., 04; + + s' + Uj. with the probability PB(s|u:)Ppp PP (s(s'(s) (p,ylss) Therefore, each entry of the cost matrix C can be calculated in practice by approximating as follows (note that we abuse the notation of "=" in Eqn. g instead of " for easier viewing):
0.
if U, U6
Cyyy min (- (o]((s PF(s Pr) s')), 1 (og(( u.))), else if Uq + Uソ E A -log(P Pr(s' (s) PF)u 8'),
otherwise.
(8)
Definition 3.2 (Optimal Transport Formulation of the Path Regularization) For any complete + 8n), we define the path regularization via OT as follows:
trajectory (80 + 81
Cor(T) .TC.. (PF(-1s::0) Pp((s
(9)
t=0
where Cr. is the cost matrix where each entry is defined in g
Te is the training policy usually that given by Pr(-1::0 or a modified version of - then the
trajectory loss is updated along trajectories sampled from ne, i.e., with stochastic gradient:
Ex-xo (C(CTB(7) + XCor(T)).
(10)
where A E R, A 0 indicates that want to minimize the path regularization and vice versa.
4

---

Under review as a conference paper at ICLR 2023
Figure 1: OT distance between Pr(-|s) and Pr(-1s'). The for- ward policy P=(-)s) is discrete probability measure supported by Child(s) = {ur, U2 == u,).. Similarly, P=((|s') ) is a discrete probability measure supported by Child(s') = {ur = u1, 22, v2) The cost matrix is a 3 matrix. For example, C11 = d(u1, 01) = (hecause U1 = v1). There exist many possible paths
move from Uy to V3. First, going directly from uy to V/ with a dis- lance len(us -> v3) = - log( (ees(Pr(sslus) (u))). Second, we can move from ug to ug along back-and-forth trajectory, ie.. - 8 s' va, with distance log( PB Pa (slus)) - log(Pr(s' (ss)Pr) log(P (P(eass)))) Because the transportation cost from Ug va the length of the shortest path from us to us. Caa d(us, Va) m min(- (= (s
1 log (p (Palaa)))))
The effect of minimizing the path regularization: Following Eqn. 5 minimizing OTc (Pr(-)s), Pr((ss)) makes (rr('ls) closed to Pr(uus), where c(u,v) v) small. This is because all probability mass from u can be transferred to v with a smaller cost than other places, reducing the transportation costs' expectation. Also, Pp(uls') and Pr(uls) will increase where c(u, v) is small, inducing the similarity of the forward policies and making the flow focus on specific direc- tions. Note that the cost matrix C is not a constant and depends on the forward policies. There- fore, minimizing the OT distance affects not only the forward policies but also the transportation cost. However, the effect on the forward policies is also consistent with the effect on the trans- portation costs. As in Eqn. 14 of Theorem 3.2. increasing Pr(u[s) leads to increase (s(s))). Besides, - Loog(P((ss)))) (oosssu)Ppr PF [s)PF(u)s')) is an upper bound of c(u, v). so increasing PF s) and PF(ols') also makes c(u,e) smaller. Also, the OT distance is equal to zero when all probability mass Pr(u|s')= = (u((ss concentrate on u, U where c(u, v) = O. For examples, (r 1 is special case. When the GFlowNet visits high reward object, the prior by design helps the GFlowNet quickly adapt its flow to this high reward terminal state. More discussion about relation with entropy is provided in Theorem 3.1
and (s(5(+1(ss)
The effect of maximizing the path regularization: In contrast, maximizing makes the forward policies different, so more diverse actions are chosen, leading more diverse and novel candidates. As in the Theorem 3.1 the upper bound of the OT distance contains the entropy of forward policy H(Ps(-)s) Moreover, when the training policy is given by the forward policy Pr. the upper bound also contains the entropy of the path (P(())). Thus, maximizing the OT distance is expected to increase the upper bound, which increases the entropy as well. This means more diversity and exploration. Recall that in terms of probabilistic interpretation, we can rewrite the OT distance as the minimum expectation of the transportation costs min,--1((a,3) C,,,,,,,, v). Thus, maximizing the OT distance means maximizing the cost c(u, v). Because uu, v) an inverse function of transition probability from to e, maximizing the OT distance means minimizing P(u - v). Consequently, the flow is distributed to more states, so more diverse actions are chosen.
3.2 UPPER BOUND AND EFFICIENT IMPLEMENTATION OF THE PATH REGULARIZATION
The cost of computing the OT distance OT (Pr()Is), Pr(+ls") scales at least in O (d log(l))))) where d is the number of support points. By using the Sinkhorn algorithm (Cuturi (2013) 3), which solves
OT with entropic regularization, we can reduce the computational complexity to (id)) (Altschuler ctt al. 2017: Lin et al. 2019; 2022). However, our path regularization's 's definition requires computing the OT distances for all edges in the trajectory т, which imposes heavy burden on the computing resources and capacity. To overcome this problem, in Theorem 3.1 we propose the upper bound of the OT distance OT (Pr(--s), Pr(-ls') The upper bound provides an efficient implementation and can explain the path regularization's behaviors. The reason is that we can decompose the OT distance into entropy and other terms. Moreover, when the GFlowNet's settings satisfy certain conditions (Section 3.3). we can solve the OT problem with a closed-form solution. Our upper bound and closed-for form formulation both have the computational complexity of O(VV) where L is the maximal length of constructed sequences and V is the action space size.
Theorem 3.1 (Upper bound of optimal transport distance) For any trajectory T = (so + s1 + + sw). The path regularization via OT Lor(T) can be upper bound by:
CTa(T) = Σ
Σ (cCChild(.))
log - log (P((()))))))))) + H(PF((
t=0
(11)
5

---

Under review as a conference paper at ICLR 2023
The proof of Theorem 3.1 is provided in Appendix D.I The entropy terms in the upper bound LUB( (7) encourage the sparsity of the GFlowNet's flow. In addition, when we minimizes the upper bound (T)( the terms> (og (u) can make (s (u) become higher when (p) is high. As result, increasing PB (s; (u) makes other flows leading to u pruned. While the meanings of the terms Evechba(s,) Pr (u] (x.)) log(P (P(()))))) and entropy terms are clear, now would like to explain the meaning of regularization terms (og(Pr(st+1)s))) Direct
log(P (s
calculation indicates that:
1og(Pp(sL+1 sx)) = log
II PF (F( - -oog( t=0
(12)
Moreover, when the training policy To is given by the forward policies Pr(-I=:0,, we have:
(- (-o((((())) = H(((T))
(13)
By taking this approach, the upper bound regularizes not only on the forward policy via H())()))++1) but also on the path via H(P(r)). Besides, we think - -Po((((F) log( (T())) = - log( (((rss)) as the length of r. Then minimizing path regularization result in shorter paths and smaller numbers of paths with high flows.
3.3
CLOSED-FORM FORMULATION FOR THE PATH REGULARIZATION
In some specific circumstances, i.e., synthetic hypergrid environment (Bengio et al. 20)1a)), dis- crete probabilistic modeling (Zhang et al. (0022), ). and biological sequence design (Jain et al. (2022) ) Theorem 3.2 allows us to compute OT loss by a closed-form formulation, where the cost matrix C is re-defined by its approximation tion was derived by taking advantage of the following observations. First, in these cases, each two neighbor states and s' on each sampled trajectory, such that state. In other words, there doesn't 't exist any state s" satisfying s + s" E A and s' + s" E A. This finding is resulted from the property that each action can not decompose into the composition of others in the action space. Second, if there exists an edge connecting from a child state u of 8, such that u * 8', to child state of s' then the action of transition from s to u must be the same action of of transition from s' to v and the transition action from u to must be exactly the transition action from s to s',
in Eqn. g Specifically, the closed-form formula-
s' E ... don't 't share any child
Theorem 3.2 (Closed-form solution for optimal transport distance) Let the OT cost between the forward policies of two neighbor states be defined as in Eqn. where the cost matrix C re- defined by its approximation that 8 > s' E A, let a, be an action so that the state-action pair (8, a;) leads to uq and (s', a.))
in Eqn. [ For each non-terminal neighbor states and 8' such
leads to Us, where us Child(s) and v, Child(s s"). Let A A be the set of non- non-terminal valid actions at state s, s' and at be the action of moving from s to s'. If the following conditions satisfied: (1) a # O% + On. dai,ak, Oh A, and (2) if 0i: On, Omg On E A such that a, + an = Om + On, Gg F Om then Oj = an; Oh = ami the following result holds:
oT (T(()()))), Pr(()'))
Σ
log())
H(Pr())s))
wEChild())
+ (log)s) (s) + (P (s')ss)
(14)
+
Σ
min(Pr(u:ls), Pr(2)15))c
4E n A ta
where we define:
min(0, log( (s(lu))) + log( (Pr(s(ss)) + (pr((1))))))) - (og(P((15(ua))) if # s'
& 0,
u; s
(15)
The proof of Theorem
3.2 and the closed-form solution for OT distance at terminal states are pro- vided in Appendix D.2. These closed-fome form solutions of the OT distance will be used in our experi- ments in Section (see Appendix D.2 for the reasons).
EXPERIMENTAL RESULTS
In this section, we numerically justify the advantage of OT regularization over the baseline GFlowNet model only trained with trajectory balance loss on a wide range of tasks: hypergrid en- vironment, discrete probabilistic modeling, and biological sequence design tasks. We aim show

---

Under review as a conference paper at ICLR 2023
that: (i) Minimizing the path regularization via OT improves the GFlowNets' generalization, and the upper bound can be used as an efficient approximation when we want to minimize the regular- ization tern hyper-grid environment, discrete probabilistic modeling); while (ii) Maximizing path regularization via OT enhances the exploration ability of the GFlowNet (biological sequence tasks).
4.1 HYPER-GRID ENVIRONMENT
Task We follow the framework of Malkin et al. (2022) with slight changes to study hyper- arper-grid en- vironment, which evaluates the generalization ability of the GFlowNet to guess and sample unvisited modes of the interested distribution. Consider D-dimensional hyper-gria environment with length of each side is H. where each cell represents non- non-terminal state of the given DAG: & = ($1...., sD) where 8a E {0, 1,-.. ...... H 1} for d E {l,..., D). The source state is (0, O..., 0). For any non- terminal state, the available actions are operations of increasing coordinate i by 1 that still satisfies 8/ S H 1 and a terminating action that moves to corresponding terminal state s", which has its reward:
D II [1s - 1) - 0.5| e (0.25, 0.5]] + 2
D II d=1
R ( Ro 0.5
n | 1) 1) - 0.5| E 0.4)
(16)
where Rg is the constant that controls the discovery challenge and Il is the indicator function. This reward function indicates that only considerable rewards exist at the environment's corners, and there are correct 2 modes. The experiment is conducted for two hyper-grid environments with the number of dimensions 4 and 8 (a higher number of dimensions means more challenging). We
consider the same side length H = 8 and Rg = 10 - for both environments. To evaluate the performance this task, we use KI. divergence and the number of modes found during training as the main evaluation metrics. More details about architectures, hyper-parameters, and evaluation criteria are provided in Appendix E.
Results We plot the mean results over 10 runs for each configuration in Fig. 2 Generally, the ob- served behaviors of the proposed regularization methods become more apparent when the hypergrid environment is sparser. Precisely, although recovering full modes, the GFlowNet model trained by minimizing the path regularization via OT discovers modes faster than the baseline, which indicates its focus on finding directions leading to states with high rewards during the training progress rather than spending time exploring the environment. This also helps the model better discover the latent structures of the interested distribution and achieve lower KL error. We can also see that the upper bound is an efficient approximation in terms of complexity when using a positive regularization co- efficient, whose performance is even better. Meanwhile, Max OT seems unsuitable because of its motivation to improve the model's 'ss exploration, while we only need the high-reward states near the corners, and the majority of states have minimal rewards.
16
TB Baseline
TB + Un-OT
25
TB Baseline
TB un-aT Max OT Min oT
TB + Min OT 12
modes found
TB
20
20
10a
10a
9-²
10a
states visited
states visited
Seseline Nax
27.5
modes found
200
Mox
25.0
150
22.5
3 20.0
100
17.5
50
15.0
10P
10c
10a
104
states visited
states visited
Figure 2: Results on the 4- D (upper) and 8 D (lower) hyper-grid environment. Left: Number of modes found during training. Right: KL divergence between the true and empirical distribution.
7

---

Under review as a conference paper at ICLR 2023
4.2 BIOLOGICAL SEQUENCE DESIGN
Task We follow the framework of Jain al. (2022) simulate the process of designing biological sequences. such as anti- microbiil peptides (AMP). DNA and protein sequences (TF Bind 8, GFP). The experiments are conducted in the multi-round active learning setting, with the goal of generating a diverse set of useful candidates after evaluation rounds. We report the performance score, diversity score, and novelty score of the Topk scoring candidates to evaluate the performance of each method. More details about task description, datasets, hyper-parametery and evaluation criteria are provided in Appendix E.2
Performance
Diversity
Novelly
DynaPPO
0.938 O. 0.009
12.12 # 1.71
9:31 + 0.69
COMs
0.761 + 0.009
19.38 + 0.14
26.47 = 1.30
GFloNNet-LL (paper)
0.932 + 0.002
28.44 1.32
GElowNet-AL
0.874 4: 0.022
31.98
2:27
23.91
1.87
GFloNNet-Min OT-AL
1 0.033
1.66
GFlwwNdd-UR OT-AL
0.828 + 0.022
29.89 1 2.80
24.16
1.75
GEloNettMax OT:AL
0.917 4: 0,003
31.56
2.43
28.86 1: 0.96
Table I: Results on the AMP task with K 100.
AMP The results for AMP design task is shown in Table I We see that the GFlowNet-AL model trained by maximizing the regularization via OT performs better than other baselines in terms of di- versity and novelty. In addition, the TopK performance of the GFLowNet-AL baseline also increases from 0.874 to 0.917 when we maximize the OT regularization and is only lower than the reported performance of DynaPPO. However, DynaPPO has much lower diversity and novelty score, which implies that it mostly generates similar candidates from the training dataset.
TF Bind 8 An interesting observation here is that initial dataset Do contains only half of all possible DNA sequences of length 8 having lower scores. Specifically, low-quality data is very common in practice, and this task, poses big challenge for all the methods to have good results. From T we can see that MINs have the highest diversity compared to the other methods. However, this method has a much lower Topk performance and novelty score, which indicates its generated samples are very similar to the low-quality training dataset. Moreover, although having slightly lower diversity, the GFlowNet-AL model trained by maximizing the path regularization via OT performs better than the others when looking at all metrics it outperforms other baselines in terms of performance and novelty score.
Performance
Diversity
Novelty
DynaPPO
0.58 0.02
5.18
COMs
0.74 0.04
4.38
1: 0.24
1.16 : 0.11
BO-qEI
D.44 + 0.05
4.78
+ 0.17
0.62 + 0.23
CbAS
0.45 0.14
5.35
0.18
2 0,04
0.40 D.14
5.57
0.15
0.36 0.00
CMA-ES
0.47 + 0.12
4.89
U.01
0.64 2 0.21
AmortizedBO
0.62 + 0.01
4.97
+ 0.08
1.00 +: 0.57
GFlowNet-AL (paper)
D.54 0.05
4.68
2.12
GelonNet-AL
0.83 0:33-:000 0.01
4.66 0.08
1.14 2 0.03
GelowNet OT-AL
0.82
0.110
NettL OT-AL
0.83 0.001
4.68 0.100
1.14 2 0.05
GeloNNettMax OT-AL
0.85 1: 0.02
4.52 0.188
1.21 2:1::00 0.10
Table 2: Results on the TF Bind task with K = 128.
GFP Lastly, the results for the GFP task are shown in Table 3 where the objective is to find protein sequences having high fluorescence. We observe that the GFlowNet-AL model trained by max- imizing OT regularization generates more diverse and novel candidates than other methods. In addition, its performance score is only lower than the best one achieved by COMs and higher than the GFlowNet-AL. baseline. However, when looking all metrics, the GFlowNet-AL model trained by maximizing the path regularization via OT still outperforms all other baselines.
Note that in all biological sequence tasks, Min OT and UB OT do not improve the performance of GelowNets since these methods make the forward policy at each state have a focused tendency, which seems not to increase the diversity and novelty of the generated candidates.
4.3 SYNTHETIC DISCRETE PROBABILISTIC MODELING TASKS
Task We follow the framework of Zhang et al. (202)). Energy- based Generative Flow Networks (EB-GFN), model seven different distributions over 32-dimensional binary data that are dis-
8

---

Under review as a conference paper at ICLR 2023
Performance
Diversity
Novelly
DynaPPO
0.794 0 0.002
206.19 19 = 0.10
203 20 = 0.47
COMS
0.831 0.003
204.14
0.14
201.64 0.42
0.003
139.89
203.60 0.06
CBAS
0.817 0.012
5.42 0.18
1.81 + 0.16
MINS
0.761 = 0.0007
6.39 + 0.00
2.42 + 0.00
CMA-ES
0.063 0.00
201.43 0.12
203. 02...22-1 0.09
AmortindRO
0.001
NelowNetAA.
0.853 0.004
211.51
0.73
210.36 0.82
NeL
0.8232 0.0001
218.54 7.88
222. 5.49
GelowNet-Min OT-AL
0.8231 + 0.0001
182 03 0.25
220. 36 2.02
GFlowNettUB OT-AL
0.8232 1: O. 0.0001
221 64 0.11
218. 02 0.79
GFoowettMMax OT-AL
0.8233 0.0001
225.00 + 3.76
1.44
Table 3: Results on the GFP task with K = 128.
cretizations of continuous distributions over the plane. The state space S of the GFlowNet consists of vectors of length D = 32 with with entries in {0; 1; @2.. The source state is 50 = (Ø, D,..., 2). For any non-terminal state, the available actions are turning void entry @ to 0 or 1. After D tions, we reach the terminal states having all entries in {0; The main evaluation metrics are NLL score and MMD score. The detailed settings about architectures, hyper-parameters, and evaluation criteria are provided in Appendix E.3
The results for synthetic discrete probabilistic modeling tasks (Synthetic EB-GFN) shown in Table 4 Training the GFlowNet with either minimizing the path regularization via OT (Min OT) or via the upper bound (UB OT) gains the better NLL and MMD scores than the baseline and Max OT. We also observe that the performance of training EB-GFN with Min OT and UB OT are quite similar. Meanwhile, Max OT not useful due to the same reasons provided in hypergrid environment modeling task. Because there is gap between our reproduce results and the baseline EB-GFN Zhang et al. (2022), we only take into account the the reproduce results of EB-GFN when comparing with our methods (Min OT, Max OT, and UB OT).
Results
Zspirals
Sgaussians
cireles
I0000S
pinwheel
swissroll
checkerboard
20.
20.565
19.763
19.599
20.172
21.314
ALOE
20.280
20:350
20.565
19:287
19.821
20.160
54.653
ALOE+
20.062
19.984
20.570
19.743
19.576
20.170
21. 1442
EB-GEN (puper)
20.050
19.982
20.546
19.752
19.554
20.146
20.896
BB-GEN
20.0679
19. 9862
20.5598
19.7324
19.5733
20. 10.11599
28.6839
EB GEN Max OT
20.0673
19. 19.9857
20.5399
19.7319
19.5714
20.1597
08.8
00.0651
19.9654
20.5600
19.7305
19.5707
20.1596
BB-GEN Min OT
20.0640
19.9855
20.5598
19.7308
19.5699
20.1595
20,6831
2.160
0.954
0.188
0.962
0.505
1.382
2.831
ALOE
21.926
107.320
0. 4.499
26.894
39 39.091
0.471
61. 562
ALOE
0.149
0.078
0.636
0.516
1.746
0.718
12.138
(poper)
0.531
0.305
1,206
EB-GEN
0.3012
0.0408
-0.1724
-0.1744
0,2058
0.1555
-0,9966
EB-GEN Max
0.3258
0.0197
-0.1919
0.0763
EB GEN + UB OT
0.2902
0.0102
0.2819
-0.1253
0.1561
0.0257
-0,9923
EB-GEN Min OT 0.1816
0.0343
-0.2755
-0.1966
0.1220
0.1354
-0.1071
Table 4: Results on the Synthetic EB-GFN tasks. The negative log-likelihood (NLL) and MMD
displayed in units of 1 10-4. ALOE+ uses 30 larger parametrization than ALOE and EB- GFN.We We only take into account the the reproduce results of EB-GFN when comparing with our methods (Min OT, Max OT, and UB OT).
5 CONCLUDING REMARKS
In this paper, we propose to train the GelowNet with an additional path regularization via Optimal Transport that places prior constraints on the underlying structure of the GFlowNet. We have empiri- cally shown that minimizing the path regularization via OT improves the GFlowNet's generalization while maximizing path regularization via OT enhances the exploration ability of the GFlowNet. In addition, we derive an efficient implementation of the regularization by finding its closed-form so- lutions in specific cases and a meaningful upper bound that can be used as an approximation when we want to minimize the regularization term. A limitation of the current method is computing the optimal transport distances for all couples of nearest neighbor states. Our proposed Dropout OT (see in Appendix C might be solution. In future works, aim to develop a more efficient path regu- larization for high dimensional discrete data or propose a new cost function to compute the optimal transport distances.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Jason Altschuler, Jonathan Weeed, and Philippe Rigollet. Near-linear time approximation algo- rithms for optimal transport via Sinkhorn iteration. Advances Neural Information Processing Systems, pp. 1964-19744 2017.
Christof Angermueller, David Dohan, David Belanger, Ramya Deshpande, Kevin Murphy, and Lucy J. Colwell. Model-based reinforcement learning for biological sequence design. In ICLR, 2020.
Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL bttps::/ net/ forum?id= OPEUP3926gE
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, Trevor Siggers, Leila Shokri, Raluca Gordån, Nidhi Sahni, Chris Cotsapas, Tong Hao, Song Yi, Manolis Kellis, Mark J. Daly, Mare Vidal, David E. Hill, and Martha L. Bulyk. Survey of variation in human transcription factors reveals prevalent dna binding changes. Science, 1552000035411 2016. doi: (0.11266science.aad22... URL http:: www. science ssciece.
David Belanger, Suhani Vora, Zelda E. Mariet, Ramya Deshpande, David Dohan, Christof Anger- mueller, Kevin Murphy, Olivier Chapelle, and Lucy J. Colwell. Biological sequence design using batched bayesian optimization. 2019.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. In NeurIPS, 2021a.
Yoshua Bengio, Tristan Deleu, Edward J. Hu, Salem Lahlou, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations. CoRR, abs/2111.09266, 2021b. URL
2111.003 09266
Nicolas Bonneel, Michiel van de Panne, Sylvain Paris, and Wolfgang Heidrich. Displacement interpolation using lagrangian mass transport. ACM Trans. Graph., 30(6):158, 2011. doi: 011114//277781.22241411
hRtpp:::://doii http///....rg/11111///////0///// http:s///ooolorrr/...11///0000.0
Charlotte Bunne, David Alvarez-Melis, Andreas Krause, and Stefanie Jegelka. Learning generative models across incomparable spaces. In ICML. 2019.
G. Chartrand and S. Tian. Distance in digraphs. Computers & Mathematics with Applica- fions, 34(1 15-2-23, 1997. ISSN 0898- 1221. doi: 00216-2. URL https: wwtp///www../. sci encedi rect. rrett..com/scieenee/articllll//
S0898122 02581221999 170029970210 122
Nicolas Courty, Rémi Flamary, and Devis Tuia. Domain adaptation with regularized optimal trans- port. In ECML/PKDD, 2014.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Christopher C. Burges, Léon Bottou, Zoubin Ghahramani, and Kilian Q. Weinberger (eds.), Advances in Neural Information Processing Systems 26: 27th Anmual Conference on Neural Information Processing Systems 2013. Proceedings of a meeting held December 5-8, 2013, Lake Tahoe, Nevada, United 2013. URL ht https:: ngs heuris.. htppscccppppe
States, pp. 2013 2011://hhh422220000/222222222322
Abst html
Dai, Liu, Hanjun Dai, Niao He, Arthur Gretton, Le Song, and Dale Schuurmans, Exponential family estimation via adversarial dynamics embedding. In NeurlPS, 2019.
Hanjun Dai, Rishabh Singh, Bo Dai, Charles Sutton, and Dale Schuurmans. Learn ing discrete energy-based models via auxiliary-variable local exploration. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien
10

---

Under review as a conference paper at ICLR 2023
Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12. 2020, wttpppppe////////p/e/eepppp/
virtual, 2020. URL /tttpss:/prooeednggg..neeer
-bsses777---bbstatatassat .tmml
Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. In James Cussens and Kun Zhang (eds.), Uncertainty in Artificial Intelligence, Proceedings of the Thirty-Fighth Conference on Uncertainty Artificial Intelligence, UAI 2022, 1-5 August 2022, Findhoven, The Netherlands, volume 180 of Proceedings of Machine Learning Research, pp. 518-528. PMLR, 2022. URL http::::prrceeeddng. mlr mlz wwtppppeseeepeppppppppppppp http://11..//dee22...aa.a.. htmmll2aammmmmm
Yilun Du Shuang Li, Joshua B. Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy based models. In ICML, 2021.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn diver- gences. In AISTATS, 2018.
Will Grathwohl, Kevin Swersky, Milad Hashemi, David Duvenaud, and Chris J. Maddison. Oops I took A gradient: Scalable sampling for discrete distributions. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18- 24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 3831 3841. PMLR, 2021. httt::::://:/ttot ht.l
ngs m wtt.pppesss./eeeeeeessp /ttpaa/s/113//ggaaahhhhh2222
Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(25):723-773, 2012. URL
http:// /mlr .rgg httppaapaaa/s// httpeers/vll///rrteeeee
html
Nikolaus Hansen. The ema evolution strategy: comparing review. In Towards a New Evolutionary
Computation, 2006.
Geoffrey E. Hinton. Training products of experts by minimizing contrastive divergence. Neural Computation, 14:1771-1800, 2002.
Geoffrey E. Hinton, Simon Osindero, and Yee Whye Teh. A fast learning algorithm for deep belief nets. Neural Computation, 18:1527-1554, 2006.
Nhat Ho, XuanLong Nguyen, Mikhail Yurochkin, Hung Hai Bui, Viet Huynh, and Dinh Phung. Multilevel clustering via Wasserstein means. In International Conference Machine Learning, pp. 1501-1509, 2017.
Moksh Jain, Emmanuel Bengio, Alex Gareía, Jarrid Rector-Brooks Bonaventure F. P. Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological sequence design with gflownets. In ICML, 2022.
Tianyi Lin, Nhat Ho, and Michael Jordan. On efficient optimal transport: An analysis of greedy and accelerated mirror descent algorithms. In International Conference on Machine Learning, pp. 3982 3991, 2019.
Tianyi Lin, Nhat Ho, and Michael Jordan. On the efficiency of entropic regularized algorithms for optimal transport. Journal of Machine Learning Research, 2022.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory bal- ance: Improved credit assignment in gflownets. CORR, abs/2201. 13359, 2022. URL https:
//arxiv.org/aae 2201. 13259
Frederic P. Miller, Agnes Vandome, and John McBrewster. Levenshtein Distance: Information Dis-
Theory, Computer Science, String ( (Computer Science), String Metric, tance, Spell Checker, Hamming Distance. Alpha Press, 2009. ISBN 6130216904.
T. Nguyen, Q.-H. Pham, T. Le, T. Pham, N. Ho. and B.-S. Hua. Point-set distances for learning representations of 3DD point clouds. In ICCV, 2021.

---

Under review as a conference paper at ICLR 2023
Ofir Pele and Michael Werman. Fast and robust earth mover's distances. 2009 IEEE 12th Interna- tional Conference on Computer Vision, pp. 460-467, 2009.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Found. Trends Mach. Learn., 11:355-607, 2019.
Malak Pirtskhalava, Anthony Amstrong, Maia Grigolava, Mindia Chubinidze, Evgenia Alim- barashvili, Boris Vishnepolsky. Andrei E. Gabrielian, Alex Rosenthal, Darrell E. Hurt, and Michael Tartakovsky. DBAASP v3: database of antimicrobial/yyttoxie activity and structure of peptides as a resource for development of new therapeutics. Nucleic Acids Res., 49(Database- Issue):D288-D2977 2021. doi: 10. 1093/aargkka9911 URL https:/ /ttt:s://eei.or/1001000000 o0t.....rr/11.00.2....
nar gkaa991
Roshan Rao, Nicholas Bhattacharya, Neil Thomas, Yan Duan, Xi Chen, John F. Canny, Pieter Abbeel, and Yun S. Song. Evaluating protein transfer learning with TAPE. In Hanna M. Wal- lach, Hugo Larochelle, Alina Beygelzimer, Florence d' A'Aché-Bue, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 9686-9698, 2019. httphhttt////otceeeeeteteee ngs neur wpr.ppp..... CC cccccc/ppp.... /pper/ 7009 Aaa.ce-bbtrrrcccc
201 2t113/hash/37650068888
Ruslan Salakhutdinov and Geoffrey E. Hinton. Deep boltzmann machines. In AISTATS, 2009.
Karen S. Sarkisyan, Dmitry Bolotin, Margarita V. Meer, Dinara R. Usmanova, Alexander S. Mishin, George V. Sharonov, Dmitry N. Ivankov, Nina G. Bozhanova, Mikhail S. Baranov, Onuralp Soyle- mez, Natalya S. Bogatyreva, Peter K. Vlasov, Evgeny S. Egorov, Maria D. Logacheva, Alexey S. Kondrashov, Dmitry M. Chudakov. Ekaterina Putintseva, Ilgar Mamedov, Dan Tawfik, Konstantin A. Lukyanov, and Fyodor A. Kondrashov. Local fitness landscape of the green fluo- rescent protein. Nature. 533:397 401, 2016.
Sam Sinai, Richard Wang, Alexander Whatley, Stewart Slocum, Elina Locane, and Eric D. Kelsic. Adalead: A simple and robust adaptive greedy search algorithm for sequence design. CoRR, abs/2010.02141, 2020. URL URl/atts::/aazzivorrr owt.:/iivrrrreeoooo abs/
12010.00111
Kevin Swersky, Yulia Rubanova, David Dohan, and Kevin P. Murphy. Amortized bayesian opti- mization over discrete spaces. In UAI, 2020.
Tijmen Tieleman. Training restricted boltzmann machines using approximations to the likelihood gradient. In /CML '08, 2008
Brandon Trabucco, Aviral Kumar, Xinyang Geng, and Sergey Levine. Conservative objective mod- cls for effective offline model-based optimization. ICML, 2021.
Cédric Villani. Topics in optimal transportation. 2003.
Gui-Song Xia, Sira Ferradans, Gabriel Peyré, and Jean-François Aujol. Synthesizing and mixing stationary gaussian texture models. SIAM J. Imaging Sci., 7:476-508, 2014.
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron C. Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference Machine Leurning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 26412 26428. PMLR, 2022. URL https://procecinggt mlr. http://www......../sssessseeeeee wwtpppewweeeppppppppppppppppeeee 2hw.e//2hhaa22 html
12

---

Under review as a conference paper at ICLR 2023
Supplement to "Improving Generative Flow Networks with Path
Regularization"
A RELATED WORK
GFlowNets The objective of GFlowNets is related to MCMC methods for sampling from given unnormalized density function, especially in discrete spaces where exact sampling is intractable (Dai al. (2020); Grathwohl et al. (2021)). However, GFlowNets amortize the complexity of iterative sampling by a training procedure that implies the data's compositional structure as its learning prob- lem. Empirically, GFlowNets' performance is better than other earlier methods in wide variety of tasks: small molecules generation (Bengio et al. (2021a), discrete probabilistic modeling (Zhang
et al. (022 Bayesian structure learning (Deleu et al. (2022) ) and biological sequence design (Jain
et al. (222)). On the theoretical side, definitions and properties of GelowNets are more investigated
Bengio et al. (2021b).
The optimal transport theory (OT) (Villani (2003)) has established a natural and useful geometric tool for comparing measures supported on metrie probability spaces. The de- velopment of OT theory has a long history, where it has been discovered in many settings and under different forms. And in recent years, another revolution in the spread of OT has been witnessed, thanks to the emergence of approximate solvers that can scale to the problem of large dimensions. As a consequence, OT being widely used to solve various problems in computer graphies (Bon-) neel et al. ( Nguuen et al. (202T)), image processing (Xia et al. (2014)) and machine learning (Courty et al. (2014), Ho al. (2017) Genevay et al. (2018). Bunne al. (2019)).
Optimal Transport
Energy-based models EBMs, or energy functions parameterized by deep neural networks, have demonstrated effectiveness generative modeling (Salakhutdinov & Hinton (2009); Hinton et al. (2066). Contrastive divergence methods (Hinton (2002 ; Tieleman (2008); Du et al. (2021)) have been proposed to handle costly MCMC processes by approximating energy gradient. Recently, has been shown that simultaneous learning of the proposal distribution can also be helpful (Dai
eet al. (2019); Arbel et al. (202T)) Then this finding has been extended to discrete spaces by using GFlowNets /Zhang al. (2022).
Biological sequence sequeeee design Various methods have been proposed to handle the biological sequence design tasks: deep model-based optimization (Trabucco et al. (202I)), Bayesian optimization (Be- langer et al. (2019); Sweessy et al. (2020)), reinforcement learning (Angermueller et al. (2020)), adaptive evolutionary methods (Hansen (2nns) Sinai al. (2020), ) and so on. Recently, GelowNets also have been proposed as a useful generator of diverse candidates for this problem in Jainn et al.
(2022)
B BACKGROUND OF GFLOWNETS
Generative Flow Networks (GFlowNets) are a recently proposed class of generative model, which aims sample structural object x with probability proportional to a given reward function R(x). From the reinforcement learning viewpoint, GFlowNets learn stochastic policy to generate object E X by applying sequence of discrete actions E where the action space. The construc- tion of an object E X defines complete trajectory initial state, 8n = x E X' is the terminal state (indicating entirely constructed object), and sy the final state. Note that the same terminal state can be formed by different sequences of actions. These states and actions correspond to the vertices and edges of a directed acyclic graph G = (S,.).. In addition, for each transition - s' E A, we call parent of s', and s' child of s. T is defined as the set of all complete trajectories.
(Sp, $1, 8n 8y) where so is the
Following Bengio et al. 20211), trajectory flow is any nonnegative function defined on the set of complete trajectories, such as F T + R+. Correspondingly, the flow through a state (state flow) is defined as F(8) ΣFET F(T) and the flow through edge (edge flow) is defined as F(s s') ErET. F(r). Additionally, the forward transition probabilities Pp and the backward transition probabilities Pp are defined as follows:
PF(S'(s) = F(s + s')
(17)
F(s)
13

---

Under review as a conference paper at ICLR 2023
Pa(s)s') F(s 5)
(18)
F(s')
Then the training objective of the GFlowNet is to learn a consistent flow (Bengio et al. (2021b);
Malkin et al. (2022)) ) that has the terminal flow F (x - 8y) approximately equal given reward function K(z) for any I E X. In addition, when the flow consistent, the forward transition probabilities Pr and the backward transition probabilities Pa correspondingly define a distribution over the children and parent of each state, which can be considered as the forward and backward
policy of Gets
Specifically, followed by Malkin et al. 2022). the GflowNet models the forward policy, backward policy and total flow of Markovian flow F by 0)) PBl.,,; 0) and Ze. The trajectory balance objective is then optimized for each complete trajectory т sampled from the training policy ITa
CTB (T, 0) = log( Ze
II Pr(st| - (og(((()) II Pa((t-1(s8;6)
(19)
[=1
which is derived from the trajectory balance constraint (Malkin et al. (0222)
Moreover, as already proved by [Bengio et al. (2021b), Ty can be chosen as any distribution on the set of complete trajectories T with full supports, or the GflowNet can be trained with offline policy as well, such as a mixture between the GelowNet's forward policy and an uniform distribution over allowed actions in each state:
7o (1 Pp 9) + a Uniform
(20)
There also exist other objectives for learning a GFlowNet, which are based on flow matching con-
straint detail balance constraint in [Bengio al. ()) However, Malkin et al. (2022) em- pirically shows that the trajectory balance objective improves the training of a GHlowNet in terms of more efficient credit assignment and faster convergence, compared the previously proposed objectives. These advantages make us choose it as the training objective in this paper.
C DROPOUT OPTIMAL TRANSPORT
A limitation of the current method is computing the optimal transport distances for all couples of nearest neighhor states, especially in high dimensional discrete data. Our proposed dropout OT might be a solution. This is because rather than sampling trajectories 7 and using all edges from them, we can separately sample edges > proportional edge flows, allowing us to efficiently compute path regularization.
Theorem C.1 For any complete trajectory T (8o + 81 + - 8n) sampled from the training policy
E Er (For(T)) x Eg-se' (OT (P)(s))))
(21)
The proof of Theorem C.I is in Appendix D.3 Here we train GFlowNets with trajectory balance objective. Therefore, when sampling trajectory т, we get set of edges from 7. We just sample uniformly a percentage of edges to compute OT loss.
To sample p percentage of edges, let sample r, Ber(p).
Eggsst~r, (OT (Pr())s), Pr(+s)))) = Brr,Ber(L)) E, (F
(22)
We approximate the path regularization loss via:
1
Σ () Pp Pp()))44))
Cor(T) p
(23)
with x, drawn independently from Ber(p) for all <t n n-1. Intuitively, if x; = then we don't tt need to calculate the corresponding optimal transport cost anymore, which reduces considerable amount of computing time and memory down to percentage.
D PROOFS
D.1 PROOF OF THEOREM 3.1
For any trajectory 7 = (50 - 81
sn), ). we first prove that for any t E 0,
OT Pr(.)s < Σ
(()((())))).
(sChild(ss)
(24)
14

---

Under review as a conference paper at ICLR 2023
Consider two neigboor states and $141 with the children sets: (s)l = {u, ...y U4 and (5+41 = (U1,..., v1). By definition 5 the optimal transportation distance between two distri-
butions (r([s,) and Pr(.(s++1) is defined as:
OTc (PP()))), PF(()()))))
min PP(( (C,x),
(25)
where the admissible couplings set is defined as:
((PP((lss) { e R x1, = Pf(+(s)), x 11 PP(-11ct1)
(26)
We have,
OT(Pp(-1s)) (r((s4+1))
Σ Ti Cy
ΣΣ Tíj log (Pa((//ll)))Prosssssssssssop
ΣΣ Tej log (Pa(s(lus) ΣΣ Tss log (PD(S141(se)
ΣΣ Tay log (Py (1y((1144))
Σ log (PB(sr(u:)) Σ
Tej log (1'F(441(ss)) Σ Túj Σ log (PF(u)]s(+1)) Σ Taf
E
Σ log (PB(s) (1))))) - log (Pp Σ log (Pp Pe (vgll0044)
Σ
logulss))ogg(plll) u)) - log( (s:) + H(PF(1)441)
(EChild(s.)
(27)
The first inequality obtained by the definition optimal transport distance in Eq. 25 the second inequality comes from Eq. 8 the fifth equality is due to the constraints of admissible couplings in Eq. 26
As consequence, the upper bound loss is obtained by summing up all inequality 24 for all t.
D.2 PROOF OF THEOREM 3.2
Recall from definition 5 the optimal transportation distance between two distributions Pr(.ls) and Pr(.(s)) is defined as:
OTc (PP(()))) =
min (C, T).
(28)
(e
Let decompose the total cost (C, m)
(C, m)
Tzz Cij
Ttyl (o((ooo(PP(s))p)) - log(PP(s'(s)) - (og(PP(())ss)))
Σ
(oggo log(Pp(5((s))
(s)(( log(P + log (P(()))s)) + Cit)
ucds' my EChildii)
Th (loo((Pp((s)u)) + log( (og(PF(sss (og
(29)
We will prove that u; Fūj i.ej.i.e Child(s) Child(s)) O. indeed.
a; * ax + ag Va,, Cz Qn A - a; a * + = s + a, # s a + -a
4s #
Vi.j
(30)
15

---

Under review as a conference paper at ICLR 2023
We have:
# 8
by Child(u:),
0: #1 d2 + djs
a, + dg
s a; + au
0, * a
(31)
a; + d;
a + au a + +aj a a
+ ag:
= aj # aa au =
As result, we can rewrite Eq. as:
(C, Σ
log (oog-oo(P((()).)) - (og((()))))))
ij
+ Σ
Tiy
(32)
=
Σ
[) (o)( + C.).
H aa=aazoa
The first term of above equation actually is the upper bound of the optimal transport distance. There- fore, we can rewrite the total transportation cost as:
(C, m)
E
PF(u)s) log(Pn(s(u)) log(Pr(s')s) H(Pr(-16))
WEChild(s)
+ PP (log())))) s)) Log(PP(s
(33)
Σ
The (log (Tg(lo(PP(s(ua))) (o)
Ci).
aa _a==ay#
From the definition of c4 in Eq. 15
we have
log(Pn(s)u.)) + log (r(s'ss) log(Pr(1,ls)) + Cu: if u, + a, = a, + T
4
(34)
0
us = s' or a; = T.
From Eq. and Eq.
33
34
we have
Tij (log( (lo((Pa()))))) log log((P(s'(s)) log(Pr(u;)s)) Cti) Σ k...y. (35)
Thus, we find that
arg min TE (1s)
(C,
arg min rEI[(P=((s)PP((19))
(C', T)
(36)
where, C' is a diagonal matrix with the diagonal df < For convenience, if action a, invalid so the cost matrix of the optimal transport distance still is square matrix with the zero cost a invalid actions, then applying the Lemma 1, we have:
state s, we assign Pr(u; s)
min (e))
(C',T) = Σ min (u.))) PF(s))))
(37)
We obtain the dosed-form formulation for optimal transport distance OT (PP())))
E
((PP H(Pp)))))))))))))))))))))
(EChild(s)
+ Pr(s' (log (s) + log (r(('(s)))
(38)
min (u/s), Pp(usss))),
REA; n
Lemma Given a squared diagonal cost matrix C' with non-positive entities in the diagonal, the solution of optimal transport problem between two distribution Pr(-)s) and Pp(-|s'). which has the same number of support points, given cost matrix C' is given by:
min (e(l(P(((b))PP((e))
(C', = Σ min(Pp(u (
(39)
16

---

Under review as a conference paper at ICLR 2023
Proof of Lemma Let define
F(T) = (C', TT,
P'.)) (r) - mim(mi / p. (m-ymminmm
ifi=j
Pej
i.
Lx min (FE
where p² = (ss and p/z = Pr(usls')-
We will prove that के E II (Pr(-1s), (Pr(())))PP(((ss)) and F(T) 2 F(T) Vn (P()) It is not difficult to show that vzz 2 0. From the definition of म. we have
(p' min (p', p'y)) (P' min(p'. p2))
Tij
min(y's (40)
Σx min P5)
j
If mininp's p( p then
1
Ray 0 min(p'. = p's Σ
(41)
else min (p., pp)) = pp then
Σ p/. - min(p'. 14)) = Σ V min(p'. = Σ min min((",,p))
1+1
Tij (p-- min min(p'.p'))) Ejst pp.. min min(ps,
(42)
n
+ min(p'.Pp) p':)) ps.
Σx min(p. p3,)
Therefore Σn Kij = p,= = Pp(uuls). Similarly, Σ Tij = p. = Pa(usss"), combining with i 2 0, we have
T (1
(43)
Moreover F(T) (C,, Σ T C Σ min(p'..'.)
-ं = () F(T) I Pp(-ss))
(44)
As consequence, we obtained the solution of optimal transport problem.
closed-form solution for optimal transport distance at terminal state. We will derive the closed- form solution for optimal transport distance in case of two neighbor states s < s', in which s' is a terminal state. In the case of Hyper-grid environment, EB-GFN experiments, and Biological Sequence Design, all terminal state have only one child that is the final state sy, and PF = 1 Vz. Thus, the admissible couplings set [I P=(-Is')) )) has only one element. That is T Pr(+.s). As a result, the optimal transportation distance between Pp(.ls) and Pr(.ls") is: OT (PF(-Is), P=(-'s))
min =<][(Pr((e).PP((l)
(C, = (C,x*))
(45)
Specially, in EB-GFN experiments, all children u; of s is terminal state so di d(u,, 8y) - log(1) = This makes C = and OT (Pr(-ls), (((()) = O. In Hyper-grid environment experiment, for terminal sate s because c, = 0, we have:
OT (p((s))) = Σ
Pp(u)s) log(P
(EChil(s))
Pr(s' (log( s)) + log((P(s')s)) (s)))
(46)
The Hyper-grid environment (Bengio et al. 2021a) (in section 4.1 and EB- GFN experi-
ments (Zhang et al.] 2022) (in section 4.3) satisfy two condition in Theorem 3.2 Biological Sequence Design (Jain et al., 2022) (in section 4.2) such as protein and DNA sequences, the action space consists of actions adding a nucleic acid in [A, T, G, U} and a amino acid respectively. Such settings satisfy former condition a; + ag + an Va,, ags an. E A. However, the later condition a, + 4/ = am + an, G; + am > a an aj = 0m, a, # am is no longer true because the order property of action space, i.e, 0; + a, a, + a, In this situation, the third terms in Eq. is zero and we can still using the formulation in Eq. [4 unique factorization.
Generally, the action space is independence and
17

---

Under review as a conference paper at ICLR 2023
D.3
PROOF OF THEOREM C.I
By definition the edge flow we have
F(T)
F(s + 8))
Σ
P(T)=
Σ
P(s 8).
(47)
Z
Tis-s'Er
T18-teEr
From that equation, we find that
Erwno (Cor(T)) = -EyrTo
Σ
(T(P)((6))P((11))
a-to'ET
= Σ
Σ
(TPPA(())
STSEET
Σ
Σ (PPP((
(48)
T:8-
Σ (TP(( P=(-15)
E P(T)
Σ OT
P(s
5
X E E (
(F
E EXPERIMENT SETTINGS
In this part, we report experiment settings, including evaluation metrics for comparing the meth- ods, hyper-parameter choices, and neural network architectures for all experiments. For biological sequence design tasks, we also give more details about the task description and datasets used for training. Note that the regularization coefficients provided in this part are task-specific. Specifically, A is chosen from predefined set of values with different scales in each task. However, because all tasks in our experiment parts do not change the target distribution between the training and test time, the reported À is chosen to have the best model's performance.
E. HYPER-GRID ENVIRONMENT
E.1.1 EVALUATION CRITERIA
To evaluate the performance, we measure the KL divergence between the actual and empirical dis- tribution of the last 2 x 10 visited states. The number of modes found during the training progress also used to measure the learned models' performance.
E.1.2 IMPLEMENTATION DETAILS
GFlowNet: For the implementation of the GFlowNet model, we also follow the framework of
Malkin et al.] (2022): an MLP with two hidden layers of 256 dimensions each. The GFlowNet policy model, which includes both Pp and PB. is trained with a learning rate 0.001 while the learning rate for total flow Zo is 0. We use mini-batch size of 16 and 62500 training steps with the trajectory balance objective.
Proposed OT regularization The regularization coefficient is 0.02 for both Min OT, UB-OT, and Max OT - hypergrid environment and 0.1 for both Min OT, UB-OT, and Max OT in D hypergrid environment.
E.2 BIOLOGICAL SEQUENCE DESIGN
E.2.1 TASK DESCRIPTION & DATASETS
These experiments simulate the process of designing biological sequences, such as anti-microbial peptides, DNA, and protein sequences... in drug discovery applications. This process often consists of an active loop with several rounds ideating molecules and multiple-stage evaluations for fil- tering candidates, with rising levels of precision and cost. This characteristic makes the diversity of proposed candidates a considerable concern in the idcation phase because many similar candidates can all fail in the later phases.
18

---

Under review as a conference paper at ICLR 2023
Specifically, consider the problem of finding objects in the space of discrete objects X' that maximize a given oracle f r> Rt. Here, we can only query this oracle N times, each with an imput batch of fixed size b. This can form N rounds of evaluation in the active learning setting, where the generative policy is initially given a dataset Do = { { vi) ........ (.% collected from the oracle, where vf- = f(x)) for 1 S sism.. n.n. n.
Because the oracle can only be called limited, we also train supervised proxy model M that predicts from x to approximate the oracle f. Specifically, in i-th round, given the current dataset D. this proxy model can be used as a reward function R collect additional observations train our generative policy to propose batch of candidates B; = - {z)....,.fj .. Then the current dataset
D, is updated for the next round of evaluation D.+1 D, U (ti. (.... (z), vt))
where
v} = f (1))-
Following the framework of Jain et al. (2022). we will conduct experiments on the biological se- quence design tasks:
Anti-Microbial Peptide Design: This task aims to generate short amino-acid sequences of length lower than 51, which have anti-microbial properties. The vocabulary has 20 amino-acids [A, C, D, E, F, G. H, I, K, L, M, N, P. Q, R, S, T, W, YJ. The active learning algorithm is eval- uated for N 10 rounds, with the number of candidates generated each round 6 1000. The initial dataset Do contains 3219 AMPs and 4611 non-AMP sequences, which is collected from the DBAASP database Pirtskkalava al. 2021h.
TFBind 8:
The goal of this task is to generate DNA sequences of length 8, which have high binding activity with human transcription factors. The vocabulary has 4 nucleobases [A, C, T, G7. The active learning algorithm is evaluated for N = 10 round, with the number of candidates generated each round 128. The initial dataset Do contains 32, 898 samples, which half of all possible DNA sequences of length having lower scores. The data and the oracle used are from Barrera et al. (2016).
The objective of this task is to generate protein sequences of length 237 that have high fluorescence. The vocabulary is similar to the one the AMP task (size 20). The active learning algorithm is evaluated for N 10 round, with the number of candidates generated each round b= 28. The initial dataset Dn contains 5, 000 samples, which is from Rao et al. (2019); Sarkisyan al. 2016) together with the oracle.
GFP:
E.2.2 EVALUATION CRITERIA
To evaluate the performance, we also use the metrics as in Jain et al. (2022). Specifically, considering a set candidates D, we have the following metrics:
Performance score:
score of the candidates in the set E(m))ee V/
Mean(D)
(49)
[D
Diversity: measurement of how well the generated candidates can capture the modes of the
distribution implied by the oracle
Diversity(D)- E(HDDED E(s,,, ((a,,))))) d (2x, 23g)
(50)
[D -1
where d is distance defined over X, such as Levenshtein distance Miller et al. (2009).
Novelty: measure of the difference between the candidates in D and Do Novelty = E(x.BDeD mins,ED, d (Z);
(51)
D
These metrics will be evaluated on the set of candidates that have top K scores Topk (Dw\Do).
E.2.3 IMPLEMENTATION DETAILS
For the implementation of the GFlowNet-AL baseline model, we use the previously published im- plementation with slight changes, which follows the training setups of Jain al. (2022):
19

---

Under review as a conference paper at ICLR 2023
Proxy model: We parameterize it as an MLP with two hidden layers, each having 2048 hidden units, and use ReLU activation. We also use ensembles of 5 models with same architecture for uncertainty estimation. For the acquisition function, with MSE loss using mini-batch size 256 and Adam optimizer with (Bo. 31) = (0.9, 0.999) and learning rate 10-1 During training, early stopping is also used by evaluating the validation set containing 10% of the data.
use UCB (μ kσ) with
0.1. The proxy is trained
GFlowNet generator: We use MLP with hidden layers of 2048 hidden units each. The model trained with trajectory balance objective as the main loss function, by using Adam optimizer with (So. 31) (0.9. 0.999). Additionally, log Z is trained with learning rate of 10 for AMP, TF Bind 8 task, and 5 x 10 for GFP task. Other prper-paraeeters are shown in the following table:
AMP
TF Bind 8
GFP
8: Uniform Policy Coefficient
0.001
0.001
0.05
Learning rate
5 x 10-1
10-5
10-3
m. : Minibatch size
32
32
32
B : Reward Exponent R(x))
3
3
3
T Training steps
10,000
5,000
50,000
Table 5: Hyper-parameters for the GelowNet.
There are some changes in the TF Bind 8 task and the GFP task compared to the original training setups of Jain et al.) (2022). However, during the experiment, we observed that these settings helped us get the closest results to the reported one in Jain al. (2022).
choices and the number of active learning rounds in
Proposed OT regularization The regularization coefficients for Min OT, UB OT, and Max OT are the same for each biological sequence design task. Specifically, the coefficients for the AMP, TF Bind 8, and GFP task are 0.025, 0.1, and 0.02 correspondingly.
E.3 SYNTHETIC DISCRETE PROBABILISTIC MODELING TASKS
E.3.1 EVALUATION CRITERIA
To evaluate the performance, keep the same evaluation criteria in /hang al. (2022) where they use the NIL of a large independent sample of ground truth data and the exponential Hamming MMD (Gretton al. (2012)) between ground truth data and generated samples as performance metrics. To measure NNL and MMD, we use 10 fixed sets, and each set consists of 4000 ground truth data samples.
E.3.2 IMPLEMENTATION DETAILS
GFlowNet: For the implementation of the GFlowNet model, we use an MLP with hidden layers of 512 dimensions each. The GFlowNet policy model, which includes both Pp and PB. is trained with a learning rate of 0.001. We use mini-batch size of 128 and 1e5 training steps with the trajectory balance objective.
EBMs: For the implementation of the Buergy-Basee Model, we use an MLP with hidden layers of 256 dimensions each. The learning rate is 0.001.
Proposed OT regularization: The regularization coefficient is 0.001 for both Min OT and UB OT and is the same for all tasks.
F ADDITIONAL EXPERIMENT RESULTS
F.1 ABLATION STUDY ABOUT VARYING À
Specifically, we will further investigate the proposed path regularization via OT with different values of the regularization coefficient À in the 8-D hyper-grid environment in Section 4.1 In addition, the regularization coefficient is selected from the set (0.001, 0.01, 0.1, 0.4). We plot the mean results over 10 runs for each configuration in Fig.
20

---

Under review as a conference paper at ICLR 2023
30.0
18 Saseine
TB
27.5
TB UB-OT
200
TB + Min OT Max
TB OT
modes found
25.0
150
22.5
20.0
100
17.5
50
15.0
12.5
10a
106
10
10
states visited
states visited
250
TB Rnseine
30.0
Raseine
27.5
UB-OT
200
TD +
18 Min of
found
TK + Max OT
25.0
TR: Max OT
150
22.5
20.0
100
17.5
50
15.0
0
12.5
10P
10s
10a
10
states visited
states visited
250
1B Bese ine
30.0
18 Saseiie
15 + Ub-01
27.5
200
T0 + Max OT
25.0
TB Max OT
150
22.5
modes
20.0
100
17.5
15.0
12.5
0
10a
states visited
states visited
250
TB Base
TB Raseine
TI + UR-OT
27.5
TB- UB-OT
200
25.0
Max aT
150
22.5
modes
20.0
100
17.5
15.0
0
12.5
10a
10a
visited
visited
Figure 3: Results on the 8 D hyper- grid environment with e (0.001, 0.01, 0.1,0.4) (from top to bottom). Left: Number of modes found during training. Right: KL divergence between the true and empirical distribution.
Note that the good range of values for the regularization coefficient is observed to highly depend on the specific setting of the experiment task. Here, can see that when is relatively small, such as À E (0.001, 0.01). the performance of GFlowNets trained additionally with our proposed regu- larization via OT does not seem to be significantly different from the baseline model's performance, which holds for both UB OT, Min OT, and Max OT. This may be resulted from the small contribu- tion of regularization to the regularized training objective, which is caused by the not large enough value of A. Specifically, when À 0.01, by minimizing the upper bound is slightly better than the baseline's 's result.
still see that the performance of GFlowNets trained
In addition, when À is relatively large (A 0.4), the result of learned GFlowNets is even worse than the baseline, which may be due to the large value of À forces the model's 's learning focus on the regularization part more than necessary, which badly affects the optimization of the main training objective (trajectory balance objective). Specifically, this can observed in the lower KL diver- gence of both UB OT, Min OT, and Max OT compared to the baseline.
21

---

Under review as a conference paper at ICLR 2023
Meanwhile, when À 0.1, GFlowNets trained by minimizing the OT regularization and its upper bound clearly perform better than the baselines regarding the number of modes found and KL diver- gence between the actual and empirical distribution, which proves our motivation that minimizing the proposed path regularization is more beneficial in this circumstances.
F.2
ADDITIONAL RESULTS THE HYPERGRID ENVIRONMENT
We also plot the mean results over 10 runs for each configuration with variance in Fig.
17.5
TB Base Inc
TB Baseline
25
TB Min OT
TB Max OT
modes found
12.5
TB Max OT
20
TB Min OT
10.0
15
10
5.0
2.5
10a
10a
90
10*
states visited
states visited
30.0
250
Baseline
TB Baseline
Ta UB:OT
modes found
200
T8 Max OT
25.0
+ Min OT
Max
22.5
20.0
17.5
15.0
12.5
10a
10
10a
109
states visited
states visited
Figure 4: Results with variance on the D (upper) and 8 (lower) hyper-gre -eird environment. Left: Number of modes found during training. Right: KL divergence between the true and empirical distribution.
22


Reviews:
[{"review_id": "qhkHbxdipp7", "reviewer": "Reviewer_86A8", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This work proposes three different ways to regularize GFlowNet training with Wasserstein distance. Implementations of the regularization term and an upper bound of it are derived for tractable training. The path regularization is claimed to be capable of generating more diverse candidates via maximizing the Wasserstein distance or improve the generalization by minimizing the Wasserstein distance. The proposed methods are evaluated on a series of GFlowNet-related tasks, including hypergrid, biological sequence design and synthetic probabilistic modeling. ", "summary_of_the_review": "This work is with reasonable writing and experiments. However, I have some concerns about the motivation, and also feel more experiments could be conduct to amplify the story of this paper. I would definitely consider raising the score if some of the main weaknesses could be reasonably addressed :)", "strength_and_weaknesses": "## Strengths\nThe proposal of using optimal transport theory to regularize GFlowNet training is novel and interesting. Three different regularization terms are proposed accordingly.\n1. The authors propose the definition of directed distance given the \"back-and-forth trajectory\" idea in EB-GFN paper.\n2. Tractable calculation methods are derived with theoretical insights to achieve the proposed algorithms.\n3. A series of experiments are conduct to verify the advantage of the path regularization method.\n\n## Weaknesses\n1. In Sec. 1, it is written that the motivation of minimizing path regularization is to enhance flow on high flow trajectories. This is true, but what I do not understand is the connection with generalization. First, GFlowNet is a way to amortize the MCMC sampling problem, and since there is no change of the target distribution between the training and test time, I do not think \"generalization\" is a proper notion to discuss the performance of GFlowNet. Second, such motivation is contrary to the one of GFlowNet, who aims to sample *proportional* to a given reward function, rather than finding the high-score objects. Such difference is the core motivation for why people need GFlowNets instead of directly using reinforcement learning or other related methods.\n\n2. In Sec. 1 and Sec. 3.1, the authors also say maximizing the path regularization term gives more diverse effects. I agree with this point, but since minimizing and maximizing the path regularization term gives different regularization effects, I wonder which effect is needed by GFlowNet? It seems the authors only simply testify the two methods (Min OT & Max OT) without further analysis. Maybe different tasks would require different behaviors (and thus need different regularization effects)? For example, in the toy hypergrid experiment, if we make the reward function to be more sparse, would the Max OT method show more benefit? I think it would be great if the author could provide more analysis on this point.\n\n3. Ablation about the hyperparameter $\\lambda$, namely the coefficient of the regularization term, is missing. It would be great to have a plot / table of the performance of different methods with different values of $\\lambda$ (e.g. both positive and negative values) on hypergrid (since it is the easiest task, but it would be better to have ablation results on other tasks).\n\n4. The forward and backward policy are in symmetric positions, which says, either of them is able to determine the whole GFlowNet distribution. If the method proposed in this work is effective, then one thing worth trying is to put a similar technique to regularize the distance of backward policy on intermediate states.\n\n5. The result of \"Max OT\" is missing in the probability modeling tasks. Is there any particular reason?\n\n### Minors\n1. In Sec. 4.3, there is a gap between the MMD values in EB-GFN paper and this paper. I feel this is about the number of samples used to calculate the MMD. Theoretically, MMD value should always be positive. However, it could be negative if estimated with a small number of samples, for example in original ALOE paper they use a fixed 4000 samples. In EB-GFN paper, the MMD is calculated with 10 times average, each time with a different set of 4000 samples (not fixed samples), and the results are positive.  It is OK to compare with a consistent metric like in the current paper, but it would be great if the numbers of this part could be re-calculated for better reliability. \n\n2. One of the advantage of Wasserstein distance is that it does not require the distributions to have the same support, as mentioned in Sec. 1. In molecule tasks, this unequal support would really matter as the action space is more structured in graph generation. In this sense, the molecule task should be a better playground for the proposed method (otherwise, why not directly use KL or other simple divergences? They are much easier to compute than OT). \n\n3. There are a couple of typos and misuses of notations. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "g0-Pwt-0UJ", "reviewer": "Reviewer_ayzy", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes a way to regularize GFlowNet policies using an optimal transport cost that can encourage either similarity or dissimilarity of policies at states that are close in the state space. This cost is converted into several proxy regularization terms that can be added to the trajectory balance objective in on-policy training of GFlowNets. Experiments are performed on three domains from recent work.", "summary_of_the_review": "There is a potential for a good contribution here: optimal transport is an interesting way to answer the question mentioned in \"Originality\" above. However, I vote to reject because the weaknesses listed above are dominating: in particular, the lack of (1) experiments that give insight about hypothesized benefits of the prosed regularization and (2) consideration of simpler algorithms, like entropy regularization on $P_F$ or $P_B$, that may have a similar effect.\n\nPost-rebuttal update: After reading all reviews and responses, I increase the recommendation from 3 to 5; see comments.", "strength_and_weaknesses": "Comments on experiments:\n- (++) Breadth of experiments: there is hypergrid, sequence design, energy-based modeling -- which all use different codebases in the original implementations from the papers that introduced them. I appreciate the engineering work and the variety of applications considered (simple reward matching, active learning, energy-based modeling).\n- (--) Many of the results are unconvincing. \n  - For the hypergrid environment, the difference between algorithms is almost invisible, and it could be compensated by small changes in learning rate. I would suggest trying  larger values of H and D, where sparsity may become more important, and showing more clearly the dependence on the most important hyperparameters (regularization weight -- how was it chosen?).\n  - In the EB-GFN experiments, there is practically no difference between the proposed algorithm and baselines. Furthermore, the values highlighted in Table 4 are sometimes not actually the best one in their columns, which is misleading.\n- (+/-) It is interesting that the Max OT approach performs best in the active learning settings. What is the interpretation of this result? Are more entropic policies beneficial in the AL setting? This is a potential strong point that should be investigated further.\n- (-) In practice, what is the computational overhead of UB and Min/Max OT?\n  - There are simpler (and less computationally expensive) ways to regularize a GFlowNet to either encourage or discourage sparsity. The simplest is just to penalize the entropy of the backward policy. It would be interesting to see how these compare to the OT-motivated regularization proposed here.\n\nMajor comments on writing and theory:\n- (++) The motivation for regularizing the GFlowNet policy makes sense and, as far as I can tell, the math is right.\n- (-) Large parts of the background on GFlowNets, especially Appendix B, closely follow or even copy [Malkin et al., 2022], abbreviated [TB], in the flow of exposition and in the notation. For example, the first sentence of \"Learning GFlowNets\" (p.14) is identical to the first sentence of 2.2 in [TB], and all choices of notation are the same. Borrowing the structure and notation is fine, but it should be attributed, with a sentence such as \"this exposition closely follows that of ...\".\n  - In addition, in the section \"Learning GFlowNets\" on p.3, the proof of correctness of TB is misattributed (in fact, it appears in [TB], not [Bengio et al.]). \n- (-) In Definition 3.1, the distance between states is only defined for states that are ancestors of one another. (By the way, in that definition, one should not call a sequence of actions a \"transition\".) However, on p.4, it is generalized to distances between any pair of states. \n  - This generalization should be explicitly stated, and it should also be stated in what sense the resulting d(s,s') is a \"distance\". For example, it is not symmetric, and it is possible that the distance between two distinct states is 0 if the transition probability between them is 1. Thus all we can say about d is that it is a pseudoquasimetric (since it is nonnegative, d(s,s)=0, and it satisfies the triangle inequality).\n  - The cost matrix in equation (9) is only an approximation to the true distance, since the true distance may be the NLL of a trajectory with more than three edges and may not just go \"back\" and then \"forth\". This should also be stated.\n\nMinor comments:\n- This sentence on p.2 does not make sense to me: \"Compared with KL divergence, the biggest problem of KL divergence is that it was infinite for a variety of distributions with unequal support\". Does this mean that KL divergence cannot be used as the distance metric because KL(p||q) is infinite if (support of p)-(support of q) has nonzero measure under p?\n- The \"pseudo backward policy\" requires some more details. It is not actually a conditional distribution (may not sum to 1). So what exactly is meant by the cross-entropy in (17)?\n- In (19), $\\pi_\\theta$ was not defined. (Is it the untempered forward policy?)\n- The paper would benefit from some proofreading for small writing bugs, which nonetheless do not interfere with understanding of the paper. Here are a few:\n  - Inconsistent capitalization and grammatical number in the name of the algorithm throughout the text, such as  \"GFlownet\", \"GflowNet\", \"GFLowNet\" instead of \"GFlowNet\". GFlowNet is singular and GFlowNets are plural.\n  - Similar comment about \"trajectory balance objective\" (sometimes called \"trajectory balanced objective\" and inconsistently capitalized)\n  - \"natural interpreted\" -> \"naturally interpreted\" (p.2)\n  - \"shortage path\" -> \"shortest path\" (several places)\n  - \"can be define\" -> \"can be defined\" (p.4)\n  - Some equations are not in math mode (p.6)\n  - Incomplete citations:\n    - [E. Bengio et al., 2021] was published in NeurIPS 2021.\n    - [Dai et al., 2020] was published in NeurIPS 2020.\n    - [Deleu et al., 2022] was published in UAI 2022.\n    - [Grathwohl et al., 2021] was published in ICML 2021.\n    - [Malkin et al., 2022] will be published in NeurIPS 2022.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "HxkOBfsvFJ", "reviewer": "Reviewer_MjKh", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "GFlowNets learn stochastic policies to sample discrete objects proportionally to a non-negative reward. The paper proposes path regularization to improve exploration and generalization in GFlowNets. The authors begin by defining a notion of directed distance within the GFlowNet DAG, followed by the transport cost between forward policies of neighboring states. The authors then establish an upper-bound on this transport cost and discuss efficient implementation of the transport cost. The authors discuss two particular instantiations of the path regularization. Specifically, the authors claim that minimizing the transport cost results in faster adaptation to high reward regions, whereas maximizing the transport cost leads to improved coverage of the state space. This is accompanied by empirical results on various tasks studied in prior work on GFlowNets. ", "summary_of_the_review": "In summary, the paper explores an interesting direction of optimal transport to improve GFlowNet training. Some key details of the method are not well motivated and some observations seem to be incorrect. The empirical improvements are also marginal. I am thus leaning towards rejection but I encourage the authors to address the feedback during the discussion and rebuttal. ", "strength_and_weaknesses": "**Strengths**\n\n- The proposed regularization scheme is well formulated and described in sufficient detail. I think the optimal transport perspective fits training of GFlowNets quite well, and is an interesting direction to be explored. \n\n**Weaknesses**\n\n- While I believe the optimal control perspective is interesting, the particular instantiation proposed in the paper is not too clear and well motivated. In Section 3.1 the authors discuss the cases of maximizing and minimizing the path regularization. For motivating the minimization of the regularization term the authors argue that making the forward policy at neighboring states closer results in better generalization - but this point is not super clear to me (and also doesn't seem well supported by the experiments). In general choices in the method are clearly described but quite poorly motivated. \n- In Section 3.3 the authors make the observation \"any two neighbor states $s < s′$ do not have the same child state\". I do not believe this observations holds true in the DAGs studied in the paper. For example, in the EB-GFN setup, the permutations of actions in a subtrajectory result in the same state. \n- The improvement due to the regularization seems marginal at best, based on the empirical results reported in the paper. Another thing missing in the experiments is careful analysis of the effect of the regularization, i.e. how does varying $\\lambda$ affect the performance. It is also not clear to me why the Max-OT is missing in the EB-GFN experiments? (Minor - in Table 3 specifically the authors report results with a precision on 10^-5 but as the oracle used in the task is a neural network that precision seems insignificant)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FBOi2SkkLDs", "reviewer": "Reviewer_J937", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper introduces a novel path regularizer for Generative Flow Networks (GFlowNets, Bengio et al., 2021), based on concepts from Optimal Transport (OT). This regularizer, defined in terms of directed distances in the GFlowNet, is in general expensive to compute. Fortunately, the authors show that this regularizer can be upper-bounded by a term more convenient for optimization. Moreover, under some conditions on the structure of the GFlowNet that are met in practice for the applications considered in prior works, the authors show that this regularizer can be OT regularizer can even be computed in closed form. This regularizer was tested in conjunction with the GFlowNet on a synthetic environment, discrete probabilistic modeling, and biological sequence design, and shows advantages in terms of performance and diversity compared to the baseline model without the regularizer.\n\n---\n\n*Yoshua Bengio, Tristan Deleu, Edward J. Hu, Salem Lahlou, Mo Tiwari, and Emmanuel Bengio. Gflownet foundations.*", "summary_of_the_review": "This paper provides a practical solution to improve GFlowNets, and is completely novel. I am currently leaning towards acceptance, but there are some technical aspects raised here that may have significant effects on the results presented in this paper (notably Theorem 3.2). Therefore I can only recommend borderline acceptance for now.", "strength_and_weaknesses": "**Strengths**: This regularizer is a well-founded way to improve diversity of samples from a GFlowNet. The upper-bound provides an effective strategy to add path regularization to GFlowNets. The authors evaluated this regularizer thoroughly on multiple environments over multiple metrics of interest in previous works.\n\n**Weaknesses**:\n 1. Since the core of the paper is about the added benefit of the regularizer over standard GFlowNets, I would have appreciated a more thorough study of the effect of the regularization constant $\\lambda$ on the performance and diversity. What trade-offs should one expect from setting the regularization to some value? Currently, the regularization constants (provided in Appendix E) are task-specific, which is expected, but there is no mention how those regularization constants were chosen: did you choose them based on a validation set? What was the metric you optimized for?\n 2. The definition of the cost function $C_{ij} = \\min_{\\tau} (-\\log P(\\tau\\mid u_{i}))$ in Equation 9 seems incorrect (except for $u_{i} \\equiv v_{j}$): for example you consider that if $u_{i} \\not\\rightarrow v_{j}$, then $C_{ij} = -\\log (P_{B}(s\\mid u_{i})P_{F}(s'\\mid s)P_{F}(v_{j}\\mid s'))$. However, we can imagine having a trajectory $\\tau = u_{i} \\rightarrow s \\rightarrow s'' \\rightarrow s' \\rightarrow v_{j}$ that has a smaller $-\\log P(\\tau\\mid u_{i})$; intuitively, it might be more unlikely to move from $s$ to $s'$ via an intermediate step $s''$, rather than directly.\n 3. Related to 2., the proof of Theorem 3.2 might also be erroneous, for similar reasons. The proof in Appendix D.2 starts with $C_{ij} = -\\log (P_{B}(s\\mid u_{i})P_{F}(s'\\mid s)P_{F}(v_{j}\\mid s'))$, even though there is no assumption in Theorem 3.2 preventing us from having a situation as above with $\\tau = u_{i} \\rightarrow s \\rightarrow s'' \\rightarrow s' \\rightarrow v_{j}$. All we can say, and that's what you properly used in the proof of Theorem 3.1 for the upper bound, is that $C_{ij} \\leq -\\log (P_{B}(s\\mid u_{i})P_{F}(s'\\mid s)P_{F}(v_{j}\\mid s'))$.\n 4. There are some confusions regarding the notions of entropy and cross-entropy. The formula in Equation 19 only corresponds to the entropy of $P(\\tau)$ if $\\pi_{\\theta} = P$. Moreover since $P_{B}^{\\star}$ is not a probability distribution, the cross-entropy $H(P_{F}(\\cdot\\mid s_{t}), P_{B}^{\\star}(\\cdot\\mid s_{t}))$ is not properly defined. One can understand what this means from the Proof in Appendix D.1, but you should be precise in your use of these terms, especially since you are drawing conclusions about the behavior of the regularizer from maximizing entropies (Section 3.1) and minimizing cross-entropies (Section 3.2).\n 5. The behavior of Trajectory Balance with and without OT regularization seem to be very similar on the synthetic task (Figure 2). It would be informative to add shaded areas to show the variance for the 10 runs, and evaluate if the effect of the OT regularization is statistically significant on this environment.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject