Source-linked AI summary
Multi-Access Speculative Inference: Uplink or Downlink?
Chang Cai, Kaibin Huang
TL;DR
Multi-SPIN faces a communication bottleneck when rejected drafts require transmitting vocabulary-wide distributions, with fixed uplink or downlink policies poorly suited to heterogeneous multi-device resources. The proposed framework adaptively selects communication modes, controls draft lengths, and allocates power, significantly improving token goodput on Qwen2.5 and DeepSeek-R1 model pairs.
Problem
Rejected drafts require vocabulary-wide distribution transmission, creating an uplink bottleneck while assigning all devices to downlink correction can overload shared downlink resources.
Method
The framework jointly optimizes adaptive UL/DL mode selection, device-specific draft lengths, and power allocation through stochastic decomposition, prefix-structured mode search, and greedy draft-length control.
Results
The proposed framework significantly improves token goodput on Qwen2.5 and DeepSeek-R1 model pairs, while adaptive mode selection outperforms fixed UL or fixed DL policies as device competition increases.
Takeaways & Limitations
Adaptive communication-mode selection balances limited UL rates against shared DL constraints, making neither fixed UL nor fixed DL universally optimal in Multi-SPIN.
Abstract
from arXiv · showhide
Multi-access speculative inference (Multi-SPIN) extends SPIN to multi-device edge networks to accelerate cooperative token generation. It allows on-device small language models (SLMs) to autoregressively draft multiple tokens for individual generation tasks, while an edge-server large language model (LLM) verifies them in parallel. The major communication overhead arises when a drafted token is rejected by the server, in which case sampling the correction token requires access to both the SLM-output draft distribution and the LLM-output target distribution over the full token vocabulary. Existing designs typically perform correction at the server by uploading the draft distribution, but transmitting a vocabulary-wide distribution creates a critical uplink (UL) bottleneck. Alternatively, the correction can be performed at the device by downloading the target distribution, leveraging the high transmission rates available on the downlink (DL). Motivated by this insight, we introduce communication-mode selection as a new design dimension for Multi-SPIN. Specifically, each device can adaptively switch between the UL and DL modes to balance the UL bottleneck against the shared DL resource constraint, thereby relieving the overall communication burden. We formulate a sum-token-goodput maximization problem that jointly accounts for mode selection, draft-length control, and power allocation. For mode selection, we reveal a simple optimal structure that enables efficient search over the number of UL devices, with the corresponding transmit powers optimized accordingly. For draft-length control, we develop a greedy-search algorithm that adapts device-specific draft lengths to heterogeneous computation and communication capabilities. Experimental results on Qwen2.5 and DeepSeek-R1 model pairs demonstrate that the proposed framework significantly improves token goodput.
I. INTRODUCTION
Multi-SPIN addresses token-correction communication bottlenecks by adaptively selecting uplink or downlink correction modes while jointly controlling draft lengths and power. The resulting framework uses structured optimization and greedy draft-length search to improve token goodput across heterogeneous edge conditions.
- Communication bottleneck: Rejected drafts create the main communication bottleneck because correction sampling requires full vocabulary-wide SLM and LLM distributions.Acceptance checking needs only scalar probabilities, whereas correction requires both distributions at the rejection position.
- Communication-mode selection: Multi-SPIN lets each rejected device correct through UL or DL, balancing limited uplink rates against competition for shared downlink resources.All-device DL assignment can overload the shared downlink, while adaptive switching exploits heterogeneous wireless capabilities.
- Joint optimization: The framework jointly optimizes expected sum token goodput through proactive draft-length control and reactive mode selection with power allocation.Verification outcomes determine reactive correction decisions, while draft lengths are selected before those outcomes are known.
- Reactive optimization: Sorting devices by minimum UL delay yields a prefix-structured optimal mode assignment and reduces mode selection to a one-dimensional search over the UL/DL split.UL devices occupy the prefix, while remaining devices use DL; powers are then optimized through maximum-power UL and delay-equalizing DL allocation.
- Draft-length control and evaluation: A greedy-search algorithm selects device-specific draft lengths under heterogeneous computation and communication conditions, with sample average approximation avoiding exhaustive outcome enumeration.Experiments use Qwen2.5 and DeepSeek-R1 model pairs across code, mathematics, conversation, and instruction-following tasks.
- SPIN background: SPIN combines local SLM autoregressive drafting with parallel server-side LLM verification to generate multiple tokens per round.The LLM evaluates the drafted block through a single forward pass, preserving the target distribution while accelerating inference.
B. Key Insights for Edge Deployment
SPIN communication is negligible when drafts are accepted but becomes costly after rejection, when correction requires a full distribution. Multi-SPIN therefore chooses UL or DL correction per device to balance asymmetric link rates against shared downlink contention.
- Communication overhead: Acceptance checking uses only scalar draft and target probabilities, whereas rejected-token correction requires both full distributions at the rejection position.The distinction explains why token correction, rather than initial verification, dominates communication overhead.
- UL/DL alternatives: Correction can use UL draft-distribution upload or DL target-distribution download because the two distributions are available at opposite endpoints.Both distributions share the vocabulary, so either direction carries the correction information.
- Multi-access trade-off: Although DL rates are often higher than UL rates, assigning every correcting device to DL can overload shared server transmit power.Moving a suitable subset to UL can relieve DL contention and better exploit device-side uplink capability.
- Multi-device architecture: Multi-SPIN serves independent device tasks while batching their drafted sequences for parallel LLM verification.The architecture exploits both parallelism across draft positions and parallelism across device requests.
A. Mode-Selection-Based Multi-SPIN Protocol
The protocol uses SLMs to draft token blocks, server-side batched LLM verification, and outcome-dependent bonus-token or correction handling. Devices with rejected drafts are assigned UL or DL correction modes with associated transmit powers.
- Common Steps: Each device autoregressively drafts γk tokens locally, uploads token–probability pairs, and the server verifies all device drafts and positions in parallel.The server also produces target distributions for drafted positions and the next position.
- Devices with Accepted Drafts: When all drafted tokens are accepted, the server samples and downloads one bonus token, then both sides append it and begin the next round.
- Devices with Rejected Drafts: For rejected drafts, the server selects a communication mode and transmit power for each affected device.The mode indicator uses 0 for UL and 1 for DL; the power variable denotes device UL power or server-allocated DL power accordingly.
- Devices with Rejected Drafts: In UL mode, the device uploads the draft distribution at the first rejection position, enabling server-side residual sampling before correction-token download.
- Devices with Rejected Drafts: In DL mode, the server downloads the target distribution at the first rejection position, the device samples the correction token, and then uploads that token.
- Round Timeline: One synchronized round selects draft lengths first, verifies after all device drafts arrive, and transmits distributions only for devices with rejected drafts.Lightweight signaling steps are omitted from the timeline illustration.
B. Computation Throughput and Latency
The computation model accounts for device-side autoregressive drafting and batched server verification. Drafting latency depends on device-specific draft lengths and per-token generation speeds, while verification exploits parallelism across positions and devices.
- Local Drafting: Device k generates γk draft tokens autoregressively, with drafting latency determined by its draft length and average per-token generation latency.
- Total Computation Latency: Total computation latency combines synchronized local drafting with the server-side batched verification latency.Because the server waits for all devices before verification, the round is governed by the coordinated computation stages.
- Batched Verification: The server verifies drafted positions in parallel within each request and across device requests using a batched LLM forward pass.
- Batched Verification: Verification latency depends weakly on moderate draft length and primarily on batch size, modeled through fixed and per-additional-device latency terms.
C. Communication Overhead and Latency
Multi-SPIN communication focuses on the vocabulary-wide distribution exchange required after rejection. The model captures UL/DL transmission latency, shared-resource constraints, and the two-stage timing of proactive draft-length and reactive communication decisions.
- Communication Model: Distribution transmission is treated as the dominant communication overhead, while other lightweight signaling operations are neglected.
- Communication Model: UL and DL correction transmit distributions over the same vocabulary and therefore incur the same distribution payload size, despite using different links.Qwen2.5 has 152,064 vocabulary entries and GPT-4o has 200,019 entries.
- Communication Latency: Under OFDMA, each device receives fixed orthogonal bandwidth for either UL or DL transmission, with channel-aware achievable rates and power constraints.
- Communication Latency: For rejected devices, each distribution-transmission latency depends on its mode and power, while simultaneous transmissions make the slowest device determine overall communication latency.
- Joint Optimization: The optimization jointly selects proactive device-specific draft lengths and reactive modes and powers to maximize expected sum-token goodput.
- Joint Optimization: Draft lengths determine verification-outcome distributions, and each realized outcome determines the rejected devices and reactive communication problem, enabling decomposition into two subproblems.
1) Reactive Mode Selection and Power Allocation:
For fixed draft lengths, optimizing reactive communication reduces to minimizing communication latency separately for each verification outcome. The resulting mixed-integer problem jointly chooses UL/DL modes and powers under a shared DL power constraint.
- Reactive Optimization: With draft lengths fixed, expected output tokens and computation latency do not depend on the reactive policy, so reactive optimization minimizes expected communication latency.
- Reactive Optimization: The reactive minimization separates across verification outcomes, allowing an optimal communication action to be selected independently for each realized outcome.
- Reactive Optimization: The per-outcome problem jointly optimizes binary mode assignments and continuous transmit powers, with a shared DL power constraint coupling DL devices.
- Reactive Optimization: Because possible mode assignments grow exponentially with device count, exhaustive search is computationally prohibitive.
2) Proactive Draft-Length Control:
The section formulates proactive draft-length control around expected communication latency and develops the reactive optimization components needed to evaluate each draft-length choice.
- Proactive Draft-Length Control: Draft-length control is evaluated through the minimum expected communication latency over verification outcomes.The reactive communication problem is solved for each realized outcome before evaluating a draft-length vector.
- Reactive Communication Optimization: The resulting reactive optimization jointly determines communication modes and powers under UL and shared-DL constraints.The formulation is mixed-integer because mode selections are binary while power allocations are continuous.
- Reactive Communication Optimization: For a fixed mode selection, UL devices transmit at maximum available power to minimize transmission latency.The DL side instead minimizes the maximum device delay under a shared sum-power constraint.
- Reactive Communication Optimization: DL power allocation has a closed-form solution parameterized by the resulting DL delay and can be computed efficiently by bisection.The result is stated for any nonempty DL user set.
B. Optimal Mode Selection and Power Allocation
Optimal mode selection has a prefix structure after devices are sorted by minimum UL delay, reducing joint mode and power optimization to a one-dimensional split search.
- Delay Comparison: The optimized communication delay for split n is the maximum of the UL and DL delays.The candidate split therefore balances the two mode-specific delays.
- Prefix Structure: An optimal mode assignment uses a UL prefix and assigns the remaining devices to DL after sorting by minimum UL delay.The UL set is U(n) = {1, 2, ..., n}, while the DL set is D(n) = {n + 1, ..., |R|}.
- Delay Comparison: As more devices move to UL, UL delay is non-decreasing while DL delay is non-increasing.The DL set shrinks and shares the same total DL power budget among fewer devices.
- Bisection Search: The optimal split is found by comparing the two candidate integers around the first non-positive delay difference.This split-point search can be obtained efficiently by bisection.
- Power Allocation Coupling: Each split evaluation embeds optimized power allocation rather than using fixed-power delay values.UL delay uses maximum-power transmission, while DL delay is optimized for the current DL set.
V. PROACTIVE DRAFT-LENGTH CONTROL
The paper controls device-specific draft lengths with a two-phase greedy search, while using sample average approximation to evaluate an otherwise exponentially costly expected-goodput objective.
- Optimization Challenge: Draft-length optimization is highly non-convex because each candidate requires solving the underlying reactive communication problem.The objective depends on minimum expected communication latency and lacks a closed-form solution.
- Greedy Search: The greedy algorithm first applies synchronized all-ones updates to exploit global draft-length improvements.An update is accepted when it increases expected goodput.
- Greedy Search: After synchronization converges, coordinate-wise refinement selectively increases individual draft lengths to capture residual gains.The device with the largest expected-goodput improvement is selected at each iteration.
- Expected Goodput Evaluation: Direct expected-goodput evaluation requires averaging over 2^K verification outcomes and solving a reactive problem for each outcome.The minimum communication latency lacks a closed-form expression because it depends on mode selection and power allocation.
- Expected Goodput Evaluation: Sample average approximation replaces exhaustive outcome enumeration with N_SAA sampled outcomes and provides a consistent estimate as N_SAA increases.The estimate is used throughout the greedy algorithm.
VI. EXPERIMENTAL RESULTS
The experiments evaluate Multi-SPIN under heterogeneous model, task, computation, and wireless settings, comparing adaptive mode selection and draft-length control with fixed and computation-only baselines.
- Experimental setup: Experiments use Qwen2.5 and DeepSeek-R1 model pairs across code generation, mathematical reasoning, conversation, and instruction-following tasks.Each device is independently assigned one of four representative tasks, while the server performs batched verification on an NVIDIA A100 GPU.
- Experimental setup: The evaluation models heterogeneous device computation by scaling each device’s drafting latency independently by a factor uniformly drawn from [1, 1.15].Reference per-token drafting latencies are 15.2 ms for Qwen2.5-0.5B and 19.6 ms for DeepSeek-R1-Distill-Qwen-1.5B.
- Experimental setup: The wireless evaluation uses OFDMA with 20 devices by default, 10 MHz bandwidth, equal bandwidth allocation among rejected devices, and quantized effective vocabulary distributions.Channels follow independent Rayleigh small-scale fading, and devices share specified uplink power budgets.
- Benchmarks: Mode-selection comparisons include all-UL, all-DL, and random-selection policies, with powers optimized according to each policy’s communication direction.Draft lengths are separately optimized under each communication policy to isolate mode-selection gains.
- Benchmarks: Draft-length comparisons include optimal uniform lengths and computation-only control, while both use the proposed optimal reactive mode-selection and power-allocation policy.The uniform length is exhaustively selected from {0, 1, ..., 20}; computation-only control ignores communication delay.
B. Adaptive Draft-Length and Mode-Selection Behaviors
The proposed controller adapts draft lengths and correction directions to DL power, device heterogeneity, and shared-resource competition. Across the evaluated model pairs and power-budget settings, adaptive mode selection and draft-length control outperform fixed or uniform alternatives.
- B. Adaptive Draft-Length and Mode-Selection Behaviors: As DL power increases, sum token goodput rises and gradually saturates when local drafting and server verification become the dominant bottlenecks.The average optimized draft length generally increases because higher DL power reduces the communication penalty of rejection.
- B. Adaptive Draft-Length and Mode-Selection Behaviors: With limited DL power, most correction traffic uses UL; as DL power increases, rejected devices progressively shift toward DL correction.Some devices remain in UL mode at high DL power to avoid competition for the shared DL budget.
- C. Performance Gain of Optimal Mode Selection: Under both DL-to-UL power-budget ratios, optimal mode selection achieves the highest goodput by assigning only a suitable subset of devices to DL.At ratio 4, all-DL suffers shared-server-power competition and all-UL is constrained by per-device transmit power; at ratio 16, all-DL becomes more competitive but remains inferior.
- C. Performance Gain of Optimal Mode Selection: Neither a fixed UL nor a fixed DL policy is universally optimal as the number of devices increases and shared DL competition intensifies.The results support adapting the transmission direction to system conditions rather than relying on a predetermined mode.
- C. Performance Gain of Optimal Mode Selection: Proactive draft-length control consistently outperforms optimal uniform and computation-only draft-length baselines across DL-to-UL power-budget ratios.The gains reflect adaptation to device-specific latency heterogeneity and the importance of accounting for communication effects.
- VII. Concluding Remarks: Future designs could jointly select transmission direction, exchanged representation granularity, model collaboration, and draft length.The proposed perspective extends communication-mode selection toward broader co-design of computation placement, information representation, and transmission direction.
APPENDIX A PROOF OF LEMMA 2
The proof shows that an optimal mode selection can use a prefix UL set, reducing the search to the number of UL devices. Monotone UL and DL delay behavior then enables efficient localization of the optimum.
- Prefix structure: Any non-prefix UL assignment can be transformed into a prefix assignment without increasing communication delay.Moving an earlier device from DL to UL does not increase minimum UL delay and cannot increase minimum DL delay because the DL user set shrinks.
- Prefix structure: The resulting optimal UL set has the form U(n) = {1, 2, ..., n} for some n in {0, 1, ..., |R|}.Devices are ordered by their minimum UL delays, so the transformation repeatedly removes inversions from the UL set.
- Efficient search: Because UL delay is non-decreasing and DL delay is non-increasing in n, their delay difference is non-increasing and its first nonpositive crossing can be found by bisection.The global minimizer lies around this crossing because the dominant delay switches from DL to UL there.