PaperlayerView exact revision on arXiv

Source-linked AI summary

X-CoSD: Communication-Efficient Cross-Vocabulary Collaborative Speculative Decoding

Jaeduk Lee, Wan Choi

arXiv:2609.09166v1cs.CLcs.DCcs.LG

TL;DR

Existing CoSD assumes shared SLM–LLM vocabularies, and residual resampling can require substantial distribution exchange. X-CoSD and X-CoSD-E redesign resampling for heterogeneous vocabularies, preserving the server LLM distribution while improving communication efficiency and token throughput without sacrificing generation quality.

  • Problem

    Existing CoSD assumes shared SLM–LLM vocabularies, while heterogeneous token spaces make conventional verification and resampling inapplicable; naive extension can require full LLM-distribution transmission.

  • Method

    X-CoSD uses hybrid resampling across common-vocabulary and LLM-only regions, while X-CoSD-E adds server resampling with device verification.

  • Results

    Both methods preserve the server LLM distribution, and experiments show improved token throughput without sacrificing the server LLM's generation quality.

  • Takeaways & Limitations

    X-CoSD and X-CoSD-E provide lossless, communication-efficient CoSD for heterogeneous vocabularies under limited communication resources.

  • Takeaways & Limitations

    The framework assumes that device and server share tokenizer specifications in advance.

Abstract

from arXiv · show

This paper investigates collaborative speculative decoding (CoSD), a distributed large language model (LLM) inference framework in which an on-device small language model (SLM) drafts candidate tokens and a server LLM verifies them. Existing CoSD methods assume a shared vocabulary between the SLM and the LLM and incur substantial communication load because residual resampling requires token distribution exchange between the user device and the edge server. To address these limitations, we propose cross-vocabulary CoSD (X-CoSD), a lossless and communication-efficient CoSD framework for heterogeneous SLM-LLM vocabularies. X-CoSD is built on hybrid resampling (HR), which splits residual resampling across the common-vocabulary region on the device and the LLM-only region on the server, so that distribution transmission is required only for the common-vocabulary region. We further propose X-CoSD-E, an enhanced variant based on server resampling with device verification (SR-DV), in which the server sends only replacement candidates sampled from the server LLM and their corresponding probabilities for local verification at the device. We prove that both X-CoSD and X-CoSD-E preserve the server LLM distribution, and experiments show that they significantly improve token generation speed while maintaining generation quality comparable to that of the server LLM.

1 Introduction

Collaborative speculative decoding distributes drafting and verification between an on-device SLM and a server LLM, but existing methods face communication bottlenecks and assume shared vocabularies. X-CoSD addresses heterogeneous vocabularies by redesigning rejection resampling to reduce distribution transmission while preserving exact decoding.

  • Collaborative speculative decoding: CoSD uses an on-device SLM for drafting and a server LLM for parallel verification, enabling distributed speculative decoding.The framework retains lossless decoding with respect to the server LLM while distributing computation across device and server.
  • Communication bottleneck: Repeated exchange of high-dimensional token distributions during rejection resampling can dominate latency over wireless networks.The communication burden is not limited to token IDs because rejected candidates require collocated SLM and LLM distributions.
  • Prior limitations: Prior communication-efficient CoSD methods reduce communication through selective verification or distribution compression but may degrade quality or retain uplink distribution transmission.U-HLM skips verification for some candidates and may compress SLM distributions using top-k truncation.
  • Vocabulary mismatch: Existing CoSD frameworks assume identical SLM and LLM vocabularies, making their verification and resampling procedures inapplicable to heterogeneous token spaces.Different model families or vendors may use incompatible tokenizers, so their token distributions cannot be directly compared.
  • X-CoSD: X-CoSD uses hybrid resampling to split residual sampling between the common-vocabulary region on the device and the LLM-only region on the server.Only the common-region portion of the server LLM distribution is transmitted instead of the full distribution.

2 Preliminaries

CoSD iteratively drafts candidate tokens on the device, verifies them on the server, and resamples rejected tokens, while heterogeneous vocabularies require token-space alignment. A naive TLI-based extension enables alignment but remains impractical because rejection can require transmitting the full LLM distribution.

  • CoSD procedure: Each CoSD verification round consists of device drafting, server verification, and resampling of a rejected token on the device.The device generates candidate tokens autoregressively, then requests server-side verification.
  • CoSD procedure: The device uploads candidate tokens and their generation probabilities rather than the full SLM distribution.The server runs the LLM on the concatenated prefix and candidate sequence to obtain verification distributions.
  • CoSD procedure: When a candidate is rejected, the server sends the accepted length and full LLM distribution at the rejected position for device-side residual resampling.If all candidates are accepted, the server instead samples and sends a bonus token.
  • Shared-vocabulary CoSD: Shared-vocabulary CoSD preserves the server LLM distribution while using asymmetric communication that limits routine uplink transmission.The server sends its distribution over the downlink only upon rejection.
  • Heterogeneous vocabularies: With heterogeneous vocabularies, the SLM and LLM distributions are not directly comparable, so conventional CoSD cannot be applied without modification.The mismatch arises because the distributions are defined over incompatible token spaces.
  • Naive heterogeneous extension: TLI restricts drafting to the common vocabulary, while LLM-only finalized tokens may be decoded into strings and retokenized by the device.The device and server share tokenizer specifications in advance, and this metadata exchange does not add generation-time communication.
  • Naive heterogeneous extension: The naive heterogeneous extension is impractical because every rejected candidate can require transmitting the full LLM distribution over the LLM vocabulary.This creates substantial downlink latency and motivates lossless resampling methods that avoid full-distribution transmission.

3 X-CoSD

X-CoSD uses hybrid resampling to support heterogeneous SLM–LLM vocabularies while reducing the distribution information exchanged after rejection. Its SR-DV variant further limits communication by sending server-sampled candidates for device verification and falls back to HR when needed, while both variants preserve the server LLM distribution.

  • Hybrid resampling: X-CoSD uses hybrid resampling to split residual sampling between the common-vocabulary region on the device and the LLM-only region on the server.The residual distribution is decomposed into region-wise components, allowing the resampling region to be selected before sampling from its conditional distribution.
  • Hybrid resampling: HR transmits only the server LLM distribution over the common vocabulary and the scalar θo, rather than the full distribution.The device computes θc and selects between the common and LLM-only regions using this information.
  • Hybrid resampling: Server resampling for the LLM-only region sends a single detokenized token string instead of a probability distribution, limiting its additional communication latency.The device retokenizes the received string with the SLM tokenizer for subsequent drafting.
  • Hybrid resampling: HR preserves exact decoding: the token finalized by X-CoSD follows the server LLM distribution q(x).This establishes lossless decoding for the HR-based framework.
  • Server resampling with device verification: SR-DV sends K server-generated replacement candidates with their probabilities for device-side verification and invokes HR only if all KM candidates are rejected.The procedure repeats for at most M iterations, while fallback to HR preserves exact residual resampling.
  • Server resampling with device verification: SR-DV reduces communication overhead by probabilistically avoiding transmission of the common-region distribution, and X-CoSD-E preserves the server LLM distribution q(x).The common-region distribution is transmitted only in the worst-case HR fallback.

4 Experiments

Experiments evaluate heterogeneous-vocabulary CoSD methods across generation quality, communication load, throughput, and latency. X-CoSD and X-CoSD-E preserve server-LLM-level quality while reducing communication and improving efficiency, with X-CoSD-E generally strongest.

  • Generation quality: X-CoSD and X-CoSD-E achieve generation quality comparable to the server LLM, while heuristic and lossy baselines show unstable or degraded quality.Quality is measured using BLEU, ROUGE-2, or accuracy depending on the dataset.
  • Communication efficiency: X-CoSD substantially reduces downlink communication versus Naive, and X-CoSD-E reduces it further through SR-DV with fallback to HR.Communication load is measured in bits per generated token; UL instead incurs substantial uplink load.
  • Token throughput: X-CoSD and X-CoSD-E consistently outperform Naive and achieve higher token throughput than server LLM inference as downlink rates vary.X-CoSD-E maintains high throughput across the tested downlink-rate range, whereas X-CoSD remains more dependent on downlink transmission rate.
  • Per-token latency: X-CoSD and X-CoSD-E show lower per-token latency than lossless baselines, with X-CoSD-E consistently lower than X-CoSD.This trend is reported for WMT-DeEn and is also observed on XSum, CNN/DailyMail, GSM8K, and MMLU.

5 Conclusions

X-CoSD and X-CoSD-E preserve the server LLM distribution while reducing communication for heterogeneous vocabularies. The frameworks improve throughput without sacrificing server-LLM generation quality.

  • X-CoSD and X-CoSD-E preserve the server LLM distribution through HR and SR-DV, respectively.X-CoSD-E uses server-sampled replacement candidates with device-side verification and can fall back to HR.
  • The proposed frameworks improve token throughput without sacrificing the server LLM's generation quality.The conclusion presents this as the overall experimental outcome.
  • The latency analysis models token-wise acceptance events as i.i.d. and accounts for uplink candidate transmission, downlink bonus tokens, and SR-DV iterations.If all M iterations fail, the analysis includes fallback communication for HR.
  • A larger replacement-candidate acceptance probability reduces fallback frequency and the associated communication overhead.The acceptance probability is constrained by α(x) ≤ 1, and the largest feasible value minimizes fallback communication.

A.4 Proof of Theorem 2

The proof decomposes the output probability into accepted and rejected cases, then shows their sum equals the server LLM distribution.

  • The output probability is decomposed into contributions from candidate acceptance and rejection.This decomposition is written as P[Y = x] = P[Y = x, A] + P[Y = x, R].
  • Under rejection, the conditional output distribution is the residual distribution, while the acceptance contribution is min(pc(x), q(x)).The rejected contribution is max(q(x) − pc(x), 0).
  • min(pc(x), q(x)) + max(q(x) − pc(x), 0) = q(x), proving exact preservation of the server LLM distribution.

B Implementation details

The implementation uses specified SLM and server LLM models, constructs a directly mapped common vocabulary, and compares communication against several baselines.

  • Experiments use vicuna-68m as the on-device SLM and Qwen2-7B and Llama-3.1-8B as server LLMs.All experiments run on a single NVIDIA A100 40GB GPU, with on-device time scaled using TITAN RTX measurements.
  • The common vocabulary contains tokens admitting a one-to-one direct mapping between the SLM and server LLM.
  • The baseline comparison includes Naive, UL, GR, TR, U-HLM, and the server LLM.
  • U-HLM measures candidate-token uncertainty using E temperature samples and verifies only tokens whose uncertainty exceeds uth.The experiments set E = 20, ηmax = 2.0, and uth = 0.5.
  • Communication load per verification round is summarized in Table 4, excluding the full-acceptance case's bonus-token transmission from the rejection comparison.

C.1 Communication load per token on additional datasets

Across additional datasets and server LLMs, X-CoSD and X-CoSD-E achieve lower downlink communication load than Naive and avoid the large uplink overhead of UL and U-HLM.

  • Table 5 reports communication load per generated token for CNN/DailyMail, GSM8K, and MMLU.The table uses N = 2 and reports load in bits.
  • X-CoSD and X-CoSD-E achieve substantially lower downlink communication load than Naive across all additional datasets and server LLMs.The reported quantity is communication load per generated token.
  • UL and U-HLM continue to incur large uplink communication overhead in these comparisons.These results support using X-CoSD and X-CoSD-E under limited wireless resources and heterogeneous vocabularies.

C.2 Token throughput on additional datasets

X-CoSD and X-CoSD-E achieve higher token throughput than the compared baselines across CNN/DailyMail, GSM8K, and MMLU by reducing downlink communication overhead. X-CoSD-E remains consistently strong across downlink transmission rates.

  • X-CoSD and X-CoSD-E achieve higher token throughput than Server LLM, Naive, and UL across CNN/DailyMail, GSM8K, and MMLU.The passage attributes this improvement to reduced downlink communication overhead.
  • X-CoSD-E maintains high token throughput regardless of downlink transmission rates through SR-DV.
  • U-HLM also attains high token throughput, but it does not guarantee lossless decoding.

C.3 Per-token latency on additional datasets

X-CoSD and X-CoSD-E show lower per-token latency than the lossless baselines across XSum, CNN/DailyMail, GSM8K, and MMLU. X-CoSD-E further reduces latency relative to X-CoSD by reducing communication overhead.

  • X-CoSD and X-CoSD-E show lower per-token latency than other baselines across XSum, CNN/DailyMail, GSM8K, and MMLU, except U-HLM.U-HLM achieves lower per-token latency but does not guarantee lossless decoding.
  • X-CoSD-E consistently shows lower per-token latency than X-CoSD by significantly reducing communication overhead.

D Limitations

The method’s TLI-based formulation restricts on-device drafting to the vocabulary shared with the server LLM, making performance dependent on common-vocabulary size. Experiments cover two server LLMs, despite diverse-task evaluation in low-overlap settings.

  • Under TLI, the on-device SLM can generate candidate tokens only from the vocabulary shared with the server LLM.
  • Performance may depend on common-vocabulary size because it affects speculative-decoding efficiency and device-server communication load.
  • Experiments are limited to two server LLMs, Qwen2-7B and Llama-3.1-8B.The evaluation nevertheless spans diverse tasks under challenging settings where the common vocabulary is less than 20% of the full LLM vocabulary.
Loading 2609.09166v1…