首个Mamba系列模型量化方案MambaQuant解读-电子发烧友网

MambaQuant在Mamba系列模型上实现了W8A8/W4A8量化的方法，精度逼近浮点，超过Quarot等SOTA方法。该工作已被人工智能顶会ICLR-2025接收。

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设计，为其进一步的应用发展奠定了基础。

•论文链接：https://arxiv.org/abs/2501.13484

•单位：后摩智能、哈尔滨工业大学（深圳）、南京大学、东南大学

Introduction

为了建立一套针对Mamba模型的综合量化方法，我们首先研究其中涉及的潜在限制和挑战：

Mamba模型的权重和激活值中都存在显著的异常值。我们观察到，线性层的权重中存在异常值，尤其是在用于语言任务的 Mamba-LLM 的门投影层（图1a）中。我们还发现，线性层的某些输入在通道维度上表现出显著的方差。这种情况在用于视觉任务的 Vim 的输出投影层（图1b）中尤为明显。

并行扫描（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量化的工作，主要贡献包括：

a. 在离线模式下，我们提出基于 Karhunen - Loève 变换（KLT）的增强旋转。此技术将 Hadamard 矩阵与 KLT 矩阵相乘，使旋转矩阵能够适应不同的通道分布。

b. 在在线模式下，我们引入平滑融合旋转。这种方法在 Hadamard 变换之前进行平滑处理。额外的平滑参数被灵活地整合到 Mamba 模块的权重中，以避免额外的内存空间和推理步骤成本。

因此，量化数据的最大值和方差在通道维度上都得到了充分对齐，方法效果如图3所示：

实验表明，MambaQuant 能够高效地将权重和激活值量化为8位，且在基于Mamba的视觉和语言任务上，准确率损失均小于1%。

Method

对Hadamard旋转效果不佳的分析

我们发现，该方法无法对齐量化变量的通道方差，从而忽略了不同通道之间的分布一致性。详细来说，给定一个中心化的数据矩阵（矩阵的列均值为零）X（权重或激活值），其维度为(n, m)，以及维度为(m, m)的Hadamard变换矩阵H，变换后的矩阵XH的协方差矩阵可以表示为：

协方差矩阵的第l个对角元素可以表示为：

对于给定的l值，公式(2)表示第l个通道的方差。由于向量随l变化，在大多数情况下无法证明通道方差在数值上接近。此外，考虑到H是一个固定矩阵，而K和λ都取决于输入，在所有情况下，Hadamard变换都不可能统一调整通道方差。Hadamard变换的这一特性不可避免地为每个通道形成了不同的分布，从而导致次优的量化效果。

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模块和层间连接（其中输出投影、门投影和状态投影被变换）。

Smooth对齐旋转

为了在在线旋转中实现通道方差对齐，我们在执行在线Hadamard旋转之前引入了平滑(smooth)技术。采用这种方法的动机是通过一个平滑向量来使通道方差均匀化。通常，平滑因子可以被吸收到量化层的相邻层中例如SmoothQuant, OmniQuant。这种操作有效地避免了因引入额外参数而产生的额外内存分配和计算开销需求。然而，这种方法在Mamba模块中并不完全适用，这是由于非逐元素的SiLU操作以及PScan的复杂循环结构。为此，我们分别针对输出投影和矩阵乘法提出了两种不同的设计。