Your GFlowNet Secretly Learns an Optimal Transport Plan
Ian Maksimov 1 Nikita Morozov 1 Denis Belomestny 1 2 Sergey Samsonov 1
arXiv:2606.06272v1 [cs.LG] 4 Jun 2026
Abstract
GFlowNet define a generation process that follows a sequence of stochastic transitions of discrete nature, inducing a graph structure G. The generation starts at a special initial state s0 , and then moves along the edges of the graph until reaching a special sink state sf . The last state x from which the process transitioned into sf is then considered the final sample. While GFlowNets were initially developed to operate in acyclic graph environments (Bengio et al., 2021), further works generalized them to non-acyclic graphs (Brunswic et al., 2024; Morozov et al., 2025). In this setting, the expected length of a trajectory sampled in the graph becomes a key quantity controlling sampling efficiency. To effectively minimize it in practice, Brunswic et al. (2024) and Morozov et al. (2025) proposed flow regularization techniques. Moreover, non-acyclic GFlowNets have been recently applied to tasks that lie outside the domain of sampling, e.g., imitation learning (Brunswic et al., 2025a), multi-agent learning (Brunswic et al., 2025b) and shortest path problems (Morozov et al., 2026), further demonstrating the generality of the framework.
Generative Flow Networks (GFlowNets) are a framework for sampling structured objects via stochastic trajectories in a directed graph. In this work, we establish a theoretical connection between non-acyclic GFlowNets and optimal transport (OT). We show that fixing the initial flow distribution in a minimum-flow GFlowNet reduces its objective to a Kantorovich OT problem with graph-induced shortest path costs. At the optimum, the learned GFlowNet policy therefore encodes an optimal transport plan from the source distribution to the target distribution: we show that sampling trajectories from the minimum-flow GFlowNet recovers the corresponding optimal coupling. Our formulation enables applying the GFlowNet learning framework to OT problems on large graphs via edge flows and neural parameterization. Experiments confirm agreement with exact OT solvers and demonstrate that GFlowNets can learn high-quality transport plans.
Optimal transport, on the other hand, provides a principled way to compare probability distributions by taking into account the geometry and structure of the underlying space (Villani, 2008; Peyré & Cuturi, 2019). The central object is a transport plan: a coupling that specifies how much mass should be moved from each source point to each target point while minimizing a prescribed transportation cost. This formulation is especially attractive in machine learning because many data objects - images, documents, point clouds, empirical samples, and learned representations can naturally be viewed as discrete distributions. Unlike purely pointwise divergences, optimal transport distances remain meaningful and can exploit semantic or geometric structure through the ground cost. This has led to applications in image retrieval via Earth Mover’s Distance (Rubner et al., 2000), natural language processing via Word Mover’s Distance (Kusner et al., 2015), domain adaptation (Courty et al., 2017), and generative modeling via Wasserstein objectives (Arjovsky et al., 2017), as well as representation learning and barycenterbased methods (Tolstikhin et al., 2018; Agueh & Carlier, 2011).
1. Introduction Generative Flow Networks (GFlowNets, Bengio et al., 2021) are a class of methods for sampling structured discrete objects from probability distributions given by unnormalized probability mass function. A GFlowNet constructs objects sequentially through stochastic transitions induced by a forward policy. GFlowNets have been successfully used in a range of applications, including molecule generation (Bengio et al., 2021; Shen et al., 2024; Koziarski et al., 2024; Cretu et al., 2025), biological sequence design (Jain et al., 2022; Kim et al., 2024), combinatorial optimization (Zhang et al., 2023a;b; Kim et al., 2025), and the fine-tuning of large language models and diffusion models (Hu et al., 2023; Venkatraman et al., 2024; Uehara et al., 2024; Zhang et al., 2025; Kwon et al., 2024; Lee et al., 2025; Zhu et al., 2025). The theoretical foundations of GFlowNets were established in (Bengio et al., 2023; Lahlou et al., 2023). 1 HSE University 2 Duisburg-Essen University. Correspondence to: Ian Maksimov <[email protected]>.
In this paper, we uncover a connection between the minimum-flow formulation of non-acyclic GFlowNets and discrete optimal transport. The main idea of this connec-
Preprint. June 5, 2026.
1
Your GFlowNet Secretly Learns an Optimal Transport Plan
of the combinatorial space where exact solutions become intractable. We thus expand the generality of the GFlowNet framework, showing that it extends into the computational optimal transport domain.
2. Background 2.1. Non-Acyclic GFlowNets Figure 1. Visualization of the solution to the GFlowNet LP problem and Kantorovich OT optimal plan on a hypergrid environment. OT induces optimal coupling that connects source and target states and how much mass is transfered between them. GFlowNet on the other hand builds a policy that samples target states starting from source states, thus inducing a coupling with respect to the shortest path metric. Both formulations achieve the same transport cost of 4.351.
Brunswic et al. (2024) and Morozov et al. (2025) extend the theoretical construction of GFlowNets (Bengio et al., 2023) to non-acyclic environments. We use theory and notation of Morozov et al. (2025) which is summarized below as a starting point. The sequential sampling process is described by a directed graph G = (S, E), where S is a finite state space and E ⊆ S × S is a finite set of edges (or transitions). We will write out(s) for the set of children of s, and in(s) for the set of parents of s.
tion is based on two facts: 1) if an additional constraint is added to probabilities of transitions from the initial state s0 , GFlowNet effectively induces a transport plan between two probability distributions; 2) minimizing the total flow is equivalent to minimizing the graph distance transport cost of the aforementioned transport plan. This equivalence provides a novel view on the GFlowNet learning problem, as well as allows for the plethora of existing GFlowNet learning algorithms to be applied to OT problems.
Let T be a set of all finite trajectories τ = (s0 → s1 → . . . → snτ → sf ) from s0 to sf , where we use nτ to denote the length of the trajectory τ . We use a standard convention that snτ +1 = sf . We will say that τ stops at s if its last transition is s → sf . The edges of the form (s → sf ) are called terminating transitions, and the states s that have an outgoing edge into sf are called terminal states. The set of terminal states is denoted by X , and the probability distribution of interest R(x)/Z is defined on it, Pwhere R(x) > 0 is called GFlowNet reward and Z = x∈X R(x) is an unknown normalizing constant.
The main contributions of the paper can be summarized as follows: 1. We show that the minimum-flow non-acyclic GFlowNet learning problem (Morozov et al., 2025) can be equivalently formulated as a linear programming problem. We then theoretically demonstrate that fixing the initial edge-flow distribution in it effectively turns the minimum-flow objective into a Kantorovich optimal transport problem where the cost is defined as a distance in the graph, which is also equivalent to the discrete graph-based formulation of the Beckmann problem (Beckmann, 1952; Essid & Solomon, 2018). The optimal GFlowNet forward policy in this formulation samples optimal paths between the initial and the terminal distributions, thus inducing an optimal coupling. As a particular case, this construction recovers the shortest path theory of Morozov et al. (2026).
GFlowNet is essentially a pair of bidirectional policies, where a forward policy PF (s′ | s) is a distribution over children of each state, and a backward policy PB (s | s′ ) is a distribution over parents of each state. GFlowNet is trained in a way to find such pair of policies that the distributions over trajectories they induce coincide: P(τ ) =
Qnτ
t=0 PF (st+1 | st ) =
Qnτ
t=0 PB (st | st+1 ) .
(1)
The main goal for sampling is to have such P that for any x ∈ X, probability that τ ∼ P terminates in x coincides with R(x)/Z. This property is called the reward matching condition, and w.r.t PB it can be simply written as
2. We show that the GFlowNet learning framework provides an effective way to approximate solutions to graph optimal transport problems. Based on a neural parameterization, GFlowNet objectives allow to learn a policy that aims to recover an optimal coupling. We experimentally demonstrate the ability of the framework to recover exact OT solutions in cases when they are tractable, as well as its scalability, exhibited by the ability to efficiently approximate solutions with the growth
PB (x | sf ) = R(x)/Z ∀x ∈ X.
(2)
If both (1) and (2) are satisfied, PF can be used to sample terminal states from the reward distribution (see Bengio et al. (2023) and Morozov et al. (2025) for more details). Morozov et al. (2025) defines state and edge flows through 2
Your GFlowNet Secretly Learns an Optimal Transport Plan
the expected number of visits of a certain state "n +1 # τ X F(s) = Z · Eτ ∼P I{st = s} ,
balance (Bengio et al., 2023) or variants of trajectory balance loss (Madan et al., 2023; Malkin et al., 2022; Morozov et al., 2026), where λ is a regularization coefficient.
t=0 ′
F(s → s ) = Z · Eτ ∼P
"n τ X
# (3)
2.2. Kantorovich OT via Linear Programming
′
I{st = s, st+1 = s }
While GFlowNet framework aims to sample from some target distribution R(x)/Z, optimal transport problem on the other hand seeks an optimal coupling of two given source and target distributions, that moves probability mass from the source distribution to the target with respect to the defined cost function (Villani, 2008). Definition 2.1. Let µ and ν be two probability measures on a finite discrete spaces X and Y. A coupling of µ and ν is a joint probability distribution Π on space X × Y such that:
t=0
It is clear from the definition that the flow functions satisfy the following conditions (Proposition 3.6 of Morozov et al. (2025)): X X F(s → s′ ) = F(s′′ → s) (4) s→s′
s′′ →s
and F(s0 ) = F (sf ) = Z.
n X
The main problem arising when extending standard GFlowNet framework to non-acyclic is that the nτ is unbounded, since the policy can go in cycles. E[nτ ] exactly represents the mean trajectory length needed to obtain a sample from target distribution following actions sampled from forward policy, and it also signifies the practical efficiency of the sampler. Thus, it is natural to formulate the problem of training a non-acyclic GFlowNet that minimizes E[nτ ]. Using the above definition, one check that the expected trajectory length P coincides with the normalized total flow E[nτ ] = Z1 s∈S\{s0 ,sf } F(s), thus Morozov et al. (2025) gives a theoretical ground for learning a model with the smallest E[nτ ] which is equivalent to learning a model with the smallest total flow.
i=1
s.t.
Πi,j = µi
j=1
where d(xi , yj ) is some cost of moving mass form xi to yj . Usually this linear programming problem is solved via classical methods e.g. simplex-method (Dantzig, 1951).
3. GFlowNets and Optimal Transport Consider a non-acyclic GFlowNet environment graph G = (S, E). We impose the following assumptions on the structure of the graph G, similarly to Essid & Solomon (2018): Assumption 3.1. • There is a special initial state s0 with no incoming edges and a special sink state sf with no outgoing edges; • There are sets of states U and X s.t. s0 has outgoing edges only to U and sf has incoming edges only from X; • There exists a finite-length path between any u ∈ U and x ∈ X; • We are given two distributions L(u) and R(x) on U and X, respectively, that is, X X L(u) = R(x) = 1 .
s∈I ′
F(s)PF (s | s) = F (s′ )PB (s | s′ ) , ∀(s, s′ ) ∈ E, F(sf )PB (x | sf ) = R(x),
m X
Suppose we have two finite set of states {x1 , . . . xn } and {y1 . . . ym } with probability distributions µ, ν. Then Kantorovich OT problem is formulated as follows: Pn,m min i,j=1 d(xi , yj )Πi,j Π≥0 (6) P P s.t. i Πi,j = νj , j Πi,j = µi .
Morozov et al. (2025; 2026) formulates the non-acyclic GFlowNet problem as the following constraint optimization problem. Define the set of internal states I := S \ {s0 , sf }., then: X min F(s) F ,PF ,PB
Πi,j = νj ,
x ∈ X.
(5) The constraint F(s)PF (s′ | s) = F (s′ )PB (s | s′ ) is referred to as the detailed balance condition. As shown in (Bengio et al., 2023; Malkin et al., 2022; Morozov et al., 2025), satisfying detailed balance for every adjacent pair of states is equivalent to the trajectory balance condition (1). This optimization problem naturally encodes all necessary constraints needed for training GFlowNet. Equality constraints satisfy Equations 1 & 2, while the objective minimizes total flow, which is equivalent to minimizing E[nτ ] as per Morozov et al. (2025).
u∈U
x∈X
We adopt the convention that both L and R are extended by zero outside their supports and thus defined on all internal states I. Next, define a cost function d(u, x) for pairs of u ∈ U and x ∈ X as d(u, x) = |τu,x | , (7)
The simplest approach to this problem proposed in Brunswic et al. (2024) and further explored in Morozov et al. (2025) is to add λFθ (s) to the utilized GFlowNet loss e.g detailed 3
Your GFlowNet Secretly Learns an Optimal Transport Plan
where τu,x is a shortest path from u to x and |τu,x | denotes its length. Thus, one can consider a Kantorovich OT problem (6) between distributions L and R with the cost function d(u, x):
programming problem writes as X min ′ F(s) F (s),F (s→s )
X
s.t.
s∈I
F(s → v) = F (s),
s ∈ S \ {sf },
F(u → s) = F (s),
s ∈ S \ {s0 },
v∈out(s)
min
Pn,m
s.t.
P
Π≥0
X
i,j=1 d(ui , xj )Πi,j i Πi,j = R(xj ) ,
(8)
(10)
u∈in(s)
P
j Πi,j = L(ui ) .
F(x → sf ) = R(x), F(s0 → u) = L(u), ′
F(s → s ) ≥ 0,
Where Π ∈ Rn×m is a matrix which defines some coupling ≥0 between the distributions L(u) on U and R(x) on X. Here |U | = n, and |X| = m. It is also worth noting that this graph optimal transport with shortest-path ground costs has a classical min-cost-flow formulation (Essid & Solomon, 2018). We next show that the minimum-flow objective arising in non-acyclic GFlowNets, once the initial edgeflow distribution is fixed, reduces exactly to this graph OT problem.
x ∈ X, u ∈ U, ′
(s, s ) ∈ E.
From GFlowNets theoretical construction, it is unclear whether such problems with additional equality constraints has a feasible solution. However, we further show that a solution always exists (Theorem 3.2) Extended primal problem can be reduced by omitting the flow variable F(s). Define E ◦ := {s → s′ ∈ E : s ̸= s0 , s′ ̸= sf },
3.1. Minimum-flow LP formulation
The reduced primal formulation is: X min′ F(s → s′ )
We start with the GFlowNet minimum-flow problem (5) that was introduced in Morozov et al. (2025). The detailed balance constraints in (5) are bilinear in the decision variables, since
F (s→s )≥0
X
s.t.
s→s′ ∈E ◦
F(s → v)
v: s→v∈E ◦
− F(s → s′ ) = F(s)PF (s′ | s) = F (s′ )PB (s | s′ ) .
X
F(u → s) = L(s) − R(s),
s ∈ I.
u: u→s∈E ◦
(9)
(11) Note that when reducing the initial problem, we have the relation X X X F(s0 → u) + F(s → s′ ) F(s) =
The problem becomes linear after policies PF and PB are eliminated and replaced by independent state and edge-flow variables. We will now treat every state-flow F(s), s ∈ S, and edge-flow F(s → s′ ), s → s′ ∈ E, as an independent variable.
s∈I
s→s′ ∈E ◦
s0 →u
X
=1+
F(s → s′ ) .
s→s′ ∈E ◦
In this case, the detailed balance condition (9) from the optimization problem (5) can be substituted with the flow matching constraint (4), allowing for an equivalent reformulation of (5) as a linear program. We provide the proof of this equivalence in Appendix A.1.
We will omit this constant, as it does not change minimum of the problem. Note that if one sums up flow-matching constraint for all internal states, one would have: X X R(x) = L(u).
We fix the distribution L(u) defined on U which we will call Leward (left reward). Then we fix the first-step distribution induced by the underlying GFlowNet similarly to the reward matching condition (2):
x∈X
u∈U
This equality imposes a mass-balance constraint: the two functions must either have equal total mass or be normalized as probability distributions. We adopt the latter formulation, as it eliminates the need to handle an unknown normalizing constant. This is advantageous both in the theoretical analysis and in practical implementations, where the constant would otherwise need to be learned or estimated. Accordingly, we include this condition in Assumption 3.1.
F(s0 → u) = L(u) .
After adding this constraint, the extended GFlowNet linear 4
Your GFlowNet Secretly Learns an Optimal Transport Plan
3.2. GFlowNet-OT equivalence
is feasible, thus the solution set is non-empty. Now let us compute the value of (11): X X F(s → s′ ) = Π⋆u,x I[s → s′ ∈ τu,x ]
We are now ready to establish the equivalence between the Kantorovich OT problem (8) and the GFlowNet problem (11). Denote by GFlow⋆ the optimal value of the functional (11) and by OT⋆ the optimal value of the functional (8) in the Kantorovich problem. Then the following equivalence holds:
X
=
X
=
X
Π⋆u,x
I[s → s′ ∈ τu,x ]
s→s′ ∈E ◦
u∈U,x∈X
Theorem 3.2. Under Assumption 3.1, the reduced primal problem (11) is equivalent to the OT Kantorovich problem (8) between L on U and R on X, that is, GFlow⋆ = OT⋆ .
s→s′ ∈E ◦ u∈U, x∈X
s→s′ ∈E ◦
Π⋆u,x |τu,x | ,
u∈U,x∈X
(14) where |τu,x | is the length of the path from u to x. Since τu,x is a shortest path, |τu,x | = d(u, x), and X X F(s → s′ ) = d(u, x)Π⋆u,x . (15)
(12)
Moreover, if P ⋆ is the trajectory distribution induced by the solution of GFlow⋆ , it induces an optimal coupling Π⋆u,x := P ⋆ τ :u⇝x P (τ ) , u ∈ U , x ∈ X in (8).
s→s′ ∈E ◦
u,x
The flow we constructed is feasible, but might not be optimal. Hence, we obtain the inequality
Proof. 1. Suppose that we are given an optimal coupling Π⋆u,x which is the solution of the minimization problem (8). Existence of Π⋆u,x is guaranteed by Theorem 4.1 from (Villani, 2008, p. 43). We construct a feasible flow in a sense of (11) as follows. For every u ∈ U and x ∈ X, we choose an arbitrary shortest path τu,x from u to x. Then we set XX F(s → s′ ) = Π⋆u,x · I[s → s′ ∈ τu,x ]
GFlow⋆ ≤ OT⋆ .
(16)
2. Note that the optimal solution exists under Assumptions 3.1 and by the fact that our problem is bounded below. Take this optimal solution to the problem (11), and denote it by F ⋆ (s → s′ ). This flow follows equivalence establised in (A.1), this means that we now may introduce forward policy: X PF⋆ (s′ |s) = F ⋆ (s → s′ )/ F ⋆ (s → s′′ )
u∈U x∈X
Let us check the feasibility of this flow. Define for s ∈ I, s→s′′ u ∈ U , and x ∈ X, the auxiliary function and the respective probability over trajectories τ = (v0 = X X ∆u,x = I[s → v ∈ τ ]− I[v → s ∈ τ ] . s u,x u,x 0 , v1 . . . vnτ −1 , vnτ = sf ): s v:s→v∈E
v:v→s∈E
P ⋆ (τ ) =
Since τu,x is a directed path, it has only outgoing edge at its start, only incoming edge at its end, and equal incoming/outgoing counts at every intermediate node. So we can equivalently write
nτ Y i=0
Then the definition (3) from Morozov et al. (2025) implies that F ⋆ (s → s′ ) = Z
∆u,x = I[s = u] − I[s = x] s
X
P ⋆ (τ )
τ ∈T
I[s = u]L(u) −
u∈U
I[st = s, st+1 = s′ ] .
t=0
τ :u⇝x
u∈U, x∈X
=
nτ X
Assumption 3.1 implies that Z = 1, so we omit it in further derivations. Now define the joint distribution Π on U × X by X Πu,x := P ⋆ (τ ) , u ∈ U , x ∈ X .
Let ∆F(s) denote the left-hand side of the equality constraint in (11). Using ∆u,x = I[s = u] − I[s = x] and the s marginal constraints of Π⋆ , we have X ∆F(s) = Π⋆u,x I[s = u] − I[s = x] X
PF⋆ (vi+1 |vi ) .
X
It is clear that X X X Πu,x = P ⋆ (τ )
I[s = x]R(x) (13)
x∈X
= L(s) − R(s).
x→sf τ :u⇝x
x∈X
=
The above equations implies the equality constraint from (11), therefore we have checked that the flow we defined
X τ ∋s0 →u
5
P ⋆ (τ ) = PF⋆ (u|s0 ) = L(u) ,
Your GFlowNet Secretly Learns an Optimal Transport Plan
Theorem 3.3. Consider the dual to the primal problem (10) without first-step condition. X max R(x)πx
and, similarly, X
Πu,x = R(x) .
u∈U
Thus, the joint distribution Π defined above is a coupling of R(x) and L(u). Let us write down the value of primal problem using our flow: X GFlow⋆ = F ⋆ (s → s′ )
π
s.t.
X
X
s→s′ ∈E ◦
τ
=
=
X
P ⋆ (τ )
=
I[st = s, st+1 = s′ ]
t=0 nτ X
X
I[st = s, st+1 = s′ ]
Moreover, complementary slackness conditions in the optimum gives us the following relation:
|τ |P ⋆ (τ ) .
F ⋆ (s → s′ )(πs⋆′ − (1 + πs⋆ )) = 0
τ
(17) Recall that |τu,x | ≥ d(u, x), since d(u, x) is the length of the shortest path. Thus: X X X X |τ |P ⋆ (τ ) |τ |P ⋆ (τ ) =
which is also a primal-dual gap. The complementary slackness conditions provide a characterization of the optimal flow: it can be non-zero only on the tight subgraph, where (πs⋆′ − (1 + πs⋆ )) = 0. Since πx⋆ coincides with the shortest-path metric, it follows that the flow can be supported only on a subgraph of shortest paths. This recovers the corresponding claim of Morozov et al. (2026). Detailed proof of Theorem 3.3 can be found in Appendix A.5
s0 →u x→sf τ :u⇝x
τ
≥
X X X
d(u, x)P ⋆ (τ )
s0 →u x→sf τ :u⇝x
=
X X
d(u, x)
s0 →u x→sf
=
XX
X
P ⋆ (τ )
(18)
τ :u⇝x
d(u, x)Πu,x .
3.3. Learning-based algorithm
u∈U x∈X
To approximate the solution to (11) with neural networks, we utilize non-acylic GFlowNet training methodology. While one can apply various existing GFlowNet losses with flow regularization to train the model (Brunswic et al., 2024; Morozov et al., 2024), Morozov et al. (2026) recently showed the effectiveness of trajectory balance (TB) objective (Malkin et al., 2022) for training non-acyclic GFlowNets in pathfinding problems, which we also found the most suitable in our setup.
Hence, we constructed a coupling Π such that XX GFlow⋆ ≥ d(u, x)Πu,x . u∈U x∈X
For the optimal coupling Π⋆ , therefore, it holds that GFlow⋆ ≥ OT⋆ .
∀(s → s′ ) with s′ ̸= sf .
Then: 1. every feasible solution π satisfies: πs ≤ d(s) ∀s ̸= sf ; 2. the function d(·) = |τs0 ,· | is a feasible dual solution, and it is dual-optimal; if R(x) > 0 for every x ∈ X, then every optimal solution π ⋆ satisfies: πx⋆ = d(x) ∀x ∈ X.
s→s′ ∈E ◦ t=0
τ
X
P ⋆ (τ )
πs0 = 0, πs′ − πs ≤ 1,
s→s′ ∈E ◦ nτ X
x∈X
(19)
Combining (16) and (19), we obtain that the GFlow optima is the same as the OT optima, therefore (12) is proved.
For each state in the graph, the model predicts transition probabilities for forward and backward policies PF and PB . During training, we sample a batch of trajectories using the forward policy PF , and compute the regularized TB objective: Qnτ 2 PF (st+1 |st , θ) L(s1 ) t=1 LTB (θ, τ ) = log Qnτ −1 R(snτ ) t=0 PB (st |st+1 , θ) R(snτ ) +λ . PF (sf |snτ , θ) (20) The difference from the usual TB objective is the appearance of leward L due to the additional conditions on the initial distribution in (11). In terms of the policies, reward
This theorem give us characterization of a GFlowNet as an implicit way to solve Kantorovich optimal transport, as a corollary we obtain dual equivalence, this essentially follows from the strict duality of linear programming optimization problems. Detailed derivation of the dual to GFlowNet problem in terms of OT functional can be found in Appendix A.4. Furthermore, under appropriate structural assumptions, the optimal transport problem admits a reduction to a shortestpath problem. We show that, in this regime, non-acyclic GFlowNets implicitly recover such shortest-path solutions. The result is stated formally below. 6
Your GFlowNet Secretly Learns an Optimal Transport Plan d is the empirical total variation distance between samples from Table 1. Comparison of the method on the Hypergrid environment. TV the model and the true distribution, TV⋆ is the total variation obtained by a perfect sampler, E|τ | is the average trajectory length of the learned sampler, and OT⋆ is the optimal OT cost obtained via the POT solver (Flamary et al., 2021). L = Ball H
d↓ TV
⋆
TV ↓
L = Moon E|τ |
OT
⋆
d↓ TV
TV⋆ ↓
E|τ |
OT⋆
10 0.024 ±0.0004 0.024 3.990 ±0.015 3.997 0.022 ±0.018 0.024 4.352 ±0.015 4.351 15 0.036 ±0.0009 0.033 6.325 ±0.011 6.303 0.023 ±0.006 0.033 6.907 ±0.010 6.868 20 0.037 ±0.0006 0.040 8.326 ±0.021 8.325 0.032 ±0.008 0.040 9.001 ±0.112 9.059
Table 2. Results obtained on permutations environment with ablation of λ. Reference OT⋆ cost obtained via the POT solver (Flamary et al., 2021) for n = 4 is 0.567 and is n = 8 is 1.008, for larger permutations OT⋆ is intractable. C(k)L1 is the empirical L1 error between fixed points probabilities obtained from model and the true ones (Morozov et al., 2025) E|τ | is the expected trajectory length. All results are averaged over 3 seeds. n=4 C(k) L1 ↓
λ
n=8 E|τ |
C(k) L1 ↓
n = 20 E|τ |
C(k) L1 ↓
E|τ |
−1
λ = 10 0.012 ±0.001 0.445 ±0.002 0.011 ±0.005 0.645 ±0.013 0.016 ±0.000 2.313 ±0.018 λ = 10−2 0.002 ±0.000 0.557 ±0.001 0.001 ±0.002 1.010 ±0.011 0.002 ±0.000 4.436 ±0.014
matching can be written as PB (s | sf ) = R(s), and leward matching as PF (s | s0 ) = L(s), which explains the form of the loss in (20). The regularizer λR(s)/PF (sf |s) corresponds to the state flow F(s) in terminal states due to the reward matching conditions in (5). Morozov et al. (2026) also showed that if each internal state in G is terminal, i.e. X = I, the objective (20) can be computed for every prefix of the sampled trajectory, resulting in a more sampleefficient approach. We utilize this option in our experiments.
Since these environments remain small enough for exact LP optimization, they allow direct comparison between the learned sampler and the optimal OT solution. For additional details see Appendix B.1. States are lattice points in {0, . . . , H −1}D augmented with source and terminal states s0 , sf . Sampling starts from U with distribution L and transitions modify a single coordinate by ±1 while remaining inside the grid. Every state also admits a terminating transition to sf . We assume that leward L admits tractable sampling, and focus on measuring two quantities: 1) what is the average transport cost of the trained model and 2) how close is the distribution of terminal states sampled from the model to the reward distribution R.
4. Experiments In this section, we experimentally evaluate the proposed GFlowNet-based method for solving graph OT tasks. We consider the tasks based on the environments studied in Brunswic et al. (2024); Morozov et al. (2025). On a smaller hypergrid task, we empirically demonstrate that the exact solution to the GFlowNet LP formulation (11) is equivalent to Kantorovich OT formulation (8), and can be obtained both by the exact LP solver and learning-based neural network approximation described in Section 3.3. We then validate our learning-based method on larger-scale permutation environments. In particular, we recover the optimal OT solution for permutations of length n = 4 and n = 8, and demonstrate convergence for n = 20. Finally, we ablate the effect of the state-flow regularization parameter λ, illustrating its impact on the optima reached by the method.
We consider multimodal rewards with modes near the grid corners (Bengio et al., 2021) and use moon-shaped and ball-shaped leward distributions. Sampling quality is measured via the total variation distance between the reward distribution and the empirical distribution of 2 · 105 model d compares it with samples. Table 1 reports the resulting TV, the perfect-sampler reference TV⋆ (it is non-zero because and empirical distribution is considered), and evaluates convergence to the OT optimum OT⋆ using the mean trajectory length. One can clearly see that our learned policy recovers optimal OT⋆ without biasing the sampling across all selected leward distributions and all hypergrid sides. All results are averaged over three seeds. Figure 1 additionally visualizes the OT solution for H = 10 and L = Moon, as well as the exact solution of the equivalent GFlowNet LP formulation (11) obtained via scipy.linprog solver (Virtanen et al., 2020).
4.1. Hypergrids We evaluate our method on OT tasks on non-acyclic hypergrid environments proposed in Brunswic et al. (2024). 7
Your GFlowNet Secretly Learns an Optimal Transport Plan
ularization coefficient λ that must be carefully tuned to balance OT optimality and sampling quality. In addition, existing GFlowNet methodology could be potentially applied to learn transport maps between unnormalized probability distributions.
4.2. Permutations Next, we consider the environment induced by the Cayley graph of the symmetric group Sn , the group of permutations of n elements 1, 2, . . . , n, proposed in Morozov et al. (2025). Each internal state s ∈ S \ s0 , sf corresponds to a permutation of fixed length, (s(1), . . . , s(n)), Transitions are defined by adjacent transpositions: from a given state, one may swap two neighboring entries s(k) and s(k + 1). We take L to be the uniform distribution over all permutations, distribution as Pn and use the reward R(s) = exp 12 k=1 I{s(k) = k} /Z (Morozov et al. (2025) derived closed form expression for its normalizing constant). We adopt the diagnostic protocol of Morozov et al. (2025); a detailed description of the corresponding metric is provided in Appendix B.2. Table 2 reports results on permutation environments, measuring both the reward sampling quality and transport cost similarly to Section 4.1. All results are averaged over three random seeds.
Acknowledgements This research was supported in part through computational resources of HPC facilities at HSE University (Kostenetskiy et al., 2021). We would like to thank FORTUNA 812 for his inspirational music that kept us sane during this research.
References Agueh, M. and Carlier, G. Barycenters in the wasserstein space. SIAM Journal on Mathematical Analysis, 43(2): 904–924, 2011. Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein generative adversarial networks. In Proceedings of the 34th International Conference on Machine Learning, pp. 214–223, 2017.
For small permutation sizes, n = 4 and n = 8, our method recovers solutions very close to optimal, achieving an average trajectory length almost equal to OT⋆ . For larger permutations, such as n = 20, computing the exact optimal OT solution is no longer tractable; nevertheless, our method produces a reasonable approximation. We further ablate the effect of the state-flow regularization coefficient λ. Consistently with Morozov et al. (2025), we observe that λ controls a trade-off between the trajectory length of the learned sampler and sampling optimality. The larger λ brings bias to the sampling but reduces the trajectory length, on the contrary, the lower value induces accurate sampling but larger trajectories.
Beckmann, M. A continuous model of transportation. Econometrica: Journal of the Econometric Society, pp. 643–660, 1952. Bengio, E., Jain, M., Korablyov, M., Precup, D., and Bengio, Y. Flow network based generative models for noniterative diverse candidate generation. Advances in Neural Information Processing Systems, 34:27381–27394, 2021. Bengio, Y., Lahlou, S., Deleu, T., Hu, E. J., Tiwari, M., and Bengio, E. Gflownet foundations. Journal of Machine Learning Research, 24(210):1–55, 2023.
5. Conclusion We extended the theoretical framework of non-acyclic GFlowNets by establishing an equivalence with optimal transport on graphs. Unlike standard OT approaches that only learn a coupling, our framework learns a stochastic policy that transports samples through feasible local moves with near-optimal expected path cost, thereby addressing both where and how transport occurs. We further provided a novel interpretation of the shortest-path formulation of GFlowNets from Morozov et al. (2026). The equivalence established in Theorem 3.2 opens several promising directions, including backward policy optimization algorithms (Jang et al., 2024; Gritsaev et al., 2025) and reinforcement learning techniques motivated by the connection of GFlowNets to entropy-regularized RL (Tiapkin et al., 2024; Deleu et al., 2024; Mohammadpour et al., 2024; Lau et al., 2024; Morozov et al., 2024; He et al., 2025). Future work could also address practical challenges highlighted by our experimental evaluation. While presenting a computationally efficient training method, non-acyclic GFlowNets include the reg-
Brunswic, L., Li, Y., Xu, Y., Feng, Y., Jui, S., and Ma, L. A theory of non-acyclic generative flow networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 11124–11131, 2024. Brunswic, L. M., Clémente, M., Yang, R. H., Sigal, A., Rasouli, A., and Li, Y. Ergodic generative flows. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pp. 5649–5668. PMLR, 2025a. Brunswic, L. M., Wang, H., Luo, S., Hao, J., Rasouli, A., and Li, Y. A theory of multi-agent generative flow networks. arXiv preprint arXiv:2509.20408, 2025b. Courty, N., Flamary, R., Tuia, D., and Rakotomamonjy, A. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 39 (9):1853–1865, 2017. 8
Your GFlowNet Secretly Learns an Optimal Transport Plan
Cretu, M., Harris, C., Igashov, I., Schneuing, A., Segler, M., Correia, B., Roy, J., Bengio, E., and Lio, P. SynFlowNet: Design of diverse and novel molecules with synthesis constraints. In The Thirteenth International Conference on Learning Representations, 2025.
Kim, H., Kim, M., Yun, T., Choi, S., Bengio, E., HernándezGarcı́a, A., and Park, J. Improved off-policy reinforcement learning in biological sequence design. arXiv preprint arXiv:2410.04461, 2024. Kim, M., Choi, S., Kim, H., Son, J., Park, J., and Bengio, Y. Ant colony sampling with GFlowNets for combinatorial optimization. In International Conference on Artificial Intelligence and Statistics. PMLR, 2025.
Dantzig, G. B. Maximization of a linear function of variables subject to linear inequalities. In Koopmans, T. C. (ed.), Activity Analysis of Production and Allocation, pp. 339–347. John Wiley & Sons, New York, 1951.
Kostenetskiy, P., Chulkevich, R., and Kozyrev, V. Hpc resources of the higher school of economics. In Journal of Physics: Conference Series, volume 1740, pp. 012050. IOP Publishing, 2021.
Deleu, T., Nouri, P., Malkin, N., Precup, D., and Bengio, Y. Discrete probabilistic inference as control in multi-path environments. In The 40th Conference on Uncertainty in Artificial Intelligence, 2024.
Koziarski, M., Rekesh, A., Shevchuk, D., van der Sloot, A., Gaiński, P., Bengio, Y., Liu, C., Tyers, M., and Batey, R. Rgfn: Synthesizable molecular generation using gflownets. Advances in Neural Information Processing Systems, 37:46908–46955, 2024.
Essid, M. and Solomon, J. Quadratically regularized optimal transport on graphs. SIAM Journal on Scientific Computing, 40(4):A1961–A1986, 2018. doi: 10.1137/17M1132665. URL https://epubs.siam. org/doi/10.1137/17M1132665.
Kusner, M. J., Sun, Y., Kolkin, N. I., and Weinberger, K. Q. From word embeddings to document distances. In Proceedings of the 32nd International Conference on Machine Learning, pp. 957–966, 2015.
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Boisbunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., Gautheron, L., Gayraud, N. T. H., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V., Sutherland, D. J., Tavenard, R., Tong, A., and Vayer, T. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1–8, 2021.
Kwon, O. J., Matsunaga, D. E., and Kim, K.-E. Gdpo: Learning to directly align language models with diversity using gflownets. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 17120–17139, 2024.
Gritsaev, T., Morozov, N., Samsonov, S., and Tiapkin, D. Optimizing backward policies in GFlownets via trajectory likelihood maximization. In The Thirteenth International Conference on Learning Representations, 2025.
Lahlou, S., Deleu, T., Lemos, P., Zhang, D., Volokhova, A., Hernández-Garcıa, A., Ezzine, L. N., Bengio, Y., and Malkin, N. A theory of continuous generative flow networks. In International Conference on Machine Learning, pp. 18269–18300. PMLR, 2023.
He, H., Bengio, E., Cai, Q., and Pan, L. Random policy evaluation uncovers policies of generative flow networks. In International Conference on Machine Learning, pp. 22351–22367. PMLR, 2025.
Lau, E., Lu, S., Pan, L., Precup, D., and Bengio, E. Qgfn: Controllable greediness with action values. Advances in neural information processing systems, 37:81645–81676, 2024.
Hu, E. J., Jain, M., Elmoznino, E., Kaddar, Y., Lajoie, G., Bengio, Y., and Malkin, N. Amortizing intractable inference in large language models. In The Twelfth International Conference on Learning Representations, 2023.
Lee, S., Kim, M., Cherif, L., Dobre, D., Lee, J., Hwang, S. J., Kawaguchi, K., Gidel, G., Bengio, Y., Malkin, N., and Jain, M. Learning diverse attacks on large language models for robust red-teaming and safety tuning. In The Thirteenth International Conference on Learning Representations, 2025.
Jain, M., Bengio, E., Hernandez-Garcia, A., Rector-Brooks, J., Dossou, B. F., Ekbote, C. A., Fu, J., Zhang, T., Kilgour, M., Zhang, D., et al. Biological sequence design with gflownets. In International Conference on Machine Learning, pp. 9786–9801. PMLR, 2022.
Madan, K., Rector-Brooks, J., Korablyov, M., Bengio, E., Jain, M., Nica, A. C., Bosc, T., Bengio, Y., and Malkin, N. Learning gflownets from partial episodes for improved convergence and stability. In International Conference on Machine Learning, pp. 23467–23483. PMLR, 2023.
Jang, H., Jang, Y., Kim, M., Park, J., and Ahn, S. Pessimistic backward policy for GFlowNets. In Advances in Neural Information Processing Systems, volume 37, pp. 107087– 107111, 2024. 9
Your GFlowNet Secretly Learns an Optimal Transport Plan
Malkin, N., Jain, M., Bengio, E., Sun, C., and Bengio, Y. Trajectory balance: Improved credit assignment in gflownets. Advances in Neural Information Processing Systems, 35:5955–5967, 2022.
diffusion models for vision, language, and control. Neural Information Processing Systems (NeurIPS), 2024. Villani, C. Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008. ISBN 9783540710493. URL https:// books.google.ru/books?id=NZXiNAEACAAJ.
Mohammadpour, S., Bengio, E., Frejinger, E., and Bacon, P.-L. Maximum entropy gflownets with soft q-learning. In International Conference on Artificial Intelligence and Statistics, pp. 2593–2601. PMLR, 2024.
Virtanen, P., Gommers, R., Oliphant, T. E., Haberland, M., Reddy, T., Cournapeau, D., Burovski, E., Peterson, P., Weckesser, W., Bright, J., van der Walt, S. J., Brett, M., Wilson, J., Millman, K. J., Mayorov, N., Nelson, A. R. J., Jones, E., Kern, R., Larson, E., Carey, C. J., Polat, İ., Feng, Y., Moore, E. W., VanderPlas, J., Laxalde, D., Perktold, J., Cimrman, R., Henriksen, I., Quintero, E. A., Harris, C. R., Archibald, A. M., Ribeiro, A. H., Pedregosa, F., van Mulbregt, P., and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261–272, 2020. doi: 10.1038/s41592-019-0686-2.
Morozov, N., Tiapkin, D., Samsonov, S., Naumov, A., and Vetrov, D. Improving gflownets with monte carlo tree search. arXiv preprint arXiv:2406.13655, 2024. Morozov, N., Maksimov, I., Tiapkin, D., and Samsonov, S. Revisiting non-acyclic GFlowNets in discrete environments. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pp. 44887–44910. PMLR, 2025. Morozov, N., Maksimov, I., Tiapkin, D., and Samsonov, S. Learning shortest paths with generative flow networks. arXiv preprint arXiv:2603.01786, 2026.
Zhang, D., Dai, H., Malkin, N., Courville, A. C., Bengio, Y., and Pan, L. Let the flows tell: Solving graph combinatorial problems with gflownets. In Advances in Neural Information Processing Systems, volume 36, pp. 11952– 11969, 2023a.
Peyré, G. and Cuturi, M. Computational optimal transport: With applications to data science. Now Foundations and Trends, 2019.
Zhang, D., Zhang, Y., Gu, J., Zhang, R., Susskind, J. M., Jaitly, N., and Zhai, S. Improving gflownets for textto-image diffusion alignment. Transactions on Machine Learning Research, 2025. ISSN 2835-8856.
Rubner, Y., Tomasi, C., and Guibas, L. J. The earth mover’s distance as a metric for image retrieval. International Journal of Computer Vision, 40(2):99–121, 2000.
Zhang, D. W., Rainone, C., Peschl, M., and Bondesan, R. Robust scheduling with gflownets. In The Eleventh International Conference on Learning Representations, 2023b.
Shen, T., Seo, S., Lee, G., Pandey, M., Smith, J. R., Cherkasov, A., Kim, W. Y., and Ester, M. TacoGFN: Target-conditioned GFlowNet for structure-based drug design. Transactions on Machine Learning Research, 2024. ISSN 2835-8856.
Zhu, X., Cheng, D., Zhang, D., Li, H., Zhang, K., Jiang, C., Sun, Y., Hua, E., Zuo, Y., Lv, X., et al. Flowrl: Matching reward distributions for llm reasoning. arXiv preprint arXiv:2509.15207, 2025.
Tiapkin, D., Morozov, N., Naumov, A., and Vetrov, D. P. Generative flow networks as entropy-regularized rl. In International Conference on Artificial Intelligence and Statistics, pp. 4213–4221. PMLR, 2024. Tolstikhin, I., Bousquet, O., Gelly, S., and Schoelkopf, B. Wasserstein auto-encoders. In International Conference on Learning Representations, 2018. URL https:// openreview.net/forum?id=HkL7n1-0b. Uehara, M., Zhao, Y., Biancalani, T., and Levine, S. Understanding reinforcement learning-based fine-tuning of diffusion models: A tutorial and review. arXiv preprint arXiv:2407.13734, 2024. Venkatraman, S., Jain, M., Scimeca, L., Kim, M., Sendera, M., Hasan, M., Rowe, L., Mittal, S., Lemos, P., Bengio, E., Adam, A., Rector-Brooks, J., Bengio, Y., Berseth, G., and Malkin, N. Amortizing intractable inference in 10
Your GFlowNet Secretly Learns an Optimal Transport Plan
A. Full theoretical results A.1. Correctness of LP formulation Consider original minimum-flow formulation (5), we will now proof that it could be indeed reformulated as linear programming problem. X
min
F ,PF ,PB
F(s)
s∈I
F(s)PF (s′ | s) = F (s′ )PB (s | s′ ),
s.t.
(21)
′
s → s ∈ E, F(sf )PB (x | sf ) = R(x),
x ∈ X.
Proposition A.1 (Elliminating forward policy). Fix s ̸= sf . The following are equivalent: 1. There exists a probability distribution PF (· | s) ∈ ∆(out(s)) such that F(s → v) = F (s)PF (v | s),
∀v ∈ out(s).
2. The variables F(s → v) satisfy the linear conditions F(s → v) ≥ 0
X
∀v ∈ out(s),
F(s → v) = F (s).
v∈out(s)
Proof. Assume (i). Since PF (· | s) is a probability distribution, X X X F(s → v) = F(s)PF (v | s) = F(s) PF (v | s) = F(s), v∈out(s)
v∈out(s)
v∈out(s)
and clearly F(s → v) ≥ 0 for all v. Now assume (ii). If F(s) > 0, define PF (v | s) :=
F(s → v) , F(s)
v ∈ out(s).
Then PF (v | s) ≥ 0 and X
PF (v | s) =
v∈out(s)
X 1 F(s → v) = 1. F(s) v∈out(s)
Hence PF (· | s) ∈ ∆(out(s), and by construction F(s → v) = F (s)PF (v | s). If F(s) = 0, then ∆(out(s)) works.
P
v F(s → v) = 0 and F(s → v) ≥ 0 imply F(s → v) = 0 for all v, so any choice of PF (· | s) ∈
Proposition A.2 (Eliminating the backward policy). Fix s′ ̸= s0 . The following are equivalent: 1. There exists a probability distribution PB (· | s′ ) ∈ ∆(in(s′ )) such that F(u → s′ ) = F(s′ )PB (u | s′ ),
∀u ∈ in(s′ ).
2. The variables F(u → s′ ) satisfy the linear conditions F(u → s′ ) ≥ 0
∀u ∈ in(s′ ),
X u∈in(s′ )
11
F(u → s′ ) = F (s′ ).
Your GFlowNet Secretly Learns an Optimal Transport Plan
Proof. The proof is identical to the previous one, with incoming edges in place of outgoing edges. Corollary A.3 (Lifted LP). After eliminating PF and PB , the minimum-total-flow problem becomes X (P) min ′ F(s) F (s),F (s→s )
s∈I
X
s.t.
F(s → v) = F (s), ∀s ∈ S \ {sf },
v∈out(s)
X
F(u → s) = F (s),
∀s ∈ S \ {s0 },
u∈in(s)
F(x → sf ) = R(x), ′
F(s → s ) ≥ 0,
∀x ∈ X, ∀(s, s′ ) ∈ E.
This is a linear program. Proof. The objective is linear in (F(s), F(s → s′ )), and all constraints are affine. A.2. Duality proofs For notational convenience define the cost vector c by ( cs :=
s ∈ I, s ∈ {s0 , sf }.
1, 0,
Then the primal objective is cT F =
X
F(s) .
s∈I
Introduce dual variables: • αs ∈ for the outgoing-flow constraints, s ∈ S \ {sf }; • βs ∈ for the incoming-flow constraints, s ∈ S \ {s0 }; • ηx ∈ for the terminal constraints, x ∈ X. We also define padded vectors ᾱ, β̄ by ( ᾱs :=
( βs , s ̸= s0 , β̄s := 0, s = s0 .
αs , s ̸= sf , 0, s = sf ,
Proposition A.4. The dual of (A.3) is X
max
α,β,η
s.t.
R(x)ηx
x∈X
c = ᾱ + β̄,
(22)
αs + βs′ ≥ 0,
∀s → s′ with s′ ̸= sf ,
ηx ≤ αx + βsf ,
∀x ∈ X.
Moreover, this dual can be reduced: max π
s.t.
X
R(x)πx
x∈X
(23)
πs0 = 0, πs′ − πs ≤ 1,
∀(s, s′ ) ∈ E with s′ ̸= sf . 12
Your GFlowNet Secretly Learns an Optimal Transport Plan
Proof. The Lagrangian is L(F(s), F(s → s′ ), α, β, η) = cT F +
X
αs
s̸=sf
+
X
X
F(s → v) − F (s)
v∈out(s)
βs
X
s̸=s0
X F(u → s) − F (s) + ηx R(x) − F (x → sf ) . x∈X
u∈in(s)
Now collect all terms involving F(s) and F(s → s′ ). Since each edge s → s′ contributes αs through the outgoing constraint at s, and βs′ through the incoming constraint at s′ , we obtain X X L(F(s), F(s → s′ ), α, β, η) = R(x)ηx + cs − ᾱs − β̄s F(s) x∈X
s∈S
X
+
(αs + βs′ ) F(s → s′ )
(s,s′ )∈E s′ ̸=sf
+
X
(αx + βsf − ηx ) F(x → sf ).
x∈X
The dual function is g(α, β, η) :=
inf
F (s), F (s→s′ )≥0
L(F(s), F(s → s′ ), α, β, η).
We now determine when this infimum is finite. Step 1: minimization over F(s). The variables F(s) are unrestricted in sign inside the Lagrangian. Therefore the infimum over F(s) is finite if and only if each coefficient of F (s) is zero: cs − ᾱs − β̄s = 0,
∀s ∈ S.
Equivalently, c = ᾱ + β̄. Step 2: minimization over F(s → s′ ). Each variable F(s → s′ ) is constrained only by F(s → s′ ) ≥ 0. Hence the infimum over F(s → s′ ) is finite if and only if each coefficient of F(s → s′ ) is nonnegative: αs + βs′ ≥ 0,
∀(s → s′ ) with s′ ̸= sf ,
and αx + βsf − ηx ≥ 0,
∀x ∈ X.
The latter condition is equivalent to ηx ≤ αx + βsf ,
∀x ∈ X.
Under these conditions, g(α, β, η) =
X
R(x)ηx .
x∈X
Therefore the dual problem is exactly (22). Since R(x) > 0 for every x ∈ X, the optimal choice of ηx saturates the inequality ηx ≤ αx + βsf , hence ηx = αx + βsf . 13
Your GFlowNet Secretly Learns an Optimal Transport Plan
Therefore (22) is equivalent to max
X
s.t.
c = ᾱ + β̄,
α,β
R(x) αx + βsf
x∈X
∀(s → s′ ) with s′ ̸= sf .
αs + βs′ ≥ 0, From c = ᾱ + β̄, we obtain αs0 = 0,
β s = 1 − αs
βsf = 0,
∀s ∈ I.
Substituting into the edge inequalities gives αs + βs′ ≥ 0
⇐⇒
αs + (1 − αs′ ) ≥ 0
⇐⇒
αs′ − αs ≤ 1.
Renaming αs as πs yields (23). Note that the derivation of dual to the extended problem 10 is exactly the same up to an additional equality constraint from the fixed first step. A.3. Shortest Path via Dual Potentials Theorem A.5 (Shortest paths dual characterization). Consider the dual to the primal problem A.4 X R(x)πx subject to πs0 = 0, πs′ − πs ≤ 1 ∀(s → s′ ) with s′ ̸= sf . max π
x∈X
For every non-sink state s, define d(s) := min{k ≥ 0 : ∃ s0 = v0 → v1 → · · · → vk = s}. Then: 1. every feasible solution π satisfies πs ≤ d(s)
∀s ̸= sf ;
2. the function d(·) is a feasible dual solution, and it is dual-optimal; 3. if R(x) > 0 for every x ∈ X, then every optimal solution π ⋆ satisfies πx⋆ = d(x)
∀x ∈ X.
Moreover, complementary slackness conditions in the optimum gives us the following relation: F ⋆ (s → s′ )(πs⋆′ − (1 + πs⋆ )) = 0 which is also a Primal-Dual gap. Proof. We prove the three statements one by one. Step 1: every feasible π satisfies πs ≤ d(s). Let π be any feasible solution of the dual problem, and fix a state s ̸= sf . Take any directed path from s0 to s: s0 = v0 → v1 → · · · → vk = s. For every edge vi → vi+1 in this path, feasibility gives πvi+1 − πvi ≤ 1. 14
Your GFlowNet Secretly Learns an Optimal Transport Plan
Now sum these inequalities over i = 0, . . . , k − 1: k−1 X
X k−1 πvi+1 − πvi ≤ 1 = k.
i=0
i=0
The left-hand side is a telescoping sum, so it simplifies to πvk − πv0 = πs − πs0 . Since πs0 = 0, we obtain πs ≤ k. This holds for every path from s0 to s. Therefore, taking the minimum over all such paths, πs ≤ d(s). This proves part (i). Step 2: the function d(·) is feasible. Define πsmax := d(s). We show that π max satisfies all dual constraints. First, πsmax = d(s0 ) = 0. 0 Next, take any edge s → s′ with s′ ̸= sf . By definition of d(s), there exists a shortest path from s0 to s of length d(s): s0 = v0 → v1 → · · · → vd(s) = s. Appending the edge s → s′ gives a path from s0 to s′ of length d(s) + 1. Since d(s′ ) is the length of the shortest path from s0 to s′ , we must have d(s′ ) ≤ d(s) + 1. Equivalently, d(s′ ) − d(s) ≤ 1. Thus πsmax − πsmax ≤ 1. ′ So π max is feasible. Step 3: d(·) is dual-optimal. Let π be any feasible solution. By Step 1, for every terminal state x ∈ X, πx ≤ d(x). Since the rewards are nonnegative, X
R(x)πx ≤
x∈X
X
R(x)d(x).
x∈X
But Step 2 shows that the choice πsmax = d(s) is feasible, and for this choice the objective value is exactly X R(x)d(x). x∈X
Therefore no feasible solution can have a larger objective value, and d(·) is dual-optimal. This proves part (ii). 15
Your GFlowNet Secretly Learns an Optimal Transport Plan
Step 4: Every optimal solution agrees with d on terminals. Let π ⋆ be any optimal solution. By Step 1, πx⋆ ≤ d(x)
∀x ∈ X.
Suppose, for contradiction, that there exists some terminal state x0 ∈ X such that πx⋆0 < d(x0 ). Because all rewards are strictly positive, this would imply X X R(x)πx⋆ < R(x)d(x). x∈X
x∈X
But the right-hand side is the optimal value, since it is attained by the feasible solution d(·). This contradicts the optimality of π ⋆ . Hence no such x0 can exist, and therefore πx⋆ = d(x)
∀x ∈ X.
This proves part (iii).
A.4. Proof of dual equivalence of OT and GFlowNet Proof. It is easy to see that dual to the extended primal problem is of the following nature (see Appendix A.4 for derivation) X X max R(x)πx + L(u)(−πu ) π
s.t.
u∈U
x∈X
πs′ − πs ≤ 1, ∀s → s′ : s ̸= s0 , s′ ̸= sf
(24)
πs0 = 0 if we now redefine ax = πx , bu = −πu , then taking any path p: ax + bu = πx − πu =
|p| X
πpi+1 − πpi ≤ d(u, x)
i=0
Hence: max a,b
s.t.
X
R(x)ax +
x∈X
X
L(u)bu
u∈U
(25)
ax + bu ≤ d(u, x)
Which has exactly the same functional form in decision variables ax , bu as Dual to Problem 6
B. Experimental details All models are parameterized by MLP with 2 hidden layers and 128 hidden size, which accept a one-hot encoding of s as input. Fθ (s), PF (s′ | s, θ), PB (s | s′ , θ) share the same backbone, with different linear heads predicting the logits of the forward policy and the logits of the backward policy. We train our models on-policy, using batch size 512. We use AdamW optimzier with lr = 10−3 for all our experiments and weight decay 10−4 . Our implementations are based on the published code of (Morozov et al., 2025). All experiments were performed on CPUs. B.1. Hypergrid distribution definitions We define moon shaped distribution by the following formula: ∥z − b∥2 2 2 L(s) ≜ I ∥z − c∥22 ≤ rout 1 − I ∥z − (c − δe1 )∥22 ≤ rin 0.5 + 2 min 1, max 0, 1 − + ε, rout 16
Your GFlowNet Secretly Learns an Optimal Transport Plan
Figure 2. Visualization of the distributions used in the experiments: moon-shaped (left), corner-shaped (center), and ball-shaped (right).
Where: z=
s , M −1
c=
1 1 ,..., , 2 2
e1 = (1, 0, . . . , 0),
b=c+
rout e1 . 2
and rin , rout are inner and outer balls radius. The ball-shaped distribution defined as follows: ∥z − b∥2 2 2 L(s) ≜ I ∥z − c∥2 ≤ rout 0.5 + 2 min 1, max 0, 1 − + ε, rout The corner-shaped distribution definition follows previous works (Madan et al., 2023; Malkin et al., 2022): D Y R(s) ≜ R0 + R1 I 0.25 < i=1
si − 0.5 H −1
D Y + R2 I 0.3 < i=1
si − 0.5 < 0.4 , H −1
See Figure 2 for visualization. B.2. Permutations Denote C(k) as the probability that a permutation sampled from the reward distribution has k fixed points. We compute the L1 error between the vector (C(0), C(1), . . . , C(n)) and its empirical estimate over the last 105 samples seen in training for convergence diagnostics. Suppose that x1 , . . . , xm is a set of GFlowNet samples (terminal states of trajectories sampled from PF ). Then, the empirical L1 error of fixed point probabilities is defined as: N X k=0
m
C(k) −
1 X I{xi (k) = k} , m i=1
We use the exact formula from (Morozov et al., 2025) for exact computation of C(k).
17