Your GFlowNet Secretly Learns an Optimal Transport Plan

2026-06-04

AI Summary

This work investigates efficient optimal transport (OT) computation on large-scale graph structures. By fixing the initial flow distribution of an acyclic generative flow network (GFlowNet), the minimum-flow objective is reformulated as a Kantorovich OT problem under a graph-induced shortest-path cost, establishing for the first time a theoretical connection between GFlowNets and optimal transport. The authors prove that the learned policy implicitly encodes the optimal coupling between source and target distributions at optimality. Leveraging an edge-flow-based GFlowNet framework that combines neural parameterization with shortest-path metrics on graphs, the resulting transport plans exhibit strong alignment with those produced by exact OT solvers, enabling the efficient generation of high-quality optimal transport solutions on large-scale graphs.

Abstract

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.

Problem

• Optimal Transport
• GFlowNets
• Graph-based OT
• Stochastic Trajectories
• Large-scale Graphs

Innovation

• GFlowNets
• Optimal Transport
• Minimum-flow
• Graph-induced costs
• Neural parameterization