Not All Tokens Matter: Towards Efficient LLM Reasoning via Token Significance in Reinforcement Learning

Titel

Niet alle tokens tellen mee: Naar efficiënter redeneren van LLM's via token-significantie in versterkend leren

1. Het Probleem

Grote Taalmodellen (LLM's) vertonen sterke redeneervermogens, vaak via "Chain-of-Thought" (CoT), maar genereren vaak onnodig lange en redundante uitleggen. Dit leidt tot inefficiëntie in termen van rekenkosten, latentie en opslag.

• Huidige aanpak: Bestaande methoden voor versterkend leren (RL) om redenering te verbeteren, focussen voornamelijk op nauwkeurigheid en gebruiken vaak uniforme lengte-straffen (rewards/penalties).
• De beperking: Een uniforme straf voor elke token in een antwoord is suboptimaal. Het behandelt alle tokens alsof ze even belangrijk zijn. Hierdoor kan het model essentiële redeneerstappen verwijderen om de lengte te verkorten, wat de nauwkeurigheid schaadt. Omgekeerd kan het model redundantie behouden omdat de straf niet discrimineert tussen nuttige en overbodige tokens.

2. Methodologie: Het BINGO Framework

De auteurs introduceren BINGO (Boosting Efficient ReasonING in Policy Optimization), een RL-framework dat twee kernmechanismen combineert om zowel nauwkeurigheid als efficiëntie te optimaliseren.

A. Token Significantie (Token Significance)

Het paper stelt dat niet alle tokens in een CoT-redenering even waardevol zijn.

• Definitie: Tokens worden ingedeeld in significant (essentieel voor het antwoord) en insignificant (redundant, overbodige tussenstappen of herhalingen).
• Bepaling: Een contextuele scorefunctie (geïmplementeerd via LLMLingua-2) beoordeelt elke token op basis van de volledige sequentie om een significantie-score $S(y_i)$ te geven. Tokens onder een drempelwaarde $\tau$ worden als insignifiant beschouwd.
• Significantie-bewuste Lengte Reward: In plaats van de totale lengte te straffen, straft deze reward alleen de insignificante tokens af.
  • Voor correcte antwoorden: Er wordt een zachte, cosinus-gebaseerde straf toegepast op de lengte van de insignificante tokens, terwijl essentiële stappen behouden blijven.
  • Voor incorrecte antwoorden: Er wordt gestraft op insignificante tokens, maar er wordt ook een stimulans gegeven voor het genereren van meer significante redeneerinhoud.

B. Dynamische Lengte Control

Om de balans tussen exploratie (uitproberen van lange redeneringen voor moeilijke problemen) en compressie (kort en krachtig zijn) te vinden, wordt een dynamische reward gebruikt die evolueert tijdens het trainingsproces.

• Fase 1 (Vroege training / Exploratie): Voor incorrecte samples wordt een positieve beloning (bonus) gegeven voor het gebruik van significante tokens. Dit moedigt het model aan om diep te graven en complexe oplossingsstrategieën te ontdekken.
• Fase 2 (Late training / Compressie): Zodra de nauwkeurigheid stabiliseert (gemeten via een dynamische schalingsfactor $k$), schakelt de reward over naar het straffen van te lange significante tokens. Dit dwingt het model om zijn redenering te comprimeren zonder de kern van het antwoord te verliezen.

C. Integratie in PPO

Deze reward-mechanismen worden geïntegreerd in het standaard Proximal Policy Optimization (PPO) algoritme. De totale reward $R_{BINGO}$ combineert de nauwkeurigheidsindicator (Exact Match) met de bovenstaande lengte-rewards, waardoor het beleid (policy) wordt geoptimaliseerd voor een optimale afweging tussen correctheid en beknoptheid.

3. Belangrijkste Bijdragen

1. Inzicht in Token Significantie: Het onderscheid maken tussen significante en insignificante tokens in redeneertraces. Dit leidt tot een gerichte straal die redundantie verwijdert zonder de logische structuur te beschadigen.
2. Dynamische Lengte Control: Een strategie die de reward-prikkels aanpast gedurende de training: eerst exploratie stimuleren voor moeilijke vragen, en later focussen op beknoptheid.
3. Het BINGO Framework: Een nieuw RL-framework dat deze twee strategieën combineert en toont dat het mogelijk is om de response-lengte drastisch te verminderen terwijl de nauwkeurigheid behouden of zelfs verbeterd wordt.

4. Resultaten

Het framework werd getest op diverse benchmarks (MATH500, GSM8K, TheoremQA, AIME2024) met modellen van verschillende groottes (DeepSeek-R1-Distill-Qwen-1.5B/7B en Qwen2.5-Math-7B).

• Efficiëntie vs. Nauwkeurigheid: BINGO presteert consequent beter dan bestaande methoden (zoals Vanilla PPO, O1-Pruner, Kimi-k1.5) op de lengte-genormaliseerde nauwkeurigheid (L-Acc).
  • Voorbeeld: Op MATH500 reduceerde BINGO-A de gemiddelde response-lengte met 68% (van 2.771 naar 894 tokens) ten opzichte van Vanilla PPO, terwijl de nauwkeurigheid zelfs steeg (van 81,4% naar 82,2%).
• Behoud van Kwaliteit: In tegenstelling tot methoden die simpelweg alles inkorten (wat vaak leidt tot fouten), behoudt BINGO de essentiële redeneerstappen. De verhouding van significante tokens in de output nam toe (van ~71% naar ~75%).
• Robuustheid: De resultaten waren consistent over verschillende modelgroottes (1.5B en 7B parameters) en verschillende soorten redeneringsopgaven (wiskunde, logica, commonsense).
• Ablatiestudies: Het verwijderen van zowel de significantie-aware reward als de dynamische reward leidde tot een duidelijke daling in prestaties, wat aantoont dat beide componenten essentieel zijn voor het succes.

5. Betekenis en Conclusie

Dit paper biedt een fundamenteel nieuw perspectief op het optimaliseren van LLM-redenering. Het toont aan dat "korter" niet automatisch betekent "minder goed", zolang de verkorting gericht is op het verwijderen van irrelevante informatie in plaats van het uniform verkorten van de output.

Door token-significantie te modelleren en de lengte-prikkels dynamisch aan te passen, slaagt BINGO erin om de "overdenking" (overthinking) van LLM's te elimineren. Dit heeft grote implicaties voor de schaalbaarheid en kostenefficiëntie van LLM's in de praktijk, omdat het de rekenkosten verlaagt en de responsietijden verkort zonder in te leveren op de kwaliteit van de oplossing. De methode biedt een solide basis voor toekomstige RL-toepassingen waarbij efficiëntie en nauwkeurigheid gelijktijdig moeten worden geoptimaliseerd.