Your GFlowNet Secretly Learns an Optimal Transport Plan

Technical Summary: Your GFlowNet Secretly Learns an Optimal Transport Plan

Problem Statement
Generative Flow Networks (GFlowNets) are designed to sample structured discrete objects via stochastic trajectories in a directed graph. While initially developed for acyclic graphs, recent work has extended GFlowNets to non-acyclic environments. In these settings, minimizing the expected trajectory length is crucial for sampling efficiency, leading to the formulation of "minimum-flow" GFlowNets. Concurrently, Optimal Transport (OT) provides a principled framework for comparing probability distributions by minimizing a transportation cost, typically solved via linear programming (Kantorovich formulation). However, solving OT on large graphs is often computationally intractable using exact solvers.

This paper addresses the gap between these two fields by investigating whether the minimum-flow formulation of non-acyclic GFlowNets inherently solves an Optimal Transport problem.

Methodology
The authors establish a theoretical equivalence between non-acyclic GFlowNets and discrete Optimal Transport on graphs. The core methodology involves:

1. Theoretical Reformulation: The authors start with the minimum-flow optimization problem for non-acyclic GFlowNets, which minimizes the total flow (equivalent to minimizing expected trajectory length) subject to detailed balance and reward matching constraints. By fixing the initial edge-flow distribution (the flow leaving the source state $s_0$) to match a specific source distribution $L(u)$, they transform the GFlowNet objective into a linear programming (LP) problem.
2. Equivalence Proof: They demonstrate that this constrained LP is mathematically equivalent to the Kantorovich Optimal Transport problem. In this formulation:
   • The source distribution is $L(u)$ defined on a set of initial states $U$.
   • The target distribution is the reward distribution $R(x)$ defined on terminal states $X$.
   • The transport cost $d(u, x)$ is defined as the shortest path distance between $u$ and $x$ in the graph.
   • The optimal GFlowNet policy implicitly learns the optimal coupling $\Pi$ that minimizes the total transport cost.
3. Learning Algorithm: To solve this on large graphs where exact LP is infeasible, the authors employ a neural parameterization of the GFlowNet. They utilize the Trajectory Balance (TB) objective with flow regularization. A key modification is the inclusion of a "leward" (left reward) term in the loss function to enforce the fixed initial distribution $L(u)$, alongside the standard reward matching for the target distribution $R(x)$.

Key Contributions

1. Theoretical Equivalence: The paper proves that fixing the initial flow distribution in a minimum-flow GFlowNet reduces its objective exactly to a Kantorovich OT problem with graph-induced shortest path costs. Consequently, the optimal GFlowNet forward policy encodes an optimal transport plan, and sampling trajectories from it recovers the optimal coupling.
2. Connection to Beckmann Problem: The authors show that this formulation is equivalent to the discrete graph-based formulation of the Beckmann problem, linking GFlowNets to classical network flow theory.
3. Scalable OT Solver: The work demonstrates that the GFlowNet learning framework can be applied to approximate solutions for graph-based OT problems. By using neural networks and stochastic trajectory sampling, the method scales to combinatorial spaces where exact OT solvers become intractable.
4. Recovery of Shortest Path Theory: As a corollary, the construction recovers the shortest-path theory previously established for non-acyclic GFlowNets, showing that optimal flows are supported only on the subgraph of shortest paths.

Experimental Results
The authors validate their theory and method on two environments:

• Hypergrid Environment: On small hypergrids ($H=10, 15, 20$), the learned GFlowNet policy achieves transport costs ($OT^\star$) that match the exact solutions obtained via classical LP solvers (POT). The method successfully recovers the optimal coupling across different source distributions (ball and moon shapes) without biasing the sampling.
• Permutation Environment: On the Cayley graph of the symmetric group $S_n$ (permutations), the method recovers near-optimal solutions for $n=4$ and $n=8$, where exact solutions are tractable. For $n=20$, where exact OT is intractable, the method produces reasonable approximations with low error in fixed-point probabilities.
• Ablation: Experiments on the regularization coefficient $\lambda$ confirm the trade-off between trajectory length (sampling efficiency) and sampling optimality (accuracy of the transport plan). Lower $\lambda$ values yield more accurate transport plans but longer trajectories.

Significance and Claims
The paper claims to uncover a "secret" connection: GFlowNets, when formulated to minimize flow, are implicitly solving Optimal Transport problems. The significance of this work lies in:

• Unification: It provides a novel theoretical view unifying GFlowNet learning and Optimal Transport, suggesting that GFlowNets can be viewed as a method for learning transport maps.
• Practical Application: It enables the application of GFlowNet learning algorithms (which scale via neural networks and sampling) to large-scale graph OT problems that are otherwise computationally prohibitive for exact solvers.
• Framework Extension: It extends the generality of the GFlowNet framework beyond sampling and generation tasks into the domain of computational optimal transport, offering a stochastic policy that not only determines where mass is transported but how it is transported via feasible local moves with near-optimal path costs.

The authors maintain modesty regarding future implications, noting that while the method is computationally efficient, the regularization coefficient $\lambda$ requires careful tuning to balance optimality and efficiency. They suggest future work could explore applying these methods to learn transport maps between unnormalized distributions.