Your GFlowNet Secretly Learns an Optimal Transport Plan

Technische Samenvatting: Uw GFlowNet leert in het geheim een optimaal transportplan

Probleemstelling
Generative Flow Networks (GFlowNets) zijn ontworpen om gestructureerde discrete objecten te samplen via stochastische trajecten in een gerichte graaf. Hoewel ze oorspronkelijk zijn ontwikkeld voor acyclische grafen, heeft recent werk de GFlowNets uitgebreid naar niet-acyclische omgevingen. In deze settings is het minimaliseren van de verwachte trajectlengte cruciaal voor de efficiëntie van het samplen, wat leidt tot de formulering van "minimum-flow" GFlowNets. Gelijktijdig biedt Optimal Transport (OT) een principieel kader voor het vergelijken van waarschijnlijkheidsverdelingen door een transportkosten te minimaliseren, wat doorgaans wordt opgelost via lineaire programmering (Kantorovich-formulering). Het oplossen van OT op grote grafen is echter vaak computationeel onuitvoerbaar met exacte solvers.

Dit artikel adresseert de kloof tussen deze twee velden door te onderzoeken of de minimum-flow formulering van niet-acyclische GFlowNets inherent een Optimal Transport probleem oplost.

Methodologie
De auteurs vestigen een theoretische equivalentie tussen niet-acyclische GFlowNets en discrete Optimal Transport op grafen. De kernmethodologie omvat:

1. Theoretische Herformulering: De auteurs beginnen met het minimum-flow optimalisatieprobleem voor niet-acyclische GFlowNets, dat de totale flow minimaliseert (gelijk aan het minimaliseren van de verwachte trajectlengte) onderworpen aan detailed balance en reward matching restricties. Door de initiële edge-flow verdeling (de flow die de bronstaat $s_0$ verlaat) vast te leggen om overeen te komen met een specifieke bronverdeling $L(u)$, transformeren zij de GFlowNet doelstelling in een lineair programmeerprobleem (LP).
2. Equivalentiebewijs: Zij demonstreren dat deze beperkte LP wiskundig equivalent is aan het Kantorovich Optimal Transport probleem. In deze formulering:
   • Is de bronverdeling $L(u)$ gedefinieerd op een verzameling initiële staten $U$.
   • Is de doelverdeling de rewardverdeling $R(x)$ gedefinieerd op terminale staten $X$.
   • De transportkosten $d(u, x)$ worden gedefinieerd als de kortste pad afstand tussen $u$ en $x$ in de graaf.
   • De optimale GFlowNet policy leert impliciet de optimale koppeling $\Pi$ die de totale transportkosten minimaliseert.
3. Leeralgoritme: Om dit op grote grafen op te lossen waar exacte LP onhaalbaar is, gebruiken de auteurs een neurale parametrisatie van de GFlowNet. Zij maken gebruik van de Trajectory Balance (TB) doelstelling met flow regularisatie. Een belangrijke modificatie is de inclusie van een "leward" (links reward) term in de loss functie om de vaste initiële verdeling $L(u)$ af te dwingen, naast de standaard reward matching voor de doelverdeling $R(x)$.

Belangrijkste Bijdragen

1. Theoretische Equivalentie: Het artikel bewijst dat het vastleggen van de initiële flow verdeling in een minimum-flow GFlowNet de doelstelling exact reduceert tot een Kantorovich OT probleem met graaf-geïnduceerde kortste pad kosten. Bijgevolg codeert de optimale GFlowNet forward policy een optimaal transportplan, en het samplen van trajecten uit deze policy herstelt de optimale koppeling.
2. Verbinding met het Beckmann Probleem: De auteurs tonen aan dat deze formulering equivalent is aan de discrete graaf-gebaseerde formulering van het Beckmann probleem, waarmee GFlowNets worden gelinkt aan klassieke netwerkstroomtheorie.
3. Schaalbare OT Solver: Het werk demonstreert dat het GFlowNet leerframework kan worden toegepast om benaderende oplossingen te vinden voor graaf-gebaseerde OT problemen. Door gebruik te maken van neurale netwerken en stochastische traject sampling, schaalt de methode naar combinatorische ruimtes waar exacte OT solvers onuitvoerbaar worden.
4. Herstel van Kortste Pad Theorie: Als een bijgerecht herstelt deze constructie de kortste pad theorie die eerder is vastgesteld voor niet-acyclische GFlowNets, waarbij wordt aangetoond dat optimale flows alleen worden ondersteund op de subgraaf van kortste paden.

Experimentele Resultaten
De auteurs valideren hun theorie en methode in twee omgevingen:

• Hypergrid Omgeving: Op kleine hypergrids ($H=10, 15, 20$) bereikt de geleerde GFlowNet policy transportkosten ($OT^\star$) die overeenkomen met de exacte oplossingen verkregen via klassieke LP solvers (POT). De methode herstelt succesvol de optimale koppeling over verschillende bronverdelingen (bal en maan vormen) zonder de sampling te verstoren.
• Permutatie Omgeving: Op de Cayley graaf van de symmetrische groep $S_n$ (permutaties), herstelt de methode bijna-optimale oplossingen voor $n=4$ en $n=8$, waar exacte oplossingen hanteerbaar zijn. Voor $n=20$, waar exacte OT onuitvoerbaar is, produceert de methode redelijke benaderingen met een lage fout in fixed-point waarschijnlijkheden.
• Ablatie: Experimenten met de regularisatiecoëfficiënt $\lambda$ bevestigen de trade-off tussen trajectlengte (sampling efficiëntie) en sampling optimaliteit (nauwkeurigheid van het transportplan). Lagere $\lambda$ waarden leveren nauwkeurigere transportplannen op maar langere trajecten.

Betekenis en Claims
Het artikel claimt een "geheime" connectie te hebben ontdekt: GFlowNets, wanneer geformuleerd om flow te minimaliseren, lossen impliciet Optimal Transport problemen op. De betekenis van dit werk ligt in:

• Unificatie: Het biedt een nieuw theoretisch perspectief dat GFlowNet leren en Optimal Transport verenigt, wat suggereert dat GFlowNets kunnen worden gezien als een methode voor het leren van transportkaarten.
• Praktische Toepassing: Het maakt het mogelijk om GFlowNet leeralgoritmen (die schalen via neurale netwerken en sampling) toe te passen op grootschalige graaf OT problemen die anders computationeel onmogelijk zijn voor exacte solvers.
• Framework Extensie: Het breidt de algemeenheid van het GFlowNet framework uit voorbij sampling en generatieve taken naar het domein van computationele optimale transport, waarbij een stochastische policy wordt geboden die niet alleen bepaalt waar massa wordt getransporteerd maar ook hoe het wordt getransporteerd via haalbare lokale bewegingen met bijna optimale padkosten.

De auteurs blijven bescheiden over toekomstige implicaties en merken op dat hoewel de methode computationeel efficiënt is, de regularisatiecoëfficiënt $\lambda$ zorgvuldige afstemming vereist om de balans tussen optimaliteit en efficiëntie te bewaren. Zij suggereren dat toekomstig werk de toepassing van deze methoden kan verkennen om transportkaarten te leren tussen ongenormaliseerde verdelingen.