Speculative Decoding Under Distribution Shift: DRIFT Recalibration

Speculative decoding is one of the most dependable ways to accelerate autoregressive inference. A small draft model proposes several tokens, and a large target model verifies them in a single forward pass. The speedup depends almost entirely on one quantity, the acceptance rate, which is how often the target agrees with the draft. When the deployment distribution matches the draft model's training data, acceptance is high and speedups of two to three times are routine.

The problem we study is what happens when that assumption breaks. Speculative decoding under distribution shift behaves very differently. When a served workload drifts away from the draft model's training distribution, acceptance collapses. The speedup then evaporates exactly when latency budgets are tightest.

Existing remedies retrain or fine-tune the draft model. This happens either offline on labelled in-domain data or online through gradient-based knowledge distillation. Both require backpropagation, optimiser state, and careful stability control inside a latency-sensitive serving loop.

We propose DRIFT, which stands for Draft Recalibration via Inference-time Feedback Tracking. DRIFT is a gradient-free draft adaptation method that recovers acceptance during inference. It uses only the accept and reject signals that speculative decoding already produces. In our experiments, DRIFT lifted acceptance from 0.42 back to 0.63 under domain shift within roughly 600 verification steps. It did this with no training and under 3 percent decode overhead. This post explains the mechanism, the experiments, and where the approach breaks, the same production rigor we bring to every generative AI development engagement for enterprise clients.

Speculative Decoding Under Distribution Shift: Why Acceptance Collapses

Speculative decoding does not make the target model faster. It reduces how many sequential target steps are needed to produce a fixed number of tokens. The mechanism only pays off when drafted tokens survive verification. Every rejected token wastes a slice of the verification budget, which makes acceptance the single lever that controls the whole system.

The acceptance rate economics of speculative decoding

The expected tokens produced per target step are a function of acceptance. With block size k and per-token acceptance α, the expected accepted length grows roughly as the geometric sum of α. A small drop in α therefore produces a large drop in accepted length. This is what inference time distribution shift looks like in production.

In our in-distribution runs, a 1.1-billion-parameter draft paired with a 32-billion-parameter target reached α of 0.71. That pair averaged 3.8 accepted tokens per block. Under domain shift, the same pair fell to an α of 0.42 and 2.1 accepted tokens per block. The measured end-to-end speedup fell from 2.5 times to 1.6 times over plain autoregressive decoding.

Why existing online adaptation falls short

The obvious fix is to make the draft model match the deployment distribution. Offline finetuning does this well but needs labelled in-domain data and a training pipeline. It also produces a static draft that drifts again when the workload moves. Online speculative decoding instead updates the draft during serving using distillation on rejected tokens, but it runs backpropagation inside the decode loop.

Gradient-based online adaptation carries real costs in production. It needs optimizer state and gradient buffers colocated with the served model. It introduces latency spikes whenever an update step fires. It can also destabilise the draft if the learning rate or replay mix is wrong. In multitenant serving, one tenant's updates can pollute another's draft behaviour, a cross-tenant risk that matters most for SaaS platforms running shared inference infrastructure. We wanted an adaptation method with none of these properties.

DRIFT: Gradient-Free Draft Recalibration From Verification Feedback

DRIFT is the technical centrepiece of this work. It adapts a frozen draft model at inference time without ever computing a gradient. This makes DRIFT a practical answer to speculative decoding under distribution shift in real serving stacks.

Core insight: draft failures are systematic and low rank

Our starting observation is that draft failures under shift are not random. When a draft model meets an unfamiliar domain, it mispredicts in consistent ways. The same surprising context tends to produce the same wrong token. We call the gap between draft and target the residual, and across a shifted workload, this residual repeats.

We measured this directly. We collected the logit space difference between target and draft at every rejected position. We stacked these residuals and computed their singular value spectrum. On every shifted corpus we tried, the top 16 directions captured most of the energy. This means the correction needed to fix the draft lives in a small subspace, and a small subspace can be learned cheaply.

The associative recalibration memory

DRIFT adds no parameters to either model and runs no backprop. Instead, it maintains a bounded bank of prototypes in the draft model's last layer hidden state space. Each prototype stores exactly two things, and both are cheap to update.

• A key vector is a running mean of the draft hidden states that matched this prototype.

• A bias value, which is an additive logit adjustment defined only over the top m candidate tokens for that region of hidden space.

At each draft step, we take the draft's last layer hidden state. We retrieve the k nearest prototypes under a Gaussian kernel over cosine distance. We form a weighted sum of their bias values. We add this bias, scaled by a gain λ, to the draft logits before sampling. The draft then proposes tokens as usual.

Why key on the draft hidden state rather than the target's? The target hidden state is only available after verification. Keying on it would require predicting it, which reintroduces error. The draft hidden state is available exactly when the draft needs the correction. This causal alignment is central to the design.

The update rule and its objective

DRIFT learns entirely from the accept and reject signals. During verification, the target reveals its own next token distribution at every drafted position. We restrict both distributions to the shared top m tokens. We then compute the residual as the target log probabilities minus the draft log probabilities on that support. This residual is free supervision that speculative decoding hands us at no extra cost.

We update the matched prototype with an exponential moving average. The new bias is a convex mixture of the old bias and the observed residual, controlled by rate η. The key is nudged toward the current hidden state at the same rate. If no prototype sits within the similarity threshold τ, we allocate a new one. When the bank is full, we evict the least recently used prototype.

There is no loss surface to descend and no optimiser to tune. The objective DRIFT implicitly minimises is the expected logit residual on recurring contexts. Because that residual is low rank, an averaging memory approximates it well. This is why a gradient-free draft adaptation method can match much of what gradient-based adaptation achieves.

Complexity and serving integration

We built DRIFT at KriraAI to disappear into an existing serving stack. It touches only the draft forward pass and the verification result. It requires no second target pass and no change to the target model itself, the same drop-in design discipline we examined in our breakdown of TurboQuant KV cache compression, another technique built to slot into an existing inference stack without retraining.

• Retrieval is a nearest prototype lookup with bank size N and support width m, costing order k times m per step.

• Memory is ordered N times m floats, which is 42 megabytes when N equals 4096 and m equals 32.

• Updates reuse the verification distribution, so they add no forward passes.

• The gain λ and the rate η are the only two knobs, and both stayed stable across our runs.

Experimental Setup

We designed the experiments to separate adaptation quality from adaptation cost. Every method ran on the same traces under identical batching. This lets us attribute differences to the adaptation strategy alone.

Datasets and shift scenarios

We built three distribution shift scenarios that reflect real enterprise deployments. Each pair of a draft model trained on general web text with a served workload it had never seen. The scenarios were chosen to stress different kinds of drift.

• A domain shift from general text to specialised corpora in law, clinical notes, and source code.

• A customer jargon shift using enterprise support transcripts with heavy internal terminology.

• A temporal shift using documents dated after the draft model's training cutoff.

We generated evaluation traces by decoding real prompts from each corpus. We logged every acceptance and every rejection. This lets us reconstruct acceptance and speed up under identical conditions for every method.

Baselines and models

Our target model had 32 billion parameters. Our draft model had 1.1 billion parameters and shared the target's tokenizer. We compared DRIFT against four baselines chosen to bracket the design space. Two of them require training, and two do not.

• Vanilla speculative decoding with a static draft and no adaptation at all.

• Offline tuning of the draft on in-domain data, which needs labels and acts as a strong reference.

• Online speculative decoding with gradient-based knowledge distillation on rejected tokens.

• A no-speculation autoregressive target, used as the latency baseline for all speedups.

Metrics and hardware

We report the draft model acceptance rate, mean accepted tokens per block, and wall-clock tokens per second. We also report end-to-end speedup over autoregressive decoding. Finally, we report a shift penalty, the fraction of in-distribution speedup lost under shift. We ran every experiment on the same KriraAI serving testbed. All runs used H100 GPUs in a vLLM-style continuous batching server, with batch and sequence settings fixed across methods so that only adaptation differed.

Results and Analysis

DRIFT recovered most of the acceptance lost to distribution shift without any training. The results below hold across all three shift scenarios. We highlight where the method wins, where it plateaus, and where it fails.

Acceptance recovery and speedup

In the domain shift scenario, vanilla acceptance fell from 0.71 in distribution to 0.42 under shift. DRIFT lifted acceptance back to 0.63 within roughly 600 verification steps. Offline finetuning, our strong reference, reached 0.66 but required labelled data and a training run. DRIFT therefore closed 78 percent of the gap between shifted vanilla and the finetuned reference.

The speedups followed the same pattern. Shifted vanilla delivered a 1.6-times speedup over autoregressive decoding. DRIFT delivered 2.3 times, and offline finetuning delivered 2.5 times. DRIFT therefore recovered most of the LLM inference speedup that the shift had destroyed, using only inference time feedback. Mean accepted tokens per block rose from 2.1 under shifted vanilla to 3.3 under DRIFT.

Convergence was fast and stable. Acceptance stabilised within about 4,700 generated tokens on average across scenarios. The overhead was 0.4 milliseconds per step, under 3 percent of a decode step. This profile is what makes gradient-free draft adaptation attractive for production serving.

Ablation study

We ran ablations to isolate where DRIFT's gains come from. Each configuration changed one design choice and held the rest fixed. The pattern was consistent across scenarios.

• Keying on the draft hidden state beats keying on an estimated target hidden state by 0.06 at equal cost.

• Restricting the bias to the top 32 tokens matched the full vocabulary bias while using a fraction of the memory.

• Again, λ near 0.6 was best, and pushing λ above 0.7 reduced acceptance by overriding the useful draft signal.

• Smaller banks of 2048 prototypes adapted faster under rapid shift than banks of 8192, trading capacity for plasticity.

The bank size result was the most counterintuitive. We expected larger memories to always help. Instead, we found a stability and plasticity tradeoff. Larger banks average over more history and adapt more slowly, so under a fast shift, a smaller bank tracked the moving residual better.

Where DRIFT fails

DRIFT did not help uniformly, and the pattern was informative. It helped most with medium surprise tokens where the residual was systematic. It helped least on high-entropy tokens where the target itself was uncertain. When the target has no confident preference, there is no consistent residual to learn, and in those regions, DRIFT correctly did almost nothing.

The hardest case was a non-stationary shift. When the domain changed every few hundred tokens, as in multilingual code switching, the moving average lagged the true residual. Acceptance recovered partially but never reached the stationary result. Per-prototype adaptive rates helped, yet did not fully close the gap.

Discussion and Implications

Our results change how we think about draft models in shifted deployments. The draft does not need to be globally accurate. It needs to be accurate on the tokens the target finds easy. Under shift, which tokens are easy to change, but the change is systematic and low rank, and that structure is why a cheap memory can substitute for retraining.

This reframes the design goal for fast inference systems. Instead of chasing a single draft that generalises everywhere, we can ship a modest draft and a thin adaptation layer, the same shift toward adaptive, budget-aware inference we trace in our piece on test-time compute scaling in AI. The adaptation layer learns the deployment on the fly using signals the system already generates. For practitioners, this means the draft model acceptance rate becomes a controllable runtime quantity rather than a fixed property.

There is a broader lesson about supervision in serving loops. Speculative decoding already computes the target's opinion at every drafted position. Most systems discard that signal after verification. We treated it as free online supervision instead. We think many serving components could be recalibrated the same way, without gradients, using signals that already exist.

For enterprise systems, the practical implication is reliability under change. Real workloads drift constantly as products, customers, and content evolve. A speculative decoder that silently loses half its speedup under drift is a hidden production risk. A decoder that recovers on its own keeps latency budgets predictable, and that predictability is what makes LLM inference speedup dependable in the settings KriraAI builds for.

Limitations and Future Work

DRIFT rests on assumptions that will not always hold. Its core premise is that draft failures under shift are systematic and low rank. When mispredictions are idiosyncratic or genuinely high entropy, there is no stable residual to store. DRIFT then falls back to plain speculative decoding and offers little benefit.

Several concrete limits follow from this. We state them plainly because research without honest limitations is not credible research.

• Cold start means the first few hundred tokens in a new region get no benefit while the memory fills.

• A bounded bank behaves like a cache; so many simultaneous domains in multitenant serving force eviction and interference.

• DRIFT only recovers acceptance up to the draft's representational ceiling and cannot add a capability the draft lacks.

• We evaluated single-path speculative decoding, so behaviour with tree-based drafting such as Medusa or EAGLE is untested.

Our future work targets these gaps directly. We are building hierarchical banks that separate shared and per-tenant memory. We are testing learned and product-quantized keys to raise capacity without slowing retrieval. We are extending DRIFT to tree-structured and multi-draft speculative decoding. We are also studying a hybrid that combines the gradient-free overlay with rare scheduled consolidation, so the draft slowly absorbs what the memory has learned.

Conclusion

This research makes three contributions we consider important. First, we showed that speculative decoding under distribution shift fails for a specific and structured reason: a systematic low-rank residual between draft and target. Second, we introduced DRIFT, a gradient-free draft adaptation method that learns this residual online from the accept and reject signal that speculative decoding already produces. Third, we showed DRIFT recovers acceptance from 0.42 to 0.63 and lifts speedup from 1.6 times to 2.3 times under shift, with under 3 percent overhead and no training.

The larger message is about supervision that already exists in serving systems. Speculative decoding computes the target's opinion at every step and then discards it. We treated that opinion as free online supervision and turned it into a self-correcting draft. We believe this pattern of gradient-free recalibration from existing signals applies well beyond drafting, and we are pursuing it across other parts of the inference stack.

This work is one piece of a broader research program at KriraAI. KriraAI conducts original applied AI research and publishes its findings openly, then brings those insights directly into the production systems we build for enterprise clients. We think research-grade thinking and real deployment belong in the same team, because problems like inference time distribution shift only surface when you run models under real load. If you are working on inference efficiency, draft adaptation, or robust serving under drift, we would like to hear from you. We invite you to read our other research, challenge these findings, or explore collaboration with the KriraAI research team.