AttnPO: Attention-Guided Process Supervision for Efficient Reasoning
Paper · 2026 · Shuaiyi Nie, Dingsiyu, Wenyuan Zhang et al.
Summary
Reasoning models trained with RLVR tend to overthink, producing long chains of thought with little benefit. Trajectory-level length penalties treat every step the same, so they often fail to shorten outputs and can hurt accuracy. Process-supervised alternatives need extra Monte Carlo sampling or a separate reward model, and they still assign credit imprecisely. The authors find that some attention heads, which they call Key-Focus Heads (KFHs), already separate useful steps from wasted ones. While the model writes its final answer after </think>, these heads attend strongly to essential reasoning steps and largely ignore redundant ones. AttnPO uses the attention scores of these heads to rescale GRPO-style advantages for each step of correct responses. It shrinks the positive advantage given to redundant steps and softens the negative advantage applied to essential steps. This adds almost no training overhead. Across 9 benchmarks, AttnPO shortens reasoning considerably and also raises accuracy. This makes it a cheap route to more efficient reasoning models.
Key points
- Probing data: 300 easier Deepscaler questions were used. Steps were labeled as essential or redundant by GPT-5, Claude-4-sonnet and Gemini-2.5-pro, keeping only labels all three agreed on.
- Head quality: the best KFHs reach Step Ranking Accuracy of about 95% in DeepSeek-R1-Distill-Qwen-1.5B and 96% in the 7B model. Most other heads score near chance (~0.5).
- Where KFHs sit: they are concentrated in middle-to-late layers and stay stable during length-penalized RL. A few heads are enough, since ensembling more of them saturates quickly.
- Main result: on DeepSeek-R1-Distill-Qwen-1.5B, AttnPO gains +7.3 points of average accuracy while cutting reasoning length by about 60% across six math benchmarks.
- Generalization and scope: training used only math data, yet the method transfers to LiveCodeBench, GPQA and MMLU. Experiments are limited to 1.5B and 7B models, and the code is at github.com/NieSYsc20/AttnPO.
Institutions
Fields, topics and concepts
📚 Shelves 5The 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 worksComputer Vision and Pattern Recognition · 3,486 works
🧭 Topics 1Research subjects and languages the work is about.Reasoning · 1,685 works
Cite
Cite as Shuaiyi Nie, Dingsiyu, Wenyuan Zhang, Linhao Yu, Tianmeng Yang, Yao Chen, Weichong Yin, Yu Sun, Hua Wu and Tingwen Liu. 2026. AttnPO: Attention-Guided Process Supervision for Efficient Reasoning. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).
@inproceedings{nie2026attnpo,
title = {{AttnPO: Attention-Guided Process Supervision for Efficient Reasoning}},
author = {Shuaiyi Nie, Dingsiyu, Wenyuan Zhang, Linhao Yu, Tianmeng Yang, Yao Chen, Weichong Yin, Yu Sun, Hua Wu and Tingwen Liu},
booktitle = {Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
year = {2026},
doi = {10.18653/v1/2026.acl-long.1845},
url = {https://doi.org/10.18653/v1/2026.acl-long.1845}
}References
- Pranjal Aggarwal and Sean Welleck. 2025. L1: Controlling how long a reasoning model thinks with reinforcement learning. arXiv preprint arXiv:2503.04697.
- Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzieh Fadaee, Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. 2024. Back to basics: Revisiting REINFORCE-style optimization for learning from human feedback in LLMs. In Proceedings of the 62nd Annual Meeting of the Association for Compu… find
- Zhengyang Ai, Zikang Shan, Xiaodong Ai, Jingxian Tang, Hangkai Hu, and Pinyan Lu. 2026. Shape: Stage-aware hierarchical advantage via potential estimation for llm reasoning. arXiv preprint arXiv:2604.06636. find
- Shengnan An, Xunliang Cai, Xuezhi Cao, Xiaoyu Li, Yehao Lin, Junlin Liu, Xinxuan Lv, Dan Ma, Xuanlin Wang, Ziwen Wang, and 1 others. 2025. Amo-bench: Large language models still struggle in high school math competitions. arXiv preprint arXiv:2510.26768. find
- Daman Arora and Andrea Zanette. 2025. Training language models to reason efficiently. In The Thirty-ninth Annual Conference on Neural Information Processing Systems.
- Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, and 39 others. 202… — Evaluating Large Language Models Trained on Code
- Weize Chen, Jiarui Yuan, Tailin Jin, Ning Ding, Huimin Chen, Zhiyuan Liu, and Maosong Sun. 2025a. The overthinker’s diet: Cutting token calories with difficulty-aware training. NeurIPS. find
- Xingyu Chen, Jiahao Xu, Tian Liang, Zhiwei He, Jianhui Pang, Dian Yu, Linfeng Song, Qiuzhi Liu, Mengfei Zhou, Zhuosheng Zhang, Rui Wang, Zhaopeng Tu, Haitao Mi, and Dong Yu. 2025b. Do not think that much for 2+3=? on the overthinking of o1-like llms. Preprint, arXiv:2412.21187.
- Yao Chen, Yilong Chen, Yinqi Yang, Junyuan Shang, Zhenyu Zhang, Zefeng Zhang, Shuaiyi Nie, Shuohuan Wang, Yu Sun, Hua Wu, and 1 others. 2026. Sparse growing transformer: Training-time sparse depth allocation via progressive attention looping. arXiv preprint arXiv:2603.23998. find
- Xiaoxue Cheng, Junyi Li, Zhenduo Zhang, Xinyu Tang, Wayne Xin Zhao, Xinyu Kong, and Zhiqiang Zhang. 2025a. Incentivizing dual process thinking for efficient large language model reasoning. arXiv preprint arXiv:2505.16315. — [arXiv:2505.16315](https://arxiv.org/abs/2505.16315) find
- Zhengxiang Cheng, Dongping Chen, Mingyang Fu, and Tianyi Zhou. 2025b. Optimizing length compression in large reasoning models. arXiv preprint arXiv:2506.14755.
- Daewon Choi, Jimin Lee, Jihoon Tack, Woomin Song, Saket Dingliwal, Sai Muralidhar Jayanthi, Bhavana Ganesh, Jinwoo Shin, Aram Galstyan, and Sravan Babu Bodapati. 2025. Think clearly: Improving reasoning via redundant token pruning. In Findings of the Association for Computational Linguistics: EMNLP… find
- Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems. Preprint, arXiv:2110.14168.
- Muzhi Dai, Chenxu Yang, and Qingyi Si. 2025. S-grpo: Early exit via reinforcement learning in reasoning models. NeurIPS. find
- Jingcheng Deng, Liang Pang, Zihao Wei, Shichen Xu, Zenghao Duan, Kun Xu, Yang Song, Huawei Shen, and Xueqi Cheng. 2025. Latent reasoning in llms as a vocabulary-space superposition. arXiv preprint arXiv:2510.15522. — [arXiv:2510.15522](https://arxiv.org/abs/2510.15522) find
- Haonan Dong, Kehan Jiang, Haoran Ye, Wenhao Zhu, Zhaolu Kang, and Guojie Song. 2026. Neureasoner: Towards explainable, controllable, and unified reasoning via mixture-of-neurons. Preprint, arXiv:2604.02972.
- Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Peiyi Wang, Qihao Zhu, Runxin Xu, Ruoyu Zhang, Shirong Ma, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, and 175 others. 2025. DeepSeek-R1: Incentivizing reasoning capability in llms…
- Hasan Abed Al Kader Hammoud, Hani Itani, and Bernard Ghanem. 2025. Beyond the last answer: Your reasoning trace uncovers more than you think. Preprint, arXiv:2504.20708. — [arXiv:2504.20708](https://arxiv.org/abs/2504.20708) find
- Chaoqun He, Renjie Luo, Yuzhuo Bai, Shengding Hu, Zhen Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, Jie Liu, Lei Qi, Zhiyuan Liu, and Maosong Sun. 2024. OlympiadBench: A challenging benchmark for promoting AGI with olympiad-level bilingual multimodal scientific problems. In Proce…
- Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.
- Bairu Hou, Yang Zhang, Jiabao Ji, Yujian Liu, Kaizhi Qian, Jacob Andreas, and Shiyu Chang. 2025. Thinkprune: Pruning long chain-of-thought of llms via reinforcement learning. arXiv preprint arXiv:2504.01296.
- Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang, Armando Solar-Lezama, Koushik Sen, and Ion Stoica. 2025. Livecodebench: Holistic and contamination free evaluation of large language models for code. In The Thirteenth International Conference on Learning Representatio…
- Kehan Jiang, Haonan Dong, Zhaolu Kang, Zhengzhou Zhu, and Guojie Song. 2026. Foe: Forest of errors makes the first solution the best in large reasoning models. arXiv preprint arXiv:2604.02967.
- Shuyang Jiang, Yusheng Liao, Ya Zhang, Yanfeng Wang, and Yu Wang. 2025. Overthinking reduction with decoupled rewards and curriculum data scheduling. arXiv preprint arXiv:2509.25827. — [arXiv:2509.25827](https://arxiv.org/abs/2509.25827) find
- Renren Jin, Pengzhi Gao, Yuqi Ren, Zhuowen Han, Tongxuan Zhang, Wuwei Huang, Wei Liu, Jian Luan, and Deyi Xiong. 2025. Revisiting entropy in reinforcement learning for large reasoning models. arXiv preprint arXiv:2511.05993. find
- Qiming Li, Zekai Ye, Xiaocheng Feng, Weihong Zhong, Libo Qin, Ruihan Chen, Baohang Li, Kui Jiang, Yaowei Wang, Ting Liu, and 1 others. 2025a. Cai: Caption-sensitive attention intervention for mitigating object hallucination in large vision-language models. arXiv preprint arXiv:2506.23590. — [arXiv:2506.23590](https://arxiv.org/abs/2506.23590) find
- Ruoran Li, Xinghua Zhang, Haiyang Yu, Shitong Duan, Xiang Li, Wenxin Xiang, Chonghua Liao, Xudong Guo, Yongbin Li, and Jinli Suo. 2026. Mempo: Self-memory policy optimization for long-horizon agents. arXiv preprint arXiv:2603.00680. find
- Zhong-Zhi Li, Duzhen Zhang, Ming-Liang Zhang, Jiaxin Zhang, Zengyan Liu, Yuxuan Yao, Haotian Xu, Junhao Zheng, Pei-Jie Wang, Xiuyi Chen, Yingying Zhang, Fei Yin, Jiahua Dong, Zhiwei Li, Bao-Long Bi, Ling-Rui Mei, Junfeng Fang, Xiao Liang, Zhijiang Guo, and 2 others. 2025b. From system 1 to system 2:…
- Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2024. Let’s verify step by step. In The Twelfth International Conference on Learning Representations.
- Guoming Ling, Zhongzhan Huang, Yupei Lin, Junxin Li, Shanshan Zhong, Hefeng Wu, and Liang Lin. 2026. Neural chain-of-thought search: Searching the optimal reasoning path to enhance large language models. arXiv preprint arXiv:2601.11340. find
- Tianci Liu, Ran Xu, Tony Yu, Ilgee Hong, Carl Yang, Tuo Zhao, and Haoyu Wang. 2026. Openrubrics: Towards scalable synthetic rubric generation for reward modeling and llm alignment. Preprint, arXiv:2510.07743.
- Wei Liu, Ruochen Zhou, Yiyun Deng, Yuzhen Huang, Junteng Liu, Yuntian Deng, Yizhe Zhang, and Junxian He. 2025a. Learn to reason efficiently with adaptive length-based reward shaping. Preprint, arXiv:2505.15612. — [arXiv:2505.15612](https://arxiv.org/abs/2505.15612) find
- Yule Liu, Jingyi Zheng, Zhen Sun, Zifan Peng, Wenhan Dong, Zeyang Sha, Shiwen Cui, Weiqiang Wang, and Xinlei He. 2025b. Thought manipulation: External thought can be efficient for large reasoning models. arXiv preprint arXiv:2504.13626. — [arXiv:2504.13626](https://arxiv.org/abs/2504.13626) find
- Zheng Liu, Mengjie Liu, Siwei Wen, Mengzhang Cai, Bin Cui, Conghui He, and Wentao Zhang. 2025c. From uniform to heterogeneous: Tailoring policy optimization to every token’s nature. arXiv preprint arXiv:2509.16591. — [arXiv:2509.16591](https://arxiv.org/abs/2509.16591) find
- Haotian Luo, Li Shen, Haiying He, Yibo Wang, Shiwei Liu, Wei Li, Naiqiang Tan, Xiaochun Cao, and Dacheng Tao. 2025a. O1-pruner: Length-harmonizing fine-tuning for o1-like reasoning pruning. In 2nd AI for Math Workshop @ ICML 2025. find
- Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Y. Tang, Manan Roongta, Colin Cai, Jeffrey Luo, Li Erran Li, Raluca Ada Popa, and Ion Stoica. 2025b. Deepscaler: Surpassing o1-preview with a 1.5b model by scaling rl. https://pretty-radio -b75.notion.site/DeepScaleR-Surpassing-O 1-Preview-… find
- Xinyin Ma, Guangnian Wan, Runpeng Yu, Gongfan Fang, and Xinchao Wang. 2025. Cot-valve: Length-compressible chain-of-thought tuning. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 6025–6035.
- Mathematical Association of America. 2025a. American invitational mathematics examination 2024. find
- Mathematical Association of America. 2025b. American invitational mathematics examination 2025. find
- OpenAI. 2024. Learning to reason with LLMs. https://openai.com/index/learning-to-reason-wit h-llms/. ↗
- 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.
- Matthew Renze and Erhan Guven. 2024. The benefits of a concise chain of thought on problem-solving in large language models. arXiv preprint arXiv:2401.05618. — [arXiv:2401.05618](https://arxiv.org/abs/2401.05618) find
- John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017. Proximal policy optimization algorithms. Preprint, arXiv:1707.06347. find
- Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. Preprint, arXiv:2402.03300.
- Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. 2025. Hybridflow: A flexible and efficient rlhf framework. In Proceedings of the Twentieth European Conference on Computer Systems, pages 1279–1297.
- Zelin Tan, Hejia Geng, Xiaohang Yu, Mulei Zhang, Guancheng Wan, Yifan Zhou, Qiang He, Xiangyuan Xue, Heng Zhou, Yutao Fan, and 1 others. 2025a. Scaling behaviors of llm reinforcement learning post-training: An empirical study in mathematical reasoning. arXiv preprint arXiv:2509.25300.
- Zezhong Tan, Hang Gao, Xinhong Ma, Feng Zhang, and Ziqiang Dong. 2025b. Towards flash thinking via decoupled advantage policy optimization. arXiv preprint arXiv:2510.15374. — [arXiv:2510.15374](https://arxiv.org/abs/2510.15374) find
- Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, Chuning Tang, Congcong Wang, Dehao Zhang, Enming Yuan, Enzhe Lu, Fengxiang Tang, Flood Sung, Guangda Wei, Guokun Lai, and 75 others. 2025. Kimi k1.5: Scaling reinforcement l…
- Songjun Tu, Jiahao Lin, Qichao Zhang, Xiangyu Tian, Linjing Li, Xiangyuan Lan, and Dongbin Zhao. 2025. Learning when to think: Shaping adaptive reasoning in r1-style models via multi-stage rl. arXiv preprint arXiv:2505.10832. — [arXiv:2505.10832](https://arxiv.org/abs/2505.10832) find
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
- Minzheng Wang, Yongbin Li, Haobo Wang, Xinghua Zhang, Nan Xu, Bingli Wu, Fei Huang, Haiyang Yu, and Wenji Mao. 2026. Adaptive social learning via mode policy optimization for language agents. In The Fourteenth International Conference on Learning Representations. find
- Shenzhi Wang, Le Yu, Chang Gao, Chujie Zheng, Shixuan Liu, Rui Lu, Kai Dang, Xionghui Chen, Jianxin Yang, Zhenru Zhang, and 1 others. 2025. Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcement learning for llm reasoning. arXiv preprint arXiv:2506.01939.
- Zihao Wei, Liang Pang, Jiahao Liu, Wenjie Shi, Jingcheng Deng, Shicheng Xu, Zenghao Duan, Fei Sun, Huawei Shen, and Xueqi Cheng. 2025. The evolution of thought: Tracking llm overthinking via reasoning dynamics analysis. arXiv preprint arXiv:2508.17627.
- Fei Wu, Zhenrong Zhang, Qikai Chang, Jianshu Zhang, Quan Liu, and Jun Du. 2026. Step potential advantage estimation: Harnessing intermediate confidence and correctness for efficient mathematical reasoning. arXiv preprint arXiv:2601.03823.
- Violet Xiang, Chase Blagden, Rafael Rafailov, Nathan Lile, Sang Truong, Chelsea Finn, and Nick Haber. 2025. Just enough thinking: Efficient reasoning with adaptive length penalties reinforcement learning. arXiv preprint arXiv:2506.05256. — [arXiv:2506.05256](https://arxiv.org/abs/2506.05256) find
- Can Xie, Ruotong Pan, Xiangyu Wu, Yunfei Zhang, Jiayi Fu, Tingting Gao, and Guorui Zhou. 2025. Unlocking exploration in rlvr: Uncertainty-aware advantage shaping for deeper reasoning. arXiv preprint arXiv:2510.10649.
- Lei Yang, Wei Bi, Chenxi Sun, Renren Jin, and Deyi Xiong. 2026. Soup: Token-level single-sample mix-policy reinforcement learning for large language models. arXiv preprint arXiv:2601.21476. — [arXiv:2601.21476](https://arxiv.org/abs/2601.21476) find
- Jingyang Yi, Jiazheng Wang, and Sida Li. 2025. Shorterbetter: Guiding reasoning models to find optimal inference length for efficient reasoning. NeurIPS. find
- Linhao Yu, Tianmeng Yang, Siyu Ding, Renren Jin, Naibin Gu, Xiangzhao Hao, Shuaiyi Nie, Deyi Xiong, Weichong Yin, Yu Sun, and Hua Wu. 2026. Knowrl: Boosting llm reasoning via reinforcement learning with minimal-sufficient knowledge guidance. Preprint, arXiv:2604.12627. — [arXiv:2604.12627](https://arxiv.org/abs/2604.12627) find
- Haohan Yuan and Haopeng Zhang. 2025. Understanding llm reasoning for abstractive summarization. arXiv preprint arXiv:2512.03503.
- Hongyuan Yuan, Xinran He, Run Shao, Bolei He, Xianwei Xue, Mengke Chen, Qiutong Pan, Haiwei Wang, and Haifeng Li. 2026. Graph-based chain-of-thought pruning for reducing redundant reflections in reasoning llms. arXiv preprint arXiv:2604.05643.
- Chuhuai Yue, Chengqi Dong, Yinan Gao, Hang He, Jiajun Chai, Guojun Yin, and Wei Lin. 2025. Promoting efficient reasoning with verifiable stepwise reward. arXiv preprint arXiv:2508.10293. — [arXiv:2508.10293](https://arxiv.org/abs/2508.10293) find
- Jiajie Zhang, Nianyi Lin, Lei Hou, Ling Feng, and Juanzi Li. 2025a. Adaptthink: Reasoning models can learn when to think. arXiv preprint arXiv:2505.13417.
- Wenyuan Zhang, Shuaiyi Nie, Xinghua Zhang, Zefeng Zhang, and Tingwen Liu. 2025b. S1-bench: A simple benchmark for evaluating system 1 thinking capability of large reasoning models. Preprint, arXiv:2504.10368. — [arXiv:2504.10368](https://arxiv.org/abs/2504.10368) find
- Wenyuan Zhang, Xinghua Zhang, Haiyang Yu, Shuaiyi Nie, Bingli Wu, Juwei Yue, Tingwen Liu, and Yongbin Li. 2026. Expseek: Self-triggered experience seeking for web agents. arXiv preprint arXiv:2601.08605.
- Ziqi Zhao, Zhaochun Ren, Jiahong Zou, Liu Yang, Zhiwei Xu, Xuri Ge, Zhumin Chen, Xinyu Ma, Daiting Shi, Shuaiqiang Wang, and 1 others. 2026. Reinforced efficient reasoning via semantically diverse exploration. arXiv preprint arXiv:2601.05053.
- Zifan Zheng, Yezhaohui Wang, Yuxin Huang, Shichao Song, Mingchuan Yang, Bo Tang, Feiyu Xiong, and Zhiyu Li. 2024. Attention heads of large language models: A survey. Preprint, arXiv:2409.03752. — [arXiv:2409.03752](https://arxiv.org/abs/2409.03752)
- Wei Zhu, Jian Zhang, Lixing Yu, Kun Yue, and Zhiwen Tang. 2026. Dissecting failure dynamics in large language model reasoning. Preprint, arXiv:2604.14528.
Updated
This page in Markdown → https://md.research.elamur.ai/free/en/papers/nie-2026-attnpo-attention-guided-process-supervision/