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

1. Problem Statement

The success of Large Language Models (LLMs) in mathematical reasoning relies heavily on generating diverse and valid solution paths during the rollout phase of Reinforcement Learning (RL). However, current exploration techniques face a fundamental trade-off:

• Token-Level Sampling: Standard methods often yield redundant trajectories that differ only in rephrasing or commutative re-orderings (e.g., "1+2" vs. "2+1") without altering the underlying reasoning logic.
• Embedding-Level Noise: Methods that inject random noise (e.g., Gaussian noise) into continuous embeddings frequently disrupt semantic consistency. Preliminary experiments showed that unconstrained noise pushes representations off the semantic manifold, transforming valid tokens into unrelated ones and derailing rollout trajectories.

Existing attempts to bridge this gap, such as Soft Thinking or HRPO, often rely on discrete token sampling for their randomness or fail to maintain semantic coherence when perturbing the continuous space.

2. Methodology: N-GRPO

The authors propose N-GRPO, a novel exploration strategy integrated into the Group Relative Policy Optimization (GRPO) framework. The core innovation is Semantic Neighbor Mixing, an embedding-level perturbation mechanism that injects diversity while strictly adhering to the local semantic manifold.

Core Mechanism: Semantic Neighbor Mixing

Instead of sampling discrete tokens or adding unconstrained noise, the method constructs input representations by mixing the embedding of an anchor token with its nearest semantic neighbors.

1. Anchor Selection: At each generation step, the model identifies the most preferred token (the semantic anchor, $v^*_t$) based on the temperature-scaled logits.
2. Neighbor Retrieval: A set of $k$ neighbor tokens ($C_t$) is retrieved from the embedding matrix based on cosine similarity to the anchor. This ensures neighbors cluster along the same semantic direction.
3. Weighted Mixing: The temperature-scaled logits are normalized over this neighbor set to compute mixing weights ($\alpha_t$).
4. Continuous Embedding Construction: A continuous mixed embedding ($\tilde{e}_{t+1}$) is generated as the weighted sum of the neighbor embeddings. This serves as the input for the next time step.

Integration with GRPO

N-GRPO integrates this mechanism into the GRPO rollout phase via a stochastic gating mechanism:

• Mixing Mask: A Bernoulli mask ($m_{i,t}$) with a fixed probability $\rho$ (e.g., 0.1) determines whether to activate Semantic Neighbor Mixing or follow standard discrete token sampling.
• Stability: By keeping the majority of steps as discrete tokens, the method preserves semantic stability while introducing necessary intermittency for exploration.
• Training: During the training phase, the candidate sets and weights are recorded to ensure reproducible trajectory reconstruction for advantage calculation. The final reward is computed based on the discrete anchor tokens corresponding to mixed steps.

3. Key Contributions

• Controllable Embedding-Level Exploration: The paper introduces a method that injects semantically guided perturbations during generation, bridging the gap between discrete token-level reasoning and continuous embedding-level exploration.
• N-GRPO Framework: The authors integrate Semantic Neighbor Mixing into the GRPO framework, creating a system that utilizes local semantic manifolds to expand the exploration space without sacrificing coherence.
• Empirical Validation: Extensive evaluations across different model scales (1.5B to 7B) and architectures (DeepSeek-R1-Distill-Qwen, Llama-3.2, Qwen3) demonstrate consistent improvements over strong baselines.

4. Experimental Results

The method was evaluated on challenging mathematical reasoning benchmarks (AIME25, AMC23, MATH500) and out-of-distribution (OOD) tasks (GPQA-Diamond).

• Mathematical Reasoning: N-GRPO consistently outperformed baselines, including standard GRPO, Soft Thinking, and the STHT baseline (which uses unconstrained Gaussian noise).
  • On the DeepSeek-R1-Distill-Qwen-1.5B model, N-GRPO improved the average Pass@32 from 77.41 (GRPO) to 79.17.
  • On the 7B model, the average Pass@32 improved from 81.94 (GRPO) to 84.20.
  • Significant gains were observed on the challenging AIME25 benchmark, where N-GRPO achieved a Pass@32 of 50.28 (1.5B) and 61.18 (7B), surpassing GRPO and STHT.
• Out-of-Distribution Generalization: On the GPQA-Diamond benchmark, N-GRPO demonstrated robust generalization, achieving the highest Pass@32 scores (92.87 for 1.5B and 92.80 for 7B), indicating the method enhances general scientific reasoning capabilities, not just math-specific skills.
• Transferability: The mechanism successfully transferred to the GSPO framework (N-GSPO), showing a +7.66 gain on AIME25, suggesting the perturbation mechanism is not tightly coupled to GRPO's specific advantage normalization.
• Ablation Studies:
  • Mixing Rate: Performance remained stable across varying mixing rates (0.1 to 0.2), though excessive mixing (0.2) on smaller models for AIME25 caused a drop, highlighting the need for a balance between diversity and stability.
  • Distance Metric: Cosine distance significantly outperformed L1 and L2 metrics for neighbor selection, attributed to the directional nature of semantic information in high-dimensional embedding spaces.
  • Inference: The authors note that applying Semantic Neighbor Mixing during the inference phase (without the GRPO filtering mechanism) degrades performance, confirming its role as a training-time exploration tool rather than a decoding strategy.

5. Significance and Claims

The paper claims that N-GRPO offers a promising avenue for advancing LLM reasoning capabilities by exploring reasoning dynamics in the embedding space. Its significance lies in:

• Resolving the Trade-off: It successfully injects diversity into the rollout process without the semantic drift caused by naive noise injection or the redundancy of token-level sampling.
• Semantic Grounding: By leveraging the local semantic manifold, the method ensures that exploration remains within valid reasoning regions.
• Generalization: The approach is not limited to specific model families or reasoning-distilled policies, showing robustness across different scales and base models.

The authors conclude that while the method introduces computational overhead during rollout and is currently validated primarily on math and science tasks, it establishes a new paradigm for policy optimization that balances exploration and semantic consistency.