N-GRPO: Embedding-Level Neighbor Mixing for Enhanced Policy Optimization
Paper · 2026 · Xukun Zhu, Hang Yu, Peng Di et al.
Summary
RL fine-tuning methods such as GRPO depend on the rollout phase producing varied reasoning paths. Existing approaches have drawbacks. Ordinary token sampling tends to produce near-duplicate solutions that differ only in wording or operand order. Adding Gaussian noise to embeddings, by contrast, often moves tokens to semantically unrelated words, which the authors show with a PCA visualization. N-GRPO addresses this with Semantic Neighbor Mixing:
- A Bernoulli mask with rate ρ decides at which generation steps mixing is used.
- At those steps, the most likely token (argmax) serves as an anchor.
- Its k nearest neighbors in the embedding matrix are retrieved.
- Their embeddings are combined using logit-based weights, re-normalized over the neighbor set, and the mixture is fed back as the next input.
This adds randomness while staying near the anchor's meaning. The method was evaluated on DeepSeek-R1-Distill-Qwen at 1.5B and 7B. It outperforms strong baselines on math reasoning benchmarks in Pass@16 and Pass@32, and it also generalizes to out-of-distribution scientific tasks. The result suggests that constrained, embedding-space exploration is a practical way to enrich RL training for reasoning.
Key points
- Comparison with prior methods: N-GRPO is contrasted with HRPO, SofT-GRPO, and Gaussian noise injection (Butt et al., 2025). The authors argue that HRPO and SofT-GRPO still draw their randomness from discrete tokens. A baseline that adds Gaussian noise to embeddings improves on the vanilla model but scores below N-GRPO on average.
- Distance metric for neighbor retrieval (1.5B model, Mean@32 / Pass@16 / Pass@32):
- Cosine distance performs best: 54.11 / 76.82 / 79.17.
- L2 distance: 53.77 / 75.26 / 77.68.
- L1 distance: 53.07 / 72.98 / 75.16.
- The authors attribute this to semantic information being encoded mainly in vector direction rather than magnitude.
- Mixing rate matters: Applying mixing to every token ("w/o rate" ablation) lowers performance. The mixing rate therefore acts as a stability filter.
- Mixing is for training only: Using mixing at inference time reduces accuracy compared with standard temperature sampling.
- 7B model: Pass@32 falls from 84.20 to 81.45.
- 1.5B model: Pass@32 falls from 79.17 to 77.05.
- Limitations:
- Neighbor retrieval and weighted mixing add computational overhead during rollouts.
- The method has not been tested on code generation.
Institutions
Fields, topics and concepts
📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 works
Cite
Cite as Xukun Zhu, Hang Yu, Peng Di and Linchao Zhu. 2026. N-GRPO: Embedding-Level Neighbor Mixing for Enhanced Policy Optimization. Findings of the Association for Computational Linguistics: ACL 2026.
@inproceedings{zhu2026n,
title = {{N-GRPO: Embedding-Level Neighbor Mixing for Enhanced Policy Optimization}},
author = {Xukun Zhu, Hang Yu, Peng Di and Linchao Zhu},
booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
year = {2026},
doi = {10.18653/v1/2026.findings-acl.1443},
url = {https://doi.org/10.18653/v1/2026.findings-acl.1443}
}References
- Chenjia Bai, Yang Zhang, Shuang Qiu, Qiaosheng Zhang, Kang Xu, and Xuelong Li. 2025. Online preference alignment for language models via count-based exploration. arXiv preprint arXiv:2501.12735. — [arXiv:2501.12735](https://arxiv.org/abs/2501.12735) find
- Natasha Butt, Ariel Kwiatkowski, Ismail Labiad, Julia Kempe, and Yann Ollivier. 2025. Soft tokens, hard truths. arXiv preprint arXiv:2509.19170. — [arXiv:2509.19170](https://arxiv.org/abs/2509.19170) find
- Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, and Tieyan Liu. 2019. Representation degeneration problem in training natural language generation models. In International Conference on Learning Representations.
- Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, and 1 others. 2025. Deepseekr1 incentivizes reasoning in llms through reinforcement learning. Nature, 645(8081):633–638. find
- Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian. 2024. find
- Training large language models to reason in a continuous latent space. arXiv preprint arXiv:2412.06769. find
- Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the math dataset. In Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks, volume 1.
- Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, and 1 others. 2024. Openai o1 system card. arXiv preprint arXiv:2412.16720.
- Austin C Kozlowski, Callin Dai, and Andrei Boutyline. 2025. Semantic structure in large language model embeddings. arXiv preprint arXiv:2508.10003. — [arXiv:2508.10003](https://arxiv.org/abs/2508.10003) find
- Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Y Tang, Manan Roongta, Colin Cai, Jeffrey Luo, Tianjun Zhang, Li Erran Li, and 1 others. 2025. Deepscaler: Surpassing o1-preview with a 1.5 b model by scaling rl. Notion Blog. find
- Minh Nhat Nguyen, Andrew Baker, Clement Neo, Allen Roush, Andreas Kirsch, and Ravid Shwartz-Ziv. 2024. Turning up the heat: Min-p sampling for creative and coherent llm outputs. arXiv preprint arXiv:2407.01082. — [arXiv:2407.01082](https://arxiv.org/abs/2407.01082) find
- Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, and 1 others. 2022. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744.
- David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. 2024. GPQA: A graduate-level google-proof q&a benchmark. In First Conference on Language Modeling.
- Adriaan MJ Schakel and Benjamin J Wilson. 2015. Measuring word significance using distributed representations of words. arXiv preprint arXiv:1508.02297. — [arXiv:1508.02297](https://arxiv.org/abs/1508.02297) find
- Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, and 1 others. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300.
- Zhenyi Shen, Hanqi Yan, Linhai Zhang, Zhanghao Hu, Yali Du, and Yulan He. 2025. Codi: Compressing chain-of-thought into continuous space via self-distillation. arXiv preprint arXiv:2502.21074.
- Vaishnavi Shrivastava, Ahmed Awadallah, Vidhisha Balachandran, Shivam Garg, Harkirat Behl, and Dimitris Papailiopoulos. 2025. Sample more to think less: Group filtered policy optimization for concise reasoning. arXiv preprint arXiv:2508.09726. — [arXiv:2508.09726](https://arxiv.org/abs/2508.09726) find
- Wenhui Tan, Jiaze Li, Jianzhong Ju, Zhenbo Luo, Jian Luan, and Ruihua Song. 2025. Think silently, think fast: Dynamic latent compression of llm reasoning chains. arXiv preprint arXiv:2505.16552.
- Jianwei Wang, Ziming Wu, Fuming Lai, Shaobing Lian, and Ziqian Zeng. 2025a. Synadapt: Learning adaptive reasoning in large language models via synthetic continuous chain-of-thought. arXiv preprint arXiv:2508.00574. — [arXiv:2508.00574](https://arxiv.org/abs/2508.00574) find
- Xinyi Wang, Jinyi Han, Zishang Jiang, Tingyun Li, Jiaqing Liang, Sihang Jiang, Zhaoqian Dai, Shuguang Ma, Fei Yu, and Yanghua Xiao. 2025b. Hint: Helping ineffective rollouts navigate towards effectiveness. arXiv preprint arXiv:2510.09388. — [arXiv:2510.09388](https://arxiv.org/abs/2510.09388) find
- Xilin Wei, Xiaoran Liu, Yuhang Zang, Xiaoyi Dong, Yuhang Cao, Jiaqi Wang, Xipeng Qiu, and Dahua Lin. 2025. Sim-cot: Supervised implicit chain-of-thought. arXiv preprint arXiv:2509.20317. — [arXiv:2509.20317](https://arxiv.org/abs/2509.20317) find
- An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, and 1 others. 2025. Qwen3 technical report. arXiv preprint arXiv:2505.09388.
- Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, and 1 others. 2025. Dapo: An open-source llm reinforcement learning system at scale. arXiv preprint arXiv:2503.14476.
- Zhenrui Yue, Bowen Jin, Huimin Zeng, Honglei Zhuang, Zhen Qin, Jinsung Yoon, Lanyu Shang, Jiawei Han, and Dong Wang. 2025. Hybrid latent reasoning via reinforcement learning. arXiv preprint arXiv:2505.18454. — [arXiv:2505.18454](https://arxiv.org/abs/2505.18454) find
- Jintian Zhang, Yuqi Zhu, Mengshu Sun, Yujie Luo, Shuofei Qiao, Lun Du, Da Zheng, Huajun Chen, and Ningyu Zhang. 2025a. Lightthinker: Thinking step-by-step compression. arXiv preprint arXiv:2502.15589.
- Zhen Zhang, Xuehai He, Weixiang Yan, Ao Shen, Chenyang Zhao, Shuohang Wang, Yelong Shen, and Xin Eric Wang. 2025b. Soft thinking: Unlocking the reasoning potential of llms in continuous concept space. arXiv preprint arXiv:2505.15778.
- Chujie Zheng, Shixuan Liu, Mingze Li, Xiong-Hui Chen, Bowen Yu, Chang Gao, Kai Dang, Yuqiong Liu, Rui Men, An Yang, and 1 others. 2025. Group sequence policy optimization. arXiv preprint arXiv:2507.18071.
- Zhi Zheng and Wee Sun Lee. 2025. Soft-grpo: Surpassing discrete-token llm reinforcement learning via gumbel-reparameterized soft-thinking policy optimization. arXiv preprint arXiv:2511.06411. — [arXiv:2511.06411](https://arxiv.org/abs/2511.06411) find
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/zhu-2026-n-grpo-embedding-level-neighbor/