輔助式文本生成 - 文本生成新范式

輔助式文本生成 - 文本生成新范式

@(NLP)[文本生成]

引言

LLM 要跨越從早期采用者到大眾市場(chǎng)的鴻溝，其必要條件是價(jià)格大眾化，也就是降低每詞元的價(jià)格。



這種降低最好是“免費(fèi)”的降低，而不是“討價(jià)還價(jià)”的降低。拿電價(jià)打比方，這種降低最好是直接的電費(fèi)降低，而不是使用“省電小妙招”帶來(lái)的降低，這樣才能真正做到大眾化。在這種語(yǔ)境下，prompt 壓縮、模型量化這類優(yōu)化就是“討價(jià)還價(jià)”型優(yōu)化，它們要么需要用戶具備某些專家性的知識(shí)、要么需要用戶放棄一些生成質(zhì)量的一致性要求，總之是需要用戶付出一些成本。

模型量化不是“免費(fèi)”的降價(jià)
以 BLOOM-176B 為例，對(duì)提示 Once upon a time, there existed a little girl, who liked to have adventures. She wanted to go to places and meet new people and have fun.
原始 BF16 模型的生成為：She wanted to see the world, and she wanted to be a part of it. She wanted to make a difference, and she wanted to make the world
直接量化后的 INT8 模型生成為：She wanted to learn new things. She wanted to do things she had never done before. She wanted to do things she had never done before. She wanted to
INT8 量化模型結(jié)巴了！

輔助式文本生成帶來(lái)了這種可能性，不需要“頭禿”去做 prompt 壓縮，也不需要小心翼翼地對(duì)模型進(jìn)行量化以使質(zhì)量損失控制在可接受的范圍內(nèi)，它通過(guò)引入一個(gè)“起草 + 審閱”的新范式，引入了一個(gè)無(wú)損的模型加速范式。其原理如下圖所示：


假設(shè)要對(duì) BLOOM-176B 模型進(jìn)行加速，在輔助式文本生成的范式中，我們把它叫做 target generator（藍(lán)色框）。我們給它配一個(gè)小弟，就叫它 draft generator（黃色框），顧名思義，就是負(fù)責(zé)給它起草回答的。這樣，班子就搭起來(lái)了，工作模式就變成了：小弟負(fù)責(zé)起草，大哥負(fù)責(zé)審閱。這種范式的好處很直白：

  • 速度快。我們知道，目前的 LLM 幾乎全部都是 decoder-only 架構(gòu)的，其本質(zhì)是自回歸模型，也就是輸出需要一個(gè)字一個(gè)字地生成，這也是 LLM 生成很難加速的本質(zhì)原因，因?yàn)橛?step-wise dependency。輔助式文本生成將 step-wise dependency 乾坤大挪移給了開銷小的 draft model，而 target model 只需要負(fù)責(zé)驗(yàn)證的部分，而驗(yàn)證是一個(gè)典型的 CLM(Causal Language Modeling) 工作負(fù)載，一把頭就可以全部驗(yàn)證完。舉個(gè)實(shí)際的例子： 用 LLaMA-2-13B 在 A100 GPU 上生成 128 個(gè)詞元的時(shí)間所花的時(shí)間是對(duì)同樣長(zhǎng)度的序列做 CLM 前向所花時(shí)間的約 100 倍^{[1]}。因此可以想見，step-wise dependency 的轉(zhuǎn)移，釋放了速度的潛能。
  • 生成一致。因?yàn)樽罱K輸出結(jié)果是由 target model 審定的，在使用“精確匹配”的審定策略時(shí)，可以做到輸出一致，無(wú)需生成質(zhì)量的任何折衷。

完美切中我們的訴求。
我們以上圖為例，我們過(guò)一遍輔助式生成范式的工作流程：

  • 第一步：由 draft model 為提示 "The orange cat" 生成 K 個(gè)補(bǔ)全詞。本例中, K 為 3，draft model 生成了 3 個(gè)補(bǔ)全詞 ate the fish。
  • 第二步：將 The orange cat ate the fish 送給 target model 執(zhí)行前向 logit 計(jì)算。target model 算得第 2 個(gè)補(bǔ)全詞 the 的 logit 并非最大，而是 my 最大，因此拒絕了 draft model 生成的第 2 個(gè)詞及其之后的詞，也即其接受了 \frac{1}{3} 的補(bǔ)全詞，最終本輪的輸出補(bǔ)全詞為 ate my，詞數(shù)為 2。
  • 第三步： 確定是否滿足結(jié)束條件，如滿足結(jié)束生成；如不滿足，提示變成 "The orange cat ate my" 返回第一步，繼續(xù)生成。

在上述流程中：K 為草稿窗口，即每次打草稿生成多少個(gè)詞，也叫 look ahead window; \frac{1}{3} 為接受率，我們用 \alpha 表示；最終經(jīng) target model 審定后輸出的詞數(shù)為 \alpha K+1，我們叫它壓縮率（compression ratio）。

由此，我們可以得到一個(gè)輔助式文本生成的經(jīng)驗(yàn)加速公式：
acc\_ratio=\frac{l_t \times (\alpha K+1)}{Kl_d + l_t}=\frac{\alpha K+1}{K\frac{l_d}{l_t}+1}
其中 l_d 為 draft model 的每詞元延遲，l_t 為 target model 的每詞元延遲。
可以看到，想要獲得好的加速需要提高接受率 \alpha、降低 \frac{l_d}{l_t}，這兩者有時(shí)是正相關(guān)的，因此需要折衷。

What-If 分析

  • 設(shè) K=4，??=0.75, 如需獲得 2 倍加速，需有 \frac{l_d}{l_t} \le \frac{1}{4};
  • 設(shè) K=4，??=0.5, 如需獲得 2 倍加速，需有 \frac{l_d}{l_t} \le \frac{1}{8};

注意這里 K 可以不是固定值，而是根據(jù) \alpha 聯(lián)動(dòng)的動(dòng)態(tài)值；審定策略也未必一定是精確匹配，也可以是軟匹配。但這些戰(zhàn)術(shù)上的微調(diào)，都不影響范式層面的大結(jié)論。

范式確立以后，下面就是看基于這個(gè)范式有哪些方法了。

現(xiàn)有方法

vanilla

最直觀的方法就是使用一個(gè)小模型（SLM，Small Language Model）為 LLM 打草稿了，如下圖所示。該方法也是 Hugging Face transformers 默認(rèn)的方法。


我們做一個(gè)實(shí)驗(yàn)看一下效果如何，實(shí)驗(yàn)設(shè)置如下：

  • target model
    opt-66b (AutoTP 跑在 8 張 A100)
  • draft model
    opt-125m (單張 A100 或 單路 SPR CPU)
  • IO
    input tokens: 32
    new_tokens: 32
    data type: bf16
  • 其他設(shè)置
    K：動(dòng)態(tài)取值。初始值為1，如所有的 target model 接受了 draft model 所有的草稿，則 K 加 1；否則，K 減 1。

結(jié)果如下：


可以看到，在 Hugging Face transformers 下 opt-66b 的每詞元延遲是 43 ms，使用了 opt-125m 作為 draft model 并采用輔助式文本生成架構(gòu)后，加速比達(dá) 1.37。甚至可以做異構(gòu)輔助式文本生成，將 opt-125m 放到機(jī)頭的 CPU 上運(yùn)行，在我們的測(cè)試環(huán)境中，加速比達(dá) 1.5，還實(shí)現(xiàn)了整機(jī)硬件的充分利用。
在嘗試 vanilla 方法的過(guò)程中，我們還發(fā)現(xiàn)了一些有趣的事情：

  • 不同模型在相同的 K 下的 compression ratio 有較大的差異。下圖是 bloom-560m + bloom-176b 的 compression ratio，可以看到其比上圖的 opt-125m + opt-66b 低不少。我們的直覺是覺得這是由于 bloom 系列模型比 opt 系列模型詞表大很多的緣故（bloom 為 250880 vs opt 為 50272）。

  • dynamic K 比較有用，免去設(shè)超參的煩惱，拯救頭皮！下圖給出了 bloom-560m + bloom-176b 在固定 K 和 dynamic K 下的 compression ratio 表現(xiàn)，可以看出效果基本與 K=2 相當(dāng)。

  • 軟匹配對(duì)提高接受率有明顯的效果。對(duì)于使用 top-K 或者 top-P 這類帶有隨機(jī)性的生成策略的場(chǎng)景而言，可以不使用精確匹配，而使用軟匹配。

軟匹配策略
對(duì) draft model 生成的每個(gè)詞元，比較 draft model 計(jì)算的 logit q 以及 target model 計(jì)算的 logit p:
若 q < p，則接受該詞元
否則：

  • 若 1 - \frac{p}{q} > \beta，則接受該詞元
  • 否則，拒絕該詞元

看下軟匹配策略的效果，以 bloom-560m + bloom-176b 和 bloom-1b7 + bloom-176b 為例，如下圖，可見其顯著提高了接受率。

自投機(jī)解碼

起初，在選擇 draft model 時(shí)，一般選擇同家族中的小尺寸模型，這些模型是基于同一個(gè)訓(xùn)練集訓(xùn)練出來(lái)的，且共享同一個(gè)詞表，因此不難假設(shè)其與 target model 的對(duì)齊性較好。自投機(jī)方法的想法就是為何不直接在 target model 中選一些層組成小模型呢，反正 LLM 也就是一組同規(guī)格的 transformer block 的堆疊，跳過(guò)一些層并不會(huì)導(dǎo)致數(shù)據(jù)尺寸不兼容。由此，作者使用了貝葉斯優(yōu)化方法對(duì)每個(gè)注意力層和 MLP 層選擇跳過(guò)或不跳過(guò)，以最終決定抽取哪些層組成 draft model。這一整套路數(shù)復(fù)用的是之前 NAS 的套路，以在 l_d 和接受率之間求得最優(yōu)解。

以上是事實(shí)，以下是觀點(diǎn)。

這一套雖然很好看，但如果從工程的角度看，個(gè)人并不十分感冒。貝葉斯這一頓裁是基于某個(gè)校準(zhǔn)數(shù)據(jù)集的，因此裁剪出的模型不管是性能（如提示長(zhǎng)度、生成長(zhǎng)度等）還是生成質(zhì)量其實(shí)都是基于這個(gè)數(shù)據(jù)集的最優(yōu)，泛化性始終都是打問(wèn)號(hào)的。招子比較花，但是實(shí)用性有限。自投機(jī)的方向，個(gè)人還是覺得如果能不跳層，選前 N 層作為 draft model 是最好的，這樣：

  1. 可輕易地實(shí)現(xiàn)端云協(xié)同，既拆分，又能聯(lián)合在線學(xué)習(xí)。
  2. 前 N 層的 KV-cache 可復(fù)用，進(jìn)一步加速 target model。
  3. 在端云場(chǎng)景，原生支持隱私計(jì)算，可不將用戶數(shù)據(jù)傳給云端，只傳 N 層的輸出給云端。
    以上僅為個(gè)人觀點(diǎn)。

SpecInfer

基本架構(gòu)已經(jīng)搭成，下面就是如何進(jìn)一步加速了。從上文的估算公式可以看出，從算法上來(lái)看，提速的關(guān)鍵在于提高接受率 \alpha。那么如何提高接受率呢？關(guān)鍵在于審定機(jī)制，上述方法的做法是：draft model 在每個(gè)時(shí)間步都只起草一個(gè)詞元供 target model 審定。這種方法對(duì) draft model 和 target model 的 top-1 對(duì)齊能力有很高的要求。我們知道，一旦當(dāng)前時(shí)間步 top-1 不對(duì)齊，后面所有的預(yù)測(cè)都作廢了，因此極大限制了接受率的提高。而解脫之道就在其中！我們知道：1）大小模型的 top-1 對(duì)齊肯定是不能指望太多的，否則小模型的能力其實(shí)就跟大的差不多了，不需要大的了；2）但我們也知道，這倆能力相差也不會(huì)太大，大模型的輸出在小模型的 top-k 中的概率肯定不會(huì)低。至此，SpecInfer 就應(yīng)運(yùn)而出了！

如上圖所示，SpecInfer 采用 beam search 的方法（expansion based）或 ensemble 方法（merge based） 生成草稿，其目的就是為每個(gè)時(shí)間步生成多個(gè)候選詞，以提高接受率。他還做了一件事情，就是把生成的多份草稿整理成一棵樹的形式，以充分利用它們之間重疊的詞，減少計(jì)算量，加速 l_t。思路很簡(jiǎn)單，也很合理。
以下是 opt-125m + opt-13b 在固定 K=5 設(shè)定下的實(shí)驗(yàn)結(jié)果?？梢钥吹皆黾?beam width 到 8， compression ratio 有明顯提高，再 ensemble 5 個(gè) 模型，又有進(jìn)一步提高，驗(yàn)證了上述說(shuō)法。

總結(jié)一下，SpecInfer 的架構(gòu)如下：


LADE

捫心自問(wèn)一下：draft model 就必須就是個(gè) LLM 的小兄弟嗎？好像也不見得，只要能輸出接受率較高，管它是不是個(gè)有參的 LM。話說(shuō)到這兒就引入了 LADE（Look Ahead DEcoding）。LADE 做得比較絕，如下圖所示，它用了一個(gè)無(wú)參的、在線學(xué)習(xí)的 n-gram LM 作為 draft model。

其流程如下：

  1. 用提示中的最后一個(gè)詞元為鍵去從 n-gram 模型找它的 n-1 個(gè)補(bǔ)全詞。有可能找到多個(gè)，此時(shí)最多用 G 個(gè)；也有可能沒有，此時(shí)隨機(jī)補(bǔ)全
  2. LLM 審定
  3. 將 LLM 審定通過(guò)的新 n-gram 更新入 n-gram 模型
  4. 如生成完成，退出；否則，返回第 1 步

下圖以 2-gram 為例展示了該流程：

我們做個(gè)簡(jiǎn)單的實(shí)驗(yàn)，7-gram + TinyLlama-1.1B-Chat-v0.3，且 K 設(shè)為 20，輸入長(zhǎng)度為 32， 輸出長(zhǎng)度為 256 時(shí)，性能如下：

可以看到在 A100 單卡下，性能有 1.58 倍的提升。

總結(jié)

以上介紹了輔助式文本生成范式，該范式為“免費(fèi)”降低文本生成每詞元成本提供了可能性。同時(shí)還介紹了該范式下的四種典型算法。個(gè)人認(rèn)為這一范式是深具潛力的，原因之一正如開篇所講，其提供了無(wú)損的透明加速范式，使得用戶可以專注于應(yīng)用本身，而不是花很多時(shí)間、金錢和資源用于壓縮模型，權(quán)衡收益。其二是其范式構(gòu)型的靈活性，可以適應(yīng)多種部署設(shè)計(jì)：全部部署在加速器上？可以! 異構(gòu)部署，draft model 部署在 CPU 上， target model 部署在加速器上？ 也行！端云協(xié)同，draft model 部署在端側(cè)， target model 部署在云側(cè)？ 未嘗不可！想象的畫卷正在展開......

參考文獻(xiàn)

  1. Fast Inference from Transformers via Speculative Decoding
  2. Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding
  3. SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification
  4. Break the Sequential Dependency of LLM Inference Using Lookahead Decoding
最后編輯于 ：
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點(diǎn)，簡(jiǎn)書系信息發(fā)布平臺(tái)，僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容