模型量化: 訓(xùn)練時(shí)噪音帶來的極致壓縮

我們在模型量化這片文章中，介紹了在訓(xùn)練時(shí)是如何做量化的。

https://zhuanlan.zhihu.com/p/353597270

文獻(xiàn)[1]基于這個(gè)工作，進(jìn)一步引入了其他的壓縮方法，并用噪音的方式進(jìn)行了改進(jìn)，成功的將Roberta壓縮到了14M，將EfficientNetB3壓縮到了3.3M，分別達(dá)到了34x和14x的壓縮倍率。

量化噪音

這個(gè)概念其實(shí)挺容易理解，就是在訓(xùn)練時(shí)做量化，會(huì)帶來權(quán)重的變化，這個(gè)變化稱之為噪音。在模型量化: 只有整數(shù)計(jì)算的高效推理，模型在訓(xùn)練時(shí)就把權(quán)重和activation做了量化，并基于量化后的參數(shù)進(jìn)行fine-tune。

在文獻(xiàn)[1]中，提出了部分量化的做法，即每次只把一部分參數(shù)和激活值做量化，其他的值不變，按照原來的方式進(jìn)行梯度計(jì)算。這樣的方法比訓(xùn)練時(shí)就全部做量化要好一些，見下圖：

但在inference的時(shí)候，是全部都要做量化的。

部分量化是論文[1]的核心點(diǎn)，剩下的就是將訓(xùn)練時(shí)部分量化與其他的量化技術(shù)結(jié)合。

量化壓縮

在神經(jīng)網(wǎng)絡(luò)中，需要量化的參數(shù)一般都是矩陣W，這里假設(shè)矩陣的size是n * p。常見的壓縮方法是對(duì)矩陣進(jìn)行切分，切分成mxq個(gè)小矩陣。然后小矩陣上做codebook查找，即找到K個(gè)vector做成codebook，然后m * q個(gè)小矩陣，每個(gè)都表示成在codebook上查到的索引，從而達(dá)到壓縮的目的。

基于上面的介紹，量化就分成了兩種：

• Fixed-point scalar Quantization，即對(duì)值進(jìn)行量化。
• Product quantization，即對(duì)矩陣進(jìn)行量化

值量化

其中，對(duì)值進(jìn)行量化就是給當(dāng)前值找到最接近的量化后的值：

其中，

取決于量化后的bit數(shù)目，假設(shè)量化后是N位bit，那么壓縮率就是32/N。如果N=4，那么就是int4 quantization, 如果N=8，則是int8 Quantization

矩陣量化

對(duì)于矩陣量化(PQ)來說，傳統(tǒng)的方法是：

• 將大矩陣的每一列分成m個(gè)子vector，這樣就有m * p個(gè)子vector。
• 每個(gè)子vector表達(dá)成codebook中離它最近的vector。
• codebook使用K-means方法學(xué)習(xí)得到。
• 所以矩陣被表示成了索引矩陣，如果codebook中一共只有256個(gè)vector，那么索引矩陣中的每個(gè)值都是8bit的，就有了量化的效果，注意，此時(shí)codebook中的vector還可以是float值。

所以目標(biāo)函數(shù)如下所示，即讓原矩陣和量化后矩陣歐氏距離最近。

但這種矩陣量化的方式會(huì)導(dǎo)致偏差積累，從而導(dǎo)致層次比較深的時(shí)候會(huì)發(fā)生比較大的漂移。

因此，一個(gè)改進(jìn)的方法就是從底層到高層迭代的去做矩陣量化，稱之為迭代式矩陣量化(Iterative Product Quantization, iPQ)。迭代到某一層的時(shí)候，會(huì)對(duì)其上層做fine-tune。

矩陣量化和值量化是可以結(jié)合起來達(dá)到更高的壓縮率。它們還可以和上一部分介紹的噪音也進(jìn)行疊加。在訓(xùn)練時(shí)，矩陣量化可以一個(gè)epoch重新運(yùn)行k-means算法來更新。

實(shí)驗(yàn)結(jié)果

各種量化方法的對(duì)比如下圖，其中QAT方法是不加噪音的fine-tune?？梢钥吹郊恿嗽胍艉髸?huì)在更高的壓縮比上得到不錯(cuò)的效果。

另外，不同的壓縮比率和效果與其他模型的對(duì)比：

除此之外，還可以和weight sharing與pruning聯(lián)合使用。得到：

參考文獻(xiàn)

[1]. Fan, A., Stock, P., Graham, B., Grave, E., Gribonval, R., Jégou, H. and Joulin, A., 2020. Training with quantization noise for extreme model compression. arXiv e-prints, pp.arXiv-2004.

如果覺得有用，就請(qǐng)分享到朋友圈吧！