Future-KL Regularized GRPO: Process-Level Credit Assignment from $f$-Divergence Regularization

Technische Samenvatting: Future-KL Geregulariseerde GRPO

Probleemstelling
Group Relative Policy Optimization (GRPO) is een dominant paradigma geworden voor critic-vrije versterkingsleer (RL) in het natrekken van Large Language Models (LLM), met name voor wiskundig redeneren. GRPO elimineert de noodzaak van een geleerde waardenfunctie door uitkomstbeloningen te normaliseren binnen een groep van bemonsterde antwoorden. Echter, standaardimplementaties van KL-regularisatie in GRPO behandelen de straf doorgaans als een lokale, verlieszijdige tokenstraf (bijvoorbeeld met gebruik van de $k_3$-schatter).

De auteurs identificeren een fundamenteel gebrek in deze aanpak: het negeert de procesniveau-toewijzing van krediet die wordt veroorzaakt door autoregressieve generatie. In een autoregressieve trajectorie incasseert een vroege token niet alleen zijn eigen lokale regularisatiekosten, maar verandert hij ook de verdeling van toekomstige voorvoegsels, waardoor de KL-kosten van alle daaropvolgende tokens worden beïnvloed. Standaard lokale verlies termen differentiëren alleen de straf van de bemonsterde token, en laten de causale "toekomstige-regularisatie-opbrengst tot einde" (future-regularization return-to-go) weg.

Bovendien benadrukt het artikel een structurele incompatibiliteit tussen de groepsnormalisatie van GRPO en standaard KL-in-beloning formuleringen. In tegenstelling tot lineaire beloningsdoelen (bijvoorbeeld MaxEnt RL), waarbij beloning en regularisator kunnen worden samengevoegd, induceert de groepsnormalisatie van GRPO een niet-lineaire nuttigheid op prompt-niveau (specifiek, een arcsin-transformatie voor binaire beloningen). Deze niet-lineariteit betekent dat het toevoegen van KL-straffen aan beloningen vóór normalisatie het impliciete doel vervormt, token-niveau KL-informatie reduceert tot een enkel traject-niveau scalair, en faalt in het behouden van het juiste gradiëntsignaal.

Methodologie
Het artikel leidt de on-policy policy gradient af voor GRPO-achtige doelen onder algemene token-weise $f$-divergentie regularisatie. Het belangrijkste theoretische inzicht is dat de gradient decomposeert in twee distincte termen met verschillende schalingseigenschappen:

1. De Beloningsterm: Geërfd van de niet-lineaire nuttigheid $\ell(p)$, draagt deze term een prompt-afhankelijke winst $\ell'(p)$. Voor binaire beloningen herstelt dit de gestandaardiseerde GRPO-voordeel.
2. De Regularisatorterm: Deze term is onafhankelijk van $\ell'$ en decomposeert in een lokale tokenstraf plus een causale toekomstige-regularisatie-opbrengst tot einde.

Gebaseerd op deze afleiding stellen de auteurs Future-KL Geregulariseerde Policy Optimization (FRPO) voor. Het algoritme wijzigt de standaard GRPO-update op een ontkoppelde manier:

1. Bereken de groeps-genormaliseerde uitkomstvoordeel ($\hat{A}^{GRPO}$) strikt vanuit de originele verifiërbeloningen, waarbij de impliciete niet-lineaire nuttigheid behouden blijft.
2. Construeer een token-niveau toekomstige-KL correctieterm door de omgekeerde cumulatieve som van per-token log-ratio's ($\delta_{i,j} = \log \frac{\pi_\theta}{\pi_{ref}}$) te berekenen van de huidige token tot het einde van de trajectorie.
3. Voeg deze toekomstige-KL term toe aan het voordeel na de voordeelconstructie, in plaats van de beloning te vormen vóór normalisatie.

De updateregel voor een token $t$ in antwoord $i$ is:
$$\hat{A}^{FRPO}_{i,t} = \text{sg}\left( \hat{A}^{GRPO}_i - \beta \sum_{j=t}^{T_i} \delta_{i,j} \right)$$
waarbij $\text{sg}$ stop-gradient aanduidt. Deze aanpak vereist geen critic, geen extra modelpasses, en behoudt dezelfde rekenkundige overhead als standaard GRPO.

Belangrijkste Bijdragen

1. Identificatie van Niet-Lineaire Nuttigheid: Het artikel toont aan dat de groepsnormalisatie van GRPO een niet-lineaire nuttigheid induceert ($\ell(p) = 2 \arcsin \sqrt{p}$ voor binaire beloningen), waardoor beloning en KL niet samenvoegbaar zijn. Het toevoegen van KL aan beloningen vóór normalisatie vervormt het doel structureel.
2. Afleiding van On-Policy Gradiënt: De auteurs leiden de correcte on-policy gradiënt af voor token-weise $f$-divergentie regularisatie, en tonen aan dat deze een toekomstige-regularisatie-opbrengst tot einde-term vereist die afwezig is in lokale verlieszijdige implementaties.
3. FRPO Algoritme: Zij stellen FRPO voor, een eenvoudige, ontkoppelde update die het uitkomstvoordeel van GRPO behoudt terwijl de nodige toekomstige-KL krediettoewijzing wordt ingebracht. Deze methode scheidt traject-niveau ranking van procesniveau KL-regularisatie.
4. Optimaliteitsinterpretatie: De toekomstige-KL structuur blijkt te aligneren met de optimaliteitsvoorwaarden van entropie-geregulariseerde RL, waarbij het optimale beleid een proces-beloningsvorm voldoet die toekomstige KL-kosten omvat.

Experimentele Resultaten
De auteurs evalueren FRPO op wiskundige redeneertaken met behulp van Qwen3 basismodellen (1.7B, 4B, en 30B-A3B) getraind op het DAPO-Math-17k dataset.

• Prestatie: Op het 30B-A3B model verbetert FRPO pass@16 met bijna 5% absolute winst ten opzichte van de sterkste baseline (standaard GRPO met verlieszijdige KL) op de MATH500 benchmark. Het presteert ook beter dan baselines op AIME24, AIME25 en andere Olympiade-niveau benchmarks.
• Stabiliteit en Drift: FRPO behoudt een hogere uiteindelijke entropie en lagere beleid-drift (gemeten door PPO-KL) in vergelijking met verlieszijdige KL baselines. Dit suggereert dat de methode lange Chain-of-Thought (CoT) redeneerpatronen oproept zonder overmatige afwijking van het referentiebeleid.
• Ablatie: Experimenten bevestigen dat reverse KL (mode-seeking) beter presteert dan forward KL (mode-covering) voor deze taak. Het artikel merkt ook op dat het gebruik van de $k_1$-schatter (log-ratio) de voorkeur verdient boven $k_3$ voor de toekomstige-KL term, aangezien de niet-negativiteit van $k_3$ een eenzijdig strafsignaal creëert dat krediettoewijzing hindert.
• Foutmodi: Het artikel observeert dat "KL-in-beloning" benaderingen (het toevoegen van KL aan beloningen vóór normalisatie) leiden tot modelinstorting of prestatiedegradatie, wat de theoretische claim valideert dat dergelijke fusie structureel incorrect is onder de niet-lineaire nuttigheid van GRPO.

Betekenis
Het artikel betoogt dat de betekenis van dit werk ligt in het corrigeren van een subtiel maar kritiek implementatiefout in een veelgebruikt RL-kader. Door het perspectief te verschuiven van "lokale verliesstraf" naar "procesniveau krediettoewijzing", tonen de auteurs aan dat GRPO kan worden verbeterd zonder complexe waardenfuncties of extra rekenkosten in te voeren. De voorgestelde FRPO-methode biedt een theoretisch onderbouwde, lichtgewicht modificatie die redeneercapaciteiten verbetert terwijl de stabiliteit en generalisatie-eigenschappen van het basismodel behouden blijven. Het werk onderstreept dat in autoregressieve settings regularisatie inherent een procesniveau signaal is, en het negeren van de toekomstige kosten tot einde leidt tot suboptimale krediettoewijzing.