MambaQuant在Mamba系列模型上實現(xiàn)了W8A8/W4A8量化的方法，精度逼近浮點，超過Quarot等SOTA方法。該工作已被人工智能頂會ICLR-2025接收。

Abstract

Mamba是一種高效的序列模型，可與Transformer相媲美，在各類任務(wù)中展現(xiàn)出作為基礎(chǔ)架構(gòu)的巨大潛力。量化技術(shù)常用于神經(jīng)網(wǎng)絡(luò)，以減小模型大小并降低計算延遲。然而，將量化應(yīng)用于Mamba的研究尚少，現(xiàn)有的量化方法雖然在CNN和Transformer模型中效果顯著，但對Mamba模型卻不太適用（例如，即使在W8A8配置下，QuaRot在Vim-T模型上的準確率仍下降了21%）。我們率先對這一問題展開探索，并識別出幾個關(guān)鍵挑戰(zhàn)。

首先，在門投影、輸出投影和矩陣乘法中存在大量異常值。其次，Mamba獨特的并行掃描操作進一步放大了這些異常值，導致數(shù)據(jù)分布不均衡且呈現(xiàn)長尾現(xiàn)象。第三，即使應(yīng)用了Hadamard變換，權(quán)重和激活值在通道間的方差仍然不一致。為此，我們提出了MambaQuant，這是一種訓練后量化（PTQ）框架，包含：1）基于Karhunen-Loève變換（KLT）的增強旋轉(zhuǎn)，使旋轉(zhuǎn)矩陣能適應(yīng)不同的通道分布；2）平滑融合旋轉(zhuǎn)，用于均衡通道方差，并可將額外參數(shù)合并到模型權(quán)重中。

實驗表明，MambaQuant能夠?qū)?quán)重和激活值量化為8位，且基于Mamba的視覺和語言任務(wù)的準確率損失均小于1%。據(jù)我們所知，MambaQuant是首個針對Mamba系列模型的綜合性PTQ設(shè)計，為其進一步的應(yīng)用發(fā)展奠定了基礎(chǔ)。

?論文鏈接：https://arxiv.org/abs/2501.13484

?單位：后摩智能、哈爾濱工業(yè)大學（深圳）、南京大學、東南大學

Introduction

為了建立一套針對Mamba模型的綜合量化方法，我們首先研究其中涉及的潛在限制和挑戰(zhàn)：

Mamba模型的權(quán)重和激活值中都存在顯著的異常值。我們觀察到，線性層的權(quán)重中存在異常值，尤其是在用于語言任務(wù)的 Mamba-LLM 的門投影層（圖1a）中。我們還發(fā)現(xiàn)，線性層的某些輸入在通道維度上表現(xiàn)出顯著的方差。這種情況在用于視覺任務(wù)的 Vim 的輸出投影層（圖1b）中尤為明顯。

并行掃描（PScan）進一步放大了激活值的異常值。為了在每個時間戳獲得隱藏狀態(tài)，PScan算子（Smith等人，2022）會對一個固定的參數(shù)矩陣不斷進行自乘運算。在這種情況下，值較高的通道會被放大，而值相對較低的通道則會被削弱。這種通道間明顯的數(shù)值差異會直接擴展到激活值上（例如，如圖1（c）所示的矩陣乘法的輸入變量，以及圖2所示）。

最近，基于 Hadamard 的方法因其能夠使最大值均勻化以及具有等價變換特性，在 Transformer-based LLMs (T-LLMs) 的量化中取得了顯著成功。例如，使用 QuaRot 將 LLAMA2-70B 量化為 4 位時，能保持 99% 的零樣本性能。然而，將這種方法直接應(yīng)用于 Mamba 模型會導致準確率大幅下降（例如，即使在 8 位量化的情況下，在 Vim上平均準確率仍然下降超過 12%）。

為了解決上述問題，我們發(fā)表了MambaQuant這篇文章，（據(jù)我們所知）這是首個在Mamba系列模型上實現(xiàn)了高準確率W8A8/W4A8量化的工作，主要貢獻包括：

a. 在離線模式下，我們提出基于 Karhunen - Loève 變換（KLT）的增強旋轉(zhuǎn)。此技術(shù)將 Hadamard 矩陣與 KLT 矩陣相乘，使旋轉(zhuǎn)矩陣能夠適應(yīng)不同的通道分布。

b. 在在線模式下，我們引入平滑融合旋轉(zhuǎn)。這種方法在 Hadamard 變換之前進行平滑處理。額外的平滑參數(shù)被靈活地整合到 Mamba 模塊的權(quán)重中，以避免額外的內(nèi)存空間和推理步驟成本。

因此，量化數(shù)據(jù)的最大值和方差在通道維度上都得到了充分對齊，方法效果如圖3所示：

實驗表明，MambaQuant 能夠高效地將權(quán)重和激活值量化為8位，且在基于Mamba的視覺和語言任務(wù)上，準確率損失均小于1%。

Method

對Hadamard旋轉(zhuǎn)效果不佳的分析

我們發(fā)現(xiàn)，該方法無法對齊量化變量的通道方差，從而忽略了不同通道之間的分布一致性。詳細來說，給定一個中心化的數(shù)據(jù)矩陣（矩陣的列均值為零）X（權(quán)重或激活值），其維度為(n, m)，以及維度為(m, m)的Hadamard變換矩陣H，變換后的矩陣XH的協(xié)方差矩陣可以表示為：

協(xié)方差矩陣的第l個對角元素可以表示為：

對于給定的l值，公式(2)表示第l個通道的方差。由于向量隨l變化，在大多數(shù)情況下無法證明通道方差在數(shù)值上接近。此外，考慮到H是一個固定矩陣，而K和λ都取決于輸入，在所有情況下，Hadamard變換都不可能統(tǒng)一調(diào)整通道方差。Hadamard變換的這一特性不可避免地為每個通道形成了不同的分布，從而導致次優(yōu)的量化效果。

KLT增強旋轉(zhuǎn)

為了克服上述限制，我們引入了KLT來均衡通道方差。KLT識別數(shù)據(jù)中的主成分，并將數(shù)據(jù)投影到這些成分上，通過關(guān)注方差最大的方向來保留每個通道的最關(guān)鍵信息。在實際應(yīng)用中，Mamba權(quán)重和激活值的均值通常接近于零，滿足KLT的適用條件。具體而言，我們對由校準數(shù)據(jù)得到的中心化矩陣X的協(xié)方差矩陣進行特征值分解來應(yīng)用KLT：

接下來，如公式(4)所示，通過將KLT應(yīng)用于Hadamard矩陣H，可以得到KLT增強旋轉(zhuǎn)矩陣：

在公式(4)基礎(chǔ)上，公式(1)可因此轉(zhuǎn)化為公式(5)：

而公式(2)可變?yōu)楣?6)：

這樣，每個通道的方差變得相同，使得量化變得容易得多。這種變換具有雙重目的：它不僅均衡了不同通道之間的方差，還體現(xiàn)了KLT矩陣與Hadamard矩陣的獨特屬性，后者能夠平衡最大值。在實踐中，KLT是離線執(zhí)行的，以避免額外的計算成本。為了將這種KLT增強的旋轉(zhuǎn)矩陣應(yīng)用于Mamba結(jié)構(gòu)，我們修改了QuaRot中的離線變換。如圖5所示，我們將此策略應(yīng)用于LoRA模塊和層間連接（其中輸出投影、門投影和狀態(tài)投影被變換）。

Smooth對齊旋轉(zhuǎn)

為了在在線旋轉(zhuǎn)中實現(xiàn)通道方差對齊，我們在執(zhí)行在線Hadamard旋轉(zhuǎn)之前引入了平滑(smooth)技術(shù)。采用這種方法的動機是通過一個平滑向量來使通道方差均勻化。通常，平滑因子可以被吸收到量化層的相鄰層中例如SmoothQuant, OmniQuant。這種操作有效地避免了因引入額外參數(shù)而產(chǎn)生的額外內(nèi)存分配和計算開銷需求。然而，這種方法在Mamba模塊中并不完全適用，這是由于非逐元素的SiLU操作以及PScan的復雜循環(huán)結(jié)構(gòu)。為此，我們分別針對輸出投影和矩陣乘法提出了兩種不同的設(shè)計。

對于輸出投影層：我們提出S - SiLU，改進了傳統(tǒng)的SiLU激活函數(shù)，以滿足平滑融合量化的需求：

如圖6(a)所示，S - SiLU函數(shù)在門投影上的應(yīng)用可以表示為如下公式：

對于矩陣乘法層：如圖6(b)所示，平滑參數(shù)s可以被自然的吸收到權(quán)重B和權(quán)重C中，然而A矩陣會在推理時執(zhí)行多次的自乘運算，因此我們引入了計算友好的addcmul算子，僅對第一個時間步的A矩陣的運算做s參數(shù)的吸收融合，如公式(8)所示：

經(jīng)過平滑處理后，輸出投影和矩陣乘法的激活值的通道方差變得相對均勻。隨后，我們針對Mamba結(jié)構(gòu)修改并應(yīng)用了在線Hadamard旋轉(zhuǎn)，如圖7所示。Hadamard矩陣H被動態(tài)地應(yīng)用于輸出投影和矩陣乘法的輸入激活值，而轉(zhuǎn)置后的H^T可以被吸收到相應(yīng)的權(quán)重中。

Experiments

實驗結(jié)果表明，MambaQuant 在 Mamba 系列模型上都取得了驚人的效果，表現(xiàn)超過Quarot，并遠超其他的量化方案。例如其W8A8的精度在多種視覺語言的評估任務(wù)上都表現(xiàn)出小于1%的精度損失，其W4A8的量化也是實現(xiàn)了SOTA的效果。

值得一提的是，我們的通道方差對齊方法對精度有很明顯的提升，該提升的可視化效果也十分顯著，比如KLT：

這項工作首次在Mamba模型上實現(xiàn)了高精度量化，為Mamba模型的高效部署和推理開辟了新的途徑，尤其是邊緣設(shè)備上。