PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xukun Zhu, Hang Yu, Peng Di, Linchao Zhu

arXiv:2606.10768v1cs.LGcs.CL

TL;DR

Rollout exploration in GRPO can produce redundant trajectories through token sampling or disrupt semantic consistency through random embedding noise. N-GRPO mixes an anchor token’s embedding with semantic neighbors and consistently outperforms baselines across model scales while generalizing to scientific tasks.

  • Problem

    GRPO rollout exploration can yield redundant trajectories through token-level sampling or disrupt semantic consistency through direct embedding or logit noise.

  • Method

    N-GRPO integrates Semantic Neighbor Mixing into GRPO by combining an argmax anchor token’s embedding with weighted nearest semantic neighbors during generation.

  • Results

    N-GRPO consistently outperforms strong baselines on Pass@16 and Pass@32 across 1.5B and 7B models, with superior generalization on scientific tasks.

  • Takeaways & Limitations

    The findings support semantically guided embedding-level exploration as a way to expand reasoning-path diversity while maintaining semantic stability.

  • Takeaways & Limitations

    The method adds rollout computational overhead and has not been validated on code-generation tasks.

Abstract

from arXiv · show

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.

1 Introduction

N-GRPO addresses redundant token-level exploration and semantically disruptive embedding-level noise by mixing sampled-token embeddings with nearest semantic neighbors during GRPO rollouts. Across 1.5B and 7B models, it consistently outperforms strong baselines on mathematical reasoning benchmarks under Pass@16 and Pass@32.

  • Motivation: Token-level sampling often produces redundant trajectories that differ mainly through paraphrasing or commutative re-orderings while preserving the same reasoning logic.Examples include equivalent expressions such as “1+2” and “2+1”.
  • Motivation: Embedding-level exploration must adapt to local semantic context because Transformer embedding spaces typically exhibit strong anisotropy.The paper motivates semantic control partly by illustrating semantic drift from unconstrained Gaussian noise applied to token embeddings.
  • Results: N-GRPO consistently outperforms strong baselines on challenging mathematical reasoning benchmarks across 1.5B and 7B model scales.The reported improvements hold for both Pass@16 and Pass@32 metrics.
  • Method: N-GRPO integrates controllable embedding-level exploration into GRPO through Semantic Neighbor Mixing during generation.The method mixes a sampled token’s embedding with embeddings of nearest neighbors retrieved by embedding similarity, which cluster along the anchor’s semantic direction.

2 Related Work

Related work frames rollout diversity as central to reinforcement-learning alignment and surveys efforts to move reasoning or exploration into continuous latent representations. These approaches reduce reliance on explicit Chain of Thought or discrete sampling but face semantic drift, instability, or off-manifold perturbations.

  • RL alignment methods depend on rollout phases that generate diverse trajectories for accurate gradient estimation and advantage calculation.
  • Continuous Exploration: HRPO, SofT-GRPO, and Gaussian-noise methods move rollout exploration into continuous representations, but randomness remains tied to token sampling or token-level noise.
  • Limitations: Unconstrained embedding or logit noise can push representations off the semantic manifold and disrupt rollouts, while latent adaptation can cause semantic drift and instability.SIM-COT uses step-level supervision with an assistant decoder, and CODI uses self-distillation alignment between explicit and implicit reasoning.
  • Latent Reasoning: Latent reasoning carries intermediate computation in continuous hidden vectors rather than explicit text, decoupling thinking from language until the final answer.
  • Latent Reasoning: Token-based latent-reasoning methods replace or compress explicit Chain of Thought to reduce reasoning or context costs.

3 Method

N-GRPO expands rollout exploration by mixing embeddings within an anchor token’s local semantic neighborhood, with weights determined by current logits. Integrated into GRPO through stochastic gating, it preserves standard optimization and discrete reward computation while introducing continuous perturbations.

  • Semantic Neighbor Mixing: N-GRPO selects the model’s argmax token as an anchor, then forms a k-token neighbor set using cosine similarity in the embedding space.The neighbor set includes the anchor and the k−1 most similar tokens.
  • Semantic Neighbor Mixing: It normalizes temperature-scaled logits over the candidate set to obtain mixing weights and constructs a continuous mixed embedding for the next step.The resulting perturbation remains within the anchor’s local semantic neighborhood while adapting its direction and intensity to current logits.
  • N-GRPO Rollout Integration: N-GRPO applies this embedding-level sampling during GRPO rollouts while preserving GRPO’s group-relative advantage estimation and optimization objectives.For each input, GRPO continues sampling a group of outputs from the old policy and updating from group-relative advantages.
  • N-GRPO Rollout Integration: A Bernoulli(ρ) mask stochastically selects embedding mixing or standard token sampling at each step, retaining semantic stability through partial standard sampling.The mixing rate is controlled by the fixed probability ρ.
  • N-GRPO Rollout Integration: Mixed steps use their anchor tokens for discrete verification and rewards, while advantages and optimization follow standard GRPO with reconstructed rollout representations.Candidate sets and weights are recorded for reproducibility, and the objective retains PPO-style clipping with KL regularization.

4 Experiments

Experiments show that N-GRPO improves mathematical reasoning across distilled and non-distilled backbones, generalizes to out-of-distribution scientific reasoning, and transfers beyond GRPO. Ablations indicate robustness to moderate mixing-rate changes but sensitivity to excessive or inference-time mixing.

  • Experimental Setup: N-GRPO is evaluated on reasoning-distilled DeepSeek-R1-Distill-Qwen models and non-distilled Llama-3.2-1B and Qwen3-1.7B-Base backbones.The evaluation covers AMC23, AIME25, MATH500, and the out-of-distribution GPQA-Diamond benchmark.
  • In-Domain Results: 79.17: N-GRPO raises the 1.5B model’s average Pass@32 from 74.62 for the base model and 77.41 for GRPO.On AIME25, Pass@32 increases to 50.28 from 41.19 for the base model and 47.31 for GRPO.
  • In-Domain Results: 5.00 points: N-GRPO improves Qwen3-1.7B-Base AIME25 Pass@32 over GRPO and achieves the highest overall average Pass@32.On Llama-3.2-1B, it attains the highest average Pass@16 and Pass@32 and the best Pass@32 on all three math benchmarks.
  • Out-of-Distribution Generalization: 92.87: N-GRPO achieves GPQA-Diamond Pass@32 for the 1.5B model, exceeding the base model’s 90.79 and GRPO+Soft Thinking’s 91.92.The 7B model also achieves the highest performance across all reported metrics.
  • Transfer Beyond GRPO: 79.04: N-GSPO improves average Pass@32 from 77.34 over GSPO, including a +7.66 gain on AIME25.Applying Semantic Neighbor Mixing to GSPO suggests the mechanism is not tightly coupled to GRPO’s advantage normalization.
  • Ablations and Limitations: Mixing-rate performance remains relatively stable, but increasing the rate to 0.2 significantly reduces 1.5B AIME25 Pass@32.The study attributes this exception to excessive embedding-level exploration disrupting semantic stability; applying mixing at inference also degrades accuracy versus standard temperature sampling.

5 Conclusion

N-GRPO presents an embedding-level exploration strategy that improves reinforcement-learning efficiency for LLMs by generating diverse reasoning paths while preserving semantic stability.

  • 5 Conclusion: N-GRPO replaces discrete token-level randomness with continuous embedding-level perturbations to encourage diverse reasoning paths without sacrificing semantic stability.The method implements semantically grounded noise through Semantic Neighbor Mixing and integrates it into the GRPO framework.
  • 5 Conclusion: Semantic Neighbor Mixing provides the mechanism for embedding-level exploration in the proposed N-GRPO approach.It mixes embedding representations to introduce semantically grounded exploration within GRPO.

Limitations … B.1 Codebase and Framework

The paper identifies rollout-time computational overhead and limited benchmark coverage as limitations, while specifying standardized reasoning prompts and a modified verl/sglang implementation for training and evaluation.

  • Limitations: Semantic Neighbor Mixing adds rollout overhead by retrieving nearest neighbors and computing weighted embeddings at generation steps, potentially increasing inference latency.Unlike standard sampling, the method operates beyond already-computed logits.
  • Limitations: Experimental validation is primarily concentrated on mathematical and scientific reasoning benchmarks.
  • A Task Prompt: The prompt examples cover an isosceles-trapezoid math problem and a solar-neutrino GPQA-Diamond question with explicit assumptions.The GPQA example asks for a flux ratio between 700-800 KeV and 800-900 keV bands.
  • A Task Prompt: Training and evaluation prompts use explicit Chain-of-Thought instructions, placing the user query before formatting constraints.The template is intended to elicit reasoning and standardize output format.
  • A Task Prompt: Math tasks require boxed final answers, whereas GPQA multiple-choice tasks enforce a specific answer format.The prompts directly append the <think> token to initiate reasoning.
  • B.1 Codebase and Framework: The implementation modifies verl version 0.5.0 and sglang version 0.4.6.post5 to support the proposed training, evaluation, and Semantic Neighbor Mixing rollout pipelines.sglang serves as the rollout backend, with corresponding verl-interface adaptations for compatibility.

B.2 Hyperparameters for Training · B.3 Hyperparameters for Inference · C Evaluation

Training uses full-parameter fine-tuning for the 1.5B and 7B DeepSeek-R1-Distill-Qwen models, with detailed settings provided in Table 6. Inference uses sglang with specified sampling and generation-length settings aligned to training.

  • B.2 Hyperparameters for Training: Full-parameter fine-tuning is performed on DeepSeek-R1-Distill-Qwen-1.5B and DeepSeek-R1-Distill-Qwen-7B for the main experiments.
  • B.2 Hyperparameters for Training: The training setup covers both the 1.5B and 7B model sizes used in the main experiments.
  • B.2 Hyperparameters for Training: The detailed hyperparameters used during training are listed in Table 6.
  • B.3 Hyperparameters for Inference: Inference uses the sglang engine with the recommended base-model settings.
  • B.3 Hyperparameters for Inference: 0.6 is the inference sampling temperature.
  • B.3 Hyperparameters for Inference: 0.95 is used for nucleus sampling (top-p), while maximum generation length is set to 8192 to match training.

C.1 Datasets · C.2 Metrics · C.3 Answer Extraction and Grading

The evaluation spans mathematical and scientific reasoning benchmarks, using Mean@32 and Pass@k metrics over 32 generated solutions. A strict task-specific pipeline extracts and verifies mathematical answers or parses multiple-choice labels.

  • C.1 Datasets: The evaluation covers four recognized benchmarks spanning in-domain mathematical reasoning and out-of-distribution scientific reasoning.The dataset statistics are summarized in Table 7.
  • C.1 Datasets: AMC23, AIME24, AIME25, and MATH500 assess competition-level and diverse mathematical reasoning, with AIME serving as a challenging multi-step holdout.MATH500 includes algebra, geometry, number theory, and probability.
  • C.1 Datasets: GPQA-Diamond tests out-of-distribution scientific reasoning through high-difficulty expert questions in biology, physics, and chemistry.It is a subset of the GPQA benchmark.
  • C.2 Metrics: Mean@32 reports average accuracy, while Pass@16 and Pass@32 measure exploration potential from generated solutions.For each problem, the evaluation generates n = 32 candidate solutions.
  • C.2 Metrics: Mean@n estimates the expected accuracy of one sampled generation and is reported in Table 1’s “@1” column as an unbiased estimator of Pass@1.The metric is computed over n = 32 samples.
  • C.2 Metrics: Pass@k estimates whether at least one correct solution appears among k samples; for Pass@32, it reduces to checking whether c > 0.Here c is the number of correct samples among n = 32 generations.
  • C.3 Answer Extraction and Grading: Mathematical answers are extracted from the last boxed environment and checked against ground truth with Math-Verify symbolic equivalence utilities.This handles formatting variations such as equivalent fractions, decimals, and simplified expressions.
  • C.3 Answer Extraction and Grading: For GPQA-Diamond, grading parses “Answer: <Option>” and directly compares A–D labels, marking format failures or ambiguous outputs incorrect.This multiple-choice rule applies specifically to GPQA-Diamond.

D Transfer to GSPO (N-GSPO)

N-GSPO applies Semantic Neighbor Mixing to GSPO and improves average Pass@32 on DeepSeek-R1-Distill-Qwen-1.5B, including a +7.66 gain on AIME25. This transfer indicates the exploration benefit extends beyond GRPO-specific advantage estimation.

  • Transfer to GSPO (N-GSPO): N-GSPO integrates Semantic Neighbor Mixing into GSPO and evaluates the resulting method on DeepSeek-R1-Distill-Qwen-1.5B.All other training and evaluation settings follow Section 4.1.
  • Transfer to GSPO (N-GSPO): +7.66 AIME25 gain is observed for N-GSPO over GSPO in average Pass@32.The table reports that N-GSPO improves average Pass@32 over GSPO, with the particularly large gain on AIME25.
  • Transfer to GSPO (N-GSPO): Semantic Neighbor Mixing transfers to another rollout-based optimization pipeline rather than depending on GRPO-specific advantage estimation.The result supports semantically grounded embedding-level exploration during rollout beyond the GRPO framework.

E Semantic Coherence and Trajectory Diversity of Mixed Embeddings

Semantic Neighbor Mixing keeps mixed embeddings tightly aligned with anchor semantics while still producing genuinely different reasoning trajectories. On AIME24, it adds unique problem solutions without sacrificing any problems solved by standard sampling.

  • Proximity to the anchor semantics: 0.9985 average cosine similarity to the anchor, with only 1.2% of mixed steps below 0.95, demonstrates tight local semantic alignment.The minimum cosine similarity was 0.8813.
  • Meaningful trajectory variation: 10.0% of AIME24 problems were solved uniquely by Mixing, while it lost no problem solvable under standard sampling.The paired inference experiment generated 32 samples per problem with identical hyperparameters and compared Standard temperature sampling against Mixing.

F Full Ablation Results · G Rollout Throughput Overhead

The paper provides comprehensive ablations of mixing rates, strategies, distance metrics, and inference procedures across mathematical benchmarks and model scales. It also reports that Semantic Neighbor Mixing adds less than 10% rollout overhead while maintaining accuracy gains.

  • F Full Ablation Results: Full ablations report Mean@32, Pass@16, and Pass@32 across AIME25, AMC23, and MATH500; @1 denotes Mean@32.The evaluations cover the DeepSeek-R1-Distill-Qwen-1.5B and 7B scales.
  • F Full Ablation Results: Table 10 analyzes sensitivity to the mixing rate parameter ρ for the 1.5B and 7B models.Performance is evaluated using Mean@32, Pass@16, and Pass@32 across three benchmarks.
  • F Full Ablation Results: Table 11 compares N-GRPO with Gumbel Soft-Thinking and an unconstrained w/o rate variant on the 1.5B model.The unconstrained variant applies mixing to all tokens.
  • F Full Ablation Results: Table 12 compares Cosine, L2, and L1 distance metrics for latent-space neighbor selection on DeepSeek-R1-Distill-Qwen-1.5B.The metrics represent cosine, Euclidean, and Manhattan distances, respectively.
  • G Rollout Throughput Overhead: The rollout study measures average throughput in tokens/s on identical hardware using the same rollout backend.Standard Sampling denotes GRPO rollouts without mixing, whereas Semantic Neighbor Mixing denotes N-GRPO rollouts with mixing enabled.
  • G Rollout Throughput Overhead: Below 10% overhead is reported for Semantic Neighbor Mixing at both the 1.5B and 7B scales.The paper describes this as a modest practical cost relative to standard sampling.
  • G Rollout Throughput Overhead: The modest overhead arises because mixing affects few tokens, nearest-neighbor sets can be precomputed, and aggregation is a small weighted sum over k embeddings.The mixing rate ρ controls activation frequency during rollout.
Loading 2606.10768v1…