🤖 machine learning

$f$ -Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data

Dit artikel introduceert de $f$ -Trajectoriebalans, een familie van verliesfuncties die de methode van de gemiddelde kwadratische fout uitbreidt naar alle $f$ -divergenties, waardoor het trainen van generatieve modellen en LLM's met off-policy data mogelijk wordt terwijl de specifieke optimalisatie-eigenschappen (zoals het dekken van modi) van hun corresponderende on-policy $f$ -divergentie-gradiënten behouden blijven.

Oorspronkelijke auteurs: Jake Fawkes, Jason Hartford

Gepubliceerd 2026-05-18

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Jake Fawkes, Jason Hartford

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een robot leert een uitgestrekt, donker grottenstelsel te verkennen. De grot heeft vele verborgen kamers (zogenaamde "modi"), waarvan sommige goud bevatten (hoge beloningen) en anderen gewoon leeg of gevaarlijk zijn. Je doel is de robot te leren om zo veel mogelijk verschillende goudkamers te vinden, in plaats van slechts één grote goudmijn te vinden en de rest te negeren.

Dit artikel introduceert een nieuwe reeks "leermiddelen" (verliesfuncties) om deze robots te trainen, of ze nu moleculen voor nieuwe medicijnen genereren, kunst creëren of code schrijven.

Hier is de uiteenzetting van de ideeën uit het artikel, met gebruik van eenvoudige analogieën:

1. Het Probleem: De "Goudkoorts" versus de "Ontdekkingsreiziger"

In het verleden was de standaardmanier om deze robots te leren, vergelijkbaar met een Goudkoorts. De robot zou één plek met een beetje goud vinden, enthousiast worden en vervolgens al zijn tijd besteden aan het delven daar. Het zou alle andere goudmijnen in de grot negeren.

In technische termen: Dit heet "mode-seeking". Het model stort in op een paar antwoorden met hoge waarschijnlijkheid en negeert de rest van de diversiteit.
Het doel van het artikel: We willen "mode-covering". We willen dat de robot zich verspreidt en alle goudmijnen verkent, zelfs de kleinere, om een complete kaart van de grot te krijgen.

2. Het Oude Hulpmiddel: De "Kwadratische KL"-Verliesfunctie

Onlangs ontdekten onderzoekers een slimme truc om robots te leren met een methode op basis van "kwadratische fout" (het meten van de afstand tussen wat de robot denkt en wat het moet denken).

De Analogie: Stel je een leraar voor die een student beoordeelt. Als de student het antwoord fout heeft, zegt de leraar niet alleen "Fout". Ze berekenen het kwadraat van de fout.
Het Voordeel: Deze methode is zeer stabiel. Het werkt zelfs als de student leert van oude notities (off-policy data) in plaats van van live lessen.
Het Gebrek: Hoewel het stabiel is, gedraagt deze specifieke "kwadratische" methode zich nog steeds als een Goudkoorts. Het duwt de robot er natuurlijk toe om zich te focussen op slechts een paar antwoorden, waardoor de diversiteit van de grot wordt gemist.

3. De Nieuwe Oplossing: De "f-Trajectbalans"-Familie

De auteurs realiseerden zich dat de "kwadratische" methode slechts één specifieke variant was van een veel grotere familie van leermiddelen. Ze noemen deze familie f-Trajectbalans.

Beschouw deze familie als een draaiknop of een volumeknop op een radio.

De draaiknop naar de ene kant draaien (Lage getallen): Je krijgt middelen die fungeren als een Super Ontdekkingsreiziger. Deze middelen dwingen de robot om zich te verspreiden en elke mogelijke goudmijn te vinden, zelfs de moeilijk bereikbare. Dit is uitstekend voor medicijnontwikkeling, waar je elk molecuul wilt vinden dat mogelijk werkt, niet alleen het meest voor de hand liggende.
De draaiknop naar de andere kant draaien (Hoge getallen): Je krijgt middelen die fungeren als een Goudzoeker. Deze middelen vertellen de robot om de kleine mijnen te negeren en zich intens te focussen op de grootste, meest voor de hand liggende goudhoop. Dit is nuttig als je gewoon het beste enkele antwoord wilt.
Het Midden: Je kunt de draaiknop ergens in het midden instellen om een mix van exploratie en focus te krijgen.

4. Waarom Dit Een Grote Zaak Is

Het artikel bewijst twee hoofdzaakken:

De "Magische Schakelaar": Ze toonden aan dat je de standaard "kwadratische" tool kunt vervangen door een van deze nieuwe "draaiknop"-tools, en de wiskunde werkt nog steeds perfect. De robot leert even goed, maar dan met een ander persoonlijkheid (meer explorerend of meer gefocust).
Stabiliteit: Net als het oude hulpmiddel werken deze nieuwe middelen zelfs als de robot leert van oude data (off-policy). Dit is cruciaal voor toepassingen in de echte wereld zoals Large Language Models (LLM's), waar training vaak asynchroon plaatsvindt (de robot leert van data die een tijdje geleden is gegenereerd).

5. Tests in de Wereld (De Grotkaarten)

De auteurs testten deze middelen in drie verschillende "grotten":

Het Raster Spel: Een eenvoudig computerraam met vier hoeken vol goud. Het standaardmiddel vond slechts één hoek. De nieuwe "Ontdekkingsreiziger"-middelen vonden alle vier de hoeken snel.
Molecuulontdekking (SynFlowNet): Ze probeerden nieuwe chemische moleculen te genereren. Door de draaiknop op "Ontdekkingsreiziger" te zetten, genereerden ze een veel bredere variëteit aan unieke moleculen die potentieel medicijnen konden zijn, in plaats van slechts variaties van dezelfde paar chemicaliën.
Taalmodellen (LLM's): Ze stelden AI-modellen in om wiskundeproblemen op te lossen. De nieuwe middelen stonden de AI toe om verschillende manieren om problemen op te lossen te verkennen zonder vast te komen te zitten in een lus van hetzelfde antwoord herhalen, zelfs wanneer de trainingsdata vertraagd was (asynchroon).

Samenvatting

Dit artikel bedenkt geen nieuw type robot. In plaats daarvan bedenkt het een nieuwe reeks instructies voor hoe ze getraind moeten worden.

Oude manier: "Vind de grootste goudmijn en graaf daar." (Stabiel, maar saai).
Nieuwe manier: "Hier is een draaiknop. Je kunt kiezen om een Goudzoeker te zijn, een Super Ontdekkingsreiziger, of iets daartussenin. En ongeacht wat je kiest, de training zal stabiel zijn en werken met oude data."

Dit geeft ingenieurs een eenvoudige knop om te controleren hoe creatief of gefocust hun AI-modellen moeten zijn, zonder het trainingsproces te breken.

Technische Samenvatting: f-Trajectbalans

Probleemstelling

De uitlijning en fijnafstelling van generatieve modellen, waaronder Large Language Models (LLMs) en Generatieve Flow-netwerken (GFlowNets), vertrouwen doorgaans op het minimaliseren van het verschil tussen de verdeling van een model en een doelverdeling. Bij Reinforcement Learning (RL)-fijnafstelling wordt dit vaak geformuleerd als het maximaliseren van de verwachte beloning, onderworpen aan een Kullback-Leibler (KL)-divergentiestraf.

Een aanzienlijke uitdaging in dit domein is de afweging tussen mode-zoekend en mode-dekkend gedrag. Standaarddoelstellingen, zoals de omgekeerde KL-divergentie (vaak geïmplementeerd via Trajectbalans of een kwadratische KL-verliesfunctie), neigen naar mode-zoekend gedrag, waardoor modellen instorten op een paar hoogstwaarschijnlijke modi. Dit is suboptimaal voor taken die exploratie vereisen, zoals geneesmiddelenontwikkeling of het genereren van diverse kandidaten, waarbij "mode-dekking" (het vastleggen van de volledige diversiteit van de doelverdeling) de voorkeur heeft.

Bovendien worstelen bestaande methoden vaak met off-policy validiteit. Hoewel schatters zoals REINFORCE of PPO gebruikelijk zijn, lijden ze onder hoge variantie. Recentere benaderingen die gebruikmaken van de kwadratische KL-verliesfunctie ( $KL_{sq}$ ) bieden lage-variantie gradiënten on-policy en validiteit off-policy, maar ze zijn inherent gebonden aan de eigenschappen van de omgekeerde KL-divergentie, wat de mogelijkheid beperkt om de afweging tussen exploratie en exploitatie te controleren.

Methodologie

De auteurs stellen f-Trajectbalans voor, een familie van surrogate-verliesfuncties die de kwadratische KL-verliesfunctie generaliseert tot de volledige familie van f-divergenties.

Theoretische Grondslag

De kerntheoretische bijdrage is het vaststellen van een één-op-één-correspondentie tussen translatie-invariante verliesfuncties op log-kansen en f-divergenties.

Van Divergentie naar Verlies: Voor elke convexe functie $f$ die een f-divergentie $D_f(p_\theta \| p^*)$ definieert, leiden de auteurs een verliesfunctie $L_f$ af die is gedefinieerd op het verschil van log-kansen $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$ :
$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$
Dit verlies heeft twee kritieke eigenschappen:
- On-Policy Gradiënt-equivalentie: Wanneer steekproeven worden getrokken uit het modelbeleid ( $\mu = p_\theta$ ), komen de verwachte automatisch gedifferentieerde gradiënten van $L_f$ overeen met de gradiënten van de bijbehorende f-divergentie $D_f$ .
- Off-Policy Validiteit: Wanneer geëvalueerd op off-policy data (geprobeerd uit een gedragsverdeling $\mu$ ), blijft het verlies een geldig doel met dezelfde globale minimalizer ( $p_\theta = p^*$ ), mits de ondersteuning van het doel is opgenomen in de ondersteuning van het gedragsbeleid.
Van Verlies naar Divergentie: Omgekeerd komt elke strikt convexe, translatie-invariante verliesfunctie $\ell$ op log-kansen overeen met het minimaliseren van een specifieke f-divergentie. Het artikel bewijst dat deze koppelingen elkaars inverse zijn.

Praktische Implementaties

Om praktische uitdagingen bij het trainen aan te pakken, introduceren de auteurs verschillende uitbreidingen:

DevGrad-verlies (Batchgewijze Normalisatie): Om onoplosbare partitiefuncties aan te pakken (vaak voorkomend bij GFlowNets en genormaliseerde doelen), generaliseren de auteurs de VarGrad-schatter. In plaats van batchvariantie te gebruiken, maken ze gebruik van de batchgewijze gegeneraliseerde afwijking van de log-kansverschillen. Dit houdt in dat de normalisatieconstante $\log Z$ wordt geschat door het verlies over de batch te minimaliseren, wat effectief de coëfficiënten van de scorefunctie centreert en de variantie verlaagt.
Getemperd Verlies: Om numerieke overloop te voorkomen bij het werken met kleine temperatuurparameters $\beta$ (waarbij de doelverdeling $p^* \propto \pi_{ref} \exp(r/\beta)$ is), stellen de auteurs een getemperd verlies voor. Dit schaalt het verlies met $1/\beta$ en past het toe op de getemperde verdeling, waardoor gradiënten begrensd en stabiel blijven.
f-Trajectbalans voor GFlowNets: Het kader wordt toegepast op GFlowNets door trajectbalansbeperkingen af te dwingen met behulp van de $L_f$ -verliesfunctie. Dit generaliseert de standaard Trajectbalans (die overeenkomt met omgekeerde KL) naar willekeurige f-divergenties.

Belangrijkste Bijdragen

Gegeneraliseerde Verliesfamilie: Afleiding van een algemene familie van translatie-invariante surrogate-verliezen ( $L_f$ ) waarvan de on-policy gradiënten overeenkomen met die van de bijbehorende f-divergentie, terwijl off-policy validiteit behouden blijft.
Theoretische Equivalentie: Bewijs van de inverse relatie tussen translatie-invariante log-kansverliezen en f-divergenties, wat een systematisch ontwerp van nieuwe verliesfuncties met gewenste eigenschappen (bijv. mode-dekking) mogelijk maakt.
Uitbreiding tot Ongenormaliseerde Doelen: Generalisatie van de VarGrad-techniek naar willekeurige f-divergenties via het DevGrad-verlies, waarmee het probleem van de onoplosbare partitiefunctie bij RL-fijnafstelling wordt aangepakt.
Controle van Exploratie: Demonstratie dat de keuze van de f-divergentie (geparametriseerd door $\alpha$ in $\alpha$ -divergenties) direct het mode-zoekende versus mode-dekkende gedrag controleert. Lagere $\alpha$ -waarden (bijv. Forward KL, Hellinger) bevorderen mode-dekking, terwijl hogere waarden (bijv. Pearson $\chi^2$ ) mode-zoekend gedrag bevorderen.

Experimentele Resultaten

De auteurs evalueren de verliesfamilie over vier verschillende taken:

Synthetisch Rooster (GFlowNets): Op een 2D-roostertaak ontworpen om modedekking te testen, ontdekten verliezen met lagere $\alpha$ (Forward KL, Hellinger) alle vier modi van de doelverdeling aanzienlijk sneller en betrouwbaarder dan standaard Trajectbalans (Reverse KL). Omgekeerd vertoonde Pearson $\chi^2$ ( $\alpha=2$ ) sterk mode-zoekend gedrag en bleef het vastzitten in de eerste gevonden modus.
Molecuulontdekking (SynFlowNet): Toegepast op het SynFlowNet-kader voor het genereren van synthetisch toegankelijke moleculen. De auteurs ontdekten dat het afstemmen van $\alpha$ controle bood over de diversiteit van gegenereerde moleculen. Een afkoelschema (beginnen met lagere $\alpha$ voor exploratie en toenemen naar $\alpha=1$ voor exploitatie) leverde de beste afweging op, waarbij moleculen met hogere beloning werden gegenereerd met grotere diversiteit in vergelijking met standaard Trajectbalans.
Diffusiemodellen: Bij conditionele steekproeven van MNIST-cijfers (oneven versus even) faalde standaard Trajectbalans om gelijkmatig te bemonsteren over de doelmodi (oversampling van specifieke cijfers zoals 0 en 6). f-Trajectbalansverliezen, met name die met lagere $\alpha$ , bereikten een meer uniforme bemonstering over de doelmodi.
Asynchrone LLM-fijnafstelling: De verliezen werden toegepast om LLM's (Qwen- en OLMo-families) te fijnafstellen op wiskundeproblemen (GSM8k, MATH) met een asynchrone vertraging van 50 stappen. De f-divergentieverliezen toonden stabiel trainen en consistente afwegingen tussen entropie en beloning, terwijl standaard PPO-implementaties instabiliteit vertoonden onder dezelfde asynchrone omstandigheden.

Betekenis en Claims

Het artikel stelt dat de effectiviteit van de kwadratische KL-verliesfunctie geen geïsoleerd fenomeen is, maar een specifiek voorbeeld van een bredere familie van verliezen. Door de eigenschap van "on-policy gradiëntmatching" en "off-policy validiteit" uit te breiden naar de volledige familie van f-divergenties, bieden de auteurs een unificerend kader voor het fijnafstellen van generatieve modellen.

De primaire betekenis ligt in het vermogen om de afweging tussen exploratie en exploitatie expliciet te controleren door simpelweg een andere surrogate-verliesfunctie te selecteren (of een $\alpha$ -parameter af te stemmen), zonder het onderliggende trainingsalgoritme te wijzigen. Dit stelt praktici in staat om te kiezen tussen mode-zoekend (voor beloningsmaximalisatie) en mode-dekkend (voor diversiteit en exploratie) gedrag, terwijl de stabiliteit en off-policy compatibiliteit behouden blijven die nodig zijn voor moderne, grootschalige asynchrone training van LLM's en GFlowNets. Het werk suggereert dat dit kader bijzonder waardevol is voor toepassingen zoals geneesmiddelenontwikkeling en exploratieve agenten, waarbij het bestrijken van de ruimte van hoogbelonende kandidaten cruciaal is.