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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Song, Mingbo, Xia, Heming, Zhang, Jun, Leong, Chak Tou, Xu, Qiancheng, Li, Wenjie, Li, Sujian
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917210612563968
author Song, Mingbo
Xia, Heming
Zhang, Jun
Leong, Chak Tou
Xu, Qiancheng
Li, Wenjie
Li, Sujian
author_facet Song, Mingbo
Xia, Heming
Zhang, Jun
Leong, Chak Tou
Xu, Qiancheng
Li, Wenjie
Li, Sujian
contents Speculative Decoding (SD) has emerged as a widely used paradigm to accelerate the inference of large language models (LLMs) without compromising generation quality. It works by efficiently drafting multiple tokens using a compact model and then verifying them in parallel using the target LLM. Notably, Self-Speculative Decoding proposes skipping certain layers to construct the draft model, which eliminates the need for additional parameters or training. Despite its strengths, we observe in this work that drafting with layer skipping exhibits significant sensitivity to domain shifts, leading to a substantial drop in acceleration performance. To enhance the domain generalizability of this paradigm, we introduce KNN-SSD, an algorithm that leverages K-Nearest Neighbor (KNN) search to match different skipped layers with various domain inputs. We evaluated our algorithm in various models and multiple tasks, observing that its application leads to 1.3x-1.6x speedup in LLM inference.
format Preprint
id arxiv_https___arxiv_org_abs_2505_16162
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle KNN-SSD: Enabling Dynamic Self-Speculative Decoding via Nearest Neighbor Layer Set Optimization
Song, Mingbo
Xia, Heming
Zhang, Jun
Leong, Chak Tou
Xu, Qiancheng
Li, Wenjie
Li, Sujian
Computation and Language
Speculative Decoding (SD) has emerged as a widely used paradigm to accelerate the inference of large language models (LLMs) without compromising generation quality. It works by efficiently drafting multiple tokens using a compact model and then verifying them in parallel using the target LLM. Notably, Self-Speculative Decoding proposes skipping certain layers to construct the draft model, which eliminates the need for additional parameters or training. Despite its strengths, we observe in this work that drafting with layer skipping exhibits significant sensitivity to domain shifts, leading to a substantial drop in acceleration performance. To enhance the domain generalizability of this paradigm, we introduce KNN-SSD, an algorithm that leverages K-Nearest Neighbor (KNN) search to match different skipped layers with various domain inputs. We evaluated our algorithm in various models and multiple tasks, observing that its application leads to 1.3x-1.6x speedup in LLM inference.
title KNN-SSD: Enabling Dynamic Self-Speculative Decoding via Nearest Neighbor Layer Set Optimization
topic Computation and Language
url https://arxiv.org/abs/2505.16162