🤖 machine learning

Predictive Divergence Masks for LLM RL

Dit artikel stelt Predictive Divergence Masks voor, een nieuwe methode voor reinforcement learning bij LLM's die het op ratio gebaseerde richtingscriterium van PPO vervangt door een gesloten voorspelling van of de volgende gradiëntstap de waarschijnlijkheidsdivergentie zal verhogen of verlagen die wordt gebruikt voor het nabijheidscriterium, waardoor de trainingsstabiliteit en prestaties over verschillende modelschalen heen worden verbeterd.

Oorspronkelijke auteurs: Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang

Gepubliceerd 2026-07-14

📖 6 min leestijd🧠 Diepgaand

CC BY 4.0

Oorspronkelijke auteurs: Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een superintelligente robot leert om verhalen te schrijven. Je wilt dat de robot leert van zijn fouten, maar je hebt een strikte regel: "Verander je schrijfstijl niet te wild in één keer." Als de robot te enthousiast wordt en in een totaal andere stem begint te schrijven, moet je op de rem trappen. Dit is de taak van Reinforcement Learning (RL) voor Large Language Models (LLMs).

Lama tijd was de standaardmanier om op die rem te trappen een methode genaamd PPO. Het werkte als een simpel verkeerslicht gebaseerd op één enkel woord. Elke keer dat de robot een woord koos, controleerde het systeem: "Is dit woord veel waarschijnlijker dan wat de robot normaal gesproken kiest? En zegt de beloning ons dat het nóg waarschijnlijker moet worden?" Als het antwoord op beide "ja" was, zei het systeem: "Stop! Dat gaat te ver!" en negeerde het de les van dat woord.

Maar hier is het probleem: de robot kiest niet zomaar één woord; hij herschikt zijn hele persoonlijkheid (zijn waarschijnlijkheidsverdeling) om dat woord te kiezen. De oude methode keek alleen naar het specifieke woord dat hij koos, alsof je een heel orkest beoordeelt door naar slechts één viool te luisteren. Het miste hoe de rest van de muziek veranderde.

Het Nieuwe Idee: De "Predictive Divergence Mask"

Een team onderzoekers van Tencent, UIUC en NUS realiseerde zich dat dit eenzijdige kijk op één woord misleidend was. Ze stelden een nieuwe regel voor genaamd de Predictive Divergence Mask.

In plaats van alleen naar het ene woord te kijken dat de robot koos, stelt deze nieuwe methode een grotere vraag: "Als we deze stap zetten, zal de volledige persoonlijkheid van de robot dan verder afwijken van zijn oorspronkelijke zelf?"

Denk er zo over na:

De Oude Manier (PPO): Je ziet een kind naar een klif rennen. Je controleert of die specifieke stap gevaarlijk is. Zo ja, dan roep je "Stop!"
De Nieuwe Manier (Predictive Divergence): Je ziet het kind rennen. Je kijkt naar het hele pad. Zelfs als die specifieke stap er oké uitziet, kan het zijn dat de momentum van het kind er bijna voor zorgt dat het een draai maakt en van een andere klif af vliegt. De nieuwe methode voorspelt de hele draai voordat de stap wordt toegelaten.

Hoe het werkt (De Magische Wiskunde)

De onderzoekers vonden een slimme wiskundige truc om deze "hele draai" te voorspellen zonder dat ze eerst de simulatie hoeven te draaien.

De Lokale Term: Dit is het deel dat de oude methode al kende. Het is de verandering in het specifieke woord dat de robot koos.
De Globale Term: Dit is het nieuwe geheime ingrediënt. Het houdt rekening met hoe het veranderen van één woord de robot dwingt om de waarschijnlijkheden van alle andere woorden aan te passen die hij niet koos (omdat de totale waarschijnlijkheid altijd 100% moet zijn).

De oude methode negeerde de Globale Term. De nieuwe methode telt deze op. Als de wiskunde laat zien dat de hele persoonlijkheid van de robot te ver afdwaalt, blokkeert de maskering de update. Als de wiskunde laat zien dat de robot eigenlijk weer terug beweegt naar veiligheid, wordt de update toegelaten.

Het "Top-K" Probleem

Er is een addertje onder het gras. In de echte wereld toont de motor van de robot (de rollout engine) alleen de meest waarschijnlijke top 20 woorden (de "Top-K") om tijd en geheugen te besparen. Het verbergt de duizenden andere woorden in een "staart"-bak.

Om hun voorspelling werkbaar te maken met dit beperkte zicht, hebben de onderzoekers twee lichtgewicht estimators uitgevonden:

De Aggregated-Tail Estimator: Doet alsof alle verborgen woorden één grote hoop zijn.
De Uniform-Tail Estimator: Doet alsof de verborgen woorden allemaal even waarschijnlijk zijn.

Ze ontdekten dat omdat de top 20 woorden meestal bijna alle waarschijnlijkheid bevatten (meer dan 99%), beide methoden bijna exact hetzelfde werken. Het is alsof je het weer probeert te voorspellen door naar de lucht te kijken; je hoeft niet elke individuele wolk te tellen om te weten of het gaat regenen.

Wat de Experimenten Laten Zien

Het team heeft dit getest op verschillende formaten robots (modellen variërend van 4 miljard tot 30 miljard parameters) en heeft het zelfs getest wanneer ze draaien op "FP8"-precisie (een super snelle, iets minder nauwkeurige rekenmodus die ervoor zorgt dat de trainings- en schrijfmotoren anders werken).

Het Resultaat: De nieuwe maskering maakte de training stabieler.
Het Bewijs: Ze voerden 61 verschillende experimenten (seeds) uit. Ze ontdekten dat de oude methode soms updates aanhield die er eigenlijk voor zorgden dat de robot verder afweek (een "unsafe keep" ratio van 36,9%). De nieuwe methode verminderde dit naar 34,2%.
De Verbetering: Hoewel de cijfers er misschien klein uitzien, helpt het voorkomen van zelfs maar een paar slechte stappen in de wereld van het trainen van gigantische AI-modellen om de robot beter en sneller te laten leren. De nieuwe methode werkte consistent over alle modelgroottes en precisie-instellingen heen.

Wat Ze Expliciet Afwezen

Het paper is zeer duidelijk over wat niet zo goed werkt:

Ze argumenteren tegen het gebruik van de sampled-token importance ratio (de oude "enkele viool"-controle) om de richting van de update te bepalen. Ze lieten zien dat deze controle op één enkel woord vaak niet overeenkomt met de werkelijke verandering in de hele persoonlijkheid van de robot.
Ze lieten ook zien dat het simpelweg nauwer maken van de "veiligheidszone" (de trust region) (door een drempel van 0,05 te gebruiken in plaats van 0,15) de prestaties overal verslechtert, wat suggereert dat alleen maar restrictiever zijn niet het antwoord is; slimmer zijn over de richting wel.

Hoe Zeker Zijn Ze?

De auteurs zijn vol vertrouwen in hun bevindingen op basis van gemeten data uit hun experimenten. Ze hebben niet alleen in een vacuüm gesimuleerd; ze hebben daadwerkelijke modellen getraind op echte wiskundige problemen (met behulp van de DAPO-Math-17k dataset) en de nauwkeurigheid gemeten op benchmarks zoals AIME24 en AIME25.

Ze suggereren dat deze aanpak een betere afstemming van de regels is: als je de veiligheidszone definieert door naar de hele persoonlijkheid te kijken (divergentie), moet je ook de richting van de update controleren door naar de hele persoonlijkheid te kijken, niet alleen naar één enkel woord. De resultaten suggereren dat dit leidt tot stabielere training, maar ze erkennen dat het nog steeds een "lokale eerste-orde benadering" is — wat betekent dat het een zeer goede schatting is op basis van de onmiddellijke volgende stap, en geen perfecte kristallen bol voor de hele toekomst.

Kortom, de nieuwe methode is als het geven van een GPS aan de robot die naar de hele kaart kijkt, in plaats van alleen naar de stap die hij op dit moment zet. Het helpt de robot sneller te leren zonder van de rand van de wereld af te dwalen.

Technische Samenvatting: Predictive Divergence Masks voor LLM RL

1. Probleemstelling

Reinforcement Learning (RL) voor Large Language Models (LLMs) is inherent off-policy vanwege mismatches tussen de inference-engines die worden gebruikt voor rollouts en de trainingsstacks, evenals beleidsveroudering (policy staleness). Om deze off-policy updates te stabiliseren, maken standaardmethoden zoals Proximal Policy Optimization (PPO) gebruik van trust-region masks. Deze masks vertrouwen op twee criteria om te beslissen of een token's gradiënt gemaskeerd (op nul gezet) wordt:

Proximity Criterion (Nabijheidscriterium): Bepaalt of het beleid te ver is afgeweken van het gedragbeleid (behavior policy).
Direction Criterion (Richtingcriterium): Bepaalt of de update het beleid verder weg van het gedragbeleid duwt (naar buiten gericht) of het terugtrekt (corrigerend).

Recent werk, specifiek DPPO, heeft het proximity criterion verbeterd door de enkelvoudige sample importance ratio te vervangen door een distributionele divergentie (bijv. KL-divergentie) die wordt geschat over de top-K vocabulaire. DPPO behoudt echter het oorspronkelijke PPO direction criterion, dat uitsluitend rust op de importance ratio van de gesamplede token ( $r_t - 1$ ).

De auteurs identificeren een fundamentele inconsistentie in deze hybride aanpak: terwijl het proximity criterion nu gedefinieerd wordt door een globale distributionele divergentie, blijft het direction criterion een single-sample proxy. De auteurs observeren dat het teken van de verandering in de gesamplede ratio kan verschillen van het teken van de werkelijke verandering in de divergentie. Als gevolg hiervan kan een token worden behouden (omdat de ratio een corrigerende update suggereert) zelfs wanneer de update de globale divergentie feitelijk vergroot, of vice versa.

2. Methodologie

Het paper stelt de Predictive Divergence Mask voor, een methode die het direction criterion afstemt op het distributionele karakter van het proximity criterion. In plaats van te vragen "Zal de gesamplede ratio naar één toe bewegen?", vraagt de methode: "Zal de volgende gradiëntstap de divergentie vergroten?"

Theoretische Afleiding

Voor de softmax-policies die worden gebruikt in autoregressieve LLM's, leiden de auteurs een gesloten vorm af voor de directionele afgeleide van de Kullback-Leibler (KL) divergentie, aangeduid als $\dot{D}$ , langs een policy gradient-stap.
$\dot{D} = \underbrace{(\pi_k - \mu_k)}_{\text{Lokale Term}} + \underbrace{\sum_{i} \pi_i (\mu_i - \pi_i)}_{\text{Globale Term}}$
Waar:

$\pi_k$ en $\mu_k$ de waarschijnlijkheden van de gesamplede token $k$ zijn onder respectievelijk het trainings- en het gedragbeleid.
De Lokale Term komt exact overeen met de hoeveelheid die wordt gebruikt in het standaard ratio-gebaseerde direction criterion ( $r - 1$ ).
De Globale Term vangt de normalisatiekoppeling over de gehele vocabulaire op. Deze term is onzichtbaar voor een enkelvoudige gesamplede ratio, maar is cruciaal voor het bepalen van de werkelijke verandering in divergentie.

De predictive mask houdt een update alleen aan als het teken van $\hat{A}_t \cdot \dot{D}_t$ niet-positief is (wat wijst op een corrigerende of neutrale beweging), mits de token zich al buiten de trust region bevindt ( $D_t > \delta$ ).

Praktische Schatting

Productie-rollout engines exposeren doorgaans alleen de top- $K$ token-waarschijnlijkheden, niet de volledige vocabulaire-distributie. Om $\dot{D}$ onder dit beperkte zicht te berekenen, stellen de auteurs twee lichtgewicht estimators voor de ongeziene staart van de distributie voor:

Aggregated-Tail Estimator: Reduceert de gehele resterende waarschijnlijkheidsmassa tot één geaggregeerde token.
Uniform-Tail Estimator: Verdeelt de resterende massa gelijkmatig over de ongeziene vocabulaire-tokens.

Beide estimators vereisen slechts eenvoudige reducties over de behouden top- $K$ waarschijnlijkheden en staartmassa's, wat een verwaarloosbare computationele overhead toevoegt.

3. Belangrijkste Bijdragen

Identificatie van Inconsistentie: Het paper benadrukt dat het mengen van een distributioneel proximity criterion (divergentie) met een single-sample direction criterion (ratio) een mismatch creëert waarbij de direction mask mogelijk de werkelijke verandering in divergentie niet goed volgt.
Predictive Divergence Mask: Een nieuwe masking-regel die de eerste-orde verandering van de divergentie zelf gebruikt om de richting van de update te bepalen.
Gesloten Vorm Afleiding: De afleiding van de directionele afgeleide voor softmax-policies, waarbij de lokale (gesamplede) en globale (normalisatie) componenten expliciet worden gescheiden.
Truncated Vocabulary Estimators: Twee praktische, laag-overhead estimators (aggregated en uniform tail) die de methode in staat stellen te functioneren onder standaard productie-rollout restricties (top- $K$ views).

4. Experimentele Resultaten

De auteurs evalueerden de methode op de DAPO-Math-17k dataset met Qwen3-modellen (4B, 8B en 30B) met verschillende precisie-instellingen (FP8).

Trainingsstabiliteit: Over alle modelmaten en precisie-instellingen heen (inclusief stress-tests met FP8), verbeterde de Predictive Divergence Mask de trainingsstabiliteit vergeleken met de baseline DPPO-TopK-KL. In tegenstelling hiertoe vertoonde standaard GRPO met clip-higher trucs instabiliteit en collaps in grotere model-settings.
Direction Criterion Alignment: Token-niveau analyse toonde aan dat het divergentie-gebaseerde direction criterion beter is afgestemd op de gerealiseerde verandering in divergentie dan het ratio-gebaseerde criterion.
- Onder de tokens waar de twee criteria verschilden, had het divergentie-gebaseerde criterion een lager "unsafe-keep" percentage (34,2% vs. 36,9%), wat betekent dat het minder waarschijnlijk updates behield die de divergentie daadwerkelijk vergrootten.
- Omgekeerd behield het een hoger deel van de updates die succesvol de divergentie lieten krimpen (65,1% vs. 62,9%).
Robuustheid: De methode demonstreerde verbeterde robuustheid ten opzichte van de trust-region drempel hyperparameter ( $\delta$ ), en presteerde beter dan de baseline zelfs onder striktere beperkingen ( $\delta = 0,05$ ).
Estimator Equivalentie: De aggregated-tail en uniform-tail estimators produceerden bijna identieke resultaten, wat suggereert dat de methode niet gevoelig is voor de specifieke aanname van de staartmodellering.

5. Betekenis en Claims

Het paper claimt dat de Predictive Divergence Mask een structurele inconsistentie oplost in huidige divergentie-gebaseerde RL-methoden voor LLM's. Door ervoor te zorgen dat zowel het proximity als het direction criterion afgeleid zijn van dezelfde distributionele divergentie, biedt de methode een getrouwere trust-region controle.

De auteurs benadrukken dat hoewel de verbetering in "unsafe-keep" rates bescheiden is (vanwege de kleine set tokens waar de criteria verschillen), de consistentie theoretisch solide en empirisch voordelig is over verschillende schalen. Zij concluderen dat het direction criterion dezelfde distributionele behandeling verdient als het proximity criterion om de voordelen van divergentie-gebaseerde trust regions volledig te benutten. De methode wordt gepresenteerd als een drop-in vervanging voor bestaande DPPO-implementaties die geen nieuwe hyperparameters vereist buiten de bestaande divergentie-drempel.