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

2025-05-22

AI Summary

Self-speculative decoding (SSD) accelerates LLM inference via layer-skipping to construct lightweight draft models, but its fixed skip-layer strategy suffers substantial performance degradation under domain shifts. To address this, we propose a KNN-driven dynamic domain-matching mechanism—the first to integrate parameter-free K-nearest-neighbor search into SSD. Our method dynamically retrieves optimal layer-skip configurations in real time based on input representations, enabling zero-training, zero-parameter cross-domain adaptation. Crucially, it requires no architectural modification to the backbone model or additional fine-tuning. Extensive experiments across multiple LLMs (Llama-2/3, Qwen) and diverse tasks (commonsense reasoning, code generation, mathematical QA) demonstrate consistent inference speedups of 1.3×–1.6×. Moreover, our approach significantly enhances SSD’s robustness and generalization under distribution shift, establishing a new paradigm for adaptive speculative decoding.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• Improves domain generalizability of Self-Speculative Decoding
• Optimizes layer skipping for dynamic domain adaptation
• Accelerates LLM inference via KNN-based layer matching

Innovation

• Dynamic layer skipping via KNN optimization
• Domain-adaptive nearest neighbor search
• Self-speculative decoding acceleration