Technische Samenvatting: N-GRPO - Embedding-Level Neighbor Mixing voor Verbeterde Policy Optimalisatie

1. Probleemstelling

Het succes van Large Language Models (LLMs) in wiskundige redenering hangt sterk af van het genereren van diverse en geldige oplossingspaden tijdens de rollout-fase van Reinforcement Learning (RL). Huidige exploratietechnieken kampen echter met een fundamentele afruil:

• Token-Level Sampling: Standaardmethoden leveren vaak redundante trajecten op die alleen verschillen in herformulering of commutatieve herordening (bijv. "1+2" versus "2+1") zonder de onderliggende redeneerlogica te veranderen.
• Embedding-Level Ruis: Methoden die willekeurige ruis (bijv. Gaussische ruis) injecteren in continue embeddings, verstoren vaak de semantische consistentie. Voorlopige experimenten toonden aan dat ongecontroleerde ruis representaties van het semantische manifold afduwt, waardoor geldige tokens veranderen in ongerelateerde tokens en de rollout-trajecten ontsporen.

Bestaande pogingen om deze kloof te overbruggen, zoals Soft Thinking of HRPO, vertrouwen vaak op discrete token-sampling voor hun willekeur of falen om semantische coherentie te behouden wanneer ze de continue ruimte verstoren.

2. Methodologie: N-GRPO

De auteurs stellen N-GRPO voor, een nieuwe exploratiestrategie geïntegreerd in het Group Relative Policy Optimization (GRPO) framework. De kerninnovatie is Semantic Neighbor Mixing, een mechanisme voor embedding-level perturbatie dat diversiteit injecteert terwijl het strikt voldoet aan het lokale semantische manifold.

Kernmechanisme: Semantic Neighbor Mixing

In plaats van discrete tokens te samplen of ongecontroleerde ruis toe te voegen, construeert de methode inputrepresentaties door de embedding van een anker-token te mengen met zijn dichtstbijzijnde semantische buren.

1. Ankerselectie: Bij elke generatiestap identificeert het model het meest geprefereerde token (het semantische anker, $v^*_t$) op basis van de temperatuur-geschaalde logits.
2. Burenretrieval: Een set van $k$ buur-tokens ($C_t$) wordt opgehaald uit de embedding-matrix op basis van cosinus-gelijkenis met het anker. Dit zorgt ervoor dat buren clusteren langs dezelfde semantische richting.
3. Gewogen Mengen: De temperatuur-geschaalde logits worden genormaliseerd over deze verzameling buren om menggewichten ($\alpha_t$) te berekenen.
4. Continue Embedding Constructie: Een continue gemengde embedding ($\tilde{e}_{t+1}$) wordt gegenereerd als de gewogen som van de embedding van de buren. Dit dient als de input voor de volgende tijdstap.

Integratie met GRPO

N-GRPO integreert dit mechanisme in de GRPO rollout-fase via een stochastisch gating-mechanisme:

• Mixing Mask: Een Bernoulli-masker ($m_{i,t}$) met een vaste waarschijnlijkheid $\rho$ (bijv. 0,1) bepaalt of Semantic Neighbor Mixing wordt geactiveerd of dat de standaard discrete token-sampling wordt gevolgd.
• Stabiliteit: Door het merendeel van de stappen als discrete tokens te behouden, bewaart de methode semantische stabiliteit terwijl het de noodzakelijke intermittentie voor exploratie introduceert.
• Training: Tijdens de trainingsfase worden de kandidaat-sets en gewichten geregistreerd om reproduceerbare traject-reconstructie te garanderen voor de berekening van de advantage. De uiteindelijke beloning wordt berekend op basis van de discrete anker-tokens die overeenkomen met de gemengde stappen.

3. Belangrijkste Bijdragen

• Controleerbare Embedding-Level Exploratie: Het artikel introduceert een methode die semantisch geleide perturbaties injecteert tijdens de generatie, waardoor de kloof tussen discrete token-level redenering en continue embedding-level exploratie wordt overbrugd.
• N-GRPO Framework: De auteurs integreren Semantic Neighbor Mixing in het GRPO-framework, waardoor een systeem ontstaat dat gebruikmaakt van lokale semantische manifolds om de exploratieruimte uit te breiden zonder de coherentie op te offeren.
• Empirische Validatie: Uitgebreide evaluaties over verschillende modelmaten (1.5B tot 7B) en architecturen (DeepSeek-R1-Distill-Qwen, Llama-3.2, Qwen3) tonen consistente verbeteringen over sterke baselines aan.

4. Experimentele Resultaten

De methode werd geëvalueerd op uitdagende wiskundige redeneerbenchmarks (AIME25, AMC23, MATH500) en out-of-distribution (OOD) taken (GPQA-Diamond).

• Wiskundige Redenering: N-GRPO presteerde consequent beter dan de baselines, inclusief standaard GRPO, Soft Thinking en de STHT-baseline (die ongecontroleerde Gaussische ruis gebruikt).
  • Op het DeepSeek-R1-Distill-Qwen-1.5B model verbeterde N-GRPO de gemiddelde Pass@32 van 77,41 (GRPO) naar 79,17.
  • Op het 7B model verbeterde de gemiddelde Pass@32 van 81,94 (GRPO) naar 84,20.
  • Significante winsten werden waargenomen op de uitdagende AIME25 benchmark, waar N-GRPO een Pass@32 behaalde van 50,28 (1.5B) en 61,18 (7B), waarmee het GRPO en STHT overtrof.
• Out-of-Distribution Generalisatie: Op de GPQA-Diamond benchmark demonstreerde N-GRPO robuuste generalisatie, waarbij het de hoogste Pass@32 scores behaalde (92,87 voor 1.5B en 92,80 voor 7B), wat aangeeft dat de methode de algemene wetenschappelijke redeneervaardigheden versterkt, niet alleen wiskunde-specifieke vaardigheden.
• Transfereerbaarheid: Het mechanisme werd succesvol overgedragen naar het GSPO-framework (N-GSPO), met een winst van +7,66 op AIME25, wat suggereert dat het perturbatie-mechanisme niet nauw gekoppeld is aan de specifieke advantage-normalisatie van GRPO.
• Ablatie-studies:
  • Mixing Rate: De prestaties bleven stabiel bij variërende mixing rates (0,1 tot 0,2), hoewel excessieve mixing (0,2) op kleinere modellen voor AIME25 een daling veroorzaakte, wat de noodzaak benadrukt van een balans tussen diversiteit en stabiliteit.
  • Afstandsmetriek: Cosinus-afstand presteerde significant beter dan L1 en L2 metrieken voor de selectie van buren, wat wordt toegeschreven aan de directionele aard van semantische informatie in hoogdimensionale embedding-ruimten.
  • Inference: De auteurs merken op dat het toepassen van Semantic Neighbor Mixing tijdens de inference-fase (zonder het GRPO-filteringmechanisme) de prestaties verslechtert, wat bevestigt dat het een exploratietool is tijdens de training en geen decoding-strategie.

5. Betekenis en Claims

Het artikel claimt dat N-GRPO een veelbelovende weg biedt voor het bevorderen van de redeneervaardigheden van LLM's door de redeneerdynamiek in de embedding-ruimte te verkennen. De betekenis ligt in:

• Het Oplossen van de Afruil: Het injecteert succesvol diversiteit in het rollout-proces zonder de semantische drift die wordt veroorzaakt door naïeve ruisinjectie of de redundantie van token-level sampling.
• Semantische Gronding: Door gebruik te maken van het lokale semantische manifold, zorgt de methode ervoor dat exploratie binnen geldige redeneerregio's blijft.
• Generalisatie: De aanpak is niet beperkt tot specifieke modelfamilies of reasoning-distilled policies, en toont robuustheid over verschillende schalen en basismodellen.

De auteurs concluderen dat hoewel de methode computationele overhead introduceert tijdens de rollout en momenteel primair is gevalideerd op wiskunde- en wetenschapstaken, het een nieuw paradigma vestigt voor policy optimalisatie dat een balans vindt tussen exploratie en semantische consistentie.