Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity

1. Probleemstelling

Reinforcement Learning with Verifiable Rewards (RLVR) heeft de nauwkeurigheid bij één poging (Pass@1) van Large Language Models (LLM's) bij redeneertaken zoals wiskunde aanzienlijk verbeterd. Een kritiek neveneffect is echter geïdentificeerd: Diversiteitsinstorting (Diversity Collapse).

• Het Fenomeen: Hoewel RLVR de kans vergroot dat een correct antwoord wordt gegenereerd, vermindert het vaak drastisch Pass@K (de kans dat ten minste één van $K$ gesamplede antwoorden correct is). Dit geeft aan dat de oplossingsruimte van het model krimpt, waarbij de waarschijnlijkheidsmassa zich concentreert op een smal subset van correcte oplossingen, terwijl andere geldige redeneerpaden worden onderdrukt.
• De Oorzaak: De auteurs identificeren een structurele beperking in standaard RLVR-doelstellingen (zoals GRPO - Group Relative Policy Optimization). Deze doelstellingen zijn onverschillig voor hoe de waarschijnlijkheidsmassa onder correcte oplossingen is verdeeld.
  • Wiskundig gezien levert elk beleid dat 100% van zijn massa op de set van correcte uitvoer $Y^+$ plaatst, dezelfde maximale beloning op, ongeacht of die massa geconcentreerd is op één oplossing (instorting) of uniform verdeeld over alle geldige oplossingen (divers).
  • In combinatie met on-policy samplingdynamiek creëert deze onverschilligheid een zichzelf versterkende feedbacklus: oplossingen die iets vaker worden gesampled, ontvangen grotere gradiëntupdates, wat hun waarschijnlijkheid en samplingfrequentie verder verhoogt, wat leidt tot een onomkeerbare modale instorting.

2. Methodologie

A. Theoretische Analyse

De auteurs formaliseren het instortingsmechanisme en definiëren de optimale beleidsstructuur:

1. Instortingsmechanisme: Ze bewijzen dat onder GRPO samplingbias leidt tot ongelijke sampleaantallen. Gradiëntupdates schalen met deze aantallen, waardoor de waarschijnlijkheid van dominante oplossingen wordt versterkt en die van minderheidssoluties wordt onderdrukt. Zodra een oplossing onder een samplingdrempel valt ($\pi(y) \ll 1/K$), ontvangt het geen gradiëntsignaal en kan het niet herstellen.
2. Definitie van Optimaal Beleid: De auteurs stellen dat het optimale beleid het Uniform-Correct Beleid (UCP) moet zijn. Dit beleid voldoet aan twee criteria:
   • Robuustheid: Het maximaliseert de worst-case behouden waarschijnlijkheidsmassa als een subset van correcte oplossingen wordt uitgeschakeld (bijvoorbeeld door een verschuiving in de verdeling).
   • Entropie-geregulariseerde Optimaliteit: Het maximaliseert uniek het doel van het balanceren van correctheid met diversiteit (entropie) binnen de correcte set.
   • Definitie: Het UCP plaatst nul waarschijnlijkheid op incorrecte uitvoer en verdeelt de waarschijnlijkheid uniform over alle correcte oplossingen ($\pi(y|x) = 1/|Y^+|$ voor $y \in Y^+$).

B. Voorgestelde Oplossing: Uniform-Correct Policy Optimization (UCPO)

Om het Uniform-Correct Beleid te bereiken, stellen de auteurs UCPO voor, een modificatie van de GRPO-doelstelling.

• Doelfunctie: UCPO verrijkt de standaard correctheidsmaximalisatie met een voorwaardelijke uniformiteitsstraf.
  $$\mathcal{L}_{UCPO} = \underbrace{\log Z_\theta(x)}_{\text{Correctheidsmassa}} - \tau \underbrace{KL(u(\cdot|x) \parallel q_\theta(\cdot|x))}_{\text{Uniformiteitsstraf}}$$

  • $Z_\theta(x)$: Totale waarschijnlijkheidsmassa op correcte uitvoer.
  • $q_\theta(\cdot|x)$: De voorwaardelijke verdeling van het beleid over correcte uitvoer.
  • $u(\cdot|x)$: De uniforme verdeling over de set van correcte uitvoer.
  • $\tau$: Een hyperparameter die de sterkte van de uniformiteitsstraf regelt.

• Gradiëntmechanisme:

  • Standaard GRPO weegt gradiënten af op basis van de huidige samplingfrequentie (bevoordelend dominante modi).
  • UCPO herbekijkt het gradiëntsignaal om deze bias te compenseren. Het interpoleert tussen de huidige voorwaardelijke verdeling en het uniforme doel.
  • Inverse-Weigting: De gradiëntschatter wijst hogere gewichten toe aan ondervertegenwoordigde correcte oplossingen (die met een lage huidige waarschijnlijkheid) en lagere gewichten aan dominante. Dit verdeelt het "voordeelbudget" effectief om minderheidscorrecte antwoorden te versterken zonder totale correctheidsmassa op te offeren.

3. Belangrijkste Bijdragen

1. Identificatie van Onverschilligheid: Het paper bewijst formeel dat standaard RLVR-doelstellingen onverschillig zijn voor diversiteit binnen de set van correcte antwoorden, waardoor stochastische trainingsdynamiek onomkeerbare modale instorting kan veroorzaken.
2. Theoretische Karakterisering: Het vestigt het Uniform-Correct Beleid als de unieke optimale oplossing onder zowel robuustheids- als entropie-geregulariseerde criteria.
3. Algoritmeontwerp (UCPO): Het leidt UCPO af, een eenvoudige maar effectieve modificatie van GRPO die het Uniform-Correct Beleid als zijn unieke optimum afdwingt.
4. Gradiëntherverdeling: Het introduceert een mechanisme om gradiëntsignalen te herverdelen naar ondervertegenwoordigde correcte antwoorden, waardoor de zichzelf versterkende instortingscyclus wordt doorbroken.

4. Experimentele Resultaten

De auteurs evalueerden UCPO over drie modellen (1,5B, 7B parameters) en vijf wiskundige redeneerbenchmarks (AIME24/25, AMC, MATH, OlympiadBench).

• Verbetering Pass@K: UCPO verbetert consistent Pass@K (specifiek Pass@64) terwijl het een concurrerende Pass@1 behoudt.
  • AIME24 (7B model): Bereikte een absolute verbetering van +10% in Pass@64 vergeleken met standaard GRPO.
  • Gemiddelde Verbetering: Over alle modellen heen verbeterde UCPO de gemiddelde Pass@64 met +3,5 tot +4,8 punten ten opzichte van GRPO.
• Diversiteitsmetrieken:
  • UCPO bereikte tot 45% hogere diversiteit op niveau van vergelijking binnen de correcte set vergeleken met GRPO.
  • In tegenstelling tot globale entropieregularisatie (die vaak massa verspreidt naar incorrecte antwoorden), behoudt UCPO een hoge correctheidsmassa ($Z_\theta \approx 1$) terwijl het de correcte set diversifieert.
• Vergelijking met Baselines: UCPO presteerde beter dan andere diversiteitsbehoudende methoden zoals Entropieregularisatie (Ent-Reg), KL-Covariantiecontrole en Focal GRPO. Globale entropieregularisatie bleek broos en degradeerde vaak Pass@1 of slaagde er niet in de diversiteit binnen de correcte set te verbeteren.
• Efficiëntie: UCPO introduceert verwaarloosbare rekenkundige overhead (alleen scalair herbekijken van voordelen) en vereist geen extra rollouts of referentiemodellen.

5. Betekenis

Dit werk adresseert een fundamenteel gebrek in huidige RLVR-trainingsparadigma's voor redeneermodellen. Door de focus te verschuiven van het simpelweg "vinden van een correct antwoord" naar "handhaven van een diverse verdeling van alle correcte antwoorden", doet UCPO het volgende:

• Voorkomt Contractie van de Oplossingsruimte: Het zorgt ervoor dat modellen geen geldige redeneerpaden "vergeten" tijdens fine-tuning.
• Versterkt Robuustheid: Een beleid dat meerdere correcte oplossingen dekt, is robuuster tegen verschuivingen in de verdeling of ruis in de verifieerders.
• Biedt een Principieel Doel: Het gaat voorbij aan heuristische entropiebonussen naar een theoretisch onderbouwde doelstelling die expliciet de structuur van het optimale beleid target.

Kortom, UCPO biedt een schaalbare, theoretisch onderbouwde oplossing voor het probleem van diversiteitsinstorting in RLVR, waardoor LLM's een hogere dekking bij meerdere samples kunnen bereiken zonder de nauwkeurigheid bij één poging op te offeren.