UpSkill: Mutual Information Skill Learning for Structured Response Diversity in LLMs

Het artikel introduceert UpSkill, een trainingsmethode die wederzijdse informatie gebruikt om de diversiteit van antwoorden in grote taalmodellen te vergroten en zo de prestaties op meervoudige pogingen (pass@k) te verbeteren zonder de nauwkeurigheid van de eerste poging te verlagen.

De kern

UpSkill: Hoe we LLM's leren om niet steeds hetzelfde te zeggen

Stel je voor dat je een zeer slimme, maar soms een beetje stijve robot hebt die wiskundepuzzels oplost. Als je deze robot één keer vraagt: "Hoeveel was doet Raymond als Sarah 400 kilo doet?", geeft hij een goed antwoord. Maar als je hem tien keer dezelfde vraag stelt, geeft hij tien keer bijna hetzelfde antwoord. Het is alsof hij in een loop zit: hij denkt dat zijn eerste idee het allerbeste is en probeert het niet meer op een andere manier.

Dit is het probleem dat de onderzoekers van Princeton (Devan Shah en zijn team) willen oplossen. Ze hebben een nieuwe methode bedacht, genaamd UpSkill. Hier is hoe het werkt, vertaald in alledaagse taal:

1. Het Probleem: De "Eenzame Denker"

Normaal gesproken wordt een slimme taalmodel getraind om één keer het juiste antwoord te geven. Als hij dat doet, wordt hij beloond. Het gevolg? De robot wordt heel goed in één specifieke manier van denken.

• Vergelijking: Stel je voor dat je een kok hebt die alleen maar pasta kookt. Als je hem vraagt om een maaltijd te maken, maakt hij elke keer pasta, zelfs als je eigenlijk een salade of soep nodig had. Hij is goed in pasta, maar hij mist de variatie.
• In de praktijk: Bij complexe taken (zoals wiskunde of programmeren) is het vaak slim om een probleem op meerdere manieren aan te pakken. Als je model maar één manier kent, en die faalt, heb je pech. Als hij tien verschillende manieren kent, is de kans veel groter dat er één werkt.

2. De Oplossing: De "Magische Knoppen" (UpSkill)

UpSkill leert de robot om te denken met verschillende "strategieën". Ze doen dit door een klein, onzichtbaar knopje toe te voegen aan de vraag. Laten we dit knopje Z noemen.

• Z = 1: De robot denkt: "Oké, ik ga dit probleem oplossen alsof ik een algebraïsche wiskundeleraar ben."
• Z = 2: De robot denkt: "Oké, nu ga ik het oplossen alsof ik een visuele denker ben die met diagrammen werkt."
• Z = 3: "Nu ga ik het stap voor stap uitleggen als een verhaal."

De truc is: de robot leert tijdens het trainen dat deze knoppen echt verschillende manieren van denken moeten opleveren. Ze mogen niet allemaal op hetzelfde uitkomen.

3. De Beloning: "Hoe uniek ben je?"

Hoe leer je een robot om uniek te zijn? De onderzoekers gebruiken een slimme beloningssysteem gebaseerd op Mutuele Informatie (een ingewikkeld woord voor "hoe goed hangt het antwoord samen met het knopje dat je hebt gedrukt?").

• De oude manier: "Als het antwoord goed is, krijg je een sterretje." (Dit leidt tot saaie, identieke antwoorden).
• De UpSkill manier: "Als je antwoord goed is én het lijkt echt op de manier die bij knopje Z hoort, krijg je een extra grote sterretje!"

Als de robot probeert om bij knopje 1 en knopje 2 precies hetzelfde te zeggen, krijgt hij geen extra punten. Hij wordt dus beloond om echt verschillende "personages" te spelen.

4. Het Resultaat: Een Team van Experts

Na het trainen met UpSkill is de robot niet alleen slim, maar ook divers.

• Als je de robot 5 keer dezelfde vraag stelt, maar elke keer met een ander knopje (Z=1 tot Z=5), krijg je 5 verschillende oplossingsroutes.
• Zelfs als de eerste route faalt, is de kans groot dat route 3 of 5 wel werkt.
• Het mooie: De robot wordt niet minder goed in het vinden van het juiste antwoord (pass@1). Hij wordt juist beter in het vinden van een goed antwoord binnen een paar pogingen (pass@k).

Een Leuke Analogie: De Detektive

Stel je voor dat je een moord moet oplossen.

• Zonder UpSkill: Je hebt één detective die altijd dezelfde theorie heeft. Als die theorie fout is, is het raam dicht.
• Met UpSkill: Je hebt één detective die vijf verschillende "hoeden" kan dragen.
  • Met de rode hoed kijkt hij naar financiële motieven.
  • Met de blauwe hoed kijkt hij naar getuigen.
  • Met de groene hoed analyseert hij de vingerafdrukken.

Door de detective te dwingen om met elke hoed een andere conclusie te trekken, heb je een veel bredere dekking. Als de rode hoed niets oplevert, heb je misschien met de blauwe hoed het antwoord gevonden.

Samenvatting

UpSkill is een trainingstechniek die grote taalmodellen leert om niet in een "echo-kamer" van zichzelf te blijven hangen. Door ze te leren om met verschillende "strategieën" (zoals verschillende knoppen) te werken, krijgen we een model dat:

1. Minder vaak vastloopt.
2. Meer creatieve oplossingen vindt.
3. Beter presteert op taken waar je meerdere pogingen nodig hebt (zoals wiskunde en coderen).

Het is alsof je een team van vijf experts in plaats van één expert inzet, maar dan allemaal in één slimme robot verpakt.

Technische samenvatting

Probleemstelling

Grote Taalmodellen (LLMs) presteren uitstekend op verifieerbare redeneertaken, zoals wiskundige problemen oplossen en code genereren. Echter, een veelvoorkomend probleem bij het herhaaldelijk afnemen van antwoorden (sampling) is dat de modellen vaak zeer vergelijkbare of identieke output genereren. Dit fenomeen, vaak veroorzaakt door optimalisatie voor single-attempt accuracy (bijv. pass@1), leidt tot een gebrek aan diversiteit.

In scenario's waarbij meerdere pogingen worden gedaan om een probleem op te lossen (zoals bij het genereren van code met tests of formele bewijzen), is de kans op succes afhankelijk van de pass@k metric: de waarschijnlijkheid dat ten minste één van de $k$ pogingen correct is. Als de $k$ pogingen sterk gecorreleerd zijn (d.w.z. ze gebruiken dezelfde redeneerstrategie), verlaagt dit de effectieve hoeveelheid onafhankelijke pogingen. Bestaande methoden om diversiteit te vergroten, zoals temperatuur-sampling of prompt-perturbatie, zijn vaak onbetrouwbaar, vereisen handmatige tuning en leiden niet noodzakelijk tot semantisch verschillende oplossingspaden. Er is behoefte aan een trainingsmethode die gestructureerde diversiteit introduceert zonder de nauwkeurigheid van een enkele poging te verstoren.

Methodologie: UpSkill

De auteurs introduceren UpSkill, een trainingsmethode die Mutual Information Skill Learning (MISL) adapteert voor LLMs. Het doel is om een set van reproduceerbare, semantisch verschillende redeneerstrategieën te leren die geïndexeerd worden door een latente variabele $z$.

Kernconcepten:

1. Conditionering op Latente Variabelen: Tijdens het trainen wordt een discrete latente variabele $z \in \{1, \dots, N\}$ toegevoegd aan de input (bijvoorbeeld als een prefix: "Strategie {z} |"). Het model leert een beleid $\pi(\cdot | x, z)$ dat reageert op zowel de vraag $x$ als de specifieke strategie $z$.
2. Mutual Information (MI) Reward: Het centrale doel is het maximaliseren van de conditionele wederzijdse informatie $I(\tau; z | x)$, waarbij $\tau$ de gegenereerde trajecten (antwoorden) zijn. Dit wordt gedaan door een nieuwe token-level mutual information reward te definiëren:
   $$r_{TMI}(\tau_i; x, z) = \sum_{t=1}^{|\tau_i|} \left[ \log p_\pi(y_t | x, z, y_{<t}) - \log p_\pi(y_t | x, y_{<t}) \right]$$
   Hierbij is de eerste term de log-waarschijnlijkheid gegeven de specifieke strategie $z$, en de tweede term de log-waarschijnlijkheid onder een uniforme mix van alle strategieën. Deze reward moedigt het model aan om trajecten te genereren die specifiek zijn voor de gekozen $z$ (lage conditionele entropie), terwijl het de totale dekking van de oplossingsruimte behoudt (hoge marginale entropie).
3. Integratie met GRPO: De methode wordt geïmplementeerd binnen Group Relative Policy Optimization (GRPO). De totale beloning voor een traject is een combinatie van:
   • Een verifieerbaar correctheidsreward ($r_{corr}$).
   • De token-level MI reward ($r_{TMI}$).
   • Een KL-straf ($\Delta_{KL}$) om te voorkomen dat het model te ver afwijkt van het basismodel.

Inferentie:
Tijdens het gebruik (inference) worden $k$ verschillende waarden van $z$ geselecteerd en voor elk een antwoord gegenereerd. Omdat elke $z$ een distincte strategie vertegenwoordigt, zijn de $k$ antwoorden semantisch divers, wat de kans vergroot dat ten minste één antwoord correct is.

Belangrijkste Bijdragen

1. UpSkill Framework: Een nieuwe trainingsbenadering die gestructureerde response-diversiteit induceert zonder prompt-engineering, door gebruik te maken van een discrete latente variabele en een MI-gedreven reward.
2. Theoretische Link: De auteurs bewijzen theoretisch dat een verbetering in de pass@k score direct gerelateerd is aan de wederzijdse informatie $I(\tau; z | x)$. Ze tonen aan dat het maximaliseren van MI een ondergrens biedt voor de verbetering van pass@k.
3. Ongecontroleerde Diversiteit: Het bewijs dat UpSkill pass@k kan verbeteren zonder ground-truth antwoorden (alleen op basis van de MI-reward), hoewel de combinatie met correctheidsrewards het meest effectief is.
4. Empirische Validatie: Uitgebreide experimenten op drie open-weight modellen (Llama 3.1-8B, Qwen 2.5-7B, en R1-Distilled-Qwen2.5-Math-1.5B) op de GSM8K dataset.

Resultaten

De experimenten tonen significante verbeteringen aan, met name op de sterkere basismodellen:

• GSM8K (Wiskundige Probleemoplossing):
  • Voor Qwen 2.5-7B resulteerde UpSkill in een gemiddelde verbetering van +3,4% in pass@k en +9,1% in plurality@k, terwijl de pass@1 score behouden bleef.
  • Voor Llama 3.1-8B werden vergelijkbare verbeteringen waargenomen.
  • Opmerkelijk is dat de methode pass@k verbeterde zonder de pass@1 te verlagen, in tegenstelling tot wat vaak wordt verwacht bij diversiteitsmethodes.
• Rekenomgeving (Arithmetic Environment):
  • In een gecontroleerde omgeving met kleine modellen bleek dat standaard GRPO training leidt tot "entropy collapse" (alle strategieën worden identiek), terwijl UpSkill diverse strategieën behoudt. Hier steeg pass@5 van 0,793 naar 0,897, terwijl pass@1 lager bleef (0,390), wat aantoont dat de extra pogingen echt verschillende oplossingen leverden.
• Modelverschillen:
  • De methode werkte minder goed op het kleinere R1-Distilled-Qwen2.5-Math-1.5B model, waar prestaties daalden. Dit wordt toegeschreven aan de beperkte capaciteit van het model en de gevoeligheid voor de MI-objectief, wat leidt tot instabiliteit en "strategy collapse" (herhaling van nonsens).

Betekenis en Conclusie

UpSkill biedt een fundamentele verschuiving in hoe we diversiteit in LLMs benaderen. In plaats van te vertrouwen op heuristieken tijdens het genereren (inference-time), wordt diversiteit een ingebouwde eigenschap van het model tijdens het trainen.

• Praktische Impact: Voor toepassingen waar meerdere pogingen nodig zijn (zoals codegeneratie of wiskundige bewijzen), biedt UpSkill een betrouwbare manier om de succeskans te verhogen door het model te leren verschillende "denkmanieren" te gebruiken.
• Theoretische Inzichten: De paper legt een sterke theoretische brug tussen informatie-theoretische objectieven (Mutual Information) en de praktische prestatie-metric pass@k.
• Toekomst: De auteurs wijzen erop dat de methode gevoelig kan zijn voor modelgrootte en hyperparameters, en dat toekomstig werk gericht moet zijn op het stabiliseren van de training voor kleinere modellen en het verbeteren van de interpretatie van de geleerde strategieën.

Samenvattend introduceert UpSkill een robuuste, trainingsgebaseerde methode om de effectiviteit van LLMs in multi-attempt scenario's te maximaliseren door gestructureerde, reproduceerbare diversiteit te leren.