Reasoning or Memorization? Direction-Aware Diversity Exploration in LLM Reinforcement Learning

Technische Samenvatting: DiRL – Richting-bewuste Diversiteit-exploratie in LLM Reinforcement Learning

Probleemstelling
Reinforcement learning (RL) is een primair paradigma geworden voor het oproepen van complexe redeneringen in Large Language Models (LLMs). Een centrale uitdaging in RL is echter exploratie: het beleid (policy) moet diverse trajecten samplen om te voorkomen dat het instort in smalle, suboptimale modi. Bestaande exploratiemethoden stimuleren doorgaans diversiteit in semantische of gradiëntruimtes zonder het bron van die diversiteit te onderscheiden. Een traject kan nieuw lijken omdat het een werkelijk nieuw redeneerpad volgt, of simpelweg omdat het gearrangeerde patronen en shortcuts varieert. Huidige benaderingen belonen beide gedragingen gelijkelijk, wat het risico loopt de exploratie te sturen naar memorisatie in plaats van naar echte verbetering van redeneervermogen. Zelfs wanneer metrieken verschuiven van semantische embeddings naar gradiëntruimtes om optimalisatiedynamiek beter te reflecteren, kan een nieuwe update nog steeds memorisatie versterken in plaats van redeneervermogen te verbeteren.

Methodologie: DiRL-framework
Het artikel stelt DiRL (Direction-Aware Reinforcement Learning) voor, een framework dat exploratie verankert aan een interne "redeneer-memorisatie-richting" afgeleid van de eigen representaties van de policy. DiRL integreert naadloos in het standaard Group Relative Policy Optimization (GRPO) algoritme. De methodologie verloopt via vier fasen:

1. Extractie van de Redeneer-Memorizatie-richting:
   Vóór de RL-training construeert de methode een vaste richtingvector $k$ in de residual stream van de policy. Dit wordt bereikt door het gemiddelde activatieverschil te berekenen tussen twee disjuncte sets taken: $D^+$ (redeneerintensieve taken, bijv. meerstaps wiskunde) en $D^-$ (geheugenintensieve taken, bijv. feitelijke herinnering). Deze richting dient als een geometrisch anker om te onderscheiden of een respons meer overeenkomt met redeneer- of memorisatiegedrag.

2. Gradiëntkenmerken met Weging op Basis van Richting:
   Voor elke respons construeert DiRL een gradiëntkenmerk $\Phi^{(i)}$ dat samenvat hoe de respons de policy zou updaten. In tegenstelling tot standaard gradiëntkenmerken worden deze gewogen door de uitlijning van de verborgen toestand (hidden state) van elk token met de redeneerrichting $k$. Tokens die uitlijnen met de redeneerkant ontvangen hogere gewichten, terwijl tokens die uitlijnen met memorisatie lagere gewichten ontvangen. Dit zorgt ervoor dat de resulterende kenmerkvector primair diversiteit in redeneren reflecteert in plaats van memorisatie.

3. Redeneer-verankerde Exploratiescore:
   Het framework verdeelt de rollout-groep in twee subgroepen: een met redeneren uitgelijnde subgroep ($G_R$) en een met memorisatie uitgelijnde subgroep ($G_M$), gebaseerd op de projectie van de gemiddelde verborgen toestanden op $k$.

   • Voor responsen in $G_R$ meet de exploratiescore de nieuwheid ten opzichte van andere met redeneren uitgelijnde tijdgenoten.
   • Voor responsen in $G_M$ meet de score de afstand tot de met redeneren uitgelijnde update-richtingen.
     Dit ontwerp zorgt ervoor dat diversiteit alleen wordt beloond wanneer het de ruimte van geldige redeneerpaden uitbreidt, terwijl variaties in memorisatie als minder waardevol of zelfs nadelig worden beschouwd.

4. Richting-bewuste Beloningsvormgeving (Reward Shaping):
   De exploratiescore wordt gebruikt om de oorspronkelijke binaire beloningen vorm te geven. Correcte redeneer-uitgelijnde responsen ontvangen een versterkte beloning, terwijl incorrecte redeneerpogingen een verzachte straf krijgen (omdat ze nog steeds nuttige directionele updates leveren). Omgekeerd ontvangen correcte memorisatie-uitgelijnde responsen een verminderde beloning om shortcuts te ontmoedigen, terwijl incorrecte memorisatie-responsen zwaarder worden gestraft. Deze gevormde beloningen worden vervolgens ingevoerd in de standaard GRPO advantage-berekening.

Belangrijkste Bijdragen

• DiRL Framework: Introductie van een richting-bewust RL-framework dat exploratie expliciet verankert aan het onderscheid tussen redeneren en memorisatie, waardoor de policies worden gestuurd naar redeneren in plaats van gememoriseerde patronen.
• Nieuw Exploratiemechanisme: Ontwikkeling van een mechanisme dat nieuwheid meet in een richting-gewogen gradiëntkenmerkruimte, waarbij de uitlijning met redeneren selectief wordt versterkt terwijl variaties in memorisatie worden onderdrukt.
• Empirische Validatie: Uitgebreide experimenten die aantonen dat DiRL de prestaties verbetert over diverse wiskundige en algemene redeneerbenchmarks vergeleken met standaard GRPO en andere diversiteits-gebaseerde exploratiemethoden.

Experimentele Resultaten
De auteurs evalueerden DiRL op Qwen3-1.7B en Qwen3-4B backbones over meerdere benchmarks:

• Wiskundig Redeneren: Op benchmarks inclusief MATH500, AMC, AIME24 en AIME25 presteerde DiRL consequent beter dan de baselines (GRPO, Entropy Bonus, EVOL-RL en G2RL). Bijvoorbeeld, op Qwen3-4B behaalde DiRL een pass@1 van 83.0 op MATH500 (een verbetering van +5.6 t.o.v. GRPO) en 19.6 op AIME24 (een verbetering van +7.3).
• Algemeen Redeneren: Op GPQA en MMLU-Pro liet DiRL consistente verbeteringen zien, wat aangeeft dat de voordelen verder gaan dan wiskundige probleemoplossing naar algemene redeneervermogens.
• Generalisatie en Robuustheid: Met DiRL getrainde modellen vertoonden superieure prestaties op GSM-Symbolic en de geperturbeerde varianten daarvan, wat suggereert dat de modellen onderliggende redeneerprocedures hebben geleerd in plaats van te overfitten op specifieke benchmarkpatronen.
• Analyse van Rollouts: DiRL verhoogde het aandeel redeneer-uitgelijnde responsen in rollouts (bijv. van 45.8% naar 60.0% op Qwen3-4B voor AIME25).
• Efficiëntie: De methode introduceert minimale computationele overhead (ongeveer 13-18% toename in trainingstijd per stap), aangezien het vertrouwt op forward-pass activaties en een vooraf geëxtraheerde, statische richtingvector.

Betekenis en Claims
Het artikel claimt dat effectieve exploratie in LLM reinforcement learning niet enkel moet focussen op het genereren van diverse trajecten, maar op het aanmoedigen van trajecten die daadwerkelijk het redeneren uitbreiden. Door de interne geometrie van LLM's te benutten om het onderscheid tussen redeneren en memorisatie te maken, biedt DiRL een mechanisme om RL te sturen naar robuustere redeneergedragingen. De auteurs stellen dat dit werk toekomstig onderzoek naar representatie-bewuste exploratie stimuleert, waarbij interne model-signatures worden gebruikt om leerprocessen te optimaliseren. Het artikel erkent beperkingen, waarbij wordt opgemerkt dat de huidige aanpak rust op een enkele globale richting geëxtraheerd uit representatieve data, en dat toekomstig werk meerdere latente dimensies van redeneergedrag zou kunnen ontdekken.