點選下方卡片,關注“CVer”公眾號
點選下方卡片,關注“CVer”公眾號
AI/CV重磅乾貨,第一時間送達
AI/CV重磅乾貨,第一時間送達
新增微訊號:CVer2233,小助手會拉你進群!
掃描下方二維碼,加入CVer學術星球!可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料,及最前沿應用!發論文/搞科研/漲薪,強烈推薦!
新增微訊號:CVer2233,小助手會拉你進群!
掃描下方二維碼,加入CVer學術星球!可以獲得最新頂會/頂刊上的論文idea和CV從入門到精通資料,及最前沿應用!發論文/搞科研/漲薪,強烈推薦!

本文宣傳一下我們被 ICLR 2025 錄用的文章—MambaQuant。

論文標題:
MambaQuant: Quantizing the Mamba Family with Variance Aligned Rotation Methods
論文地址:
https://arxiv.org/abs/2501.13484
作者單位:
後摩智慧、哈爾濱工業大學、南京大學、東南大學
值得一提的是,該工作對 transformer-based LLM 模型的量化也有很大提升,例如在 Llama2 7B 模型上超過了 Quarot,是個很通用的方法!


Abstract
Mamba 是一種高效的序列模型,可與 Transformer 相媲美,在各類任務中展現出作為基礎架構的巨大潛力。量化技術常用於神經網路,以減小模型大小並降低計算延遲。
然而,將量化應用於 Mamba 的研究尚少,現有的量化方法雖然在 CNN 和 Transformer 模型中效果顯著,但對 Mamba 模型卻不太適用(例如,即使在 W8A8 配置下,QuaRot 在 Vim-T 模型上的準確率仍下降了 21%)。
我們率先對這一問題展開探索,並識別出幾個關鍵挑戰。首先,在門投影、輸出投影和矩陣乘法中存在大量異常值。其次,Mamba 獨特的並行掃描操作進一步放大了這些異常值,導致資料分佈不均衡且呈現長尾現象。第三,即使應用了 Hadamard 變換,權重和啟用值在通道間的方差仍然不一致。
為此,我們提出了 MambaQuant,這是一種訓練後量化(PTQ)框架,包含:1)基於 Karhunen-Loève 變換(KLT)的增強旋轉,使旋轉矩陣能適應不同的通道分佈;2)平滑融合旋轉,用於均衡通道方差,並可將額外引數合併到模型權重中。
實驗表明,MambaQuant 能夠將權重和啟用值量化為 8 位,且基於 Mamba 的視覺和語言任務的準確率損失均小於 1%。據我們所知,MambaQuant 是首個針對 Mamba 系列模型的綜合性 PTQ 設計,為其進一步的應用發展奠定了基礎。

Introduction
為了建立一套針對 Mamba 模型的綜合量化方法,我們首先研究其中涉及的潛在限制和挑戰:
1. Mamba 模型的權重和啟用值中都存在顯著的異常值。我們觀察到,線性層的權重中存在異常值,尤其是在用於語言任務的 Mamba-LLM 的門投影層(圖1a)中。我們還發現,線性層的某些輸入在通道維度上表現出顯著的方差。這種情況在用於視覺任務的 Vim 的輸出投影層(圖1b)中尤為明顯。
2. 並行掃描(PScan)進一步放大了啟用值的異常值。為了在每個時間戳獲得隱藏狀態,PScan 運算元(Smith 等人,2022)會對一個固定的引數矩陣不斷進行自乘運算。在這種情況下,值較高的通道會被放大,而值相對較低的通道則會被削弱。這種通道間明顯的數值差異會直接擴充套件到啟用值上(例如,如圖1(c)所示的矩陣乘法的輸入變數,以及圖 2 所示)。


最近,基於 Hadamard 的方法因其能夠使最大值均勻化以及具有等價變換特性,在 Transformer-based LLMs(T-LLMs)的量化中取得了顯著成功。例如,使用 QuaRot 將 LLAMA2-70B 量化為 4 位時,能保持 99% 的零樣本效能。
然而,將這種方法直接應用於 Mamba 模型會導致準確率大幅下降(例如,即使在 8 位量化的情況下,在 Vim 上平均準確率仍然下降超過 12%)。為了解決上述問題,我們發表了 MambaQuant 這篇文章,(據我們所知)這是首個在 Mamba 系列模型上實現了高準確率 W8A8/W4A8 量化的工作,主要貢獻包括:
1. 在離線模式下,我們提出基於 Karhunen – Loève 變換(KLT)的增強旋轉。此技術將 Hadamard 矩陣與 KLT 矩陣相乘,使旋轉矩陣能夠適應不同的通道分佈。
2. 在線上模式下,我們引入平滑融合旋轉。這種方法在 Hadamard 變換之前進行平滑處理。額外的平滑引數被靈活地整合到 Mamba 模組的權重中,以避免額外的記憶體空間和推理步驟成本。
因此,量化資料的最大值和方差在通道維度上都得到了充分對齊,方法效果如圖 3 所示:

實驗表明,MambaQuant 能夠高效地將權重和啟用值量化為 8 位,且在基於 Mamba 的視覺和語言任務上,準確率損失均小於 1%。


Method
3.1 對 Hadamard 旋轉效果不佳的分析
我們發現,該方法無法對齊量化變數的通道方差,從而忽略了不同通道之間的分佈一致性。詳細來說,給定一箇中心化的資料矩陣(矩陣的列均值為零)X(權重或啟用值),其維度為(n, m),以及維度為(m, m)的 Hadamard 變換矩陣 H,變換後的矩陣 XH 的協方差矩陣可以表示為:

協方差矩陣的第 l 個對角元素可以表示為:

對於給定的 值,公式(2)表示第 個通道的方差。由於向量 隨 變化,在大多數情況下無法證明通道方差在數值上接近。此外,考慮到 H 是一個固定矩陣,而 和 都取決於輸入,在所有情況下,Hadamard 變換都不可能統一調整通道方差。Hadamard 變換的這一特性不可避免地為每個通道形成了不同的分佈,從而導致次優的量化效果。
3.2 KLT增強旋轉
為了克服上述限制,我們引入了 KLT 來均衡通道方差。KLT 識別資料中的主成分,並將資料投影到這些成分上,透過關注方差最大的方向來保留每個通道的最關鍵資訊。在實際應用中,Mamba 權重和啟用值的均值通常接近於零,滿足 KLT 的適用條件。具體而言,我們對由校準資料得到的中心化矩陣X的協方差矩陣進行特徵值分解來應用 KLT:

接下來,如公式(4)所示,透過將 KLT 應用於 Hadamard 矩陣 H,可以得到 KLT 增強旋轉矩陣

在公式(4)基礎上,公式(1)可因此轉化為公式(5):

而公式(2)可變為公式(6):

這樣,每個通道的方差變得相同,使得量化變得容易得多。這種變換具有雙重目的:它不僅均衡了不同通道之間的方差,還體現了 KLT 矩陣與 Hadamard 矩陣的獨特屬性,後者能夠平衡最大值。
在實踐中,KLT 是離線執行的,以避免額外的計算成本。為了將這種 KLT 增強的旋轉矩陣應用於 Mamba 結構,我們修改了 QuaRot 中的離線變換。如圖 5 所示,我們將此策略應用於 LoRA 模組和層間連線(其中輸出投影、門投影和狀態投影被變換)。

3.3 Smooth對齊旋轉
為了在線上旋轉中實現通道方差對齊,我們在執行線上 Hadamard 旋轉之前引入了平滑(smooth)技術。採用這種方法的動機是透過一個平滑向量來使通道方差均勻化。
通常,平滑因子可以被吸收到量化層的相鄰層中例如 SmoothQuant, OmniQuant。這種操作有效地避免了因引入額外引數而產生的額外記憶體分配和計算開銷需求。然而,這種方法在 Mamba 模組中並不完全適用,這是由於非逐元素的 SiLU 操作以及 PScan 的複雜迴圈結構。為此,我們分別針對輸出投影和矩陣乘法提出了兩種不同的設計。

對於輸出投影層:我們提出 S – SiLU,改進了傳統的 SiLU 啟用函式,以滿足平滑融合量化的需求:

如圖 6(a)所示,S – SiLU 函式在門投影上的應用可以表示為如下公式:

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

經過平滑處理後,輸出投影和矩陣乘法的啟用值的通道方差變得相對均勻。隨後,我們針對 Mamba 結構修改並應用了線上 Hadamard 旋轉,如圖 7 所示。Hadamard 矩陣 H 被動態地應用於輸出投影和矩陣乘法的輸入啟用值,而轉置後的 H^T 可以被吸收到相應的權重中。

Experiments
實驗結果表明,MambaQuant 在 Mamba 系列模型上都取得了驚人的效果,表現超過 Quarot,並遠超其他的量化方案。例如其 W8A8 的精度在多種視覺語言的評估任務上都表現出小於 1% 的精度損失,其 W4A8 的量化也是實現了 SOTA 的效果。

值得一提的是,我們的通道方差對齊方法對精度有很明顯的提升,該提升的視覺化效果也十分顯著。
這項工作首次在 Mamba 模型上實現了高精度量化,為 Mamba 模型的高效部署和推理開闢了新的途徑,尤其是邊緣裝置上。同時該工作對 transformer-based LLM 模型的量化也有很大提升,值得進一步探索!
何愷明在MIT授課的課件PPT下載
何愷明在MIT授課的課件PPT下載
ECCV 2024 論文和程式碼下載
ECCV 2024 論文和程式碼下載
CVPR 2024 論文和程式碼下載
CVPR 2024 論文和程式碼下載
CV垂直方向和論文投稿交流群成立
一定要備註:研究方向+地點+學校/公司+暱稱(如Mamba、多模態學習或者論文投稿+上海+上交+卡卡),根據格式備註,可更快被透過且邀請進群
▲掃碼或加微訊號: CVer2233,進交流群
CVer計算機視覺(知識星球)來了!想要了解最新最快最好的CV/DL/AI論文速遞、優質實戰專案、AI行業前沿、從入門到精通學習教程等資料,歡迎掃描下方二維碼,加入CVer計算機視覺(知識星球),已彙集上萬人!
▲掃碼加入星球學習
▲點選上方卡片,關注CVer公眾號
整理不易,請贊和在看

▲掃碼或加微訊號: CVer2233,進交流群
CVer計算機視覺(知識星球)來了!想要了解最新最快最好的CV/DL/AI論文速遞、優質實戰專案、AI行業前沿、從入門到精通學習教程等資料,歡迎掃描下方二維碼,加入CVer計算機視覺(知識星球),已彙集上萬人!

▲掃碼加入星球學習
▲點選上方卡片,關注CVer公眾號
整理不易,請贊和在看
