Beyond Token-Level Policy Gradients for Complex Reasoning with Large Language Models

Probleemstelling

Bestaande methoden voor het fine-tunen van autoregressieve taalkundige modellen (LLMs) met behulp van beleidsgradienten (zoals PPO, GRPO en DAPO) behandelen het genereren van elk volgend token als een afzonderlijke actie. Hoewel dit effectief is voor veel generatietaken, vertoont deze aanpak een fundamenteel tekortkoming bij complexe redeneertaken (zoals wiskundige problemen oplossen of code schrijven).

• Granulariteitsmismatch: In redeneertaken worden semantische beslissingen vaak genomen over blokken van meerdere tokens heen (bijvoorbeeld het definiëren van variabelen, het opstellen van vergelijkingen of het schrijven van een functie).
• Fragmentatie: Door elke token als een geïsoleerde actie te optimaliseren, wordt de samenhang van deze semantische blokken verstoord. Dit leidt tot een mismatch tussen de optimalisatie op token-niveau en de inherente blokniveau-natuur van redenering.
• Gevolg: De modellen kunnen moeite hebben om coherentie te behouden over langere redeneerstappen, wat de prestaties beperkt bij complexe taken.

Methodologie: Multi-token Policy Gradient Optimization (MPO)

Om deze kloof te overbruggen, stellen de auteurs Multi-token Policy Gradient Optimization (MPO) voor. Dit framework integreert blokniveau-acties in het proces van beleidsgradienten.

Kerncomponenten:

1. Blokniveau Acties: In plaats van één token per keer te behandelen, behandelt MPO sequenties van $K$ opeenvolgende tokens als één geïntegreerde semantische actie.
2. Multi-Token Prediction (MTP) Modules: Het model wordt uitgebreid met extra "heads" (gebaseerd op de architectuur van DeepSeek-V3) die in staat zijn om meerdere toekomstige tokens tegelijkertijd te voorspellen. Deze modules worden eerst "opgewarmd" (warm-up) met een supervised loss om de kwaliteit van de multi-token voorspellingen te garanderen.
3. Geaggregeerde Importance Sampling Ratio:
   • Traditionele PPO gebruikt de ratio voor één token: $r_t = \frac{\pi_\theta(a_t|s_t)}{\pi_{old}(a_t|s_t)}$.
   • MPO berekent een geaggregeerde ratio over een blok van $K$ tokens. Om de variantie van deze ratio te beheersen (wat hoog zou zijn bij een puur product van kansen), gebruiken de auteurs een gewogen log-sum benadering (geïnspireerd door Log-COP-TD):
     $$\tilde{R}^{(K)}_{i,t}(\theta) = \exp \left( \sum_{n=1}^{K} \beta_n \log r_{i,t+n}(\theta) \right)$$
     Hierbij zijn $\beta_n$ gewichten die afnemen naarmate de afstand tot het huidige token toeneemt (gecontroleerd door een decay-factor $\lambda$).
4. Doelfunctie: De surrogate loss wordt aangepast om deze blokniveau-ratio's te gebruiken in combinatie met de geschatte voordelen (advantages), waardoor het model wordt gestimuleerd om coherentere patronen te leren in plaats van geïsoleerde token-bewegingen.

Belangrijkste Bijdragen

1. Nieuw Framework (MPO): Een beleidsgradienten-methode die optimalisatie uitvoert over gestructureerde redeneereenheden (blokken van tokens) in plaats van geïsoleerde tokens.
2. Integratie in Post-Training: Het toont aan dat het integreren van structurele multi-token optimalisatie in de post-training-fase (na pre-training) de coherentie van redenering verbetert, zelfs voor modellen die niet oorspronkelijk met multi-token loss zijn getraind.
3. Empirisch Bewijs: MPO overtreft consistent bestaande token-level baselines (PPO, GRPO, DAPO) op benchmarks voor wiskundig redeneren en code-generatie.
4. Stabiliteit en Variansie: Het paper analyseert hoe de integratie van multi-token informatie de variantie van de importance sampling ratio verlaagt en de training stabiliseert, waardoor een betere bias-variance trade-off wordt bereikt.

Resultaten

De auteurs hebben MPO getest op drie modellen (Llama3.2-1B, DeepSeek-Distilled-Qwen2.5-1.5B en 7B) en drie benchmarks:

• GSM8K (college-niveau wiskunde)
• MATH (competitie-niveau wiskunde)
• HumanEval (code-generatie)

Kernbevindingen:

• Prestatieverbetering: MPO behaalde hogere scores (pass@1) dan alle baselines op alle modellen en datasets. Bijvoorbeeld, op de MATH-benchmark met het 1.5B-model steeg de nauwkeurigheid van 0.553 (GRPO) naar 0.632 (MPO).
• Stabiliteit: Analyse toont aan dat MPO de variantie van de importance sampling ratio verlaagt en het percentage "clipping" (waarbij updates worden beperkt vanwege te grote veranderingen) vermindert.
• Efficiëntie: Hoewel MPO extra parameters introduceert (voor de MTP-modules), is de rekentijd per iteratie slechts ongeveer 30% hoger dan bij PPO, terwijl de geheugennood vergelijkbaar blijft met GRPO. De warm-up fase is verwaarloosbaar kort (ca. 12 minuten) vergeleken met de totale trainingstijd.
• Ablatie Studies: Er is een optimale grootte voor het blok ($K=5$) en een optimale weging van de toekomstige tokens gevonden. Te veel gewicht op multi-token informatie kan leiden tot instabiliteit door ruis in de voorspellingen.

Betekenis en Toekomstperspectief

Dit paper markeert een methodologische verschuiving in de post-training van LLMs. Het onderstreept dat voor complexe redeneertaken de granulariteit van beleidslering moet worden verlegd van het token-niveau naar het semantische blok-niveau.

• Implicatie: Het suggereert dat toekomstig onderzoek zich moet richten op het modelleren van redenering als een reeks samenhangende stappen in plaats van een opeenvolging van losse symbolen.
• Beperkingen: De auteurs erkennen uitdagingen zoals het schalen van de blokgrootte (grotere $K$ vereist meer rekenkracht) en het vinden van betere warm-up strategieën voor modellen die niet oorspronkelijk met MTP zijn getraind.
• Conclusie: MPO biedt een praktische en effectieve manier om de redeneercapaciteiten van LLMs te verbeteren door de structuur van het redeneerproces beter af te stemmen op de optimalisatiemethode.