ICLR 2022｜唯快不破！面向極限壓縮的全二值化BiBERT

魚(yú)羊 2022-02-19 13:49:48 來(lái)源：量子位

作者 秦浩桐

量子位 轉(zhuǎn)載 | 公眾號(hào) QbitAI

近年來(lái)，預(yù)訓(xùn)練語(yǔ)言模型在自然語(yǔ)言處理上表現(xiàn)出色，但其龐大的參數(shù)量阻礙了它在真實(shí)世界的硬件設(shè)備上的部署。

近日，機(jī)器學(xué)習(xí)頂會(huì)ICLR 2022接收論文結(jié)果已經(jīng)正式公布，至少有9項(xiàng)工作展示了神經(jīng)網(wǎng)絡(luò)量化方向的相關(guān)進(jìn)展。

本文將介紹首個(gè)用于自然語(yǔ)言任務(wù)的全二值量化BERT模型——BiBERT，具有高達(dá)56.3倍和31.2倍的FLOPs和模型尺寸的節(jié)省。

這項(xiàng)研究工作由北京航空航天大學(xué)劉祥龍教授團(tuán)隊(duì)、南洋理工大學(xué)和百度公司共同完成。

預(yù)訓(xùn)練語(yǔ)言模型在自然語(yǔ)言處理上表現(xiàn)出色，但其龐大的參數(shù)量阻礙了它在真實(shí)世界的硬件設(shè)備上的部署。

現(xiàn)有的模型壓縮方法包括參數(shù)量化、蒸餾、剪枝、參數(shù)共享等等。

其中，參數(shù)量化方法高效地通過(guò)將浮點(diǎn)參數(shù)轉(zhuǎn)換為定點(diǎn)數(shù)表示，使模型變得緊湊。

研究者們提出了許多方案例如Q-BERT[1]、Q8BERT[2]、GOBO[3]等，但量化模型仍舊面臨嚴(yán)重的表達(dá)能力有限和優(yōu)化困難的問(wèn)題。

幸運(yùn)的是，知識(shí)蒸餾作為一種慣用的輔助優(yōu)化的手段，令量化模型模仿全精度教師模型的特征表達(dá)，從而較好地解決精度損失問(wèn)題。

在本文中，來(lái)自北航、NTU、百度的研究人員提出了BiBERT，將權(quán)重、激活和嵌入均量化到1比特（而不僅僅是將權(quán)重量化到1比特，而激活維持在4比特或更高）。

這樣能使模型在推理時(shí)使用逐位運(yùn)算操作，大大加快了模型部署到真實(shí)硬件時(shí)的推理速度。

我們研究了BERT模型在二值化過(guò)程中的性能損失，作者在信息理論的基礎(chǔ)上引入了一個(gè)高效的Bi-Attention（二值注意力）機(jī)制，解決前向傳播中二值化后的注意力機(jī)制的信息退化問(wèn)題；提出方向匹配蒸餾（Direction-Matching Distillation）方法，解決后向傳播中蒸餾的優(yōu)化方向不匹配問(wèn)題。

△圖 1 BiBERT的整體框架圖

BiBERT首次證明了BERT模型全二值化的可行性，在GLUE數(shù)據(jù)集上的準(zhǔn)確性極大地超越了現(xiàn)有的BERT模型二值化算法，甚至超過(guò)了更高比特表示的模型。

在模型計(jì)算量和體積上，BiBERT理論上能夠帶來(lái)56.3倍和31.2倍的FLOPs和模型尺寸的減少。

方法

Bi-Attention：二值化注意力機(jī)制

我們的研究表明，在BERT模型的注意力機(jī)制中，softmax函數(shù)得到的歸一化注意力權(quán)重被視為遵循一個(gè)概率分布，而直接對(duì)其進(jìn)行二值化會(huì)導(dǎo)致完全的信息喪失，其信息熵退化為0（見(jiàn)圖2）。

△圖 2 直接對(duì)softmax函數(shù)應(yīng)用二值化導(dǎo)致完全的信息喪失

一個(gè)緩解這種信息退化的常用措施是，在應(yīng)用sign函數(shù)之前對(duì)輸入張量的分布進(jìn)行移位，可以表示為：

其中，移位參數(shù)也被認(rèn)為是二值化的閾值，希望能使二值化后的熵達(dá)到最大。

我們注意到，softmax函數(shù)是保序的，這意味著存在一個(gè)固定的閾值使二值化表示的信息熵最大化。

受到Hard Attention的啟發(fā)[4]，作者用bool函數(shù)來(lái)二值化注意力權(quán)重A：

通過(guò)應(yīng)用bool函數(shù)，注意權(quán)重中值較低的元素被二值化為0，因此得到的熵值最大的注意權(quán)重可以過(guò)濾出關(guān)鍵部分的元素。

最終二值注意力機(jī)制可以被表示為：

其中，BV是通過(guò)sign函數(shù)二值化得到的value值，BA是二值化注意力權(quán)重，是一個(gè)精心設(shè)計(jì)的Bitwise-Affine矩陣乘法?（BAMM）運(yùn)算器，由和位移組成，用于對(duì)齊訓(xùn)練和推理表征并進(jìn)行有效的位計(jì)算。

DMD: 方向匹配蒸餾

作者發(fā)現(xiàn)，由于注意力權(quán)重是兩個(gè)二值化的激活直接相乘而得。

因此，處于決策邊緣的值很容易被二值化到相反一側(cè)，從而直接優(yōu)化注意力權(quán)重常常在訓(xùn)練過(guò)程中發(fā)生優(yōu)化方向失配問(wèn)題。（見(jiàn)圖3）

△圖 3 優(yōu)化方向失配

因此，作者設(shè)計(jì)了新的蒸餾方案，即針對(duì)上游的Query、Key和Value矩陣，構(gòu)建相似性矩陣進(jìn)行對(duì)激活的蒸餾：

其中，||·||表示L2正則化。之前的研究工作表明，以這種方式構(gòu)建的矩陣被認(rèn)為能夠反映網(wǎng)絡(luò)對(duì)于特定模式的語(yǔ)義理解，并無(wú)視尺度和數(shù)值大小影響，能夠更加穩(wěn)定地表示特征之間的內(nèi)生相關(guān)性，更適合二值和全精度網(wǎng)絡(luò)之間的知識(shí)傳遞。

因此，蒸餾損失可以表示為對(duì)隱藏層、預(yù)測(cè)結(jié)果和上述激活相似性矩陣的損失之和：

其中L表示transformer的層數(shù)，

。

實(shí)驗(yàn)

作者的實(shí)驗(yàn)證明了所提出的BiBERT能夠出色地解決二值化BERT模型在GLUE基準(zhǔn)數(shù)據(jù)集的部分任務(wù)上精度崩潰的問(wèn)題，使模型能夠穩(wěn)定優(yōu)化。

表1表明所提出的Bi-Attention和DMD均可以顯著提升模型在二值化后的表現(xiàn)。

△表 1 消融實(shí)驗(yàn)

表2和表3中，作者展示了BiBERT優(yōu)于其他BERT二值化方法，甚至優(yōu)于更高比特的量化方案：

△表 2 基于BERT的二值化方法對(duì)比（無(wú)數(shù)據(jù)增強(qiáng)）

△表 3 基于BERT的二值化方法對(duì)比（有數(shù)據(jù)增強(qiáng)）

其中，50%表示要求二值化后有一半的注意力權(quán)重為0，且表中無(wú)特殊說(shuō)明均采用12層的BERT模型進(jìn)行量化。

此外，作者測(cè)量了在訓(xùn)練過(guò)程中的信息熵（見(jiàn)圖4），作者提出的方法有效地恢復(fù)了注意力機(jī)制中完全損失的信息熵。

△圖 4 訓(xùn)練過(guò)程中的信息熵

同時(shí)，作者繪制了訓(xùn)練時(shí)的loss下降曲線和準(zhǔn)確率，BiBERT相比于基線明顯更快收斂、準(zhǔn)確性更高。

△圖 5 訓(xùn)練時(shí)的Loss和準(zhǔn)確率曲線

總結(jié)

作者提出的BiBERT作為第一個(gè)BERT模型的全二值化方法，為之后研究BERT二值化建立了理論基礎(chǔ)，并分析了其性能下降的原因，針對(duì)性地提出了Bi-Attention和DMD方法，有效提高模型的性能表現(xiàn)。

BiBERT超過(guò)了現(xiàn)有的BERT模型二值化方法，甚至優(yōu)于采用更多比特的量化方案，理論上BiBERT能夠帶來(lái)56.3倍的FLOPs減少和31.2倍的模型存儲(chǔ)節(jié)省。

希望該的工作能夠?yàn)槲磥?lái)的研究打下堅(jiān)實(shí)的基礎(chǔ)。

BiBERT即將基于百度飛槳開(kāi)源深度學(xué)習(xí)模型壓縮工具PaddleSlim開(kāi)源，盡情期待。

PaddleSlim：https://github.com/PaddlePaddle/PaddleSlim

傳送門(mén)
會(huì)議論文：https://openreview.net/forum?id=5xEgrl_5FAJ

Reference
[1] Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W. Mahoney, and Kurt Keutzer. Q-BERT: hessian based ultra low precision quantization of BERT. In AAAI, 2020.
[2] Ofir Zafrir, Guy Boudoukh, Peter Izsak, and Moshe Wasserblat. Q8BERT: quantized 8bit BERT. In NeurIPS, 2019.
[3] Ali Hadi Zadeh, Isak Edo, Omar Mohamed Awad, and Andreas Moshovos. GOBO: quantizing attention-based NLP models for low latency and energy efficient inference. In MICRO, 2020.
[4] Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio. Show, attend and tell: Neural image caption generation with visual attention. In ICML, 2015.

*本文系量子位獲授權(quán)刊載，觀點(diǎn)僅為作者所有

— 完 —