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