3倍生成速度還降內存成本，超越Medusa2的高效解碼框架終于來了
傳統上，大型語言模型（LLMs）被認為是順序解碼器，逐個解碼每個token。
來自上海交通大學、加利福尼亞大學的研究團隊展示了預訓練的LLMs可以輕松地被教導成為高效的并行解碼器，并介紹了一種新的并行解碼器族，稱為一致性大語言模型（CLLMs），能夠通過在每個推斷步驟中高效地解碼一個n -token序列來降低推斷延遲。
在此篇論文中，研究表明：「模仿人類在頭腦中形成完整句子后逐字表達的認知過程，可以通過簡單地微調預訓練的LLMs來有效地學習。」
具體而言，CLLMs通過將任何隨機初始化的n -token序列映射到盡可能少的步驟中，產生與自回歸（AR）解碼相同結果，來進行并行解碼的訓練。
實驗結果表明，使用該研究團隊所提出的方法獲得的CLLMs非常有效，在生成速度上顯示出該方法獲得了2.4倍至3.4倍的改進，與其他快速推斷技術如Medusa2和Eagle相媲美甚至更好，且在推斷時不需要額外的內存成本來容納輔助模型組件。
- 論文名稱：《CLLMs：Consistency Large Language Models》
- 論文鏈接：https://arxiv.org/pdf/2403.00835
圖1:在GSM8K上，使用Jacobi解碼時，CLLM-ABEL-7B-001是baseline ABEL-7B-001大約3倍速度的演示。
Jacobi解碼
大型語言模型（LLMs）正在改變人類生活的面貌，從編程到提供法律和健康建議。
然而，在推斷過程中，LLMs使用自回歸解碼逐token生成響應，如圖1所示，這導致了較長響應的高延遲。使用自回歸解碼，通常需要進行架構修改、輔助組件或初稿模型等，以通過一次生成多個token來加快推斷速度。
圖2:傳統的自回歸（AR）解碼示意圖：一次生成一個token。
Jacobi解碼源自Jacobi和Gauss-Seidel定點迭代求解非線性方程的方法，經證明與使用貪婪解碼的自回歸生成完全相同。
Jacobi解碼將順序生成過程重新構造為一個包含n個變量的n個非線性方程組，并基于Jacobi迭代可以并行求解。
每個迭代步驟可能會預測出多個正確的token（所謂的“正確”是指在貪婪采樣策略下與自回歸解碼結果對齊），從而潛在地加速自回歸解碼。
圖3：Jacobi解碼示意圖：將n -token序列饋送到LLM中，并進行迭代直至收斂。
具體來說，Jacobi解碼方法首先從輸入提示中隨機猜測序列的下一個token（以下簡稱為n -token序列，除非另有說明）。
然后，將n -token序列連同提示一起饋送到LLM中，以進行迭代更新。這個過程會持續進行，直到n -token的序列穩定下來，不再發生變化，達到一個固定點。
值得注意的是，Jacobi解碼不需要比自回歸（AR）解碼更多的對LLM的查詢。最終，n -token的序列會收斂到在貪婪策略下由AR解碼生成的輸出。從最初的隨機猜測到最終的AR生成結果的這一過程被稱為「Jacobi軌跡」。
Jacobi解碼迭代過程和Jacobi軌跡的一個實例在圖2中進行了說明。
Jacobi解碼的局限：
然而，在實踐中，普通的Jacobi解碼對LLMs的加速效果僅有微弱的提升，例如，平均加速比只有1.05倍。這是因為當LLM在先前的token中存在錯誤時，很難產生正確的token。
因此，大多數Jacobi迭代只能為n -token的序列獲得一個校正，導致如圖3左側所示的較長軌跡。
前瞻解碼和推測解碼方法試圖緩解Jacobi解碼和傳統的自回歸解碼的低效問題，但在推斷時會產生額外的內存成本。
而CLLMs則不需要這些額外的內存成本。
一致性大型語言模型（CLLMs）
初步Jacobi解碼：
給定一個prompt x和一個預訓練的LLM p(·|x)，通常研究者會使用標準的自回歸（AR）解碼方法在貪婪策略下獲得模型的響應，即：
Jacobi解碼重新構建了LLM推斷過程，將其視為解決非線性方程組的過程，以將解碼過程轉化為可并行計算的形式。考慮到：
研究者可以將上述方程重寫為一個非線性方程組：
需要注意的是：
該過程在某個k值處退出，使得：
然后，定義
作為固定點，并且將
作為Jacobi軌跡。
為了解決該問題，該研究團隊提出調整預訓練的LLMs，使它們能夠一致地將Jacobi軌跡 J 上的任意點 y 映射到固定點 y* 。
令人驚訝的是，他們發現這樣的目標類似于一致性模型的目標——一種擴散模型的主要加速方法。
在該團隊提出的方法中，使用從目標模型收集的Jacobi軌跡來訓練模型，并使用一種損失函數，該函數鼓勵在Jacobi迭代過程中實現單步收斂。
對于每個要調整為CLLM的目標模型 p ，訓練包括兩個部分：
（1）Jacobi軌跡準備：
對于每個提示，作者按順序對每個token截斷進行Jacobi解碼，直到整個響應序列 l 被生成為止，這相當于所有連續固定點的串聯。
沿軌跡生成的每個序列都會被算作一個數據條目。
此處需要注意的是，對于包含N（N ? n）個token的 I 的長響應，這種截斷避免了對長輸入的慢速模型評估。
（2）使用一致性和AR損失進行訓練：
作者聯合優化兩個損失來調整CLLMs，一致性損失確保一次預測多個token，而AR損失則防止CLLM偏離目標LLM，以保持生成質量。
圖4: one-step收斂一致性訓練的示意圖：將目標LLM調整為在Jacobi軌跡上的任何狀態作為輸入時始終預測固定點。
一致性和AR損失：
（1） 一致性損失
設 p 表示目標LLM。
使
表示為被初始化為 p 的參數 θ 的 CLLM。
對于prompt x 和相應的Jacobi軌跡 J ，令 y 和 y* 分別表示軌跡上的隨機狀態和固定點。
可以通過最小化以下損失來促使CLLM在輸入為 y* 時輸出 y ，稱為全局一致性（GC）損失：
在此公式里，
作者大量使用符號來表示從數據集中均勻抽樣。
D(·||·) 表示兩個分布之間的距離，選擇則在GKD方法中進行了討論，在本文中主要使用前向KL。
或者，根據一致性模型中的公式，使用局部一致性（LC）損失。
其中相鄰狀態：
在Jacobi軌跡 J 中，被驅使產生相同的輸出：
（2）AR損失：
為了避免偏離目標LLM的分布，作者結合了基于目標LLM p 的生成 l 的傳統AR損失：
通過將兩種損失結合在一起，使用一些權重 ω ，訓練CLLM的總損失為：
實驗
結果：
總的來說，該實驗涵蓋了三個特定領域的任務：
（1）Spider（文本到SQL）
（2）Human-Eval（Python代碼完成）和GSM8k（數學）
（3）更廣泛的開放域會話挑戰MT-bench。
所報告的實驗使用微調過的編碼器LLM、Deepseek-coder-7B-instruct、LLaMA-2-7B或ABEL-7B-001作為目標模型，具體使用則取決于任務。
訓練和評估則都在NVIDIA A100 40GB服務器上進行。
圖5：CLLM在不同下游任務上的加速效果。結果顯示：「CLLM比預訓練模型快得多，并且與Medusa相比實現了可比較的加速，但在推斷時沒有額外的成本。」
圖6：CLLM與其他基準在特定領域任務（Spider、CSN-Python、GSM8k）以及MT-bench上的對比示意圖。CLLM在與Medusa2的比較中實現了類似或甚至更好的加速效果，同時不引入額外的推斷成本（根據FLOPS和內存消耗判斷）。
專業領域：
從圖5中，可以看到，與其他基準（包括原始目標模型、Medusa2和猜測解碼）相比，CLLMs實現了最顯著的加速。
開放域會話挑戰（MT-bench）：
使用ShareGPT數據集從LLaMA2-7B訓練的CLLM與前瞻解碼（lookahead decoding）結合使用時，可以實現與Medusa2大致相同的加速，并在MT-bench上獲得可比較的分數。
然而，CLLM具有更高的適應性和內存效率，因為它不需要對目標模型的原始架構進行修改，也不需要輔助組件。
訓練成本：
CLLMs的微調成本是適度的。
例如，對于LLaMA-7B，只需傳遞大約1M token就可以在Spider數據集上實現它的3.4倍的加速。在數據集size較大的情況下（例如對于CodeSearchNet-Python），只需要使用數據集的10%來生成訓練CLLMs的Jacobi軌跡，從而獲得大約2.5倍的加速。
可以通過以下方式估算token的總數：
N = 平均每個prompt的軌跡量 × 平均軌跡長度 × Prompt數量。
圖7：Spider上目標LLM和CLLM之間的Jacobi軌跡對比。Jacobi軌跡上的每個點都是顏色編碼的序列：與AR結果匹配的正確標記為藍色，不準確的標記為紅色。CLLM表現出增強的效率，收斂到固定點的速度比目標LLM快2倍。CLLM的這種增強效率可以歸因于一致性損失，它促進了每個給定前綴的n-token 序列的結構的學習。
圖6左側顯示，目標LLM通常在一個迭代中只生成一個正確的token。相比之下，在CLLMs中，作者發現了快速推進現象，即在單個Jacobi迭代中連續多個token被正確預測。
此外，在目標LLM中，提前正確生成的token（例如圖7左側索引6和7處的「country」和「H」），往往在隨后的迭代中被不準確地替換。
另一方面，CLLMs表現出了預測正確token的能力，即使在先前有錯誤token的情況下，也確保token保持不變。
作者將這樣的token稱為「固定token」。這兩種現象共同促成了CLLMs在Jacobi解碼中的快速收斂，從而實現了相當大的生成速度提升。
研究團隊還觀察到，通過訓練，CLLMs獲得了一個關鍵的語言概念——搭配：「一系列詞或術語，它們的共同出現頻率高于隨機機會所預期的。」
語言不僅由孤立的單詞組成，而且還嚴重依賴于特定的詞對。搭配的示例在自然語言和編程語言中都很豐富。
它們包括：
- 動詞+介詞組合（例如「talk to」，「remind ... of ...」）
- 動詞+名詞結構（例如「make a decision」，「catch a cold」）
- 許多領域特定的句法結構（例如「SELECT ... FROM ...」，「if ... else」用于編程）。
一致性生成目標使CLLMs能夠從Jacobi軌跡的任何點推斷出這樣的結構，促進CLLMs掌握大量的搭配，并因此同時預測多個詞以最小化迭代步驟。
本文轉自 機器之心 ，作者：機器之心
