Beyond Normalization: Rethinking the Partition Function as a Difficulty Scheduler for RLVR

Technische Samenvatting: Voorbij Normalisatie: Het Herdenken van de Partitiefunctie als een Moeilijkheidsplanner voor RLVR

Probleemstelling
Versterkt Leren met Verifieerbare Beloningen (RLVR) heeft bewezen effectief te zijn voor het verbeteren van de redeneercapaciteiten van Grootte Taalmodellen (LLM's). Echter, standaardmethoden voor beloningsmaximalisatie zoals PPO en GRPO leiden vaak tot overmatig scherpe uitverdelingen, wat de generatiediversiteit vermindert. Hoewel recente benaderingen die GFlowNets gebruiken dit aanpakken door beleidsplannen te trainen om een door beloningen geïnduceerde doelpool te volgen, behandelen ze de geleerde partitiefunctie ($Z$) doorgaans uitsluitend als een normalisatieconstante die nodig is voor trainingsstabiliteit. Dit artikel betoogt dat dit perspectief de informatie die in $Z$ is gecodeerd, onderbenut, waardoor een kans wordt gemist om de steekproefficiëntie te verbeteren zonder bijkomende rekenkosten.

Methodologie: PACED-RL
De auteurs stellen Partition Function-Guided RL (PACED-RL) voor, een post-training framework dat de GFlowNet-partitiefunctie niet alleen interpreteert als een normalisator, maar als een signaal voor online nauwkeurigheid per prompt. De methodologie bestaat uit drie kerncomponenten:

1. Theoretische Herinterpretatie van de Partitiefunctie:
   Het artikel legt een theoretisch verband tussen de optimale partitiefunctie $Z^*(x)$ en de online nauwkeurigheid $p_{old}(x)$ van een vraagprompt $x$ onder het voor-update-beleid. Door de Trajectory Balance (TB)-verliesfunctie te modificeren zodat de KL-divergentie wordt verankerd bij het huidige beleid ($\pi_{old}$) in plaats van bij een vast referentiemodel, leiden de auteurs af dat $Z^*(x)$ direct nauwkeurigheidsinformatie codeert. Specifiek geldt dat $p_{old}(x) \approx \beta \log Z^*(x)$, waarbij de term voor KL-divergentie empirisch wordt aangetoond verwaarloosbaar te zijn tijdens standaard trainingsdynamiek. Hierdoor kan de partitiefunctie, die al wordt geleerd voor normalisatie, dienen als een gratis, online schatter voor nauwkeurigheid.

2. Adaptieve Promptselectie:
   Door gebruik te maken van de nauwkeurigheidsschattingen afgeleid van $Z_\phi(x)$, implementeert PACED-RL een moeilijkheidsbewuste promptselectiestrategie. In plaats van prompts uniform te bemonsteren, selecteert de methode gretig de top-$m$ prompts waarvan de geschatte nauwkeurigheden het dichtst bij een doelwaarde $\tau$ liggen (doorgaans 0,5). Dit richt de training op monsters van "intermediaire moeilijkheidsgraad", wat eerdere studies suggereren de meest informatieve leersignalen bieden, waardoor de steekproefficiëntie wordt gemaximaliseerd.

3. Prioritering van Replay op Basis van Fout in Nauwkeurigheidschatting:
   Om de efficiëntie verder te verhogen en de partitiefunctie te kalibreren, introduceert het framework een replay-mechanisme. Een buffer slaat prompt-uitgangsparen op met correcte beloningen ($r=1$). Binnen deze buffer worden monsters geprioriteerd op basis van de grootte van hun fout in de nauwkeurigheidschatting: $| \frac{N_{correct}}{N} - \hat{p}_{old}(x) |$. Dit zorgt ervoor dat het model monsters opnieuw bezoekt waar de voorspelling van de partitiefunctie het sterkst afwijkt van de waargenomen realiteit, waardoor de nauwkeurigheidschattingen en de normalisatie van de doelpool worden verfijnd.

Cruciaal is dat zowel de promptselectie als de prioritering van replay informatie hergebruiken die wordt gegenereerd tijdens standaard GFlowNet-training, waardoor de rekenoverhead wordt afgeschreven op het bestaande optimalisatieproces.

Belangrijkste Bijdragen

• Theoretisch Inzicht: Het artikel toont aan dat de GFlowNet-partitiefunctie voor LLM-post-training intrinsiek verbonden is met online nauwkeurigheden, waardoor deze direct kan worden gebruikt als nauwkeurigheidschatter zonder bijkomende mechanismen.
• PACED-RL Framework: De introductie van een nieuw trainingsframework dat deze nauwkeurigheidschattingen gebruikt voor adaptieve promptselectie en replay met foutprioritering, waardoor de steekproefficiëntie wordt verbeterd terwijl de diversiteitsvoordelen van GFlowNets behouden blijven.
• Empirische Validatie: Uitgebreide evaluaties op benchmarks voor wiskundig redeneren en codegeneratie tonen aan dat PACED-RL aanzienlijke prestatiewinst behaalt ten opzichte van sterke baselines (GRPO, FlowRL) en adaptieve bemonsteringsmethoden (DS, LILO, MoPPS), zonder de bijbehorende rekenoverhead van over-bemonstering.

Experimentele Resultaten
Experimenten werden uitgevoerd op codegeneratie (DeepCoder-dataset, DeepSeek-R1-Distill-Qwen-1.5B) en wiskundig redeneren (DeepScaleR-dataset, Qwen2.5-Math-1.5B/7B).

• Prestatiewinst: Op de AIME 24-benchmark verbeterde PACED-RL de gemiddelde pass@1-prestatie met 29,1% ten opzichte van GRPO en met 40,0% ten opzichte van FlowRL, gebruikmakend van het Qwen2.5-Math-7B-model.
• Steekproefficiëntie: PACED-RL bereikte de beste prestaties van GRPO en FlowRL in aanzienlijk minder tijd (bijvoorbeeld 0,42× de trainingstijd op LiveCodeBench). In tegenstelling tot methoden die over-bemonstering vereisen (DS, LILO), behield PACED-RL een rollout-ratio van 1,0× ten opzichte van standaard training, waardoor de latentieboetes die gepaard gaan met het genereren van extra monsters voor nauwkeurigheidschatting worden vermeden.
• Diversiteit: Gemeten via pass@k-metrics, presteerde PACED-RL consistent beter dan baselines, met verbeteringen tot 14,2% ten opzichte van GRPO en 9,1% ten opzichte van FlowRL, wat bevestigt dat de methode de outputdiversiteit en het exploratievermogen behoudt.
• Ablatiestudies: De auteurs verifieerden dat het weglaten van de KL-divergentieterm in de nauwkeurigheidschatter een verwaarloosbaar effect heeft op de bemonsteringsdynamiek, terwijl er aanzienlijke rekenkracht wordt bespaard. Ze toonden ook aan dat de methode de training op natuurlijke wijze verschuift naar moeilijkere vragen naarmate het model verbetert, waardoor een instorting naar een smal dataset-subset wordt voorkomen.

Betekenis en Beweringen
Het artikel beweert dat PACED-RL een "principiële en praktische manier" biedt om trainingsmethoden voor distributiematching voor LLM's te verbeteren. Door de partitiefunctie te herinterpreteren als een signaal voor online nauwkeurigheid, maakt de methode efficiëntere training mogelijk zonder de rekenlast van bijkomende nauwkeurigheidschatters of overmatige rollout-generatie. De auteurs positioneren dit werk als een stap naar efficiëntere RLVR-post-training, met de suggestie dat de rol van de partitiefunctie verder reikt dan normalisatie en fungeert als een dynamische moeilijkheidsplanner. Voor toekomstig werk wordt gesuggereerd inferentiestrategieën (bijvoorbeeld adaptieve zelfconsistentie) te verkennen en het framework uit te breiden naar asynchrone GFlowNet-omgevingen.