Coupling Without Communication and Drafter-Invariant Speculative Decoding

Authors:

Abstract
Suppose Alice has a distribution $\mathcal{P}$ and Bob has a distribution $\mathcal{Q}$. Alice wants to draw a sample $a \sim \mathcal{P}$ and Bob a sample $b \sim \mathcal{Q}$ such that $a=b$ with as high of probability as possible. It is well-known that, by sampling from an optimal coupling between the distributions, Alice and Bob can achieve $\operatorname{Pr}[a=b]=1-D_{T V}(\mathcal{P}, \mathcal{Q})$, where $D_{T V}(\mathcal{P}, \mathcal{Q})$ is the total variation distance between $\mathcal{P}$ and $\mathcal{Q}$. What if Alice and Bob must solve this same problem without communicating at all? Surprisingly, with access to public randomness, they can still achieve $\operatorname{Pr}[a=b] \geq \frac{1-D_{T V}(\mathcal{P}, \mathcal{Q})}{1+D_{T V}(\mathcal{P}, \mathcal{Q})}$ using a simple protocol based on the Weighted MinHash algorithm. This bound was shown to be optimal in the worst-case by Bavarian, Ghazi, Haramaty, Kamath, Rivest, and Sudan [ToC 2020]. In this work, we revisit the “communication-free coupling” problem. We provide a simpler proof of the optimality result from [Bavarian et al., 2020]. Moreover we show that, while the worst-case success probability of Weighted MinHash cannot be improved, an equally simple protocol based on Gumbel sampling offers a Pareto improvement: for every pair of distributions $\mathcal{P}$ and $\mathcal{Q}$, Gumbel sampling achieves an equal or higher value of $\operatorname{Pr}[a=b]$ than Weighted MinHash. Importantly, this improvement translates to practice. We demonstrate an application of communicationfree coupling to speculative decoding, a recent method for accelerating autoregressive large language models. We show that communication-free protocols can be used to construct DrafterInvariant Speculative Decoding schemes, which have the desirable property that their output is fixed given a fixed random seed, regardless of what drafter is used for speculation. In experiments on language generation, Gumbel sampling outperforms Weighted MinHash. Finally, we study the coupling problem in the setting where communication is bounded, rather than completely eliminated. We describe a protocol that uses just $O(\log (n / \epsilon))$ bits of communication to achieve $\operatorname{Pr}[a=b]=1-D_{T V}(\mathcal{P}, \mathcal{Q})-\epsilon$, i.e. to essentially match optimal coupling.

This publication has 11 references indexed in Scilit: