N-GRPO: Embedding-Level Neighbor Mixing for Enhanced Policy Optimization

2026-06-09

AI Summary

This work addresses the challenge that existing large language models struggle to balance exploration diversity and semantic consistency in mathematical reasoning: token-level sampling often yields redundant outputs, while embedding-level noise tends to disrupt semantic structure. To overcome this, the authors propose N-GRPO, a novel method that integrates a semantic neighborhood mixing mechanism into the Group Relative Policy Optimization (GRPO) framework. By dynamically fusing anchor tokens with their nearest-neighbor semantic representations in embedding space, N-GRPO enables efficient exploration within local semantic manifolds. Evaluated on the DeepSeek-R1-Distill-Qwen model series, the approach significantly outperforms strong baselines, demonstrating superior performance and enhanced generalization on both in-distribution mathematical reasoning benchmarks and out-of-distribution tasks.

Abstract

The success of Large Language Models in mathematical reasoning relies heavily on the generation of diverse and valid solution paths during the rollout phase. However, current rollout techniques face a fundamental trade-off: token-level sampling often yields redundant trajectories that differ only in rephrasing, while embedding-level methods utilizing random noise frequently disrupt semantic consistency. To resolve this, we introduce N-GRPO, a novel exploration strategy integrated into the Group Relative Policy Optimization (GRPO) framework. Rather than relying on token-level sampling or native embedding-level noise, our approach leverages Semantic Neighbor Mixing. This mechanism dynamically constructs input representations by mixing the embeddings of an anchor token and its nearest semantic neighbors, thereby injecting diversity while strictly adhering to the local semantic manifold. Experimental evaluations on the DeepSeek-R1-Distill-Qwen models across different sizes show that N-GRPO not only achieves consistent improvements over strong baselines on math reasoning benchmarks but also exhibits robust generalization capabilities on out-of-distribution tasks.

Problem

• mathematical reasoning
• rollout diversity
• semantic consistency
• policy optimization
• embedding-level sampling

Innovation

• Semantic Neighbor Mixing
• Embedding-Level Exploration
• Group Relative Policy Optimization
• Mathematical Reasoning
• Diversity-Preserving Generation