Mitchell Stern 等人於 2018 年提出了投機取樣的原型概念。這種方法後來被各種工作進一步發展和完善，包括 Lookahead Decoding、REST、Medusa 和 EAGLE，投機取樣顯著加快了大型語言模型 (LLM) 的推理過程。
一個重要的問題是：LLM 中的投機取樣會損害原始模型的準確性嗎？先說答案：不會。
標準的投機取樣演算法是無損的，本文將透過數學分析和實驗來證明這一點。
投機取樣公式可以定義如下：
其中：
- 𝑟 是從均勻分佈中取樣的實數。
是要預測的下一個token。- 𝑝(𝑥) 是草稿模型給出的下一個token分佈。
- 𝑞(𝑥) 是基礎模型給出的下一個token分佈。
為簡單起見，我們省略了機率條件。實際上，𝑝 和 𝑞 是基於字首token序列 
的條件分佈。
以下是 DeepMind 論文中關於該公式無損性的證明：
如果你覺得閱讀數學方程式太枯燥，接下來我們將透過一些直觀的圖表來說明證明過程。
這是草稿模型 𝑝 和基礎模型 𝑞 的分佈示意圖：
圖1：草案模型p和基礎模型q輸出分佈的機率密度函式
需要說明的是，這只是一個理想化的圖表。在實踐中，我們計算的是一個離散分佈，它看起來像這樣：
圖2：語言模型預測詞彙集中每個token的離散機率分佈，藍色條來自草稿模型，紅色條來自基礎模型。
然而，為了簡單和清晰起見，我們使用它的連續近似來討論這個問題。
現在的問題是：我們從分佈 𝑝 中取樣，但我們希望最終結果就像我們從 𝑞 中取樣一樣。一個關鍵思想是：將紅色區域的機率搬運到黃色區域：
圖3：接受和拒絕取樣的區域
目標分佈 𝑞 可以看作是兩部分的總和：
I. 校驗接受
該分支中有兩個獨立事件：
- 在草稿分佈 𝑝 上取樣產生特定token 𝑥。機率為 𝑝(𝑥)
- 隨機變數 𝑟 接受該token 𝑥。機率為：
將這些機率相乘：
圖4：將藍線和紅線相乘，結果是圖6中的綠線
II. 校驗拒絕
該分支中也有兩個獨立事件：
- 𝑟 拒絕了 𝑝 中的某個token，機率為：
這是一個積分值，數值與特定token x 無關
- 在分佈 𝑞−𝑝（的正數部分）上取樣產生特定token 𝑥，機率為：
其分母作用是對機率分佈進行歸一化，以保持機率密度積分等於 1。
兩項相乘，第二項的分母被約掉：
max(0,𝑞(𝑥)−𝑝(𝑥))
圖5. 該圖中的紅線與綠線對應函式相乘，結果等於圖6中的紅線
為什麼拒絕機率恰好可以歸一化max(0,𝑞−𝑝) ？看起來似乎是巧合，這裡一個重要的觀察是，圖 3 中紅色區域的面積等於黃色區域的面積，因為所有機率密度函式的積分都等於 1。
將I, II兩部分相加： 
最後，我們得到了目標分佈 𝑞。
圖6. 綠色區域與紅色區域相加，正好等於黃線以下的面積
而這正是我們的目標。
儘管我們已經從原理上證明了投機取樣是無損的，但演算法的實現仍然可能存在bug。因此，實驗驗證也是必要的。
我們對兩種情況進行了實驗：貪婪解碼的確定性方法和多項式取樣的隨機方法。
我們要求 LLM 生成兩次短篇故事，首先使用普通推理，然後使用投機取樣。兩次取樣溫度均設定為 0。我們使用了 Medusa 中的投機取樣實現。模型權重為 medusa-1.0-vicuna-7b-v1.5 及其基礎模型 vicuna-7b-v1.5。
測試執行完成後，我們得到了兩個精確相同的結果。生成的文字如下：
User: Please tell a tale about a cyperpunk knight named PPIO keeps a charming llama as his pet. Assistant: Once upon a time, in a world where technology had advanced beyond our wildest dreams, there lived a cyberpunk knight named PPIO. PPIO was a skilled warrior, with advanced cybernetic enhancements that gave him incredible strength and speed. He was also a man of great honor, always striving to do what was right and just. One day, while out on a mission, PPIO stumbled upon a beautiful llama. Despite being a cybernetic being, PPIO had a soft spot for animals and he couldn't resist rescuing the llama from whatever danger it may have been in. From that day on, the llama followed PPIO wherever he went, and the two became fast friends. The llama, whose name was Llama-ella, was a special llama. It had a unique ability to sense when someone was in danger and would let out a loud, distinctive call to alert them to the danger. PPIO was grateful for Llama-ella's ability and the two became an unstoppable team. Together, PPIO and Llama-ella took on all sorts of challenges, from battling cybernetic villains to rescuing innocent people from harm. They were a force to be reckoned with, and their bravery and honor were an inspiration to all who knew them. |
在隨機取樣的情況下，情況更加複雜。大多數在隨機程式中重現結果的方法都使用固定的隨機種子來利用偽隨機生成器的確定性。但是，這種方法不適合我們的場景。我們的實驗依賴於大數定律：如果有足夠的樣本，則實際分佈與理論分佈之間的誤差將收斂於零。
我們編制了四個提示文字，對LLM在每個提示下生成的首個token進行了 1,000,000 次投機取樣迭代。使用的模型權重為 Llama3 8B Instruct 和 EAGLE-LLaMA3-Instruct-8B。統計結果如下所示：
|
|
|
|
- 藍色：來自基礎模型輸出 logits 的 softmax
- 綠色：來自草稿模型輸出 logits 的 softmax
- 橙色：來自投機取樣（1,000,000 次）的首個token計數佔比
基礎模型輸出結果與取樣分佈之間的標準差為 9.694e-5。符合預期。
投機取樣不會損害大型語言模型的推理精度。透過嚴格的數學分析和實際實驗，我們證明了標準投機取樣演算法的無損性。其中數學證明說明了投機取樣公式如何保留基礎模型的原始分佈。我們的實驗，包括確定性貪婪解碼和機率多項式取樣，進一步驗證了這些理論發現。貪婪解碼實驗在使用和不使用投機取樣的情況下產生了相同的結果，而多項式取樣實驗表明，在大量樣本中，token分佈的差異可以忽略不計。
這些結果共同證明，投機取樣可以顯著加快 LLM 推理速度，而不會犧牲準確性，為未來更高效、更易於訪問的 AI 系統鋪平了道路。