DuQuant:通過正交變換分散 LLM 離群值,實現 SOTA 4bit 量化
一、背景
本文中我們介紹一下最近被 NeurIPS 2024 接收為 Oral Presentation 的 LLM 量化工作 DuQuant。這是一種令旋轉變換和排列變換來更有效地緩解 Massive Outlier 和 Normal Outlier 的新方法。其達到了 4bit 量化的新 SOTA。具體來說:
- 作者發現 LLM 的 FFN 模塊中,在 down_proj 的輸入中存在明顯的Massive Outlier;這種 Outlier 不同于以往發現的 Outlier,表現為絕對值大于幾百的個別異常值;這些 Outlier 會顯著限制 4bit 權重激活量化方法的性能。
- 作者提出了基于先驗知識(Outlier Channel)的旋轉變換來將Massive Outlier旋轉至臨近 Channel,考慮效率平衡,作者使用分塊旋轉矩陣。而分塊旋轉矩陣會導致部分 Block 的平均激活值的絕對值大于其他 Block,所以作者進一步提出通道置換變換,重新平衡不同 Block 之間的 Outlier,然后使用第二次旋轉變換進一步最小化 Outlier。
- DuQuant 的訓練非??焖?,可以在 50s 內完成 7B 模型的量化過程,即插即用。
對應的 Paper 為:[2406.01721] DuQuant: Distributing Outliers via Dual Transformation Makes Stronger Quantized LLMs [1]
對應的代碼庫為:https://github.com/Hsu1023/DuQuant [2]
二、 引言
2.1 量化方法
在每個 Transformer Block 中,多頭自注意力(MSA)和前饋網絡(FFN)基本上都由線性層組成,將其表示為:,其中, 是激活輸入, 表示權重矩陣。
模型量化通過將浮點數(FP16, BF16)表示的模型權重或激活值轉換為低比特的浮點數或整數(INT8、FP8、INT4)來減少內存占用,權重激活量化可以通過使用低比特乘法的算子進一步加快模型推理的速度。本文中,作者重點關注低比特整型量化,旨在獲得更好的硬件支持。具體來說, 位量化過程將 FP16 張量 映射為低位整數:
符號 表示最接近的舍入操作, 是量化步長, 表示零點。遵循主流量化方法,作者對激活 采用逐 Token 量化,對權重 采用逐 Channel 量化,這意味著為 的每個 Token 分配不同的步長(),為 的每個輸出 Channel 分配不同的步長()。
2.2 動機
作者首次在實驗中發現,LLM FFN 模塊中的 down_proj layer 存在明顯的 [2402.17762] Massive Activations in Large Language Models [3],表現為絕對值大于幾百的異常值并局限于個別的 Token 中。這些 Massvie Outlier 造成 [2211.10438] SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models [4] 和 [2308.13137] OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models [5] 等算法在 4bit WA 量化中表現較差。在 down_proj layer 中發現 Massive Outlier 現象為該工作首次提出。
如下圖 Figure 1 (a)(b) 對比了常見的 Normal Outlier,和在 FFN 中出現的 Massive Outlier。
SmoothQuant 中將激活值除以每 Channel 平滑因子并將其乘回權重矩陣,嘗試將量化難度從激活值轉移到權重。具體來說,SmoothQuant 使用逐 Channel 的平滑對角矩陣,記為 ,將原始線性層重新表達為:,對角矩陣 中的元素 的計算方式為 ,其中 是一個超參數,表示遷移強度。
然而,作者觀察到在輸入端進行這種轉移可能會導致權重矩陣也出現明顯的難以量化的 Outlier(如上圖 Figure 1 (d) 所示),這一問題的根源在于 Massive Outlier 使平滑因子 變得異常大。
此外,極大的 Outlier 還可能導致基于優化的方法出現梯度爆炸的問題,所以基于梯度優化的 OmniQuant 和 [2403.12544] AffineQuant: Affine Transformation Quantization for Large Language Models [6] 等工作會直接跳過 down_proj layer,直接退化為 SmoothQuant 的處理方式。這都表明需要一種更新、更好的方式處理兩種 Outlier,特別是平滑掉 down_proj 輸入側的 Massive outlier。
三、本文方案
3.1 概覽
DuQuant 提出通過學習旋轉和置換變換矩陣,在 Activation 矩陣內部將 Outlier 轉移到其他通道,最終得到平滑的激活矩陣,從而大幅度降低量化難度。如下圖 Figure 2 所示為 DuQuant 算法說明:
- a. 逐步展示了 DuQuant 算法對 Normal Outlier 的處理過程。
- b. DuQuant 顯著降低了 Massive Outlier。
- c. 一個 Example 說明 DuQuant 有效降低了激活矩陣的量化難度。
簡單來說,DuQuant 算法包含三個步驟:
- 旋轉矩陣的構造有效利用了特定 Outlier Channel 的位置索引,作者使用分塊對角的旋轉矩陣,在每個 Block 內部通過貪心算法將 Outlier 平攤到其他的 Channel 中。
- 由于 Block Size 的限制,可能導致某些 Block 在旋轉之后組內的平均值大于其他 Block,因此作者進一步使用通道置換技術重新分配 Activation Channel,利用 zigzag 順序排列使各組均值的方差大幅減小。
- 進一步執行一次旋轉變換達到更均勻的 Activation 分布,從而大幅降低量化難度。
3.2 旋轉矩陣
作者希望應用旋轉矩陣 進行行或列變換,減輕 Normal Outlier 和 Massive Outlier 的影響。由于 Massive Outlier 通常隨機分布在激活空間中,直接找到能夠通過單次旋轉變換減輕 Outlier 的最優旋轉矩陣 是具有挑戰性的。為了解決這一問題,作者采用帶有先驗知識的貪心搜索方法來計算旋轉矩陣 ,從而近似理想的旋轉矩陣 。具體來說, 的計算包括以下步驟:
- 識別 Outlier 主要集中的特征維度 ,即:, 其中, 表示 中第 行和第 列的元素。
- 基于搜索到的維度 ,構建旋轉矩陣如下: 是用于交換激活值的第 1 列和第 列的交換矩陣, 表示一個正交初始化的旋轉矩陣,其第 1 行均勻分布。這樣做的目的是通過 變換后減輕第 1 列中的 Outlier。為了進一步增加隨機性,作者保留了減輕 Outlier 后的第 1 列,并通過與隨機正交矩陣 相乘,隨機旋轉其他列。
- 設 為貪心搜索的步數,則近似的旋轉矩陣為 ,其中 。每個 根據上述公式和識別到的特征維度 構建。
通過這種構建方式,可以確保近似的最優旋轉矩陣 能夠有效減輕具有較大幅度 Outlier,而不僅僅是使用隨機選擇的正交旋轉矩陣。然而,直接構建整個旋轉矩陣非常耗時,并且會導致大量的內存開銷。為了實現快速矩陣乘法,參考 [2306.11987] Training Transformer with 4ibts [6],作者選擇以分塊的方式近似旋轉矩陣 ,
其中, 表示第 個塊的方陣,該矩陣按照上述三步構建。塊的數量 通過 計算得出。
3.3 通道置換矩陣
盡管采用了塊對角旋轉矩陣 來提高時間和存儲效率,但其專注于局部信息的特性帶來了進一步減少 Outlier 的潛在限制。由于在每個小塊內進行的旋轉變換無法整合跨不同塊的信息,一個塊中可能有相對較大的 Outlier,而另一個塊中則有較小的 Outlier,導致不同塊之間存在較高的方差。因此,作者提出利用通道置換矩陣平衡不同 Block 之間 Outlier 的幅度。
具體來說,在每個小塊中,將維度 中最大的 Outlier 記為 。同時, 表示第 個塊中所有 的平均值,其中 ,各個塊之間激活幅度的方差可以表示為:
作者引入了之字形置換矩陣 。具體來說,通過生成一個之字形序列,首先將激活值最高的通道分配給第一個塊,接著,將激活值次高的通道按遞減順序分配到后續的塊,直到第 個塊。在到達最后一個塊后,順序反轉,從下一個激活值最高的通道開始,按遞增順序分配。這個往復模式貫穿所有塊,確保沒有單個塊持續接收最高或最低激活值的通道。通過使用之字形置換,DuQuant 實現了不同塊之間 Outlier 的均衡分布,從而能夠使用額外的旋轉變換進一步平滑 Outlier,如圖 Figure 2 所示。
需要說明的是:
- 通道置換其實是非常重要的一步,也很簡單快速(對于推理效率的影響很小,可見后面實驗部分),即可以避免像 [2405.16406] SpinQuant: LLM quantization with learned rotations [7] 那樣復雜的訓練流程,也比 [2404.00456] QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs [8] 的 Hadamard 旋轉性能表現更好。
- 旋轉矩陣和置換變換矩陣都是正交矩陣,保證了 輸出的不變性,作者還通過嚴謹的理論推導證明了兩種變換有效降低了量化誤差,具體證明可以閱讀 Paper 里面的 Appendix。
四、實驗和結果
4.1 主要結果
作者基于提出的方案驗證了 LLaMA、Vicuna、Mistral 系列模型,其在 PPL、QA、MMLU 和 MT-Bench 等任務上都明顯提升了量化模型的性能。
如下圖 Table 1 所示,相比幾種常見的 W4A4 量化方案,DuQuant 在 LLaMA-1 和 LLaMA-2 模型上的量化達到了 SOTA 效果:
如下圖 Table 4 所示,DuQuant 在 LLaMA3 模型的量化效果,4-bit 下同樣達到了 SOTA:
如下圖 Table 3 所示,DuQuant 在 Vicuna-v1.5-13B 模型上的量化效果,0-shot 和 5-shot 任務同樣都打到了 SOTA:
4.2 推理速度
作者在 NVIDIA RTX 3090 GPU 上進行了推理速度測試,其 W4A4 Kernel 來自 Quarot。如下圖 Table 9 所示,結果表明 DuQuant 在 Prefiling 階段相比 FP16 加速比可以達到 2.08x(LLaMA2-7B)、2.34x(LLaMA2-13B)。同時,如下圖 Figure 5 所示,DuQuant 相比 INT4 推理速度會降低 10% 左右,略微高于 QuaRot,但帶來更多精度提升。
4.3 與 QuaRot 對比
DuQuant 與使用 Hadamard 旋轉矩陣的 QuaRot 主要有以下兩點不同:
1. DuQuant 構造的旋轉矩陣利用了先驗知識(具體的 Outlier Channel 索引),因此 DuQuant的 旋轉矩陣可以比 QuaRot 更好地平滑激活空間,如下圖 Figure 4 展示了 DuQuant 單次旋轉和 Hadamard 旋轉對 LLaMA2-7B Attention Key_proj 輸入的變換效果。
2. QuaRot 依賴于耗時的 GPTQ 算法來提升性能,而作者引入的通道置換矩陣可以幫助 DuQuant 在極短時間內進一步平衡 Outlier 的分布,從而取得了更好的量化效果,如下圖 Table 7 所示。
五、參考鏈接
1. https://arxiv.org/abs/2406.01721
2. https://github.com/Hsu1023/DuQuant
3. https://arxiv.org/abs/2402.17762
4. https://arxiv.org/abs/2211.10438
5. https://arxiv.org/abs/2308.13137
6. https://arxiv.org/abs/2306.11987
7. https://arxiv.org/abs/2405.16406
8. https://arxiv.org/abs/2404.00456
