CoChat
Sign Up
PreprintCited by 1Natural Language Processing TechniquesTopic ModelingAdvanced Neural Network Applications

LK Losses: Direct Acceptance Rate Optimization for Speculative Decoding

Open access
Samarin, Alexander, Krutikov, Sergei, Shevtsov, Anton, Skvortsov, Sergei, Fisin, Filipp, Golubev, Alexander·2026·arXiv
DOI 10.48550/arxiv.2602.23881arXiv 2602.23881
Open this paper in CoChat — free accountOpen PDF

Abstract

Speculative decoding accelerates autoregressive large language model (LLM) inference by using a lightweight draft model to propose candidate tokens that are then verified in parallel by the target model. The speedup is significantly determined by the acceptance rate, yet standard training minimizes Kullback-Leibler (KL) divergence as a proxy objective. While KL divergence and acceptance rate share the same global optimum, small draft models, having limited capacity, typically converge to suboptimal solutions where minimizing KL does not guarantee maximizing acceptance rate. To address this issue, we propose LK losses, special training objectives that directly target acceptance rate. Comprehensive experiments across four draft architectures and six target models, ranging from 8B to 685B parameters, demonstrate consistent improvements in acceptance metrics across all configurations compared to the standard KL-based training. We evaluate our approach on general, coding and math domains and report gains of up to 8-10% in average acceptance length. LK losses are easy to implement, introduce no computational overhead and can be directly integrated into any existing speculator training framework, making them a compelling alternative to the existing draft training objectives.

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2602.23881

Cites

1 of 1 references verified0 citation errors · 0 not checkable

Statistics

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

  • KDA extends the delta-rule recurrence with a channel-wise forget gate.

    “KDA extends the delta-rule recurrence [105,138] with a channel-wise forget gate [63].”

    From Kimi K3: Open Frontier Intelligence (2026) · cites this paper as [105]

    The cited paper’s §2.1 discusses draft-model architectures, and §5.3 discusses training losses; neither addresses KDA’s recurrence or a channel-wise forget gate. Searches for “KDA channel-wise forget gate delta-rule recurrence” and “channel-wise forget gate delta rule” found no passage supporting the attributed claim. (LLM, conf=0.78)

    Experimental — accuracy not yet measured

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Cites

Flagged first, then by citations

Research works printed but not linked: 23

  1. Amari, S.-i. and Nagaoka, H. Methods of Information Ge- ometry, volume 191 of Translations of Mathematical Monographs. American Mathematical Society and Ox- ford University Press, Providence, RI, 2000.
  2. Cai, T., Li, Y., Geng, Z., Peng, H., Lee, J. D., Chen, D., and Dao, T. MEDUSA: Simple LLM inference acceleration framework with multiple decoding heads. In Proceedings of the 41st International Conference on Machine Learn- ing, volume 235 of ICML'24, pp. 5209-5235. PMLR, 2024.
  3. Cai, Y., Liang, X., Wang, X., Ma, J., Liang, H., Luo, J., Zuo, X., Duan, L., Yin, Y., and Chen, X. FastMTP: Accelerating LLM inference with enhanced multi-token prediction. arXiv preprint arXiv:2509.18362, 2025. · arXiv
  4. Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., and Jumper, J. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023. · arXiv
  5. Chen, M., Tworek, J., Jun, H., Yuan, Q., Pinto, H. P. d. O., Kaplan, J., Edwards, H., Burda, Y., Joseph, N., Brockman, G., Ray, A., Puri, R., Krueger, G., Petrov, M., Khlaaf, H., Sastry, G., Mishkin, P., Chan, B., Gray, S., Ryder, N., Pavlov, M., Power, A., Kaiser, L., Bavarian, M., Winter, C., Tillet, P., Such, F. P., Cummings, D., Plappert, M., Chantzis, F., Barnes, E., Herbert-Voss, A., Guss, W. H., Nichol, A., Paino, A., Tezak, N., Tang, J., Babuschkin, I., Balaji, S., Jain, S., Saunders, W., Hesse, C., Carr, A. N., Leike, J., Achiam, J., Misra, V., Morikawa, E., Radford, A., Knight, M., Brundage, M., Murati, M., Mayer, K., Welinder, P., McGrew, B., Amodei, D., McCandlish, S., Sutskever, I., and Zaremba, W. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374, 2021. · arXiv
  6. Cobbe, K., Kosaraju, V., Bavarian, M., Chen, M., Jun, H., Kaiser, L., Plappert, M., Tworek, J., Hilton, J., Nakano, R., Hesse, C., and Schulman, J. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168, 2021. · arXiv
  7. DeepSeek-AI, Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. DeepSeek-V3 technical report. arXiv preprint arXiv:2412.19437, 2024. · arXiv
  8. Hu, Y., Guo, J., Feng, X., and Zhao, T. AdaSPEC: Selective knowledge distillation for efficient speculative decoders. arXiv preprint arXiv:2510.19779, 2025. · arXiv
  9. Ji, H., Ke, P., Hu, Z., Zhang, R., and Huang, M. Tailoring language generation models under total variation distance. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview. net/forum?id=VELL0PlWfc.
  10. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with PagedAttention. In Proceedings of the 29th ACM Symposium on Operating Systems Principles, SOSP '23, pp. 611-626. ACM, 2023.
  11. Leviathan, Y., Kalman, M., and Matias, Y. Fast inference from transformers via speculative decoding. In Proceed- ings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
  12. Li, J., Du, L., Zhao, H., Zhang, B.-w., Wang, L., Gao, B., Liu, G., and Lin, Y. Infinity Instruct: Scaling instruc- tion selection and synthesis to enhance language models. arXiv preprint arXiv:2506.11116, 2025a. · arXiv
  13. Li, Y., Wei, F., Zhang, C., and Zhang, H. EAGLE: Specula- tive sampling requires rethinking feature uncertainty. In International Conference on Machine Learning, 2024.
  14. Li, Y., Wei, F., Zhang, C., and Zhang, H. EAGLE-3: Scaling up inference acceleration of large language models via training-time test. In Advances in Neural Information Processing Systems, 2025b.
  15. Liu, X., Yu, J., Park, J., Stoica, I., and Cheung, A. Specula- tive decoding: Performance or illusion? arXiv preprint arXiv:2601.11580, 2026. · arXiv
  16. Sandler, J., Christopher, J. K., Hartvigsen, T., and Fioretto, F. SpecDiff-2: Scaling diffusion drafter alignment for faster speculative decoding. arXiv preprint arXiv:2511.00606, 2025. · arXiv
  17. Schulman, J., Levine, S., Abbeel, P., Jordan, M., and Moritz, P. Trust region policy optimization. In Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1889-1897. PMLR, 2015.
  18. Wertheimer, D., Rosenkranz, J., Parnell, T., Suneja, S., Ran- ganathan, P., Ganti, R., and Srivatsa, M. Accelerating production LLMs with combined token/embedding spec- ulators. arXiv preprint arXiv:2404.19124, 2024. · arXiv
  19. Yang, A., Li, A., Yang, B., Zhang, B., Hui, B., Wang, B., Li, B., Wang, C., Yu, D., Huang, F., et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388, 2025. · arXiv
  20. Yin, M., Chen, M., Huang, K., and Wang, M. A theoretical perspective for speculative decoding algorithm. In Pro- ceedings of the 38th International Conference on Neural Information Processing Systems, NIPS '24, 2024. ISBN 9798331314385.
  21. Zhao, W., Pan, T., Han, X., Zhang, Y., Sun, A., Huang, Y., Zhang, K., Zhao, W., Li, Y., Zhou, J., Zhou, H., Wang, J., Liu, Z., and Sun, M. FR-spec: Accelerat- ing large-vocabulary language models via frequency- ranked speculative sampling. In Che, W., Nabende, J., Shutova, E., and Pilehvar, M. T. (eds.), Proceedings of the 63rd Annual Meeting of the Association for Com- putational Linguistics (Volume 1: Long Papers), pp. 3909-3921, Vienna, Austria, July 2025. Association for Computational Linguistics. ISBN 979-8-89176-251- 0. doi: 10.18653/v1/2025.acl-long.198. URL https: //aclanthology.org/2025.acl-long.198/. · DOI
  22. Zheng, L., Chiang, W.-L., Sheng, Y., Zhuang, S., Wu, Z., Zhuang, Y., Lin, Z., Li, Z., Li, D., Xing, E. P., Zhang, H., Gonzalez, J. E., and Stoica, I. Judging LLM-as-a- judge with MT-Bench and Chatbot Arena. In Advances in Neural Information Processing Systems, volume 36, 2023.
  23. Zhou, Y., Lyu, K., Rawat, A. S., Menon, A., Rostamizadeh, A., Kumar, S., Kagy, J.-F., and Agarwal, R. Distill- Spec: Improving speculative decoding via knowledge distillation. In The Twelfth International Conference on Learning Representations, 2024. URL https:// openreview.net/forum?id=rsY6J3ZaTF.

Cited by

Sorted by citations

  • Kimi K3: Open Frontier Intelligence

    Kimi Team, Bai, T., Bai, Y., Bao, Y., C., M., Cai, J., Cai, X., Cao, P., Cao, Y., Chai, Z., et al. (2026) · arXiv · doi:10.48550/arxiv.2607.24653

    “…surrogate does not guarantee maximizing this rate for a capacity-limited draft model, we directly optimize the likelihood-based LK loss [105], the negative logarithm of the acceptance rate itself, with p and q evaluated at temperature 1 and no auxiliary ground-truth cross-entropy term.”

Authors
Samarin, A., Krutikov, S., Shevtsov, A., Skvortsov, S., Fisin, F., & Golubev, A.
Venue
arXiv
Publisher
arXiv
Published
2026
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2602.23881
arXiv
2602.23881
OpenAlex
W7133220763
Cited by
0
References
1
Type
posted content
Language
en
Open access
Licence
Locations
  • arXiv (Cornell University) (submittedVersion, cc-by)
  • arXiv (Cornell University) (cc-by)
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Natural Language Processing Techniques
  • Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
  • Physical Sciences › Computer Science › Computer Vision and Pattern Recognition › Advanced Neural Network Applications
Keywords
speculative decoding, Kullback-Leibler divergence, autoregressive decoding

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review