Source-linked AI summary
LiLiCorr: Lightweight Likelihood Correlation of Parallel Drafts for Speculative Decoding
Matan Rusanovsky, Yoav Miron, Roy Uziel, Omer Belhasin, Ran Zilberstein, Maor Ashkenazi, Michael Elad
TL;DR
Parallel drafters such as DFlash produce individually plausible but potentially incoherent blocks because they are trained on per-slot marginals. LiLiCorr jointly scores top-K candidates using in/out-vector compatibility in one pass and co-trains the drafter; it raises acceptance length on every benchmark and achieves the highest throughput in 70 of 72 settings.
Problem
DFlash’s per-position marginal training leaves the joint block distribution largely unconstrained, producing tokens that may be individually plausible but mutually incoherent.
Method
LiLiCorr jointly processes top-K candidates with in and out vectors, scores adjacent cosine compatibility in parallel, and trains the drafter jointly with the reranker.
Results
LiLiCorr achieves higher acceptance length than vanilla DFlash on every benchmark and the highest throughput in 70 of 72 evaluated settings.
Takeaways & Limitations
A single-pass candidate-correlation head makes coherent draft-time correction faster than correcting one draft slot at a time across the reported benchmarks, concurrencies, and input lengths.
Takeaways & Limitations
LiLiCorr cannot recover a correct token absent from the drafter’s top-K pool, so acceptance length is bounded by drafter coverage; longer-range coupling remains open.
Abstract
from arXiv · showhide
Speculative decoding accelerates language-model inference by drafting future tokens that the target model verifies in parallel. A diffusion-style block head such as DFlash is an attractive drafter, predicting an entire block of future tokens in one forward pass. However, it is trained on per-position marginals rather than the joint block distribution, so the tokens it emits are individually plausible yet jointly incoherent. We introduce LiLiCorr, a Lightweight Likelihood-based model that Correlates the per-position marginal distributions a drafter already produces. It keeps the top-k tokens at each position as candidates and processes them jointly, producing for each an in and an out vector. A pair of adjacent candidates matches when the earlier one's out vector has high cosine similarity with the later one's in vector. These matches capture the block's joint structure without ever materializing the full joint distribution. One lightweight network pass produces all the vectors, and the pairwise scores are then computed in parallel as batched matrix operations, leaving only a cheap greedy walk sequential. We further co-train the drafter with LiLiCorr, so it learns to propose candidates that correlate into longer accepted sequences. Over the vanilla DFlash drafter, LiLiCorr raises acceptance length on every benchmark by 9 to 19%, while its scoring head accounts for about 2.8% of the per-block latency. Against DFlash and two concurrent methods that also restore coherence at draft time, LiLiCorr delivers the highest throughput in 70 of 72 settings: nine benchmarks at two target sizes under greedy and temperature-one decoding, and a throughput sweep over six concurrencies, two input lengths and three entropy tiers, with all systems equally optimized on a common serving stack. Extending LiLiCorr to inputs an order of magnitude longer than it was trained on preserves that lead.
1 Introduction
LiLiCorr addresses incoherent blocks from DFlash’s independently trained per-slot marginals by correlating top-K candidates in one lightweight pass. Joint drafter training improves acceptance and throughput across broad evaluation settings.
- Motivation: DFlash predicts each slot well but leaves the block’s joint distribution largely unconstrained, producing locally plausible yet incoherent token sequences.Its per-position cross-entropy supervises marginals rather than the joint block distribution.
- Method: LiLiCorr retains top-K candidates per slot, assigns each an in and out vector, and scores adjacent compatibility with cosine similarity.The candidate interactions are computed as batched matrix operations after one network pass.
- Training: Jointly training LiLiCorr with DFlash teaches the drafter to propose candidate sets that support longer accepted sequences.The drafter is not kept frozen during training.
- Results: 9 to 19% higher acceptance length than vanilla DFlash is reported on every benchmark across target sizes and decoding regimes.The comparison covers 8B and 4B targets under greedy and temperature-1 decoding.
- Results: 70 of 72 settings have the highest throughput among compared systems, spanning benchmarks, target sizes, decoding modes, concurrencies, input lengths, and entropy tiers.The lead also survives inputs longer than the drafter’s training context.
2 Related Work
Related work addresses incoherent parallel drafts through diffusion refinement, target-side draft trees, or draft-time correlation. LiLiCorr differs from concurrent draft-time methods by avoiding one correlation-network pass per draft slot.
- Parallel drafters: Parallel and diffusion-style drafters factorize blocks into independent per-position predictions, so bidirectional attention does not guarantee joint coherence.Iterative refinement methods restore coherence by re-masking and re-predicting subsets of tokens.
- Target-side recovery: DDTree and Oda et al. construct draft trees whose candidate prefixes are verified by the target in one pass.DDTree is training-free, whereas Oda et al. learn a small autoregressive adapter first.
- Target-side recovery: DDTree and Oda et al. shift correlation work onto the target and rely on tree-structured attention, which is typically less efficient than standard causal attention.These methods therefore spend computation that speculative decoding aims to save.
3 Method
LiLiCorr reranks candidate lattices using locally normalized, prefix-aligned compatibility scores while retaining whole-block context in its representations. Training focuses on covered prefixes and jointly adapts the drafter; decoding leaves only a greedy walk sequential.
- Setup: The drafter proposes a block conditioned on verified context, and the target accepts only the longest correct prefix.Acceptance is prefix-structured under both greedy and sampling targets.
- Setup: Selecting jointly compatible candidates can lengthen accepted prefixes, whereas per-slot supervision makes tokens individually probable but not necessarily compatible.LiLiCorr therefore reranks alternatives from the drafter’s own marginals.
- Candidate lattice: LiLiCorr keeps top-K candidates at each slot and represents each with a node combining token, drafter, confidence, positional, and rank information.A small Transformer processes all candidate nodes jointly.
- Candidate lattice: Adjacent candidates are coupled by cosine similarity between the earlier out vector and later in vector, while those vectors incorporate information from the entire lattice.This preserves global context while keeping pairwise scores suitable for batched matrix products.
- Training: Local normalization conditions each slot’s candidate distribution on the candidate committed immediately before it, aligning training with prefix acceptance.This avoids optimizing a whole-chain score whose early sacrifice could invalidate a better-scoring tail.
- Training: Oracle-prefix supervision provides learning signals only when the ground-truth token is covered and the preceding prefix remains valid.A block whose first slot misses contributes zero loss, concentrating supervision on earlier slots.
- Training: The target-weighted distractor penalty complements cross-entropy by pushing down candidates that the verifier ranks substantially below the ground truth.The combined objective is reported to raise acceptance length on every short-prompt benchmark.
- Training: Joint training keeps selected candidate indices fixed while retaining gradients through draft log-probabilities and hidden states.This allows DFlash to learn candidate pools and representations that support longer accepted prefixes.
4 Experiments
Experiments compare LiLiCorr with DFlash, Domino, and DSpark across benchmark, decoding, concurrency, entropy, and input-length settings. LiLiCorr generally delivers the strongest throughput while maintaining competitive or improved acceptance lengths.
- Evaluation setup: LiLiCorr is evaluated on eight public datasets plus SPEED, using Qwen3-8B and Qwen3-4B targets under greedy and temperature-one decoding.The study also varies serving concurrency, input length, and output-entropy tier.
- Main results: LiLiCorr leads throughput in all eighteen greedy benchmark settings and improves over vanilla DFlash by 12 to 19% in acceptance length and 4 to 13% in throughput.It leads the strongest baseline by up to 4.9% on the benchmark average.
- Main results: At temperature one, LiLiCorr is fastest in seventeen of eighteen settings and fastest among single-block methods in all eighteen.Its lead over Domino widens under sampling, while DSpark’s sampled arm accepts longer prefixes but serves more slowly.
- Scaling with concurrency: LiLiCorr attains the highest throughput at 35 of 36 concurrency operating points, trailing DSpark by 0.3% only at 1K mixed with c = 32.Its throughput lead reaches 7.2% over the strongest baseline, while speedups compress as concurrency grows.
- Scaling with input length: Across 8K, 16K, and 32K blocks beyond the training range, LiLiCorr remains the highest-throughput system after identical YaRN positional extension.A hinge on the reranker’s decision gap improves robustness at these lengths.
5 Conclusion
LiLiCorr correlates per-slot marginals in one network pass, then uses batched matching and a cheap greedy walk to produce coherent drafts. Joint training improves acceptance and throughput broadly, but top-K coverage limits recoverability.
- Conclusion: LiLiCorr assigns in and out vectors to top-K candidates and scores adjacent compatibility through cosine similarity computed with batched matrix products.Only the final greedy walk remains sequential and touches no network weights.
- Conclusion: Jointly trained LiLiCorr improves vanilla DFlash acceptance on every tested benchmark and is faster almost everywhere across benchmarks, concurrencies, and input lengths.The conclusion attributes the speed advantage to comparable draft quality at lower correction cost than per-slot methods.
- Limitation: LiLiCorr cannot recover a correct token absent from the drafter’s top-K pool, so achievable acceptance length is bounded by drafter coverage.Longer-range couplings remain an open direction because their cost under prefix acceptance is unresolved.
A Implementation and Training Details
The implementation combines candidate, drafter, confidence, and positional features in a bidirectional Transformer head, followed by target-state fusion and vector prediction. Training and serving are optimized jointly and compared on a common deployment stack.
- Architecture: Each candidate representation sums projected target-token embeddings, shared drafter states, confidence features, and learned slot and rank tables.Confidence features include log-probability, probability, probability gap, normalized rank, and top-choice status.
- Architecture: The head uses a two-layer bidirectional Transformer over the candidate lattice, with learned slot-distance and same-slot attention biases.After Transformer processing, each candidate is fused with the target state before separate linear heads emit in and out vectors.
- Objective: LiLiCorr’s objective is evaluated only on the oracle prefix, while DFlash cross-entropy is applied at every slot.Slots after the first candidate-coverage miss contribute to none of the three LiLiCorr terms.
- Serving: Fused kernels, compiled matching and selection, and CUDA graphs combine candidate pooling and reranking during serving.The same optimization level is applied to Domino’s correction loop, while isolated head timing uses a separate pipeline.
- Experimental controls: LiLiCorr and baselines use the same regenerated training corpus and matched optimization effort, with Domino widened and DSpark width-matched.DSpark is evaluated under deterministic and sampled proposal contracts at temperature one.
B Additional Results
The additional results report the measurements underlying throughput comparisons, per-category SPEED outcomes, and low seed variability. LiLiCorr maintains broad throughput advantages across targets and evaluation categories.
- Underlying measurements: Block-weighted acceptance length and output throughput underpin the reported speedup comparisons for each target.Tables 5 and 6 provide these two measurements under the same protocol.
- Per-category SPEED results: LiLiCorr serves faster than every compared system in 41 of 44 SPEED category cells.It also accepts more tokens per block than Domino in 37 cells.
- Per-category SPEED results: 8.0% acceptance-length and 16.8% throughput gains are reported on Qwen3-4B under greedy decoding in the multilingual category.This is identified as the clearest category-level lead.
- Seed variability: Standard deviations of acceptance length stay below 0.28 across benchmarks and systems, with two exceptions under DSpark’s sampled contract.Variability is below 0.16 everywhere except two cases, and larger splits have standard deviations no higher than 0.14.
C Loss-Term Ablation
The distractor penalty materially improves acceptance length when its weight is reallocated within the fixed reranker budget. Its gains appear across all nine evaluation benchmarks.
- Ablation result: 1.5 to 8.9% acceptance-length gains occur on all nine benchmarks when the distractor penalty is included.The nine-benchmark average gain is 3.2%, with the largest improvement on HumanEval and the smallest on MATH-500.
- Interpretation: The distractor penalty converts ground-truth-only reranking into training against candidate confusers that consume acceptances.The comparison keeps the total reranker weight fixed rather than enlarging the objective budget.
D Derivations for the Reranker Objective
The derivation frames the distractor penalty as an expected rectified target cost over LiLiCorr’s candidate distribution. Its gradients complement cross-entropy rather than performing target-distribution matching.
- Penalty construction: The distractor penalty weights each candidate by its probability and its rectified cost relative to the ground-truth candidate.The candidate cost is defined as −log q_k, with rectification applied to cost differences.
- Rectifier: When the ground-truth candidate is target-preferred, the rectifier is inactive and the penalty reduces to an expected-cost objective.The occasional contrary ordering is attributed to numerical differences between supervision generation and teacher-forced training.
- Gradient relation: The penalty and target cross-entropy have the same gradient with respect to LiLiCorr parameters when the ground-truth offset is constant.The penalty depends on target cost differences and is unchanged by renormalizing the target over the candidate pool.
- Not distribution matching: Because the penalty is linear in LiLiCorr’s candidate probabilities, minimizing it drives commitment toward the lowest-cost candidate rather than reproducing the target distribution.Under the stated assumption, that candidate is the target’s preferred ground-truth candidate.
- Complementary gradients: Cross-entropy pulls the exact ground-truth token upward, while the penalty pushes target-rejected confusers downward.The penalty weakens when the ground-truth probability approaches zero and is indifferent among zero-cost candidates.
E Greedy Selection Versus the Highest-Scoring Path
The paper evaluates greedy left-to-right selection against exact highest-scoring path search under prefix acceptance. Greedy wins because early mismatches prevent later draft tokens from contributing to acceptance length.
- Greedy selection: The greedy walk maximizes the leading decision and then each subsequent decision conditioned on the committed prefix.This rule is tested while holding the trained model fixed.
- Selection objective: Prefix acceptance gives earlier slots greater value because later slots contribute only when every preceding slot matches.A whole-block optimizer can trade an early slot for a better tail, but that tail may never be reached.
- Ablation result: 0.90 to 4.81% lower acceptance length results from exact optimization across all nine benchmarks.The average decrease is 2.44%, while throughput differences are negligible because both rules use equally cheap fused kernels.
- Ablation result: Greedy decoding wins all eleven categories on the multi-domain SPEED suite.The exact optimizer is slower in acceptance length despite having essentially the same serving cost.
F Input-Length Sweep on the SPEED Throughput Split
The SPEED throughput sweep evaluates LiLiCorr across input lengths, concurrencies, and entropy tiers, including blocks beyond its training range. LiLiCorr maintains strong throughput and acceptance-length performance, though memory capacity constrains interpretation at the longest inputs.
- Within the training range: 1K and 2K blocks show an acceptance-length gap from 1.4% behind to 3.2% ahead of the strongest baseline, while throughput leads by up to 7.2%.Against Domino, LiLiCorr has longer blocks at all thirty-six operating points.
- Beyond the training range: 46 of 54 beyond-range operating points favor LiLiCorr in throughput, with gains up to 10.8%.The exceptions are DSpark wins on 16K and 32K blocks, costing at most 0.91% throughput.
- Beyond the training range: 38 of 54 beyond-range operating points give LiLiCorr the highest acceptance length.The eight exceptions are losses to DSpark on the 16K and 32K blocks, with at most a 4.2% acceptance-length cost.
- Positional extension: YaRN positional extension is applied identically to every system, and beyond-range acceptance remains roughly 2.8 to 5.8 tokens per block.Without the extension, every drafter loses acceptance length outside its training range.
- Memory-bound operating points: At 16K and 32K inputs, offered concurrency can exceed single-H100 key-value-cache capacity, causing some rows to represent the same resident operating point.The comparison remains valid, but the sweep no longer varies resident batch size as intended.
G Extending to Inputs Beyond the Training Range
For long inputs, the paper combines positional extension with margin regularization that targets contested decisions without changing inference. The regularized model preserves short-input behavior while improving beyond-range throughput and acceptance-length coverage.
- Margin regularization: A hinge on the decision gap is inactive after the gap clears ρ, concentrating training gradients on contested slots while leaving inference unchanged.The reported settings are µ = 0.5 and ρ = 2.0.
- Effect on the evaluation suite: Within the training range, LiLiCorr + margin differs from LiLiCorr by at most 0.52% on the unweighted average across the reported panels.It is behind by 0.10% in one panel and ahead by 0.04%, 0.10% and 0.52% in the others.
- Effect beyond the training range: LiLiCorr + margin achieves the highest throughput at all 54 beyond-range operating points, versus 46 for LiLiCorr.Its highest-acceptance-length count rises from 38 to 44.
- Effect beyond the training range: −0.1% at 8K, +1.3% at 16K and +2.8% at 32K are LiLiCorr + margin's mean acceptance-length gains over LiLiCorr.The gain increases with prompt length across the beyond-range blocks.
- Evaluation setup: Tables 14 and 17 report 8K, 16K and 32K throughput splits with positional extension applied identically to all systems.Rows marked † offer more concurrency than the key-value cache holds.