Elamur Research⌕
☰
home / papers / 2026 / Efficient KL Divergence Estimation via Truncated Top-K Integration for Large Language Models

Efficient KL Divergence Estimation via Truncated Top-K Integration for Large Language Models

Paper · 2026 · Xinyuan Wang, Luozhijie Jin, Bo Wang et al.

Summary

RL fine-tuning of LLMs (e.g., GRPO) usually includes a KL penalty that keeps the policy close to a reference model. Computing this KL exactly means summing over the whole vocabulary at every token, which uses too much GPU memory for large vocabularies and long contexts. The common workaround uses Schulman's single-token estimators, but k1 is very noisy and k2/k3 introduce bias. The authors propose TIKE (Top-k Importance-weighted KL Estimator). It relies on the observation that next-token distributions are Zipf-like, so a handful of top tokens carry nearly all the probability. TIKE computes the KL exactly over only the top-k tokens of the current policy. It then applies an importance-sampling ratio (following DeepSeek-V3.2) to correct for the gap between the rollout policy and the policy being optimized. The method was trained on DAPO-Math-17k with Qwen2.5 and DeepSeek-R1-Distill-Llama-8B, then tested on math reasoning benchmarks. TIKE beat k1, k2 and k3 and nearly matched the exact Rao-Blackwellized KL with much lower variance. The result is a cheap, stable KL regularizer for RL training of reasoning models.

Key points

  • Memory savings: For Qwen2.5 (vocabulary 151,643; batch 32; sequence length 2,048), a full-vocabulary KL tensor needs about 37 GB in FP32. Top-16 truncation needs only about 4.19 MB, a reduction of more than 99.9%.
  • Cost at k=16: Overhead is small, with about 6.9% longer step time and about 0.1% more memory. Gains level off beyond k=16, and values of k in {8, 16} worked well.
  • Evaluation setup: Models were Qwen2.5 1.5B/7B and DeepSeek-R1-Distill-Llama-8B. Benchmarks were AIME24, AIME25, MATH500, Minerva and GSM8K. Results held across different KL coefficients β and multiple seeds.
  • Accuracy of the estimate: In variance analyses, TIKE closely tracked the exact Rao-Blackwellized estimate. By contrast, k3 sometimes settled well below the true KL value.
  • Limitations: Truncation bias may grow in high-entropy settings such as creative writing or dialogue. Domains beyond math are untested. Code is available at github.com/jinluo12345/TIKE.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Fudan UniversityShanghai · China167 works · 20,987 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Xinyuan Wang, Luozhijie Jin, Bo Wang, Yuan Li, Zhangyue Yin and Xipeng Qiu. 2026. Efficient KL Divergence Estimation via Truncated Top-K Integration for Large Language Models. Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.acl-long.405

@inproceedings{wang2026efficient,
  title = {{Efficient KL Divergence Estimation via Truncated Top-K Integration for Large Language Models}},
  author = {Xinyuan Wang, Luozhijie Jin, Bo Wang, Yuan Li, Zhangyue Yin and Xipeng Qiu},
  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.405},
  url = {https://doi.org/10.18653/v1/2026.acl-long.405}
}

References

27 references, 12 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Afra Amini, Tim Vieira, and Ryan Cotterell. 2025. Better estimation of the kullback–leibler divergence between language models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems. find
  2. 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. arXiv preprint arXiv:2110.14168.
  3. Ganqu Cui, Lifan Yuan, Zefan Wang, Hanbin Wang, Yuchen Zhang, Jiacheng Chen, Wendi Li, Bingxiang He, Yuchen Fan, Tianyu Yu, and 1 others. 2025. Process reinforcement through implicit rewards. arXiv preprint arXiv:2502.01456.
  4. Jia Deng, Jie Chen, Zhipeng Chen, Daixuan Cheng, Fei Bai, Beichen Zhang, Yinqian Min, Yanzipeng Gao, Wayne Xin Zhao, and Ji-Rong Wen. 2025. From trial-and-error to improvement: A systematic analysis of llm exploration mechanisms in rlvr. arXiv preprint arXiv:2508.07534. — [arXiv:2508.07534](https://arxiv.org/abs/2508.07534) find
  5. Riccardo Di Sipio. 2025. Rethinking llm training through information geometry and quantum metrics. arXiv preprint arXiv:2506.15830. — [arXiv:2506.15830](https://arxiv.org/abs/2506.15830) find
  6. Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, and 1 others. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783.
  7. Chengqi Lyu, Songyang Gao, Yuzhe Gu, Wenwei Zhang, Jianfei Gao, Kuikun Liu, Ziyi Wang, Shuaibin Li, Qian Zhao, Haian Huang, and 1 others. 2025. Exploring the limit of outcome reward for learning mathematical reasoning. arXiv preprint arXiv:2502.06781. — [arXiv:2502.06781](https://arxiv.org/abs/2502.06781) find
  8. Mathematical Association of America. 2025. American invitational mathematics examination (AIME). https://maa.org/. Accessed: 2025-01-05. ↗
  9. 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.
  10. Sheldon M. Ross. 2002. Simulation, 3rd edition. Academic Press. find
  11. Jian Hu. 2025. Reinforce++: A simple and efficient approach for aligning large language models. arXiv preprint arXiv:2501.03262. find
  12. John Schulman. 2020. Approximating KL divergence. http://joschu.net/blog/kl-approx.html. Accessed: 2026-01-05. ↗
  13. Diederik P Kingma. 2014. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980.
  14. John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. 2015. Trust region policy optimization. In International conference on machine learning, pages 1889–1897. PMLR. find
  15. Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, and 1 others. 2022. Solving quantitative reasoning problems with language models. Advances in neural information processing systems, 35:3843–3857…
  16. Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023. Let’s verify step by step. In The Twelfth International Conference on Learning Representations.
  17. Aixin Liu, Aoxue Mei, Bangcai Lin, Bing Xue, Bingxuan Wang, Bingzheng Xu, Bochao Wu, Bowei Zhang, Chaofan Lin, Chen Dong, and 1 others. 2025a. Deepseek-v3.2: Pushing the frontier of open large language models. arXiv preprint arXiv:2512.02556.
  18. Keliang Liu, Dingkang Yang, Ziyun Qian, Weijie Yin, Yuchi Wang, Hongsheng Li, Jun Liu, Peng Zhai, Yang Liu, and Lihua Zhang. 2025b. Reinforcement learning meets large language models: A survey of advancements and applications across the llm lifecycle. arXiv preprint arXiv:2509.16679. — [arXiv:2509.16679](https://arxiv.org/abs/2509.16679) find
  19. Kezhao Liu, Jason Klein Liu, Mingtao Chen, and Yiming Liu. 2025c. Rethinking kl regularization in rlhf: From value estimation to gradient optimization. arXiv preprint arXiv:2510.01555. — [arXiv:2510.01555](https://arxiv.org/abs/2510.01555) find
  20. Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. 2025d. Understanding r1-zero-like training: A critical perspective. arXiv preprint arXiv:2503.20783.
  21. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. find
  22. Vedant Shah, Johan Obando-Ceron, Vineet Jain, Brian Bartoldson, Bhavya Kailkhura, Sarthak Mittal, Glen Berseth, Pablo Samuel Castro, Yoshua Bengio, Nikolay Malkin, and 1 others. 2025. A comedy of estimators: On kl regularization in rl training of llms. arXiv preprint arXiv:2512.21852. — [arXiv:2512.21852](https://arxiv.org/abs/2512.21852) find
  23. 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.
  24. 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.
  25. Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. 2020. Learning to summarize with human feedback. Advances in neural information processing systems, 33:3008– 3021. find
  26. Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, and 1 others. 2025. find
  27. Kimi k2: Open agentic intelligence. arXiv preprint arXiv:2507.20534. Qwen Team and 1 others. 2024. Qwen2 technical report. arXiv preprint arXiv:2407.10671, 2(3).

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/wang-2026-efficient-kl-divergence-estimation-truncated/