Riepilogo Tecnico: N-GRPO - Neighbor Mixing a livello di Embedding per un'Ottimizzazione della Policy Potenziata

1. Definizione del Problema

Il successo dei Large Language Models (LLM) nel ragionamento matematico dipende fortemente dalla generazione di percorsi di soluzione diversificati e validi durante la fase di rollout del Reinforcement Learning (RL). Tuttavia, le attuali tecniche di esplorazione affrontano un compromesso fondamentale:

• Campionamento a livello di Token: I metodi standard spesso producono traiettorie ridondanti che differiscono solo per riformulazioni o riordinamenti commutativi (ad esempio, "1+2" rispetto a "2+1") senza alterare la logica di ragionamento sottostante.
• Rumore a livello di Embedding: I metodi che iniettano rumore casuale (ad esempio, rumore Gaussiano) negli embedding continui spesso interrompono la coerenza semantica. Esperimenti preliminari hanno dimostrato che il rumore non vincolato spinge le rappresentazioni fuori dal manifold semantico, trasformando i token validi in altri non correlati e deviando le traiettorie di rollout.

I tentativi esistenti per colmare questo divario, come Soft Thinking o HRPO, si affidano spesso al campionamento discreto dei token per la loro casualità o non riescono a mantenere la coerenza semantica quando perturbano lo spazio continuo.

2. Metodologia: N-GRPO

Gli autori propongono N-GRPO, una nuova strategia di esplorazione integrata nel framework Group Relative Policy Optimization (GRPO). L'innovazione principale è il Semantic Neighbor Mixing, un meccanismo di perturbazione a livello di embedding che inietta diversità aderendo rigorosamente al manifold semantico locale.

Meccanismo Core: Semantic Neighbor Mixing

Invece di campionare token discreti o aggiungere rumore non vincolato, il metodo costruisce le rappresentazioni di input mescolando l'embedding di un token ancora con i suoi vicini semantici più prossimi.

1. Selezione dell'Anchor: Ad ogni step di generazione, il modello identifica il token più preferito (l'anchor semantico, $v^*_t$) basandosi sui logit scalati per temperatura.
2. Recupero dei Vicini: Un insieme di $k$ token vicini ($C_t$) viene recuperato dalla matrice degli embedding in base alla similarità del coseno rispetto all'anchor. Ciò assicura che i vicini si raggruppino lungo la stessa direzione semantica.
3. Mixing Pesato: I logit scalati per temperatura vengono normalizzati su questo insieme di vicini per calcolare i pesi di mixing ($\alpha_t$).
4. Costruzione dell'Embedding Continuo: Un embedding misto continuo ($\tilde{e}_{t+1}$) viene generato come somma pesata degli embedding dei vicini. Questo servirà come input per il passo temporale successivo.

Integrazione con GRPO

N-GRPO integra questo meccanismo nella fase di rollout di GRPO tramite un meccanismo di gating stocastico:

• Mixing Mask: Una maschera di Bernoulli ($m_{i,t}$) con una probabilità fissa $\rho$ (ad esempio, 0.1) determina se attivare il Semantic Neighbor Mixing o seguire il campionamento standard dei token discreti.
• Stabilità: Mantenendo la maggior parte degli step come token discreti, il metodo preserva la stabilità semantica introducendo al contempo la necessaria intermittenza per l'esplorazione.
• Addestramento: Durante la fase di addestramento, i set candidati e i pesi vengono registrati per garantire la ricostruzione riproducibile della traiettoria per il calcolo dell'advantage. Il reward finale viene calcolato sulla base degli anchor token discreti corrispondenti ai passi di mixing.

3. Contributi Chiave

• Esplorazione Controllabile a Livello di Embedding: Il paper introduce un metodo che inietta perturbazioni guidate semanticamente durante la generazione, colmando il divario tra il ragionamento discreto a livello di token e l'esplorazione continua a livello di embedding.
• Framework N-GRPO: Gli autori integrano il Semantic Neighbor Mixing nel framework GRPO, creando un sistema che utilizza i manifold semantici locali per espandere lo spazio di esplorazione senza sacrificare la coerenza.
• Validazione Empirica: Valutazioni estese su diverse scale di modelli (da 1.5B a 7B) e architetture (DeepSeek-R1-Distill-Qwen, Llama-3.2, Qwen3) dimostrano miglioramenti costanti rispetto alle strong baseline.

4. Risultati Sperimentali

Il metodo è stato valutato su benchmark di ragionamento matematico impegnativi (AIME25, AMC23, MATH500) e compiti out-of-distribution (OOD) (GPQA-Diamond).

• Ragionamento Matematico: N-GRPO ha superato costantemente le baseline, inclusi GRPO standard, Soft Thinking e la baseline STHT (che utilizza rumore Gaussiano non vincolato).
  • Sul modello DeepSeek-R1-Distill-Qwen-1.5B, N-GRPO ha migliorato la media Pass@32 da 77.41 (GRPO) a 79.17.
  • Sul modello 7B, la media Pass@32 è migliorata da 81.94 (GRPO) a 84.20.
  • Guadagni significativi sono stati osservati sul difficile benchmark AIME25, dove N-GRPO ha raggiunto un Pass@32 di 50.28 (1.5B) e 61.18 (7B), superando GRPO e STHT.
• Generalizzazione Out-of-Distribution: Sul benchmark GPQA-Diamond, N-GRPO ha dimostrato una robusta generalizzazione, ottenendo i punteggi Pass@32 più elevati (92.87 per 1.5B e 92.80 per 7B), indicando che il metodo potenzia le capacità di ragionamento scientifico generale, non solo quelle specifiche della matematica.
• Trasferibilità: Il meccanismo si è trasferito con successo al framework GSPO (N-GSPO), mostrando un guadagno di +7.66 su AIME25, suggerendo che il meccanismo di perturbazione non è strettamente accoppiato alla specifica normalizzazione dell'advantage di GRPO.
• Studi di Ablazione:
  • Mixing Rate: Le prestazioni sono rimaste stabili attraverso diversi tassi di mixing (0.1 a 0.2), sebbene un eccessivo mixing (0.2) su modelli piccoli per AIME25 abbia causato un calo, evidenziando la necessità di un equilibrio tra diversità e stabilità.
  • Metrica di Distanza: La distanza del coseno ha superato significativamente le metriche L1 e L2 per la selezione dei vicini, grazie alla natura direzionale dell'informazione semantica negli spazi di embedding ad alta dimensione.
  • Inference: Gli autori notano che applicare il Semantic Neighbor Mixing durante la fase di inference (senza il meccanismo di filtraggio GRPO) degrada le prestazioni, confermando il suo ruolo come strumento di esplorazione durante il training piuttosto che come strategia di decoding.

5. Significato e Rivendicazioni

Il paper sostiene che N-GRPO offre una strada promettente per avanzare le capacità di ragionamento degli LLM esplorando la dinamica del ragionamento nello spazio degli embedding. La sua importanza risiede nel:

• Risolvere il Compromesso: Inietta con successo diversità nel processo di rollout senza il drift semantico causato dal naive noise injection o la ridondanza del campionamento a livello di token.
• Fondamento Semantico: Sfruttando il manifold semantico locale, il metodo assicura che l'esplorazione rimanga entro regioni di ragionamento valide.
• Generalizzazione: L'approccio non è limitato a specifiche famiglie di modelli o policy distillate per il ragionamento, mostrando robustezza attraverso diverse scale e modelli base.

Gli autori concludono che, sebbene il metodo introduca un overhead computazionale durante il rollout e sia attualmente validato principalmente su compiti di matematica e scienza, esso stabilisce un nuovo paradigma per l'ottimizzazione della policy che bilancia esplorazione e coerenza semantica.