Not all tokens are needed(NAT): token efficient reinforcement learning

Het artikel introduceert NAT (Not All Tokens Are Needed), een framework dat de trainingskosten van reinforcement learning voor lange redeneringen verlaagt door alleen een subset van tokens te updaten via Horvitz-Thompson-herschaling, terwijl het prestatieniveau van volledige token-training behouden blijft.

De kern

Stel je voor dat je een zeer slimme, maar ook erg dure robot wilt trainen om wiskundige raadsels op te lossen. Deze robot denkt hard na en schrijft een heel lang verhaal (een "chain of thought") om bij het antwoord te komen.

In het verleden dachten onderzoekers dat ze elk enkel woord dat de robot schreef, moesten controleren en corrigeren om hem te leren. Het was alsof je een leraar bent die elke letter van een 100 pagina's tellend essay van een student moet nakijken, zelfs als de eerste 90 pagina's alleen maar "hallo" en "laten we beginnen" bevatten. Dit kostte enorm veel tijd, energie en computerkracht.

Deze paper introduceert een slimme nieuwe methode genaamd NAT (Not All Tokens are Needed, oftewel: "Niet alle woorden zijn nodig").

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Probleem: De "Te dure" Leraar

Stel je voor dat je een marathonloper traint. De leraar loopt de hele 42 kilometer mee en corrigeert elke stap.

• Het probleem: De leraar is moe, de batterij van zijn stopwatch gaat leeg, en het kost te veel tijd.
• De realiteit: De loper maakt alleen kritieke beslissingen op bepaalde momenten (bijvoorbeeld: "nu moet ik versnellen" of "nu sla ik een bocht"). De rest van de tijd loopt hij gewoon rechtdoor. Die "gewone" stappen zijn minder belangrijk om te corrigeren.

2. De Oplossing: NAT (De Slimme Assistent)

De auteurs van dit paper zeggen: "Waarom kijken we niet alleen naar de belangrijke momenten?"

Ze hebben een systeem bedacht dat werkt als een slimme cameraman:

• In plaats van de hele film (het hele antwoord van de robot) op te nemen en te analyseren, filmt de cameraman alleen de hoogtepunten.
• Maar hier is de magische truc: De cameraman is niet willekeurig. Hij gebruikt een wiskundige formule (de Horvitz-Thompson schatting) om ervoor te zorgen dat het beeld dat hij wel vastlegt, statistisch gezien precies hetzelfde gewicht heeft als de hele film.

Het is alsof je een hele maaltijd proeft, maar je neemt slechts één hap van elk gerecht. Als je die hap slim kiest en de smaak goed berekent, kun je toch zeggen hoe de hele maaltijd smaakt, zonder dat je de hele maaltijd hoeft op te eten.

3. Twee Manieren om te Kiezen

De paper test twee manieren om deze "belangrijke momenten" te kiezen:

• Willekeurige Steekproef (URS): De leraar pakt een willekeurige set woorden uit het verhaal.
  • Nadeel: Omdat de robot in zijn gedachten alle eerdere woorden nodig heeft om het volgende woord te begrijpen, moet de computer toch de hele tekst "lezen" (voorwaartse berekening), ook al kijkt hij maar naar een paar woorden. Het bespaart tijd, maar niet genoeg geheugen.
• Willekeurig Afsnijden (RPC) - De Winnaar: De leraar kiest een punt in het verhaal en zegt: "Ik kijk alleen naar het begin, tot hier."
  • Voordeel: Dit is als het afsnijden van de laatste pagina's van een boek. Omdat de robot nu een korter verhaal moet lezen en schrijven, hoeft de computer minder zware berekeningen te doen.
  • De magie: Zelfs als je de laatste 50% van het verhaal afsnijdt, zorgt de slimme formule ervoor dat de robot toch leert alsof hij het hele verhaal had gezien. Het is alsof je een spiegel gebruikt die de afwezigheid van de laatste pagina's compenseert.

4. Wat is het Resultaat?

De tests met een robot genaamd Qwen (een slimme AI) toonden aan:

• Snelheid: De training was tot 29% sneller.
• Geheugen: De computer had 18% minder geheugen nodig.
• Kwaliteit: De robot werd even goed in wiskunde als wanneer hij het hele verhaal had geleerd.

De Grootste Les

Deze paper leert ons dat we niet alles hoeven te doen om iets goed te doen. Net zoals je niet elke seconde van een film hoeft te bekijken om de plot te begrijpen, hoeft een AI niet elk woord te analyseren om slim te worden.

Met NAT kunnen we AI's trainen die langer en complexer nadenken, zonder dat onze computers ontploffen van de kosten. Het is een manier om de "belasting" van het trainen te verlagen, zodat we in de toekomst nog slimmere robots kunnen bouwen.

Kort samengevat:
Het is alsof je een zware rugzak vol stenen (woorden) draagt. NAT laat je de helft van die stenen weglaten, maar zorgt er met een slimme truc voor dat je toch precies evenveel kracht opbouwt als met de volle rugzak.

Technische samenvatting

Probleemstelling

Reinforcement Learning (RL) is een drijvende kracht achter de vooruitgang van grote taalmodellen (LLMs), vooral voor complexe redeneertaken zoals wiskundig redeneren via "Chain-of-Thought" (CoT). Echter, het schalen van RL naar lange CoT-trajecten wordt steeds beperkter door de noodzaak om backpropagation uit te voeren over elk gegenereerd token.

Zelfs met geoptimaliseerde generatie-engines (rollouts) vormt het trainingsproces (forward en backward passes) een bottleneck:

1. Hoge kosten: Het verwerken van lange sequenties verbruikt een groot deel van de totale trainingskosten.
2. Geheugenproblemen: Lange trajecten leiden tot een enorme toename van het activatiegeheugen (activation memory), wat vaak leidt tot "Out-of-Memory" (OOM) fouten of instabiele optimalisatie.
3. Inefficiëntie: De huidige aanpak behandelt alle tokens als even belangrijk voor het leerproces, terwijl veel tokens (zoals mechanische voortzettingen of lage-entropie tekst) weinig bijdragen aan het leer-signaal.

De kernvraag van het paper is: Hebben we echt alle tokens nodig om sterke RL-redenaars te trainen?

Methodologie: NAT Framework

Het paper introduceert NAT (Not All Tokens are Needed), een unificerend framework dat de token-begroting tot een eerste-orde optimalisatieprimitief maakt. NAT update het beleid (policy) met behulp van slechts een geselecteerd subset van gegenereerde tokens, terwijl het beloningssignaal (reward) nog steeds op de volledige respons wordt berekend.

De kern van de methode bestaat uit twee onderdelen:

1. Token Masking met Horvitz-Thompson (HT) Herweging:

   • In plaats van alle tokens te gebruiken voor de gradiëntberekening, wordt een binaire masker $m_{i,t}$ toegepast om te bepalen of token $t$ bijdraagt aan de update.
   • Om de onbevooroordeeldheid (unbiasedness) van de gradiënt te behouden ondanks dit sub-sampling, wordt de Horvitz-Thompson schatter gebruikt.
   • Elke token krijgt een opnamekans (inclusion probability) $p_{i,t}$. De bijdrage van een token aan de gradiënt wordt vermenigvuldigd met $1/p_{i,t}$.
   • Wiskundige garantie: Het paper bewijst dat deze HT-gecorrigeerde schatter een onbevooroordeelde estimator is van de originele full-sequence RL-gradiënt, mits $p_{i,t} > 0$.

2. Twee Selectie Strategieën:
   Het paper instantieert NAT met twee specifieke methoden:

   • Uniform Random Sampling (URS): Tokens worden onafhankelijk en willekeurig geselecteerd met een vaste kans $p$.
     • Nadeel: Omdat causale Transformers afhankelijk zijn van alle voorgaande tokens voor de forward pass, reduceert URS alleen de backward pass (gradiënten), maar niet de forward compute of het geheugengebruik voor de activaties.
   • Random Prefix Cutting (RPC): In plaats van losse tokens, wordt een contiguus prefix (een beginsegment) van de respons geselecteerd. Er wordt een willekeurige afkap-lengte $L_i$ gekozen.
     • Voordeel: Omdat alleen het prefix wordt verwerkt, daalt zowel de forward compute als het activatiegeheugen aanzienlijk.
     • Verschil met deterministische truncatie: Deterministisch weghalen van het einde (bijv. altijd 50% houden) introduceert systematische bias omdat late tokens (vaak cruciaal voor verificatie) nooit worden gezien. RPC voorkomt dit door ervoor te zorgen dat elke positie een niet-nul kans heeft om te worden opgenomen via de HT-herweging.

Belangrijkste Bijdragen

1. Unificerend Framework: NAT biedt een principieel raamwerk voor token-efficiënte RLVR (Reinforcement Learning from Verifiable Rewards) dat volledige respons-bewerking voor beloningen koppelt aan geselecteerde token-updates.
2. Onbevooroordeeldheid via HT: Het paper levert een theoretisch bewijs dat HT-herweging leidt tot een onbevooroordeelde schatter voor de GRPO-gradiënt, ongeacht de selectieprobabiliteit (zolang deze positief is).
3. Praktische Snelheidswinst zonder Kwaliteitsverlies: Door RPC te gebruiken, kan de backward pass worden gereduceerd met tot 50% zonder de redeneerkwaliteit te schaden. Dit resulteert in directe besparingen in GPU-geheugen en trainingsduur.

Resultaten

De methoden zijn getest op wiskundige redeneringsbenchmarks (MATH, AIME24, AIME25) met modellen zoals Qwen2.5-Math-7B en Qwen3-8B.

• Prestatie: NAT (zowel URS als RPC) bereikt prestaties die statistisch gelijk zijn aan de volledige token GRPO-baseline. In tegenstelling tot deterministische truncatie (Det. Trunc.), die vaak slechter presteert door bias, behoudt RPC de leerkracht.
• Geheugenbesparing (RPC):
  • Voor Qwen3-8B: 18% piek GPU-geheugen besparing (van ~47.7 GB naar ~39.2 GB).
  • Voor Qwen2.5-Math-7B: ~18% besparing.
• Trainingsduur (RPC):
  • Voor Qwen3-8B: 29% snellere forward/backward trainingstijd (zonder inferentie).
  • De totale tijd per stap (inclusief inferentie) daalt met ongeveer 36% voor Qwen3-8B.
• Stabiliteit: De entropie-curves van RPC convergeren naar dezelfde waarden als de volledige GRPO, terwijl deterministische truncatie leidt tot hogere entropie en instabiliteit.

Betekenis en Impact

Dit werk biedt een orthogonale aanpak voor het schalen van RL-systemen. Waar eerdere pogingen zich richtten op het versnellen van de generatie (rollouts), richt NAT zich op het efficiënter gebruiken van deze rollouts tijdens het leerproces.

• Schalbaarheid: NAT maakt het mogelijk om RL toe te passen op langere en complexere trajecten die anders te duur of geheugenintensief zouden zijn.
• Complementair: De methode is compatibel met bestaande systeemoptimalisaties (zoals vLLM, speculative decoding) en kan daarop worden gebouwd.
• Toekomstperspectief: Het framework opent de deur voor "informatie-bewuste" token-selectie, waarbij de opnamekans dynamisch wordt aangepast op basis van token-uncertainty of gradiëntgrootte, in plaats van puur willekeurig.

Kortom, NAT bewijst dat voor effectief RL-redeneren niet alle tokens nodig zijn, en biedt een wiskundig onderbouwde weg om de "verborgen belasting" van lange trajecten op te lossen.