ACL 2022丨香港大學(xué)&華為諾亞方舟新工作：生成式預(yù)訓(xùn)練語(yǔ)言模型的量化壓縮

點(diǎn)擊上方“視學(xué)算法”，選擇加"星標(biāo)"或“置頂”

重磅干貨，第一時(shí)間送達(dá)

作者丨陶超凡

生成式預(yù)訓(xùn)練語(yǔ)言模型?(pre-trained language models, PLM) 的規(guī)模不斷擴(kuò)大，極大地增加了對(duì)模型壓縮的需求。盡管有各種壓縮 BERT 或其變體的方法，壓縮生成式 PLM （例如GPT）的嘗試很少，潛在的困難仍不清楚。

在本文中，我們通過量化壓縮生成 PLM。我們發(fā)現(xiàn)以前的量化方法在生成任務(wù)上失敗，是由于模型容量減少導(dǎo)致的詞嵌入同質(zhì)化，以及不同模塊的權(quán)重分布差異。

相應(yīng)地，我們提出了一種令牌級(jí)的對(duì)比蒸餾損失函數(shù)來學(xué)習(xí)可區(qū)分的詞嵌入， 以及模塊級(jí)動(dòng)態(tài)適應(yīng)的量化器，以使量化器適應(yīng)不同的模塊。

各種任務(wù)的實(shí)證結(jié)果表明我們提出的方法優(yōu)于最先進(jìn)的壓縮生成式 PLM 方法的優(yōu)勢(shì)明顯。憑借與全精度模型相當(dāng)的性能，我們?cè)?GPT-2 和 BART 上分別實(shí)現(xiàn)了 14.4倍 和 13.4倍的壓縮率。

我們把提出的量化GPT模型、BART模型分別取名「QuantGPT」與「QuantBART」。

基于 Transformer 的生成式預(yù)訓(xùn)練語(yǔ)言模型表現(xiàn)出強(qiáng)大的多任務(wù)和小樣本學(xué)習(xí)能力，并在各種任務(wù)上取得顯著成績(jī)。但是，由于大量參數(shù)以及逐個(gè)令牌（token）的生成過程，它們?cè)谟?jì)算和內(nèi)存方面通常都很昂貴。目前學(xué)術(shù)界提出了許多方法來壓縮 PLM，但主要集中在理解任務(wù)上，比如使用 BERT 進(jìn)行句子分類。

圖1：不同量化方法的2bit GPT-2在PTB數(shù)據(jù)集的語(yǔ)言建模表現(xiàn)

在本文中，我們首先從量化角度來探索壓縮生成式PLM。我們發(fā)現(xiàn)，將以前為 BERT 或計(jì)算機(jī)視覺任務(wù)設(shè)計(jì)的量化方法直接應(yīng)用于生成式PLM 會(huì)導(dǎo)致性能不佳。如圖1所示，隨著權(quán)重比特?cái)?shù)的減小，模型的性能急劇下降。這促使我們?nèi)パ芯苛炕墒絇LM的難點(diǎn)，并提出相應(yīng)的解決策略。

2.1難點(diǎn)1：詞嵌入同質(zhì)化 (Homogeneous Word Embedding)

圖2：在 PTB 數(shù)據(jù)集上訓(xùn)練的全精度和不同 2bit的最高頻率的 500 個(gè)詞嵌入的 T-SNE 可視化。不同量化方法的詞嵌入顯示出不同程度的同質(zhì)化。

可以看出，全精度模型的詞嵌入是分散可區(qū)分的，而之前的量化方法 PACT、LSQ 和 LAQ 學(xué)習(xí)的是同質(zhì)的詞嵌入，是聚集的且不易區(qū)分，特別是對(duì)于 PACT 和 LSQ?？紤]BERT做句子分類問題，量化后的BERT只需要讓句子整體的表征是準(zhǔn)確的。但是對(duì)于GPT這種生成問題，GPT 以從左到右的順序計(jì)算每個(gè)令牌，并且前一個(gè)令牌中產(chǎn)生的量化誤差將傳遞給未來的令牌。所以量化GPT模型需要盡可能對(duì)每個(gè)量化后的詞嵌入都學(xué)習(xí)準(zhǔn)確。文章附錄提供了不同量化方法的生成Demo，同質(zhì)化的詞嵌入會(huì)導(dǎo)致模型生成重復(fù)且無邏輯的句子。

2.2難點(diǎn)2：不同模塊的權(quán)重分布差異 (Varied Distribution of Weights)

這里簡(jiǎn)單補(bǔ)充一下量化的背景知識(shí)。對(duì)一個(gè)模型的權(quán)重w,我們要學(xué)習(xí)一個(gè)量化器Q( )把全精度權(quán)重w量化為低比特的權(quán)重。其中α是量化器的一個(gè)參數(shù)，用于把權(quán)重的數(shù)值范圍限定在[-α，α]。α的選取對(duì)量化器的效果是重要的。過小的α會(huì)導(dǎo)致很大一部分的權(quán)重被估計(jì)成-α或者α，損失了這些權(quán)重的全精度信息；而過大的α會(huì)導(dǎo)致過多的全精度數(shù)值被考慮到量化器的有效范圍，導(dǎo)致量化器的量化間隔(quantization level)較大，量化的結(jié)果不準(zhǔn)。

圖3：在一個(gè)預(yù)訓(xùn)練的GPT中，同一個(gè)transformer block里兩個(gè)不同模塊的權(quán)重分布。豎線是PACT和我們方法學(xué)出的量化器的clipping factor α

從圖3可以看出，不同模塊的權(quán)重分布差異很大，這對(duì)α的估計(jì)造成了較大的挑戰(zhàn)。

3.1令牌級(jí)對(duì)比學(xué)習(xí)

我們采取了全精度模型作為老師，量化模型作為學(xué)生的蒸餾學(xué)習(xí)框架。為了讓學(xué)生模型可以在令牌粒度上學(xué)習(xí)到老師模型的信息。我們?cè)O(shè)計(jì)了令牌級(jí)對(duì)比學(xué)習(xí)：在一個(gè)句子里，同一個(gè)令牌在老師模型和學(xué)生模型中產(chǎn)生的表征應(yīng)該被拉近，視為一對(duì)正樣本；而同一個(gè)句子里不同令牌的表征視為負(fù)樣本，應(yīng)該被拉遠(yuǎn)。為了讓學(xué)習(xí)模型產(chǎn)生的令牌表征平滑過渡，我們?cè)O(shè)計(jì)了一個(gè)memory bank存儲(chǔ)詞表里每個(gè)令牌的表征，用于對(duì)比學(xué)習(xí)中負(fù)樣本的計(jì)算。

具體地，對(duì)于一個(gè)長(zhǎng)度為n的句子，令牌級(jí)的對(duì)比學(xué)習(xí)為：

其中計(jì)算了相似度，表示老師模型產(chǎn)生的令牌特征，是學(xué)生模型產(chǎn)生的令牌特征的滑動(dòng)平均版本（最后一個(gè)decoder的輸出再經(jīng)過一個(gè)線性映射）。的更新方式為

3.2 模塊級(jí)動(dòng)態(tài)適應(yīng)的量化器

我們根據(jù)每個(gè)模塊權(quán)重的統(tǒng)計(jì)數(shù)據(jù)提出了一種簡(jiǎn)單而有效的動(dòng)態(tài)縮放。具體來說，我們不是直接對(duì)所有模塊初始化一個(gè)相同的clipping factor α，而是轉(zhuǎn)而學(xué)習(xí)一個(gè)新的縮放因子 γ，它與平均權(quán)重大小相乘獲得clipping factorγ被初始化為1。這不僅簡(jiǎn)化了初始化，而且確保了初始的α不會(huì)遠(yuǎn)離全精度權(quán)重，無論不同模塊的權(quán)重分布的差異有多大。我們進(jìn)而推出了這個(gè)新的縮放因子γ的更新公式。

這個(gè)更新公式不僅考慮到了在[-α，α]范圍內(nèi)的權(quán)重對(duì)的影響，也考慮在[-α，α]范圍外權(quán)重的影響。結(jié)果證明這樣學(xué)習(xí)clipping factor,比之前工作中更新時(shí)只考慮[-α，α]范圍外權(quán)重的效果更好。

首先，我們?cè)贕PT-2上分別做了PTB，wikitext-2,wikitext-103數(shù)據(jù)集的語(yǔ)言建模任務(wù)和Persona-Chat數(shù)據(jù)集的對(duì)話理解任務(wù)。我們的方法QuantGPT在生成文本預(yù)訓(xùn)練模型的量化任務(wù)上效果提升明顯。其中W-E-A表示權(quán)重-嵌入層-激活值的量化比特?cái)?shù)。

我們也和最新的GPT-2壓縮方案對(duì)比，包括KnGPT2（基于張量分解），DistilGPT2, LightPAFF（基于蒸餾），我們的方法QuantGPT在模型大小和效果上優(yōu)勢(shì)明顯。

除了考慮decoder-only的生成模型GPT，我們也考慮了encoder-decoder的生成模型BART。下表中看出，分別在8/4/2bit的實(shí)驗(yàn)設(shè)置下，我們模型QuantBART在摘要任務(wù)上提升明顯。

本文研究了預(yù)訓(xùn)練語(yǔ)言生成模型的低比特量化。我們發(fā)現(xiàn)量化生成預(yù)訓(xùn)練語(yǔ)言生成模型的困難在于同質(zhì)化的詞嵌入和不同模塊的權(quán)重分布差異很大。為了緩解這兩個(gè)問題，我們提出了令牌級(jí)對(duì)比學(xué)習(xí)來學(xué)習(xí)可區(qū)分的詞嵌入，以及一個(gè)模塊粒度動(dòng)態(tài)縮放的量化器，以適應(yīng)不同模塊的權(quán)重分布。語(yǔ)言建模、話語(yǔ)理解和摘要的實(shí)驗(yàn)證明了我們提出的方法的有效性。

我們希望這個(gè)工作能夠?qū)ι扇蝿?wù)的壓縮有啟發(fā)作用，也可以嘗試把這個(gè)工作提出的量化器用在判別任務(wù)的壓縮里。

點(diǎn)個(gè)在看 paper不斷！