Partial Policy Gradients for RL in LLMs

Dit paper introduceert een methode voor reinforcement learning in LLMs waarbij het optimaliseren van een subset van toekomstige beloningen leidt tot betrouwbaardere schattingen en het vergelijken van verschillende policy-classes, zoals volledige planning en K-stap vooruitkijken, voor conversatieproblemen.

De kern

Samenvatting: Hoe je een AI kunt leren om niet te vergeten wie hij is

Stel je voor dat je een acteur hebt die een rol moet spelen in een heel lang toneelstuk. Hij moet de hele avond dezelfde persoon zijn: dezelfde achtergrond, dezelfde voorkeuren en dezelfde gevoelens. Het probleem met de huidige slimme computers (LLMs) is dat ze na een tijdje de draad kwijtraken. Ze beginnen te vergeten dat ze een leraar zijn, veranderen hun mening over hun favoriete eten, of zeggen dingen die in strijd zijn met wat ze eerder zeiden. Dit noemen de auteurs "persona drift" (het wegdrift van het personage).

Deze paper introduceert een slimme nieuwe manier om deze AI's te trainen, zodat ze trouw blijven aan hun rol, zelfs in gesprekken van urenlang. De techniek heet Partial Policy Gradients (Gedeeltelijke Beleidsgradienten), maar laten we het simpel houden: het gaat over hoe ver je in de toekomst moet kijken terwijl je leert.

Hier is de uitleg in drie simpele concepten:

1. Het probleem: Kijken naar alleen de volgende stap of de hele film?

Stel je voor dat je een AI traint om een gesprek te voeren.

• De oude manier (Greedy): De AI kijkt alleen naar de volgende zin. "Als ik nu deze zin zeg, krijg ik een complimentje?" Het is alsof je een film kijkt en alleen kijkt naar het volgende frame. Je weet niet wat er over 10 minuten gebeurt. Het gevolg: de AI zegt iets slimme nu, maar vergeten later dat ze dat al eerder hadden gezegd. Ze "flip-floppen" (wisselen heen en weer) en raken de draad kwijt.
• De volledige manier (Full Planning): De AI probeert de hele film van begin tot eind in gedachten te houden. "Als ik nu dit zeg, hoe beïnvloedt dat mijn rol over 50 minuten?" Dit is heel goed, maar het is ook enorm zwaar voor de computer. Het kost veel data en rekenkracht. Als je niet genoeg voorbeelden hebt, raakt de AI in de war en faalt hij compleet.

2. De oplossing: De "K-stap vooruit" bril

De auteurs zeggen: "Waarom niet een middenweg?"
Stel je voor dat je een bril opzet die je K stappen vooruit laat kijken.

• K=1: Je kijkt alleen naar de volgende stap (te kortzichtig).
• K=100: Je kijkt naar de hele film (te zwaar).
• K=3 of K=5: Je kijkt een paar stappen vooruit. "Als ik nu zeg dat ik een visser ben, moet ik onthouden dat ik over 3 zinnen nog steeds een visser ben."

Dit is de kern van hun methode: Optimaliseren voor een subset van de toekomst.
Door de AI te laten leren op basis van een beperkt aantal toekomstige stappen (bijvoorbeeld de volgende 3 of 4 zinnen), wordt het leren veel makkelijker en betrouwbaarder. De AI krijgt een duidelijker signaal over wat goed is, zonder overweldigd te worden door de hele toekomst.

3. De verrassende ontdekking: Er is geen "beste" bril voor iedereen

De paper toont aan dat de beste "K" (het aantal stappen vooruit) afhangt van het soort gesprek:

• In een therapiegesprek: Mensen hebben complexe gevoelens. Hier werkt een K=3 (kijken 3 stappen vooruit) het beste. Het is genoeg om de emotionele lijn vast te houden, zonder te ver te plannen.
• In een chatgesprek: Dit is vaak losjes en reactief. Hier werkt K=2 (kijken 2 stappen vooruit) het beste. Te veel plannen maakt het gesprek stijf en onnatuurlijk.
• In een onderwijsgesprek (tutoring): Hier moet je echt lang plannen. Een leerling bouwt kennis op over tijd. Hier werkt volledig plannen (K=alles) het beste, omdat je de hele leertraject in de gaten moet houden.

Waarom is dit belangrijk? (De "Statistische Efficiency")

De auteurs gebruiken een mooi beeld: Het is makkelijker om een klein doel te bereiken dan een groot doel.

• Als je een AI traint met weinig data (weinig voorbeelden), moet je een simpel doel geven (kijk maar 1 of 2 stappen vooruit). Dan leert hij snel en goed.
• Als je veel data hebt, kun je de AI een complexer doel geven (kijk 10 stappen vooruit).

Het is alsof je een kind leert fietsen.

1. Begin met een fiets met zijwieltjes (K=1 of K=2). Het kind leert snel balanceren zonder te vallen.
2. Als het kind goed is, haal je de zijwieltjes weg en geef je fietsles voor lange tochten (K=volledig).
   Als je een kind direct een lange tocht laat fietsen zonder zijwieltjes (te veel data, te complex), valt het kind waarschijnlijk en leert het niets.

Conclusie

Deze paper leert ons dat we AI's niet altijd moeten laten plannen voor de "eeuwigheid". Soms is het slimmer om ze te laten focussen op de korte tot middellange termijn. Door de juiste "blikrichting" te kiezen voor het juiste type gesprek, kunnen we AI's maken die trouw blijven aan hun rol, niet vergeten wie ze zijn, en veel betrouwbaarder zijn in lange gesprekken.

Kortom: Geef de AI een bril die past bij de taak, en hij zal nooit meer de draad kwijtraken.

Technische samenvatting

Titel: Partial Policy Gradients for RL in LLMs

Auteurs: Puneet Mathur, Branislav Kveton, Subhojyoti Mukherjee, Viet Dac Lai (Adobe Research)

---

1. Het Probleem

Versterkend leren (Reinforcement Learning - RL) is een raamwerk om sequentieel te handelen in een onbekende omgeving met als doel toekomstige beloningen te maximaliseren. Bij Large Language Models (LLMs) wordt RL vaak gebruikt om gedrag te optimaliseren, zoals het behouden van een consistent menselijk persona in rolspelgesprekken.

De huidige uitdagingen zijn:

• Persona-drift: Off-the-shelf LLMs wijken vaak af van hun toegewezen persona, tegenstrijdige uitspraken doen of rol-appropriaat gedrag verlaten, vooral in lange gesprekken (gemiddeld 20 tot 60 stappen).
• Statistische inefficiëntie: Bestaande algoritmen zoals PPO (Proximal Policy Optimization) en GRPO (Group Relative Policy Optimization) attribueren beloningen vaak gelijkmatig aan alle tokens of vereisen complexe per-token beloningsmodellen. Dit leidt tot hoge variantie in de gradiënt-schattingen, wat het leren van complexe beleidsstrategieën (zoals volledige planning) onbetrouwbaar maakt, vooral bij beperkte trainingsdata.
• Trade-off: Er is een fundamenteel compromis tussen de complexiteit van het geleerde beleid (hoe ver men in de toekomst pland) en de statistische efficiëntie (hoe betrouwbaar het beleid geleerd kan worden met de beschikbare data).

2. Methodologie: Partial Policy Gradients (PPG)

De auteurs stellen een natuurlijk raamwerk voor om beleidsstructuur in te voeren in policy gradients door te optimaliseren voor een subset van toekomstige beloningen.

Kernconcepten:

1. Factoring van Beloningen: De totale beloning $r(x, \tau_n)$ wordt additief ontbonden in per-stap beloningen $r_t$.
2. Attributie van Beloningen: In plaats van elke stap $a_t$ te koppelen aan de totale toekomstige beloning (zoals in standaard Policy Gradients), koppelt PPG een actie $a_t$ alleen aan een subset van toekomstige beloningen $R_t$.
   • De gradiënt wordt geschat als: $E [\sum_{t=1}^n (\sum_{\ell \in R_t} r_\ell) \nabla \log \pi(a_t | \dots)]$.
   • Hierbij is $R_t$ de subset van toekomstige stappen wier beloningen beïnvloed worden door actie $a_t$.

Beleidstypen binnen het raamwerk:

• Full Planning (PG): $R_t$ omvat alle toekomstige stappen. Dit is complex maar vereist veel data.
• Greedy (GreedyPG): $R_t$ omvat alleen de directe stap ($t$). Dit is statistisch zeer efficiënt maar neigt tot "myopisch" gedrag.
• K-Step Lookahead (K-Step-PG): $R_t$ omvat de komende $K$ stappen. Dit biedt een tussenweg: het pland voor de korte tot middellange termijn, wat leidt tot stabielere resultaten dan greedy, maar met lagere variantie dan volledige planning.

Theoretische Onderbouwing:
De auteurs bewijzen (via Hoeffding's ongelijkheid) dat het optimaliseren voor een kleinere subset van beloningen leidt tot een gradiënt-schatting die sneller convergeert (minder variantie). Dit betekent dat eenvoudigere beleidsstrategieën (kleinere $K$) betrouwbaarder kunnen worden geleerd met minder data.

Implementatie:
Het paper introduceert zowel online als offline algoritmen. Voor de experimenten wordt een offline policy gradient gebruikt, waarbij trajecten worden gesampleerd uit een gelogde dataset (uniforme verdeling) en de beloning wordt berekend op basis van persona-consistentie.

3. Belangrijkste Bijdragen

1. Universeel Raamwerk: Een generalisatie van bestaande methoden (zoals adaptive submodular policies en segment-level credit assignment) naar een uniek raamwerk voor het optimaliseren van subsets van toekomstige beloningen.
2. Statistische Efficiëntie Trade-off: Het inzicht dat het beperken van de credit assignment horizon (K) de statistische efficiëntie verhoogt, waardoor complexere beleidsstrategieën (zoals K-step lookahead) betrouwbaarder kunnen worden geleerd dan volledige planning in data-scarce omgevingen.
3. Eerste Empirische Evaluatie van K-Step Lookahead in LLMs: Het paper introduceert en evalueert K-step lookahead beleidsstrategieën specifiek voor LLMs, wat een nieuwe richting is in de RL voor taalmodellen.
4. Empirische Validatie: Uitgebreide experimenten op vier domeinen (onderwijs, therapie, chatten, generiek) met drie verschillende LLM-architecturen (Qwen, Llama, Gemma).

4. Resultaten

De experimenten zijn uitgevoerd op de Consistent-LLMs benchmark, waarbij het doel is om persona-consistentie (Persona Consistency - PC) te maximaliseren in lange dialogen.

Kernbevindingen:

• Superioriteit boven Baselines: Alle voorgestelde PPG-methoden presteerden beter dan de onbewerkte base-modellen en de state-of-the-art PPO-baseline.
• Domein-afhankelijke Optimaliteit: Er is geen "one-size-fits-all" oplossing; de beste $K$ hangt af van de complexiteit van het domein:
  • Onderwijs: Vereist Full Planning (PG). Leertrajecten vereisen lange-termijn pedagogische strategieën.
  • Therapie: 3-Step Lookahead is optimaal. Het balanceert emotionele consistentie zonder over-planning die tot onrealistische herstelverhalen leidt.
  • Chatten: 2-Step Lookahead is optimaal. Conversaties zijn reactief en vereisen minimale planning.
• Stabiliteit: K-step lookahead beleidsstrategieën (vooral $K=2$ en $K=3$) voorkomen persona-drift en tonen stabiele prestaties over lange trajecten (tot 60 stappen), terwijl greedy beleidsstrategieën oscilleren en base-modellen monotoon degraderen.
• Statistische Efficiëntie:
  • Bij weinig data (1% van de trainingset) presteert GreedyPG het beste omdat het eenvoudiger is om te leren.
  • Bij moderate data domineren K-step lookahead methoden.
  • Bij veel data kan Full Planning de beste prestaties leveren, maar alleen als er voldoende data is om de hoge variantie te overwinnen.
• Generalisatie: De bevindingen zijn consistent over verschillende modelarchitecturen (Llama, Qwen, Gemma).

5. Betekenis en Impact

Dit paper biedt een fundamenteel nieuw perspectief op het trainen van LLMs met versterkend leren:

• Praktische Richtlijn: Het biedt een praktische ontwerpprincipe voor practitioners: kalibreer de complexiteit van het beleid (de lookahead horizon $K$) op basis van de beschikbare trainingsdata en de complexiteit van de taak. Begin met greedy methoden bij beperkte data en schaal op naar K-step of full planning naarmate de dataset groeit.
• Oplossing voor Persona-drift: Het biedt een effectieve methode om het veelvoorkomende probleem van persona-drift in lange conversaties op te lossen, wat cruciaal is voor toepassingen zoals virtuele therapeuten, tutors en chatbots.
• Theoretische Verdieping: Het verbindt de theorie van concentratie-ongelijkheden met praktische RL-toepassingen in LLMs, en toont aan dat het "plannen voor de toekomst" niet altijd lineair beter is; soms is een beperkte horizon statistisch superieur.

Samenvattend introduceert dit paper een flexibele en statistisch efficiënte aanpak voor RL in LLMs, waarbij het optimaliseren van een subset van toekomstige beloningen leidt tot robuustere en consistentere modellen, afhankelijk van de data-omgeving en de specifieke toepassing.