🤖 AI

KnapSpec: Self-Speculative Decoding via Adaptive Layer Selection as a Knapsack Problem

KnapSpec is a training-free self-speculative decoding framework that reformulates adaptive layer selection as a knapsack problem to maximize inference throughput by dynamically optimizing draft model configurations based on hardware-specific latencies and context length.

Original authors: Seongjin Cha, Gyuwan Kim, Dongsu Han, Tao Yang, Insu Han

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Seongjin Cha, Gyuwan Kim, Dongsu Han, Tao Yang, Insu Han

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Slow Writer" Dilemma

Imagine you are a professional editor (the Target Model) who writes incredibly high-quality articles. However, you are very slow. For every single word you write, you have to check your grammar, fact-check every claim, and polish the sentence structure. This takes a long time.

To speed things up, you hire a fast intern (the Draft Model) to write a rough first draft. You then quickly scan the intern’s work. If the intern’s sentence is good, you keep it. If it’s bad, you rewrite it yourself. This is called Speculative Decoding.

The Catch: Usually, the intern is a different person with different skills. Sometimes, the intern writes in a style you hate, so you have to rewrite almost everything, which wastes time.

Self-Speculative Decoding (SSD) tries to fix this by making you act as your own intern. You write a rough draft using only part of your brainpower, then use your full brainpower to check it. Since it’s the same brain, the draft is usually very close to what you would have written anyway.

The Old Way: The "Blindfolded Backpacker"

Previous methods of doing this were like a backpacker trying to pack a bag for a hike, but they were blindfolded. They would just grab random items (layers of the AI model) and hope they fit.

Static Heuristics: They assumed every item weighed the same.
The Problem: In reality, some items are heavy and some are light. Worse, as the hike gets longer (the text gets longer), some items become much heavier to carry. The old methods didn’t realize this, so they kept carrying heavy items even when they slowed the hiker down.

The Solution: KnapSpec (The Smart Knapsack)

The authors propose KnapSpec. Think of it as a smart packing algorithm that solves the classic "Knapsack Problem."

In the Knapsack Problem, you have a backpack with a limited weight capacity. You want to pack items that give you the most value (in this case, "accuracy" or "good guesses") without exceeding the weight limit (in this case, "time" or "latency").

Here is how KnapSpec works differently:

1. It Separates the "Heavy" from the "Light"

In an AI model, there are two main types of layers:

Attention Layers: These look at the whole context. As the text gets longer, these get heavier (slower) to process.
MLP Layers: These process local information. Their weight stays constant (fast) regardless of text length.

Old methods treated these as one big block. KnapSpec separates them. It realizes that if you are writing a very long book, the "Attention" layers are becoming too heavy to carry every time. So, it decides to skip some of them to save time.

2. It Adapts to the Journey (Context Length)

KnapSpec doesn’t use a fixed plan. It looks at how long the current sentence or paragraph is.

Short Text: It might carry more "Attention" layers because they aren’t too heavy yet.
Long Text: It drops the heavy "Attention" layers and relies more on the lightweight "MLP" layers to keep moving fast.

It does this calculation on the fly, like a GPS recalculating your route based on traffic jams.

3. It Uses a "Cosine Compass"

How does it know which layers to keep? It uses a mathematical trick called Cosine Similarity.

Imagine the Target Model’s perfect thought is a North Star.
The Draft Model’s thought is a compass needle.
KnapSpec checks how closely the compass needle points to the North Star. If it’s close enough, it’s a good guess.

The paper provides a rigorous mathematical proof that if the compass needle is close enough to the North Star (high cosine similarity), the guess will likely be accepted. This ensures the "intern" isn’t just fast, but also accurate.

The Result: Faster Without Training

No Extra Training: You don’t need to teach the AI new tricks. It just changes how it uses its existing brain.
Plug-and-Play: It works with existing models like Llama and Qwen.
Speedup: In tests, KnapSpec was up to 1.47 times faster than the standard way of generating text.

Summary Analogy

Imagine you are a chef (the AI) cooking a complex meal.

Old Method: You chop every vegetable and sauté every spice with full precision for every single dish, even if the dish is simple.
KnapSpec: You realize that chopping onions (Attention) takes a long time if you have a huge pile of them. So, for a large order, you skip some of the fine chopping and use pre-chopped onions (skipping layers) to get the meal out faster. You check if the taste is still good (Cosine Similarity). If it is, you serve it. If not, you fix it.

KnapSpec is the smart manager that decides exactly which steps to skip based on how big the order is, ensuring the kitchen runs as fast as possible without ruining the food.

Technical Summary: KnapSpec: Self-Speculative Decoding via Adaptive Layer Selection as a Knapsack Problem

Problem Statement
Large Language Model (LLM) inference faces significant latency bottlenecks, particularly in long-context scenarios. Speculative Decoding (SD) accelerates generation by using a lightweight draft model to propose tokens that a larger target model verifies in parallel. Self-Speculative Decoding (SSD) improves upon this by using a sub-network of the target model itself as the draft model, eliminating the need for separate model maintenance and reducing distribution mismatch. However, existing SSD methods often rely on static heuristics or block-level layer skipping that ignore the dynamic computational overhead of attention mechanisms. Specifically, as context length increases, the latency of Attention layers grows linearly ( $\Theta(n)$ ), while Multi-Layer Perceptron (MLP) latency remains constant ( $\Theta(1)$ ). Static configurations fail to adapt to this shifting bottleneck, leading to suboptimal throughput.

Methodology
The authors propose KnapSpec, a training-free framework that reformulates the selection of draft model layers as a 0/1 Knapsack Problem to maximize Tokens-per-Time (TPT) throughput. The methodology consists of three core components:

Decoupled Layer Modeling and TPT Metric:
KnapSpec decouples Attention and MLP layers, treating them as distinct items with hardware-specific latencies. It introduces the Tokens-per-Time (TPT) metric, which generalizes previous metrics like Tokens-per-Layer (TPL) by explicitly accounting for length-dependent latencies. TPT is defined as the expected number of generated tokens per unit of wall-clock time:
$\text{TPT}(S, \gamma) = \frac{1 - \alpha_S^{\gamma+1}}{1 - \alpha_S} \cdot \frac{1}{\gamma t_{\text{Draft}}(S) + t_{\text{Target}}}$
where $\alpha_S$ is the acceptance rate, $\gamma$ is the draft length, and $t_{\text{Draft}}(S)$ is the latency of the selected layer subset $S$ .
Knapsack Formulation via Dynamic Programming:
The layer selection problem is decomposed into two stages:
- Stage 1 (Candidate Generation): The problem is formulated as finding a layer subset $S$ that maximizes the expected acceptance rate (proxied by cosine similarity) under a discrete latency budget. Continuous latencies are normalized into integer weights. A parallel Dynamic Programming (DP) algorithm searches for optimal hidden state representations for all possible latency budgets simultaneously. This reduces the exponential search space to a tractable set of candidates.
- Stage 2 (TPT Maximization): From the candidate sets generated in Stage 1, the algorithm performs a grid search over draft lengths to select the configuration $(S^*, \gamma^*)$ that maximizes the actual TPT metric.
Theoretical Justification of Cosine Similarity:
The paper provides the first rigorous theoretical analysis establishing cosine similarity between hidden states as a mathematically sound proxy for token acceptance rate. Lemma 4.1 proves that if the cosine similarity between the draft and target embeddings exceeds a specific margin dependent on the vocabulary embedding geometry, the greedy decoding predictions will be identical. This justifies using cosine similarity as the "value" in the knapsack optimization.

Key Contributions

Adaptive Layer Decoupling: KnapSpec is the first method to decouple Attention and MLP layers and incorporate context length into the draft model selection process, allowing it to adapt to shifting hardware bottlenecks.
Knapsack Reformulation: It reformulates draft model selection as a constrained optimization problem using a hardware-aware TPT metric, balancing latency and acceptance rate.
Theoretical Foundation: It provides a rigorous proof that cosine similarity serves as a reliable proxy for the token acceptance rate, moving beyond empirical heuristics.
Training-Free Efficiency: The approach requires no additional training or auxiliary parameters, making it a plug-and-play solution for any pre-trained LLM.

Results
Experiments were conducted on Qwen3 (reasoning tasks) and Llama3 (summarization tasks) models ranging from 1B to 70B parameters.

Performance: KnapSpec consistently outperformed state-of-the-art training-free SSD baselines (SWIFT, DEL, CLaSp). It achieved up to 1.47× wall-clock speedup on the GovReport dataset with Llama3.1-70B.
Robustness: The method maintained high speedups across various context lengths and decoding strategies (including nucleus sampling), demonstrating that optimizing for TPT yields better practical throughput than optimizing for acceptance rate alone.
Adaptivity: Ablation studies showed that KnapSpec naturally increases the ratio of skipped Attention layers as context length grows, whereas baselines like SWIFT remained static.
Efficiency: The use of cosine similarity pruning ( $\tau=0.5$ ) reduced memory consumption by 31% without sacrificing speedup.

Significance
The paper claims that KnapSpec offers a practical and efficient solution for accelerating next-generation generative AI. By directly aligning the optimization objective with decoding throughput (TPT) rather than proxy metrics like acceptance rate, and by adapting to hardware-specific latencies, KnapSpec ensures high-speed inference for long sequences without compromising the target model’s output distribution or requiring additional training.