Transformer 在自然語言處理和視覺任務中取得了令人矚目的成果，然而預訓練大模型的推理代價是備受關心的問題，華為諾亞方舟實驗室的研究者們聯合高校提出針對視覺和 NLP 預訓練大模型的後訓練量化方法。在精度不掉的情況下，比 SOTA 訓練感知方法提速 100 倍以上；量化網絡性能也逼近訓練感知量化方法。大型預訓練模型在計算機視覺和自然語言處理中展現了巨大的潛力，但是模型大、參數多的問題也給它們的商業化落地帶來了很大挑戰。模型壓縮技術是當前的研究熱點，模型量化是其中的一個重要分支。當下預訓練模型的量化為了保證性能，大多採用量化感知訓練（Quantization-aware Training, QAT）。而模型後量化（Post-training Quantization, PTQ）作為另一類常用量化方法，在預訓練大模型領域卻鮮有探索。諾亞方舟實驗室的研究者從以下四個方面對 QAT 與 PTQ 進行了詳細對比：訓練時間：QAT 由於模擬量化算子等操作，訓練耗時遠遠超出全精度訓練（FP），而 PTQ 僅僅需要幾十分鐘，大大縮短量化流程；
顯存開銷：QAT 顯存消耗大於全精度訓練（FP），使得在顯存有限的設備上難以進行量化訓練。而 PTQ 通過逐層回歸訓練，無需載入整個模型到顯存中，從而減小顯存開銷；
數據依賴：QAT 需要獲取整個訓練數據集，而 PTQ 只需要隨機採樣少量校準數據，通常 1K～4K 張 / 條圖像或者標註即可；
性能：鑑於 QAT 在整個訓練集上充分訓練，其性能在不同的量化 bit 上均領先 PTQ。因此性能是 PTQ 的主要瓶頸。
基於以上觀測，研究者提出了針對視覺和 NLP 任務的 PTQ 方法，在保持其訓練時間、顯存開銷、數據依賴上優勢的同時，大大改善其性能，使其逼近量化感知訓練的精度。具體而言，僅使用 1% 的訓練數據和 1/150 的訓練時間，即可達到 SOTA 量化方法的精度。《Post-Training Quantization for Vision Transformer》
論文鏈接：https://arxiv.org/pdf/2106.14156.pdf下圖為視覺 Transformer 後訓練量化算法框架：
自注意力層是 Transfomer 結構中十分重要的部分，也是 Transformer 與傳統卷積神經網絡不同的地方。對於自注意力層，研究者發現量化使得注意力特徵的相對順序變得不同，會帶來很大的性能下降。因此，他們在量化過程中引入了一個排序損失：
其中表示成對的排序損失函數，表示權衡係數。然後，研究者將排序損失函數與相似度損失函數相結合，得到了最終的優化目標函數：
論文當中採用了比爾森相關係數作為特徵相似度的度量，研究者認為皮爾森相關係數減去了均值，所以對特徵的分布表示更加地敏感。為了進一步減少量化帶來的誤差，他們在優化量化步長過程中採用了量化誤差補償的方法，以減小量化誤差對之後的網絡層帶來影響。因此對每個網絡層的輸出都進行了量化誤差補償。
在實現過程中，誤差的期望值可以通過校驗數據集來計算，然後在網絡層的 bias 參數中去修正。不同的 transformer 網絡層有不同的數據分布，因為有不同的量化「敏感度」。研究者提出了混合精度量化，對於更加「敏感」的網絡層分配更多的比特寬度。在論文中，研究者提出使用 MSA 模塊中注意力層特徵和 MLP 中輸出特徵矩陣的核範數來作為度量網絡層「敏感度」的方法。與 HAWQ-V2 中的方法類似，他們使用了一種帕累托最優的方式來決定網絡層的量化比特。該方法的主要思想是對每個候選比特組合進行排序，具體的計算方式如下所示：
給定一個目標模型大小，會對所有的候選比特組合進行排序，並尋找值最小的候選比特組合作為最終的混合比特量化方案。研究者首先在圖像分類任務上對後訓練量化算法進行了驗證。從下表可以看出，在 ViT（DeiT）經典 transformer 模型上，論文的量化算法均優於之前的卷積神經網絡量化算法【1】【2】。例如，在 ImageNet 數據集上，量化 Deit-B 模型也取得了 81.29% 的 Top-1 準確率。
研究者還將後訓練量化算法應用於目標檢測任務中，其中在 COCO2017 數據集上，對 DETR 進行量化，8bit 模型的性能可以達到 41.7 mAP，接近全精度模型的性能。
《Towards Efficient Post-training Quantization of Pre-trained Language Models》
論文鏈接：https://arxiv.org/pdf/2109.15082.pdf
由於 Transformer-based 的預訓練語言模型通常含有多個線性層耦合在一起，如果採用現有的逐層重構誤差優化的方法【3】，作者發現很容易陷入局部最優解。為了考慮多個線性層內部的交互，如上圖所示，研究者把預訓練語言模型切分成多個模塊，每個模塊含有多個 Transformer 層。因此該方法聚焦於逐個重構模塊化的量化誤差，即最小化全精度網絡模塊（教師模型）的輸出與量化後模型模塊（學生網絡）的輸出之間的平方損失：
與逐個模塊化重構量化誤差不同，後量化還可以並行化訓練。研究者把每個切分後的模塊可以放在不同的 GPU 上，在不同模塊之間設置輸入緩衝池（input queue）
來收集上一個模塊的輸出，同時為下一個模塊的輸入做準備。不同模塊可以通過重置抽樣從輸入池獲取輸入樣本來進行本地訓練，無需等待其前繼模塊。因此，該設計可以使並行訓練，並且實現接近理論加速比。另外一個與逐模塊訓練不同的點在於，在並行知識蒸餾訓練的初期，下一個模塊獲得的輸入是從上一個未經過充分訓練的模塊中獲得。因此，未充分訓練的模塊的輸出可能依舊含有較大的量化誤差，並且該誤差會逐層傳播，影響後續模塊訓練。為了解決該問題，研究者受教師糾正（teacher forcing) 在訓練循環網絡中的啟發，將第 n 個全精度模塊的輸出導入為第 (n+1) 個量化模塊的輸入，從而中斷在後續模塊的量化誤差傳播。然而，這種跨模塊輸入打破了與量化模型自身前繼模塊的聯繫，造成訓練和推理前向不一致。為了實現平穩過渡，他們採用了如下的凸組合：
並對連接係數
隨着迭代次數 t 進行線性縮減。研究者首先在 MNLI 數據集上進行驗證。由下表可以發現，對比逐層後量化訓練（REM）算法，提出的逐模塊量化誤差重構 (MREM-S)大大提升了後量化的準確率；同時，MREM-S 性能也可以接近量化感知訓練（QAT）的方法，對於 BERT-base 和 BERT-large 在 W4A8 的設定下僅僅比 QAT 低了 1.1% 和 0.8%，訓練時間、顯存開銷和數據消耗也有了減小。當結合併行知識蒸餾時（MREM-P），後量化訓練時間可以進一步縮短 4 倍，而且沒有明顯性能損失。例如，MREM-P 僅耗時 6 分鐘，占用 3.7GB 即可完成 BERT-base 上 2 比特權重的後量化訓練。
研究者同時在 GLUE 上與現有的其他算法進行了對比。如下表所示，本文的方法 (MREM-S 和 MREM-P) 比後量化方法 GOBO【4】取得更好的效果，甚至在多個任務上接近量化感知訓練方法 TernaryBERT。
【1】Di Wu, Qi Tang, Yongle Zhao, Ming Zhang, Ying Fu, and Debing Zhang. Easyquant: Posttrainingquantization via scale optimization. arXiv preprint arXiv:2006.16669, 2020.【2】Peisong Wang, Qiang Chen, Xiangyu He, and Jian Cheng. Towards accurate post-trainingnetwork quantization via bit-split and stitching. In International Conference on MachineLearning, pages 9847–9856. PMLR, 2020.【3】I. Hubara, Y. Nahshan, Y. Hanani, R. Banner, and D. Soudry, 「Improving post training neural quantization: Layer-wise calibration and integer programming,」 in Proceedings of the International Conference on Machine Learning, 2021.【4】A. H. Zadeh, I. Edo, O. M. Awad, and A. Moshovos, 「Gobo: Quantizing attention-based nlp models for low latency and energy efficient inference,」 Preprint arXiv:2005.03842, 2020.
©THE END
轉載請聯繫本公眾號獲得授權
投稿或尋求報道：content@jiqizhixin.com
留言列表