Fast Large Language Model Collaborative Decoding via Speculation

TL;DR

This paper introduces Collaborative decoding via Speculation (CoS), a novel framework that accelerates collaborative decoding without compromising performance and theoretically proves that CoS is never slower than standard collaborative decoding, typically achieving faster speed.

Abstract

Large Language Model (LLM) collaborative decoding techniques improve output quality by combining the outputs of multiple models at each generation step, but they incur high computational costs. In this paper, we introduce Collaborative decoding via Speculation (CoS), a novel framework that accelerates collaborative decoding without compromising performance. Inspired by Speculative Decoding--where a small proposal model generates tokens sequentially, and a larger target model verifies them in parallel, our approach builds on two key insights: (1) the verification distribution can be the combined distribution of both the proposal and target models, and (2) alternating each model as the proposer and verifier can further enhance efficiency. We generalize this method to collaboration among n models and theoretically prove that CoS is never slower than standard collaborative decoding, typically achieving faster speed. Extensive experiments demonstrate CoS is 1.11x-2.23x faster than standard collaborative decoding without compromising generation quality. Our code is available at https://github.com/Kamichanw/CoS/.

Figures (8)

• Figure 1: Comparison of (a) vanilla collaborative decoding, (b) speculative decoding, and (c) collaborative decoding via speculation. In (b) and (c), each discrete blue block represents a probability calculated by one forward pass of $\mathcal{M}_q$, while the continuous green block indicates the joint distribution requires only one forward pass of $\mathcal{M}_p$.
• Figure 2: The sketch of Alternate Proposal Framework. A continuous colored block indicates a single model invocation, with the bonus token highlighted in a red rounded box. Beginning from Step 2, $\mathcal{M}_q$ and $\mathcal{M}_p$ are invoked alternately. Each invocation involves both the verification of the current token and the generation of a bonus token. For clarity, we assume that the proposal length for each model is 1 and that all proposed tokens are accepted.
• Figure 3: The sketch of CoS in three-model collaboration scenario. The colored boxes represent the stored probability distributions, while the gray boxes represent the discarded ones. Each invocation involves scoring the current proposal tokens and generating a bonus token. For clarity, we assume that the proposal length for each model is 1 and that all proposed tokens are accepted.
• Figure 4: Comparison of speedup ratios for different $\lambda$ in WE across diverse setings. The blue and green lines represent the speedup ratios when the corresponding models serve as the proposal model, while the shaded region highlights the maximum speedup between the two.
• Figure 5: Comparison of speedup ratios for different $\mu$ in CD across different temperatures and datasets.

Theorems & Definitions (13)

• Theorem 3.1
• Theorem 3.2
• proof
• Corollary 3.3
• proof
• Corollary 3.4
• proof
• Corollary 3.5
• proof
• Theorem 3.6