Fast Inference from Transformers via Speculative Decoding

2.3 Speculative Sampling

1. State the speculative sampling rule for one token.
   Sample \(x \sim q(x)\). Keep it if \(q(x) \le p(x)\); if \(q(x) > p(x)\), reject it with probability \(1 - p(x)/q(x)\) and instead sample from \(p'(x) = \mathrm{norm}(\max(0, p(x) - q(x)))\).
   • Leviathan et al., 2023, Fast Inference from Transformers via Speculative Decoding (Proceedings of the 40th International Conference on Machine Learning (ICML 2023), PMLR 202), §2.3
2. Why does speculative sampling yield exactly \(p(x)\)?
   A token \(x'\) is guessed and accepted with probability \(q(x')\min(1, p(x')/q(x')) = \min(p(x'), q(x'))\). Rejection happens with probability \(1 - \beta\), and the residual \(p'\) gives \(x'\) with probability \((p(x') - \min(p(x'), q(x')))/(1-\beta)\); the two terms sum to \(p(x')\).
   • Leviathan et al., 2023, Fast Inference from Transformers via Speculative Decoding (Proceedings of the 40th International Conference on Machine Learning (ICML 2023), PMLR 202), §2.3
3. Why not use ordinary rejection sampling instead?
   A non-iterative rejection sampler would accept \(x\) with probability \(p(x)/(M q(x))\), where \(M = \max_x p(x)/q(x)\), and otherwise fall back to sampling \(p\) directly; its acceptance rate is much lower than speculative sampling’s \(\sum_x \min(p(x), q(x))\).
   • Leviathan et al., 2023, Fast Inference from Transformers via Speculative Decoding (Proceedings of the 40th International Conference on Machine Learning (ICML 2023), PMLR 202), §2.3