3倍生成速度還降記憶體成本,超越Medusa2的高效解碼框架終於來了

机器之心發表於2024-05-11

高效解碼n -token序列,CLLMs+Jacobi解碼框架。

傳統上,大型語言模型(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 被生成為止,這相當於所有連續固定點的串聯。

沿軌跡生成的每個序列都會被算作一個資料條目。

此處需要注意的是,對於包含NN ≫ 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掌握大量的搭配,並因此同時預測多個詞以最小化迭代步驟。

參考連結:

https://hao-ai-lab.github.io/blogs/cllm/

相關文章