Rooted Absorbed Prefix Trajectory Balance with Submodular Replay for GFlowNet Training

Dit paper introduceert RapTB en SubM, twee methoden die GFlowNets effectiever maken voor het genereren van diverse en valide moleculen door zwakke krediettoewijzing en replay-bias te verhelpen.

De kern

Stel je voor dat je een kunstenaar bent die een heel groot, complex schilderij moet maken, maar je hebt geen idee hoe het eindresultaat eruit moet zien. Je krijgt alleen een prijs (een "beloning") als het schilderij klaar is én als het mooi is.

In de wereld van kunstmatige intelligentie (AI) heet dit een GFlowNet. Het is een slimme manier om AI te leren niet alleen één perfect schilderij te maken, maar een heleboel verschillende mooie schilderijen te creëren.

Het probleem is echter: de AI wordt vaak lui of verward. Ze begint steeds hetzelfde stukje van het schilderij te maken (de "mode collapse") en stopt te vroeg of te laat. Ze vindt maar één manier om het goed te doen, in plaats van alle mooie mogelijkheden te ontdekken.

De auteurs van dit paper hebben twee nieuwe trucs bedacht om dit op te lossen: RapTB en SubM. Laten we ze uitleggen met een verhaal.

1. Het Probleem: De Verwarde Kunstenaar

Stel je voor dat de AI een lange zin schrijft (of een chemisch molecuul ontwerpt).

• Het probleem met oude methoden: De AI krijgt pas aan het einde van de zin te horen of het goed was. Als ze halverwege een fout maakt, weet ze dat pas aan het eind. Daardoor raakt ze in de war: "Moet ik nu stoppen? Moet ik doorgaan? Moet ik precies hetzelfde doen als de vorige keer?"
• Het gevolg: De AI begint steeds dezelfde korte, saaie zinnen te schrijven die net goed genoeg zijn, en vergeten de lange, creatieve, maar misschien nog betere opties.

2. De Eerste Oplossing: RapTB (De Slimme Mentor)

De eerste truc heet RapTB. Stel je voor dat je in plaats van alleen aan het eind te oordelen, een slimme mentor hebt die je bij elke stap een klein hintje geeft.

• Hoe het werkt: Normaal gesproken krijgt de AI alleen een score aan het einde. RapTB pakt die eindscore en "verspreidt" deze terug naar elke stap die de AI eerder heeft gezet.
• De Analogie: Het is alsof je een lange reis maakt. In plaats van pas bij de bestemming te horen of je de juiste route hebt gekozen, zegt de mentor bij elke afslag: "Goed zo, deze weg leidt waarschijnlijk naar een mooie plek, want we hebben daar later een mooi uitzicht gezien."
• Het resultaat: De AI krijgt veel meer feedback onderweg. Ze leert sneller welke beginwoorden (prefixen) goed zijn en welke niet, zonder in de war te raken over wanneer ze moet stoppen. Ze blijft creatief en maakt minder fouten.

3. De Tweede Oplossing: SubM (De Slimme Verzamelaar)

De tweede truc heet SubM. Dit gaat over de verzameling van voorbeelden waar de AI van leert.

• Het oude probleem: Stel je voor dat de AI een mapje met voorbeelden heeft. Als ze een keer een heel mooi schilderij maakt, slaat ze dat op. De volgende keer kijkt ze alleen naar dat ene mooie schilderij en probeert ze dat na te bootsen. Ze vergeet alle andere leuke, maar net iets minder perfecte opties. Dit heet "replay bias". De AI wordt een kloon van zichzelf.
• De nieuwe methode (SubM): SubM is als een curator in een museum. Als er nieuwe schilderijen bij komen, kiest de curator niet alleen de "mooiste" uit. Hij kijkt ook naar:
  1. Is het mooi? (Beloning)
  2. Is het anders dan wat we al hebben? (Diversiteit)
  3. Is het een ander formaat? (Lengte)
• De Analogie: Als je alleen de beste 100 foto's van een vakantie in je album plakt, maar ze zijn allemaal van hetzelfde strand, heb je geen overzicht van je hele reis. SubM zorgt ervoor dat je ook foto's van de bergen, de stad en het bos in je album plakt, zolang ze maar van goede kwaliteit zijn.
• Het resultaat: De AI leert van een veel breder scala aan voorbeelden. Ze wordt niet alleen beter in het vinden van de "beste" oplossing, maar ook in het vinden van veel verschillende goede oplossingen.

Samenvatting: Waarom is dit belangrijk?

Door deze twee methoden samen te gebruiken (RapTB voor betere feedback onderweg en SubM voor een diversere leeromgeving), kan de AI:

1. Minder vastlopen: Ze raakt niet vast in één saaie oplossing.
2. Beter leren: Ze begrijpt beter welke stappen leiden tot een goed resultaat.
3. Meer variatie: Ze kan bijvoorbeeld veel verschillende, nieuwe medicijnen ontwerpen of creatieve zinnen schrijven, in plaats van steeds hetzelfde te herhalen.

Kortom: Het is alsof je een kunstenaar niet alleen een prijs geeft aan het einde, maar haar ook een slimme gids geeft die haar onderweg helpt, en een museumcurator die zorgt dat ze van de hele wereld leert, niet alleen van één hoekje.

Technische samenvatting

1. Het Probleem: Mode Collapse in LLM-GFlowNets

Generative Flow Networks (GFlowNets) zijn ontworpen om stochastische beleidsplannen te leren waarbij de kans op het genereren van een object evenredig is met de beloning (reward). Wanneer dit wordt toegepast op Large Language Models (LLMs) voor het genereren van sequenties (zoals moleculen of wiskundige expressies), ondervinden ze echter last van mode collapse. De auteurs identificeren twee specifieke vormen van dit falen:

1. Prefix Collapse: De entropie van vroege tokens daalt scherp. Verschillende eindresultaten (terminals) delen bijna identieke beginvoorvoegsels (prefixes), wat leidt tot een gebrek aan diversiteit in de vroege generatiestappen.
2. Length Bias: Het model vertoont een systematische voorkeur voor sequenties die te kort of te lang zijn, in plaats van de optimale lengte te vinden.

De auteurs attribueren deze problemen aan twee fundamentele oorzaken:

• Zwakke Credit Assignment: Beloningen worden vaak pas aan het einde van een traject gegeven. Dit zorgt voor hoge variantie en ambiguïteit bij het toewijzen van "krediet" aan tussenliggende stappen.
• Replay Bias: Tijdens het trainen wordt vaak gebruikgemaakt van een replay buffer die alleen de hoogst beloonde trajecten bewaart. Dit leidt tot een verschuiving in de trainingsverdeling, waarbij het model over-optimaliseert op een smal subset van de zoekruimte, wat de mode collapse verergert.

2. Methodologie: RapTB en SubM

Om deze problemen aan te pakken, stellen de auteurs twee complementaire mechanismen voor: RapTB voor betere credit assignment binnen trajecten, en SubM voor het verbeteren van de diversiteit in de replay buffer.

A. RapTB (Rooted Absorbed Prefix Trajectory Balance)

RapTB is een nieuwe trainingsdoelfunctie die de bestaande Trajectory Balance (TB) verbetert door twee innovaties:

1. Rooted Prefix Supervision: In plaats van alle mogelijke subtrajecten te construeren (zoals bij Subtrajectory Balance of SubTB), focust RapTB alleen op subtrajecten die bij de wortel ($s_0$) beginnen. Dit elimineert conflicterende randvoorwaarden die ontstaan door overlappende vensters met verschillende startpunten, wat vaak leidt tot instabiliteit in de stop-kansen (termination probabilities).
2. Absorbed Suffix Rewards: Om de hoge variantie van eindbeloningen te verminderen, "absorbeert" RapTB informatie uit het toekomstige deel van het traject (suffix) om een dichter en betrouwbaarder signaal te creëren voor vroege prefixes.
   • Het berekent een doelwaarde voor een prefix $s_{0:k}$ door de beloningen van latere stappen ($j \ge k$) te aggregeren.
   • Dit gebeurt via een combinatie van een max-backup (ondergrens) en een soft-backup (gladde aggregatie in log-ruimte met een straffing voor afstand).
   • Cruciaal: De gradients van de stop-kans ($\log q_\theta(\top|s)$) worden in dit hulpdoel losgekoppeld (stop-gradient) om te voorkomen dat het model de lengte-bias oplost door simpelweg de stop-kans te verschuiven in plaats van de token-generatie te verbeteren.

B. SubM (Submodular Replay)

SubM is een strategie om de replay buffer te verversen die niet alleen kijkt naar de hoogste beloning, maar ook naar diversiteit en lengte-coverage.

• Het selecteert een subset van trajecten uit de huidige buffer en nieuwe batches door een submodulaire functie te maximaliseren.
• Deze functie bestaat uit drie termen:
  1. Kwaliteit/Validiteit: Een statische score per sample.
  2. Facility-Location Coverage: Een diversiteitsmaat die zorgt dat de buffer goed de hele ruimte van mogelijke oplossingen dekt (gebaseerd op similariteit, zoals Tanimoto voor SMILES).
  3. Length Coverage: Een concave functie die zorgt voor een evenwichtige verdeling over verschillende sequentielengtes.
• Dit voorkomt dat de buffer "verzadigt" met bijna-identieke hoog-beloonde samples (rich-get-richer dynamiek).

3. Belangrijkste Bijdragen

• Empirische Karakterisering: De auteurs tonen aan dat mode collapse in LLM-GFlowNets een reproduceerbare combinatie is van prefix collapse en length bias, gedreven door hoge variantie in credit assignment en replay-induced distribution shifts.
• RapTB: Een nieuwe objective die TB combineert met wortel-gebaseerde prefix constraints en suffix-absorptie. Dit verhoogt de dichtheid van trainingssignalen zonder de instabiliteit van SubTB te introduceren.
• SubM: Een submodulaire replay-strategie die beloning, diversiteit en lengte-coverage in één doelwit balanceren, wat leidt tot een robuustere trainingsverdeling.
• Analyse van Termination Drift: De paper identificeert dat bestaande methoden (zoals SubTB) vaak leiden tot "termination drift" (het model stopt te vroeg of te laat) omdat ze proberen overlappende vensters te voldoen door de stop-kans aan te passen. RapTB lost dit op door de stop-gradients in de hulp-branch te blokkeren.

4. Resultaten

De methoden zijn getest op drie taken:

1. Scaffold-geconditioneerde SMILES generatie (moleculaire ontdekking).
2. Expr24 (genereren van wiskundige expressies die 24 opleveren).
3. CommonGen (concept-naar-zin generatie).

Kernresultaten:

• RapTB + SubM presteert consequent beter dan de baselines (TB en SubTB) op het gebied van kwaliteit-diversiteit trade-off.
• Validiteit: RapTB behoudt een hoge validiteit (bijv. >98% voor SMILES), terwijl SubTB vaak instabiel is en veel ongeldige moleculen genereert door lengte-bias.
• Diversiteit: De combinatie van RapTB en SubM leidt tot een significante stijging in diversiteit (gemeten via Entropy en FPDiv) en een betere dekking van de zoekruimte, zelfs voor lange trajecten.
• Prefix Collapse: RapTB vermindert de concentratie op vroege prefixes aanzienlijk, wat resulteert in een bredere vertakking van generaties.
• Lengte Bias: In tegenstelling tot SubTB, dat vaak extreem lange of korte sequenties genereert, behoudt RapTB een kalibratie van de stop-kans die dicht bij de natuurlijke verdeling ligt.

5. Betekenis en Impact

Deze paper biedt een oplossing voor een fundamenteel probleem in het trainen van GFlowNets met autoregressieve modellen (LLMs). Door de instabiliteit van credit assignment en de neiging tot mode collapse aan te pakken, maakt RapTB het mogelijk om GFlowNets effectiever in te zetten voor complexe zoekproblemen zoals:

• Geneesmiddelenontdekking: Het genereren van diverse, chemisch valide moleculen met hoge eigenschappen.
• Symbolische Redenering: Het vinden van diverse oplossingen voor wiskundige problemen.

De combinatie van RapTB (voor interne stabiliteit en betere leerkracht) en SubM (voor externe exploratie en diversiteit) stelt een nieuwe standaard voor robuust autoregressief GFlowNet-training, waarbij de balans tussen exploratie en exploitatie beter wordt beheerd zonder in te leveren op de kwaliteit van de gegenereerde data.