Elamur Research⌕
☰
home / papers / 2026 / KNN-SSD: Enabling Dynamic Self-Speculative Decoding via Nearest Neighbor Layer Set Optimization

KNN-SSD: Enabling Dynamic Self-Speculative Decoding via Nearest Neighbor Layer Set Optimization

Paper · 2026 · Mingbo Song, Heming Xia, Jun Zhang et al.

Summary

Self-speculative decoding (Self-SD) speeds up LLM inference by skipping some of the model's own layers to produce draft tokens, which the full model then verifies. The authors show that a skip-layer configuration tuned for one task transfers poorly to others. A set optimized for summarization falls from about 1.35× to under 1.10× speedup on other tasks.

KNN-SSD addresses this in two stages:

  • Offline: it samples inputs from each domain, takes their last hidden vectors from the LLM, groups them with K-means, and keeps the points nearest each centroid as anchors for a KNN classifier. For each domain, Bayesian optimization searches for a dedicated skip-layer set.
  • At inference: each incoming query is routed by the KNN classifier to its domain's layer set, and standard draft-and-verify decoding follows.

The method was tested on LLaMA-2 and Qwen-2.5 models with mixed-domain sample streams covering summarization, reasoning, translation, storytelling and text-to-SQL. It achieves a lossless 1.3×–1.6× speedup and stays stable as domains change. This matters for real deployments, where query types are unpredictable.

Key points

  • Speedup: 1.3×–1.6× over autoregressive decoding, with the output distribution unchanged. Out-of-domain data still gets about 1.2×.
  • Token acceptance: above 80% for the LLaMA-2 series and above 99% for the Qwen models.
  • Evaluation setup: a "mix ratio" controls how often the input domain switches within the sample stream.
  • Alpaca with LLaMA-2-13B-Chat: speedup rises from 1.05× with 1 cluster to 1.23× with 5 clusters, then plateaus at 7 clusters. This suggests a single dataset can contain several hidden domains.
  • Code and limitations: code is at github.com/mbsong/KNN-SSD. The method does not use draft-tree verification and was tested only on moderate-size models.

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

🎓Peking UniversityBeijing · China244 works · 39,359 citations🎓Zhejiang UniversityHangzhou · China172 works · 12,988 citations🎓Hong Kong Polytechnic UniversityHong Kong · China84 works · 4,956 citations

 Fields, topics and concepts

📚 Shelves 5The 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 worksComputer Vision and Pattern Recognition · 2,869 works

Cite

Cite as Mingbo Song, Heming Xia, Jun Zhang, Chak Tou Leong, Qiancheng Xu, Wenjie Li and Sujian Li. 2026. KNN-SSD: Enabling Dynamic Self-Speculative Decoding via Nearest Neighbor Layer Set Optimization. Findings of the Association for Computational Linguistics: EACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-eacl.31

@inproceedings{song2026knn,
  title = {{KNN-SSD: Enabling Dynamic Self-Speculative Decoding via Nearest Neighbor Layer Set Optimization}},
  author = {Mingbo Song, Heming Xia, Jun Zhang, Chak Tou Leong, Qiancheng Xu, Wenjie Li and Sujian Li},
  booktitle = {Findings of the Association for Computational Linguistics: EACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-eacl.31},
  url = {https://doi.org/10.18653/v1/2026.findings-eacl.31}
}

References

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

  1. Zachary Ankner, Rishab Parthasarathy, Aniruddha Nrusimha, Christopher Rinard, Jonathan Ragan-Kelley, and William Brandon. 2024. Hydra: Sequentially-dependent draft heads for medusa decoding. Preprint, arXiv:2402.05109. — [arXiv:2402.05109](https://arxiv.org/abs/2402.05109) find
  2. Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, Ahmed Aly, Beidi Chen, and Carole-Jean Wu. 2024. LayerSkip: Enabling early exit inference and self-speculative decoding. In Proceedings of the 62nd…
  3. Elias Frantar and Dan Alistarh. 2023. SparseGPT: Massive language models can be accurately pruned in one-shot. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 10323–10337. PMLR. find
  4. Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. 2024. Break the sequential dependency of LLM inference using lookahead decoding. In Forty-first International Conference on Machine Learning.
  5. Milan Gritta, Huiyin Xue, and Gerasimos Lampouras. 2025. Dresd: Dense retrieval for speculative decoding. Preprint, arXiv:2502.15572. find
  6. Yuxian Gu, Li Dong, Furu Wei, and Minlie Huang. 2024. Minillm: Knowledge distillation of large language models. Preprint, arXiv:2306.08543.
  7. Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D. Lee, and Di He. 2024. Rest: Retrieval-based speculative decoding. Preprint, arXiv:2311.08252. find
  8. Yijin Liu, Fandong Meng, and Jie Zhou. 2024. Accelerating inference in large language models with a unified layer skipping strategy. Preprint, arXiv:2404.06954. — [arXiv:2404.06954](https://arxiv.org/abs/2404.06954) find
  9. Zechun Liu, Barlas Oguz, Changsheng Zhao, Ernie Chang, Pierre Stock, Yashar Mehdad, Yangyang Shi, Raghuraman Krishnamoorthi, and Vikas Chandra. 2023. Llm-qat: Data-free quantization aware training for large language models. Preprint, arXiv:2305.17888.
  10. Michael R. Metel, Peng Lu, Boxing Chen, Mehdi Rezagholizadeh, and Ivan Kobyzev. 2024. Draft on the fly: Adaptive self-speculative decoding using cosine similarity. Preprint, arXiv:2410.01028. find
  11. 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. Preprint, arXiv:2103.03874.
  12. Edward J Hu, yelong shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2022. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations.
  13. Yunhai Hu, Zining Liu, Zhenyuan Dong, Tianfan Peng, Bradley McDanel, and Sai Qian Zhang. 2025. Speculative decoding and beyond: An in-depth survey of techniques. Preprint, arXiv:2502.19732. — [arXiv:2502.19732](https://arxiv.org/abs/2502.19732) find
  14. Yukun Huang, Yanda Chen, Zhou Yu, and Kathleen McKeown. 2022. In-context learning distillation: Transferring few-shot learning ability of pre-trained language models. Preprint, arXiv:2212.10670. — [arXiv:2212.10670](https://arxiv.org/abs/2212.10670) find
  15. Donald R Jones, Matthias Schonlau, and William J Welch. 1998. Efficient global optimization of expensive black-box functions. Journal of Global optimization, 13:455–492. find
  16. Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. 2023. Speculative decoding with big little decoder. In Advances in Neural Information Processing Systems, volume 36, pages 39236–39256. Curran Associates, Inc. find
  17. Fangyu Lei, Jixuan Chen, Yuxiao Ye, Ruisheng Cao, Dongchan Shin, Hongjin Su, Zhaoqing Suo, Hongcheng Gao, Wenjing Hu, Pengcheng Yin, Victor Zhong, Caiming Xiong, Ruoxi Sun, Qian Liu, Sida Wang, and Tao Yu. 2025. Spider 2.0: Evaluating language models on real-world enterprise text-to-sql workflows. P… — [arXiv:2411.07763](https://arxiv.org/abs/2411.07763) find
  18. Yaniv Leviathan, Matan Kalman, and Yossi Matias. 2023. Fast inference from transformers via speculative decoding. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 19274–19286. PMLR. find
  19. Ramesh Nallapati, Bowen Zhou, Cicero dos Santos, Çağlar Gu˙lçehre, and Bing Xiang. 2016. Abstractive text summarization using sequence-to-sequence RNNs and beyond. In Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning, pages 280–290, Berlin, Germany. Association fo…
  20. Shashi Narayan, Shay B. Cohen, and Mirella Lapata. 2018. Don’t give me the details, just the summary! Topic-aware convolutional neural networks for extreme summarization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium.
  21. OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, et al. 2024. Gpt-4 technical report. Preprint, arXiv:2303.08774.
  22. Gunho Park, Baeseong Park, Minsub Kim, Sungjae Lee, Jeonghoon Kim, Beomseok Kwon, Se Jung Kwon, Byeongwook Kim, Youngjoo Lee, and Dongsoo Lee. 2024. Lut-gemm: Quantized matrix multiplication based on luts for efficient inference in large-scale generative language models. Preprint, arXiv:2206.09557. find
  23. Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. 2018. Blockwise parallel decoding for deep autoregressive models. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc. find
  24. Mingjie Sun, Zhuang Liu, Anna Bair, and J. Zico Kolter. 2024. A simple and effective pruning approach for large language models. Preprint, arXiv:2306.11695. find
  25. Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. 2023. Stanford alpaca: An instruction-following llama model. https://github.com/tatsu-lab/stanford_alpaca. ↗
  26. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, et al. 2023. Llama 2: Open foundation and fine-tuned chat models. Preprint, arXiv:2307.09288.
  27. Zhuofan Wen, Shangtong Gui, and Yang Feng. 2024. Speculative decoding with ctc-based draft model for llm inference acceleration. In Advances in Neural Information Processing Systems, volume 37, pages 92082–92100. Curran Associates, Inc. find
  28. Minghao Wu, Abdul Waheed, Chiyu Zhang, Muhammad Abdul-Mageed, and Alham Fikri Aji. 2024. LaMini-LM: A diverse herd of distilled models from large-scale instructions. In Proceedings of the 18th Conference of the European Chapter of the Association for Computational Linguistics (Volume 1: Long Papers)… find
  29. Heming Xia, Tao Ge, Peiyi Wang, Si-Qing Chen, Furu Wei, and Zhifang Sui. 2023. Speculative decoding: Exploiting speculative execution for accelerating seq2seq generation. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 3909–3925, Singapore. Association for Computation…
  30. Heming Xia, Yongqi Li, Jun Zhang, Cunxiao Du, and Wenjie Li. 2025. SWIFT: On-the-fly self-speculative decoding for LLM inference acceleration. In The Thirteenth International Conference on Learning Representations. find
  31. Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge, Tianyu Liu, Wenjie Li, and Zhifang Sui. 2024. Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding. In Findings of the Association for Computational Linguistics: ACL 2024, pages 7655–…
  32. An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, et al. 2025. Qwen2.5 technical report. Preprint, arXiv:2412.15115. find
  33. Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. 2022. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. In Advances in Neural Information Processing Systems, volume 35, pages 27168–27183. Curran Associates, Inc.
  34. Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Gang Chen, and Sharad Mehrotra. 2024. Draft & verify: Lossless large language model acceleration via self-speculative decoding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages…
  35. Lefan Zhang, Xiaodan Wang, Yanhua Huang, and Ruiwen Xu. 2025. Learning harmonized representations for speculative sampling. In The Thirteenth International Conference on Learning Representations. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/song-2026-knn-ssd-enabling-dynamic-self/