Learning Shortest Paths with Generative Flow Networks

Deze paper introduceert een nieuw leerframework dat Generative Flow Networks (GFlowNets) gebruikt om kortste paden in grafen te vinden door flow-regulering toe te passen, wat resulteert in concurrerende prestaties bij complexe problemen zoals het oplossen van de Rubik's Cube met een kleiner zoekbudget.

De kern

Stel je voor dat je in een gigantisch, compleet donker labyrint staat. Je doel is om de kortste weg te vinden naar de uitgang. Normaal gesproken zou je elke afslag moeten uitproberen, of een heel slimme kaart nodig hebben om te weten welke kant op te gaan. Dit is precies het probleem dat kunstmatige intelligentie vaak heeft bij complexe puzzels, zoals het oplossen van een Rubik's Cube of het plannen van routes voor robots.

Deze paper introduceert een nieuwe, slimme manier om dit op te lossen met een technologie die Generative Flow Networks (GFlowNets) heet. Laten we dit uitleggen met een paar leuke metaforen.

1. Het Probleem: Het Labyrint van de Rubik's Cube

Stel je een Rubik's Cube voor. Er zijn meer mogelijke combinaties dan er atomen in het heelal zijn. Als je probeert de kubus op te lossen, kun je elke draai maken, maar je kunt ook terug naar een eerdere stand gaan (je kunt een beweging ongedaan maken). Dit maakt het pad niet lineair; het is een web van kringen en lusjes.

Oude methoden proberen vaak een "kaart" te tekenen of te voorspellen hoe ver je nog van de uitgang bent. Maar in zo'n groot web is het moeilijk om de perfecte kaart te tekenen zonder het hele web eerst te verkennen.

2. De Oplossing: De "Stroom" van de Rivier

De auteurs van dit paper gebruiken een heel ander idee. In plaats van te kijken naar de afstand, kijken ze naar een rivier (de "flow").

• De Rivier: Stel je voor dat er een rivier stroomt door je labyrint. De rivier begint bij de uitgang (de oplossing) en stroomt terug naar waar je nu bent.
• De Regels: Normaal gesproken kan een rivier overal naartoe stromen, ook in kringen of doodlopende wegen. Maar de auteurs hebben een speciale regel bedacht: "De rivier moet zo kort mogelijk zijn."

Ze hebben bewezen dat als je deze regel streng toepast (de "stroom minimaliseert"), de rivier alleen de kortste, meest rechtstreekse weg neemt. Alle andere wegen, die langer zijn of in kringen lopen, drogen gewoon op. De rivier stroomt dan uitsluitend over het kortste pad.

3. Hoe werkt het in de praktijk? (De Omgekeerde Reis)

Dit is het meest creatieve deel van hun methode. Ze trainen de computer niet om van A naar B te gaan, maar van B naar A.

• De Backward Policy (Terugwaartse Politie): De AI leert hoe je vanuit de oplossing (de opgeloste kubus) terug kunt gaan naar een willekeurige, rommelige stand. Maar niet zomaar terug: de AI leert om altijd de kortste weg terug te nemen.
• De Forward Policy (Voorwaartse Politie): Dit is de tegenhanger. Als de AI weet hoe je terug moet gaan via de kortste weg, weet hij automatisch ook hoe je de kortste weg naar de oplossing vindt.

Het is alsof je een kaart tekent door te kijken hoe je vanuit de finish terugloopt naar de start, maar dan zo snel mogelijk. Als je die terugweg perfect kent, weet je ook de perfecte voorwaartse weg.

4. Waarom is dit zo slim? (De "Flow Regularization")

In de wiskunde van deze paper gebruiken ze iets dat "flow regularization" heet. In onze metafoor is dit als het geven van een beloning aan de rivier als hij kort blijft, en een boete als hij te lang duurt of in een kring loopt.

Door deze boete (de regularisatie) te maximaliseren, dwingen ze de AI om alle lange, inefficiënte paden te negeren. De AI leert dan vanzelf: "Ah, als ik hier die draai maak, kom ik in een lange lus terecht. Dat mag niet. Ik moet die andere draai kiezen om kort te blijven."

5. De Resultaten: Sneller en Slimmer

De auteurs hebben dit getest op twee dingen:

1. Een simpele puzzel (Swap Puzzle): Waar je getallen moet sorteren door ze te verwisselen.
2. Rubik's Cubes: Zowel de kleine 2x2x2 als de grote 3x3x3.

Wat bleek?

• Hun methode vond oplossingen die net zo kort waren als de allerbeste methoden die er nu zijn.
• Het grote voordeel: Ze hadden veel minder "rekenkracht" nodig om die oplossing te vinden. Terwijl andere methoden vaak heel veel mogelijke routes moeten uitproberen (zoals een zoektocht met een hele grote net), kon hun AI met een heel klein net (een kleine "beam search") al de perfecte route vinden.

Samenvattend

Stel je voor dat je een robot wilt leren een labyrint te doorlopen.

• Oude manier: "Hier is een kaart, probeer elke weg uit en onthoud welke het snelst is." (Duur en traag).
• Deze nieuwe manier: "Stel je voor dat je een rivier bent die vanuit de uitgang terugstroomt. Je mag alleen stromen als je de kortste weg neemt. Als je in een kring loopt, droog je op."

Door deze regel te volgen, leert de robot vanzelf de perfecte, kortste route, zonder dat hij het hele labyrint eerst hoeft te verkennen. Het is een elegante manier om de kunst van het "kortste pad vinden" te vertalen naar een wiskundige stroom die zichzelf optimaliseert.

Technische samenvatting

Probleemstelling

Het vinden van kortste paden in grote, discrete grafen is een fundamenteel probleem in kunstmatige intelligentie, met toepassingen in planning, routing en robotica. Klassieke methoden zoals Dijkstra's algoritme en A* zijn optimaal, maar vereisen vaak een goed ontworpen heuristiek en kunnen rekenkundig onhaalbaar worden in hoge dimensies of wanneer de staatruimte te groot is om op te slaan (zoals bij Cayley-grafen van puzzels zoals de Rubik's Cube).

Bestaande machine learning-benaderingen trainen meestal netwerken om een waardefunctie (afstand tot het doel) te schatten, die vervolgens wordt gebruikt om een zoekalgoritme (zoals beam search) te sturen. Het paper introduceert een nieuw perspectief: het direct leren van een beleid dat kortste paden volgt, zonder afhankelijk te zijn van externe zoekstrategieën tijdens het trainen.

Methodologie

De auteurs combineren Generative Flow Networks (GFlowNets) met een specifieke regularisatie om kortste paden in niet-acyclische (cyclische) omgevingen te vinden.

1. Theoretische Basis in Niet-Acyclische GFlowNets:

   • GFlowNets worden doorgaans gebruikt om objecten te genereren die evenredig zijn met een beloningsfunctie. In cyclische omgevingen kan de trajectlengte onbeperkt zijn.
   • De kerntheorie van dit paper bewijst dat het minimaliseren van de verwachte trajectlengte ($E[n_\tau]$) in een GFlowNet leidt tot een beleid dat uitsluitend kortste paden volgt.
   • Stelling 3.4: Als de verwachte trajectlengte wordt geminimaliseerd, wijst het backward-beleid ($P_B$) een kans van nul toe aan alle trajecten die geen kortste paden zijn tussen de start- en eindtoestanden. Het beleid concentreert de waarschijnlijkheidsmassa dus volledig op de kortste paden.

2. Constructie voor Padvinding:

   • Om een kortste-padprobleem in een willekeurige graaf $G$ op te lossen, wordt de graaf gemodificeerd tot een GFlowNet-omgeving:
     • De doeltoestand ($v_g$) wordt de initiële staat ($s_0$) van het GFlowNet.
     • Alle andere knopen worden terminal states.
     • De transities zijn de omgekeerde randen van de originele graaf.
     • Er wordt een uniforme beloning ($R(s)=1$) toegekend aan alle terminal states.
   • Door een backward-beleid te trainen dat de verwachte trajectlengte minimaliseert, leert het model om van elke willekeurige staat terug te lopen naar de start (het oorspronkelijke doel) via een kortste pad.

3. Trainingsalgoritme:

   • De auteurs gebruiken een geregulariseerde Trajectory Balance (TB) loss.
   • In plaats van de gebruikelijke Detailed Balance loss, kiezen ze voor Trajectory Balance omdat dit een efficiëntere credit assignment biedt, vooral omdat het doeltoestand in elke traject voorkomt.
   • Een regularisatieterm ($\lambda F_\theta(s)$) wordt toegevoegd om de totale flow (en dus de verwachte trajectlengte) te minimaliseren.
   • Het algoritme samplet gedeeltelijke trajecten van de forward-policy en berekent de loss over alle prefixes van deze trajecten.

4. Inferentie (Beam Search):

   • Hoewel het theoretische optimum exacte kortste paden garandeert, wordt bij grote grafen vaak een benadering gebruikt.
   • Tijdens de testfase wordt beam search toegepast op het geleerde backward-beleid om de meest waarschijnlijke trajecten te selecteren. Dit verbetert de oplossingkwaliteit aanzienlijk ten opzichte van puur greedy sampling.

Kernbijdragen

• Theoretisch Bewijs: Het eerste bewijs dat het minimaliseren van de verwachte trajectlengte in niet-acyclische GFlowNets equivalent is aan het vinden van exacte kortste paden.
• Constructieve Reductie: Een methode om kortste-padproblemen in willekeurige ongewogen grafen te reduceren tot het trainen van een GFlowNet met flow-regularisatie.
• Nieuw Paradigma: In tegenstelling tot methoden die waardefuncties leren om zoekalgoritmen te sturen, leert deze methode direct een beleid dat kortste paden genereert.
• Efficiëntie: Het gebruik van Trajectory Balance met flow-regularisatie voor snellere convergentie en betere prestaties dan eerdere GFlowNet-varianten.

Resultaten

De methode is getest op synthetische en complexe puzzels:

1. Swap Puzzle (Permutatie-omgevingen):

   • Getest op grafen met $n=15$ en $n=20$ elementen (respectievelijk $\approx 10^{12}$ en $\approx 10^{18}$ toestanden).
   • Het model generaliseerde uitstekend naar ongezichten toestanden. Na voldoende training vonden zowel greedy als beam search ($W=4$) exacte kortste paden voor alle testgevallen.

2. Rubik's Cube (2x2x2 en 3x3x3):

   • Vergelijking met de state-of-the-art methode CayleyPy Cube.
   • 2x2x2: Het model vond optimale oplossingen met een beam width van slechts 26, terwijl CayleyPy Cube faalde bij kleinere widths. Zelfs met greedy search ($W=1$) vonden ze geldige oplossingen.
   • 3x3x3: Beter presterend bij kleinere beam widths (tot $W=29$) en vergelijkbaar bij zeer grote widths.
   • Snelheid: Het model was aanzienlijk sneller (1.74s vs 6.19s per oplossing op een H200 GPU), ondanks een groter neuronaal netwerk (25M vs 4M parameters). Dit komt doordat het model in één forward pass de kansen voor alle buren berekent, terwijl concurrenten 12 keer meer evaluaties nodig hebben.

Betekenis en Conclusie

Dit werk positioneert niet-acyclische GFlowNets als een principieel en algemeen raamwerk voor het leren van kortste paden in discrete omgevingen. De belangrijkste doorbraak is het inzicht dat flow-minimalisatie direct leidt tot optimaliteit in padlengte.

De methode biedt een alternatief voor traditionele zoekalgoritmen en waarde-gebaseerde RL, met name in domeinen waar de staatruimte te groot is voor volledige exploratie. De resultaten op de Rubik's Cube tonen aan dat deze aanpak niet alleen competitief is in oplossingslengte, maar ook aanzienlijk efficiënter is in runtime, wat het veelbelovend maakt voor complexe planningstaken en combinatorische optimalisatieproblemen. Toekomstig werk richt zich op het uitbreiden naar gewogen grafen en nog grotere schaal.