SPO++: Stream-Aligned Policy Optimization for Asynchronous Agentic RL

In het snel evoluerende veld van kunstmatige intelligentie leren onderzoekers computerprogramma's om complexe problemen op te lossen door ze te laten proberen, falen en leren van de resultaten. Dit proces, bekend als reinforcement learning (versterkend leren), is bijzonder nuttig voor het trainen van grote taalmodellen om te fungeren als agenten die hulpmiddelen kunnen gebruiken, virtuele omgevingen kunnen navigeren of wiskundige problemen kunnen oplossen. Een centrale uitdaging in dit werk is hoe succes te meten wanneer het pad naar een antwoord lang en onvoorspelbaar is. Traditionele methoden vereisen vaak dat het systeem veel verschillende pogingen aan hetzelfde probleem tegelijkertijd genereert, waarbij wordt gewacht tot elke afzonderlijke poging is voltooid voordat er van de resultaten geleerd kan worden. Dit creëert een bottleneck: als één poging lang duurt omdat het vastloopt of veel hulpmiddelen probeert, pauzeert het hele leerproces in afwachting van die traagste poging om in te halen.

Om deze inefficiëntie op te lossen, werd een nieuwere aanpak ontwikkeld genaamd Single-stream Policy Optimization. In plaats van te wachten op een groep pogingen, stelt deze methode het systeem in staat om te leren van één poging tegelijk, gebruikmakend van een persistent geheugen van eerdere successen en mislukkingen om toekomstige beslissingen te sturen. Echter, een team van onderzoekers onder leiding van Kai Ruan en Jinghao Lin ontdekte dat hoewel deze methode de wachttijd wegnam, het een subtiele maar significante discrepantie introduceerde in de manier waarop de computer zijn voortgang berekende. Ze ontdekten dat de manier waarop het systeem zijn beloningen gemiddelde, niet overeenkwam met de manier waarop het daadwerkelijk de stappen van een oplossing verwerkte. Door deze afstemming te corrigeren, creëerden ze een verbeterde versie van de methode, die ze SPO++ noemden, waardoor kunstmatige intelligentie aanzienlijk sneller en efficiënter kan leren.

De kern van het probleem lag in de manier waarop de computer de lengte van zijn reacties afhandelde. In de oorspronkelijke methode berekende het systeem een enkele score voor een hele poging, zoals of een robot erin slaagde een kamer schoon te maken of of een wiskundige oplosser het juiste antwoord vond. Het verspreidde vervolgens deze enkele score over elk woord of "token" dat het model tijdens die poging genereerde. Dit leek logisch, maar het creëerde een vertekening. Als één poging erg lang was en een andere kort, zou de lange poging de score verduetten over veel woorden, terwijl de korte poging de score concentreerde op minder woorden. Wanneer het systeem probeerde te leren van deze scores, veranderde de lengte van de reactie deels de centrale focus van het leren, waardoor het model optimaliseerde voor de verkeerde zaken. Het was alsof het systeem een weegschaal probeerde in evenwicht te houden, maar de gewichten op de weegschaal verschoven afhankelijk van hoeveel items erop geplaatst werden, in plaats van de waarde van de items zelf.

De onderzoekers identificeerden twee specifieke gebieden waar deze afstemming niet klopte. Ten eerste hield het systeem de tijd bij waarop het de resultaten van een poging ontving, in plaats van wanneer de poging daadwerkelijk werd gegenereerd. In een asynchroon systeem waarbij taken worden verzonden en voltooid met verschillende snelheden, is de volgorde waarin resultaten arriveren vaak willekeurig en afhankelijk van de netwerksnelheid of de computerbelasting. De oorspronkelijke methode gebruikte deze aankomstvolgorde om zijn geheugen bij te werken, wat betekende dat het leersignaal werd beïnvloed door de timing van het computersysteem in plaats van door de logica van de taak zelf. Ten tweede, en nog kritieker, hield de methode van het middelen van de scores geen rekening met het feit dat het model leert van elk woord dat het genereert, en niet alleen van de uiteindelijke uitkomst. De onderzoekers realiseerden zich dat om het leerproces te herstellen, ze de scores moesten standaardiseren op basis van het aantal gegenereerde actiewoorden, om ervoor te zorgen dat het beloningssignaal overeenkwam met het werkelijke volume aan werk dat het model verrichtte.

Om deze problemen aan te pakken, introduceerde het team SPO++, wat twee belangrijke wijzigingen aanbrengt in het leerproces. Ten eerste hebben ze het geheugensysteem geherorganiseerd om de "policy event" te volgen, oftewel het specifieke moment waarop een verzoek werd verzonden, in plaats van wanneer het resultaat terugkwam. Dit zorgt ervoor dat het geheugen van het systeem over eerdere successen gekoppeld is aan de staat van het model op het moment dat de taak werd aangemaakt, waardoor het leersignaal consistent blijft, ongeacht hoe lang de taak duurde om te voltooien. Ten tweede hebben ze de manier waarop de scores worden gemiddeld gewijzigd. In plaats van elke poging als een enkele eenheid te behandelen, berekent de nieuwe methode de gemiddelde score op basis van het totale aantal gegenereerde actiewoorden over alle pogingen heen. Dit zorgt ervoor dat het leersignaal perfect is afgestemd op de manier waarop het model zijn kennis woord voor woord bijwerkt.

De resultaten van deze wijzigingen werden getest op twee verschillende soorten taken: een set van 128 gesimuleerde huishoudelijke taken waarbij een agent objecten naar specifieke locaties moet verplaatsen, en een dataset van 1.500 wiskundige problemen die het gebruik van een Python-calculator vereisen. De onderzoekers voerden experimenten uit met twee verschillende groottes van taalmodellen, één met 0,8 miljard parameters en één met 2 miljard parameters. In elke enkele test leerde de nieuwe methode, SPO++, sneller dan de oorspronkelijke aanpak. Bij de taak met de huishoudelijke klusjes was de verbetering aanzienlijk, waarbij de nieuwe methode een aanzienlijk hogere totale beloning behaalde gedurende de training. Bij de wiskundige problemen waren de winsten kleiner maar nog steeds consistent, wat aantoonde dat het model sneller een hoger prestatieniveau bereikte. De onderzoekers ontdekten dat het krachtigste deel van de nieuwe methode de verandering in de manier waarop de scores werden gemiddeld was, wat alleen al de meerderheid van de verbetering verantwoordde.

Deze bevindingen suggereren dat in de complexe wereld van het trainen van kunstmatige intelligentie-agenten, de details van hoe gegevens worden verwerkt net zo belangrijk kunnen zijn als de gegevens zelf. Door ervoor te zorgen dat de manier waarop een systeem zijn voortgang meet, overeenkomt met de manier waarop het daadwerkelijk leert, kunnen onderzoekers aanzienlijke winsten in efficiëntie boeken zonder dat daar meer rekenkracht of grotere modellen voor nodig zijn. Het werk laat zien dat zelfs kleine afwijkingen in het leerproces de vooruitgang kunnen vertragen, en dat het corrigeren van deze mismatches het systeem in staat stelt om zijn energie te richten op het oplossen van het probleem in plaats van op het compenseren voor een gebrekkige meting. Terwijl kunstmatige intelligentie steeds moeilijkere en gevarieerdere taken aanpakt, bieden methoden zoals SPO++ een duidelijker pad voorwaarts, waarbij ze ervoor zorgen dat elke stap die het model zet nauwkeurig wordt gemeten en effectief bijdraagt aan de groei.