点击下方卡片,关注「AI视界引擎」公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
状态空间模型(SSMs)正逐渐成为Transformer的引人注目的替代方案,这得益于其持续稳定的内存使用和卓越的性能。然而,在云服务或资源受限设备上扩展SSMs仍面临挑战,主要源于其存储需求和高计算功耗。
为克服这一问题,通过低比特宽度数据格式量化SSMs可以减小模型尺寸并利用硬件加速。由于SSMs容易受到量化引入的误差影响,近期研究重点在于优化特定模型或比特宽度以提高效率,同时不牺牲性能。然而,不同的场景需要不同的比特宽度配置,例如W4A8用于提升大批量解码速度,而W4A16则用于增强单用户短 Prompt 应用中的生成速度。
为此,作者提出了Quamba2,它兼容W8A8、W4A8和W4A16,适用于Mamba1和Mamba2 Backbone 网络,以满足SSM在各种平台部署日益增长的需求。基于SSMs保持通道顺序和激活持久性的特性,作者提出了一种离线方法,通过排序和聚类对输入
的线性递归输入进行8比特量化,并结合对输入相关参数
和
的每状态组量化。
为确保SSM输出中的计算不变性,作者根据聚类序列离线重新排列权重。实验表明,Quamba2-8B优于多种最先进的SSM量化方法,在预填充和生成阶段分别实现了1.3倍和3倍的速度提升,同时仅损失1.6%的平均精度,并降低了4倍的内存使用。
在MMLU上的评估展示了yinqing-Quamba2_2503的泛化性和鲁棒性。
代码和量化模型将发布在:https://github.com/enyac-group/Quamba。
unsetunset1 引言unsetunset
状态空间模型(SMs),具有恒定的内存复杂度,正逐渐成为Transformer(Vaswani 2017)在语言建模(Waleffe等人 2024;Wang等人 2024)、视觉和音频等领域的有效替代方案。一些研究扩展了模型规模,并展示了其性能与同等规模的Transformer相当。然而,状态空间模型的大规模限制了硬件选择并增加了部署成本。
PTQ(PTQ)为高效部署提供了一种有吸引力的解决方案,通过消除对大型模型进行微调的需求。PTQ将预训练权重和激活值转换为低比特格式(例如8比特),从而降低权重存储的内存使用,并利用先进的硬件单元。近期研究(Chiang等人,2025;Xu等人,2025)表明,在Transformer中有效的量化技术在状态空间模型(SSMs)上面临挑战,因为线性递归对量化引起的误差敏感。这项先前的工具新引入了针对SSMs的PTQ算法,以缩小低精度和半精度模型之间的性能差距。然而,这些方法要么没有探索不同的比特宽度(Chiang等人,2025),要么在较低比特宽度(如W4A8)上未能达到令人满意的性能(Xu等人,2025)。
特定位宽设置在某些场景中至关重要。例如,W4A8通过使用大批量输入提升了云服务吞吐量(Lin等人 2024b),而W4A16则提高了短 Prompt 应用的效率(Lin等人 2024a)。因此,当前的基于SSM的量化方法(Chiang等人 2025;Xu等人 2025)在边缘设备上可能表现不佳,或无法在云服务上最大化吞吐量。此外,最近的一项研究(Gong等人 2024;Kumar等人 2025;Zhao等人 2024a)表明,对模型权重和激活进行重度量化(例如W4A4)会损害模型在多步推理任务上的泛化能力。以往的基于SSM的研究忽略了量化模型的泛化能力。
为解决这些问题,作者提出了Quamba2,一个针对选择性状态空间模型(SSMs)的鲁棒且可扩展的模型后量化框架。如表1和表2所示,yinqing-Quamba2_2503支持Mamba1(Gu和Dao 2024)和Mamba2(Dao和Gu 2024)的端到端W8A8、W4A8和W4A16量化,满足在云和边缘平台上部署状态空间模型的需求。基于SSM计算中通道顺序保持和激活持久性的特性,如图2和图3所示,作者采用离线集群感知权重重排方法,将具有相似值范围的SSM Head 和通道分组,使它们能够共享量化缩放因子,从而提高量化精度。对于选择性SSM的输入依赖参数
,作者识别激活中的状态持久性,并按状态组进行量化。
作者的排序聚类和按状态组量化方法提高了量化精度,缩小了半精度模型的精度差距。在图1及其余实验中,作者展示了Quamba2-8B超越了多种领先的SSM量化方法,在预填充和生成方面分别实现了高达1.3倍和3倍的加速,同时内存占用减少了4倍,在六个零样本任务中仅损失了1.6%的精度。此外,作者在MMLU(Hendrycks等人 2020)这一大型多任务数据集上测试了Quamba2,展示了yinqing-Quamba2_2503的泛化性和鲁棒性。
unsetunset2 相关工作unsetunset
模型量化。使用低位宽数据类型表示模型权重和激活可以降低存储和加载参数的成本,并受益于先进的低位宽计算单元(即张量核心)。量化方法通常分为两类:量化感知训练(QAT)(Dettmers等人2024;Liu等人2024b;Tang等人2024;Yu等人2025)和训练后量化(PTQ)(Zhou等人2024;Zhu等人2024b)。QAT需要额外的GPU资源和训练工作来使模型适应低位宽。PTQ对于大语言模型(LLM)具有吸引力,因为它无需进行训练。作者的工作属于PTQ,并最小化了GPU需求。yinqing-Quamba2_2503为基于SSM的语言模型提供了W8A8、W4A8和W4A16的位宽配置,在所有目标平台上实现了通用的内存和延迟降低。
Transformer的PTQ和权重重排。预训练后量化(PTQ)技术通常分为两类:仅权重量化(例如W4A16)和权重-激活量化(例如W8A8)(Zhu等人 2024b)。仅权重量化(Frantar等人 2023;Lin等人 2024a)最小化权重存储,而权重-激活量化(Ashkboos等人 2024b;Zhao等人 2024b)通过低比特宽操作优化吞吐量。重排权重常用于提升Transformer的量化精度(Yuan等人 2024;Zhao等人 2024b)或效率(Lin等人 2024b),但其应用及其在状态空间模型(SSMs)中的后续有效性尚不明确。作者的研究表明,选择性状态空间对偶(SD)计算(Dao和Gu 2024)保留了输入和输出之间的通道顺序,激活的通道和状态随时间保持一致。
PTQ针对ssMs. Xu等人(2025)和Chiang等人(2025)的研究表明,标准Transformer量化技术对SSMs无效,并提出了专为SSMs设计的PTQ算法。尽管如此,这些策略并未提供多种位宽配置(Chiang等人2025),且在W4A8等低位宽下表现不佳(Xu等人2025)。yinqing-Quamba2_2503为Mamba1(Gu和Dao 2024)和Mamba2(Dao和
提供了W8A8、W4A8和W4A16,实现了实际速度提升和内存减少,满足了云端和边缘端SSM部署日益增长的需求。此外,Zhao等人(2024a)的研究显示4位模型会丧失泛化能力,Kumar等人(2025)指出在内存限制下6-8位宽性能最佳,而4位宽表现较差。作者进一步在大型挑战性多任务数据集MMLU(Hendrycks等人2020)上评估Quamba2,以验证yinqing-Quamba2_2503的鲁棒性。
unsetunset3 背景unsetunset
3.1 模型量化
符号说明。作者遵循Chiang等人(2025)中的符号表示。作者使用
表示浮点矩阵,使用
表示其量化矩阵及其浮点缩放因子
。对于算子,作者使用
表示函数
的量化版本(即,在函数
中权重被量化)量化。作者关注对称均匀量化,以离散
位有符号整数(即INT8或INT4)来近似浮点权重和激活,由于其硬件兼容性。一般的对称均匀量化函数定义为
其中
表示量化后的权重或激活值,
是浮点数输入矩阵,
是缩放因子(即量化步长),该因子由目标位宽
在作者的设置中)确定。静态缩放因子
在推理过程中预先校准并固定。
3.2 选择性状态空间模型
选择性SSM(Dao和Gu 2024;Gu和Dao 2024)将时不变SSM(Gu等人 2020)转换为时变系统。系统动力学由定义。
是输入相关的。
和
是
和
的离散参数。给定输入相关的
时,
和
的离散化函数定义为
是可训练参数,
是可选的残差参数。一个可选的残差分支
被应用于SSM输出,使得在输出投影之前执行
。作者遵循 Dao 和 Gu (2024) 的方法,将时间步
的选择性 SSM 计算抽象为函数。
在函数中省略了可选的
和
。作者省略下标
以表示对整个序列的计算。
SSM模块如图5所示
Mamba1. Gu和Dao(2024)提出了选择性状态空间模型(SSMs),其中参数
和
随输入变化(即时变),允许模型根据输入内容选择性地优先处理或忽略输入。与输入
的交互被定义为
,其中
和
是线性变换,将
映射到
和
。函数
包含两个连续的投影层,表示为
。
通过在时间步
的投影层从模块输入
计算得到。
Mamba2. Dao和Gu(2024)建立了选择性状态空间模型(SSMs)与自注意力机制之间的理论联系——结构化状态空间对偶(SSD)。他们还引入了一种高效的算法,该算法利用现代硬件上的矩阵乘法单元执行线性递归计算。Mamba2通过移除在
和
并行生成时使用的连续线性层来简化模块设计,通过单个投影层生成
,其中
是时间步
的模块输入。改进后的模块设计更适合在大型模型中实现张量并行(Shoeybi等人,2019)。
3.3 选择性SSMs的量化
SSM输入参数。方程3中定义的SSM以
的形式接收输入。近期研究(Chiang等人2025;Xu等人2025)表明,由于Mamba1中的线性递归机制(Gu和Dao 2024),SSM模块对
中的量化引起的误差极为敏感。作者的研究表明,这种现象在Mamba2(Dao和Gu 2024)中依然存在。为解决此问题,作者提出采用排序和聚类方法对输入
进行8位量化。yinqing-Quamba2_2503将跨头的通道按相同值范围分组,以在组内创建更平滑的景观,从而提高量化精度。
SSM异常值。先前关于Transformer的研究(Dettmers等人2022;Xiao等人2023)已检测到通道持续性异常值。一种常见的异常值消除方法是应用Hadamard变换(Ashkboos等人2024b;Liu等人2024c)。在SSM量化(Chiang等人2025;Xu等人2025)中,在线Hadamard矩阵将输入输出投影转换到更平滑的空间,提高了量化精度。尽管快速沃尔什-哈达玛变换(FWHT)可以以nlogn的复杂度并行执行(Dao 2024b;Sloane 1999),作者遵循Xu等人(2025)和Chiang等人(2025)的方法对输出投影输入进行量化,旨在最小化在线Hadamard变换的开销。
unsetunset4 提出方法:Quamba2unsetunset
4.1 SsM参数量化
yinqing-Quamba2_2503基于SSM激活中的两个发现:通道持久性和状态持久性,以及SSM的计算特性:通道顺序保持。符号遵循方程3的定义。
排序与聚类。作者观察到SSM输入
和输出
中通道幅度的持续性和通道顺序的保持,如图2所示。尽管
对Mamba2中的量化引起的误差敏感(Dao和Gu 2024),但Chiang等人(2025)的发现仍然适用,Chiang等人(2025)却忽略了SSM通道的持续特性与顺序保持特性。相比之下,作者利用这两种特性首先对 Head 通道进行排序,并将 Head 和通道进行分组。具体来说,作者首先从校准数据集中获取通道最大值。在图3(a)中,作者展示了按Mamba2-8B最后一块的离线校准通道最大值排序的
。
在在线
-token样本中仍然保持排序输入。排序后的
解耦了 Head 嵌入,从而实现 Head 分组。图4(c1-c2)显示具有相似特征的 Head 被紧密分组,导致使用无监督聚类分为
组。对于每一组 Head ,作者再次应用聚类算法将通道分组为
组。为每一组计算缩放因子,最终得到总共
个缩放因子,这些缩放因子随后用于将
量化为8位精度。详细的排序与聚类过程如图4所示。作者发现
和
在所有实验中都提供了足够好的结果。图5中的
指的是应用排序与聚类后的激活。
按状态组量化。Dao和Gu(2024)放宽了状态组的大小,并引入了一个多输入状态空间模型(SSM),其中
和
矩阵在输入
的所有通道中共享,类似于Transformer中的分组 Query 注意力(Ainslie等人,2023)。作者的研究发现,激活状态(数值较大的状态)在时间步
和输入样本中是相同的。在图3(c-f)中,作者可视化了Mamba2-8B最后一个模块中
和
的激活分布。
和
中的组数设置为8,每个组有128个通道。图3(c-d)显示只有少数几个组被激活且数值较大。例如,在图3(e-f)中,
中的第六组大部分被激活,而
和
中的第七组变化极小。因此,作者对
和
应用了按状态组量化,每个组使用一个单独的缩放因子。图5中的
和
指的是应用按状态组量化后的激活值。按状态组量化显著提高了数值范围较小的组中的量化精度,例如
和
中的第七组。作者证明,按状态组量化是减少Mamba2-8B与FP16模型性能差距的关键。
4.2 系统和框架设计
基于集群的权重重排。在排序和集群阶段,作者创建了一个新的通道和 Head 序列,将同一集群内的 Head 分组,并按照预先校准的最大值对它们的通道进行排序。为了生成具有排序和集群顺序的激活值,作者使用集群和排序索引对离线的输入投影、因果卷积、归一化和输出投影进行重排。输入投影权重的输出列和因果卷积权重的通道被重排。由于ssD计算保持通道顺序(见图2右侧),作者重排归一化权重并应用融合Hadamard量化。最后,使用相同的索引重新排列输出投影的输入行,以保持输出相同。离线基于集群的权重重排如图5所示。
离线Hadamard矩阵融合。Hadamard矩阵具有计算特性
,其中
表示
维方阵。因此,作者将Hadamard矩阵离线融合到输入和输出线性投影中。对于输出投影,Hadamard矩阵在权重矩阵的两边进行乘法运算,使得
。作者在输入投影权重的输入侧融合一个Hadamard矩阵,使得
。因此,将输入/输出投影中的Hadamard矩阵与在线Hadamard量化相结合,可以实现计算不变性(Ashkboos等人 2024a,b),从而产生相同的块输出。离线Hadamard矩阵融合如图5所示。作者在矩阵融合后对权重应用4位/8位量化。
高效的4位/8位Mamba模块。yinqing-Quamba2_2503支持W8A8、W4A8和W4A16投影核,W8A8因果卷积核,4位和8位嵌入核,以及8位选择性扫描和sD核。对于投影层,作者重新排序权重及其每组的缩放因子(Frantar等人2024;Lin等人2024b;Zhang等人2024),以最大化张量核的加载吞吐量。输出缩放因子与输入缩放因子融合,使得
,其中
。作者为Mamba1模块实现了W4A8和W4A16融合的矩阵乘-转置核。对于序列变换,作者加载8位激活值和8位缓存状态以减少内存压力,从而提高延迟,如表3所示。在正向Hadamard变换中,缩放因子
被集成,使得
,从而避免量化过程中的额外计算负载。yinqing-Quamba2_2503的高效核提供了通用的速度提升和内存减少,满足了在云和边缘部署SSM日益增长的需求。
从头到尾量化。从嵌入层到输出头(即从头到尾量化)能够带来额外的内存和延迟降低,这在内存容量有限的边缘计算平台上是必要的。如图1所示,作者的从头到尾(H2T)量化使得Mamba2-8B能够在Nano 8G上部署。具体而言,作者对嵌入层采用每token量化,对头的权重采用每组量化。如表2所示,作者实现了CUDA Kernel ,并支持4位/8位嵌入层和4位/8位输出头。因此,yinqing-Quamba2_2503实现了通用的
内存降低。
通过W4AX混合提升鲁棒性。Zhao等人(2024a)证明,将W4A4应用于所有模块会损害Transformer的泛化能力。作者将此类分析扩展至验证SSM在MMLU(Hendrycks等人2020)数据集上的鲁棒性和泛化能力。作者的研究发现,虽然完整的W4A8量化最大化了预填充加速,但它存在显著的泛化差距(在MMLU上为
,在LAMBADA上为
)。相比之下,完整的W4A16量化展示了鲁棒性,但代价是预填充延迟的增加。为此,作者在框架中引入了混合精度支持。作者根据模块的性能敏感性自动搜索显著模块,并将其分配更高的精度。作者的W4A{8/16}-mixed SSM在MMLU上实现了
的精度提升,同时预填充延迟仅增加了
。
unsetunset5 实验unsetunset
5.1 实验设置
作者在附录C中提供了框架设计细节。
评估。作者使用LM-EVAL(Gao等人,2023)评估Quamba2和 Baseline 模型在六个零样本下游任务上的表现:LAMBADA(Paperno等人,2016)、HellaSwag(Zellers等人,2019)、PIQA(Bisk等人,2020)、ARC(Clark等人,2018)和WinoGrande(Sakaguchi等人,2020),并在每个表中展示五次运行的平均准确率。为了与MambaQuant(Xu等人,2025)进行比较,作者跨五个数据集(ARC-easy、ARC-challenge、PIQA、WinoGrande和HellaSwag)计算平均准确率。完整评估在附录A中,作者遵循Mamba1(Gu和Dao,2024)中的评估协议,报告LAMBADA、WinoGrande、PIQA和ARC-easy的准确率,以及HellaSwag和ARC-challenge按序列长度归一化的准确率。为了展示泛化能力和鲁棒性,作者评估了8B模型在MMLU(Hendrycks等人,2020)上的表现,这是一个包含多个领域选择题的大型多任务测试。
Baseline 方法。在作者的W8A8设置中,作者将yinqing-Quamba2_2503与最新的SSM量化方法MambaQuant(Xu等人,2025)(W8A8,W4A8)和Quamba(Chiang等人,2025)(W8A8)在零样本下游任务上进行比较。在Quamba设置(Chiang等人,2025)中,作者将Hadamard变换应用于输出投影输入,并在输入SSM上实现了百分位数裁剪,从而建立了作者的W8A8 Mamba2 Baseline ,用于延迟和精度。作者还提供了W4A8和W4A16的延迟数据。
5.2 延迟与模型大小
作者在A5ooo上测试了所有方法用于云应用,并在Orin Nano 8G上测试用于边缘应用。对于批量大小为1,作者测量了每个输出token的时间(TPOT)和首个token生成时间(TTFT),单位为毫秒(ms)。TTFT使用1024个输入token进行性能分析。结果如表5和图1所示。在W8A8设置下,yinqing-Quamba2_2503的端到端量化将Mamba2-8B的TPOT延迟降低了1.80倍(22.73 ms vs. 12.61 ms),优于Quamba 1.61倍(Chiang et al. 2025)(22.73 ms vs. 14.12 ms)。在W4A8配置下,Quamba2在A5000上为Mamba2-8B实现了内存使用减少3.89倍、prefilling提升1.39倍、生成速度提升3.05倍。与W8A8和FP16相比,W4A8和W4A16由于去量化开销导致TTFT变慢。然而,4位权重在内存受限的生成阶段带来了延迟优势。yinqing-Quamba2_2503使得Mamba2-8B能够在Nano 8G上部署,生成速度达到每秒13个token,而FP16和W8A8无法实现,如图1和表5所示。对于SSD Kernel ,作者加载8位激活值(
)以降低内存压力,并将延迟提升1.18倍,如表3所示。
5.3 下游任务的零样本评估
作者展示了Quamba2在五个数据集上的平均准确率:ARC-easy、ARC-challenge、PIQA、WinoGrande和HellaSwag,从而可以与MambaQuant(Xu等人,2025)进行公平比较。完整的评估在附录中,作者遵循Mamba1(Gu和Dao,2024)中的评估协议。与Quamba(Chiang等人,2025)相比,当应用于Mamba1时,yinqing-Quamba2_2503利用Hadamard变换对输入和输出投影进行量化,以提高量化精度,从而提升Mamba1的准确率。如表6所示,作者的技术在Mamba2(Dao和Gu,2024)中通过排序聚类和按状态组量化超越了截断。yinqing-Quamba2_2503执行端到端量化,在延迟和内存使用方面优于Quamba(参考表5和表4),适用于W8A8 Mamba1和Mamba2。Quamba2在W4A8 Mamba1中也优于MambaQuant,并在计算平台上实现了实际加速。此外,yinqing-Quamba2_2503支持Mamba1和Mamba2的W8A8、W4A8和W4A16精度,并具有令人满意的准确率和延迟。
5.4 在大型多任务数据集上的评估
作者在MMLU数据集(Hendrycks等人,2020年)中评估了W4A16和W4A8 Quamba2-8B,这是一个包含57个不同难度 Level 学科范围的大型多任务数据集。作者的研究表明,先前的量化方法可能忽视了低比特宽度模型的泛化能力。W4A8在预填充和生成速度之间取得了平衡,但在MMLU泛化方面表现不足,而W4A16尽管预填充延迟增加,仍保持了更好的泛化能力,如表7所示。作者手工设计了两个混合精度模型,分别用W4A16替换最后14层和前14层,在表中分别 Token 为HC-last和HC-first。然而,它们在MMLU数据集上的改进微乎其微。为此,作者采用进化搜索方法来识别敏感层,并将W4A16分配给这些模块。由此产生的混合精度模型在MMLU数据集上减轻了泛化能力的损失(+2.9%),优于手工设计的混合精度和纯W4A8模型,预填充延迟仅增加了10%。
unsetunset6 消融实验unsetunset
6.1 W4A8的消融研究
作者在表8中对W4A8 Quamba2-8B进行了消融研究。在W4A8设置下,对输出投影的输入应用Hadamard变换至关重要。如果没有应用Hadamard变换,模型将无法正常工作。然而,由于SsM对量化引起的误差非常敏感,即使在对Hadamard变换应用按组量化和高阶信息GPTQ(Frantar等人,2023年)(第二阶信息),结果仍然不尽如人意。yinqing-Quamba2_2503——按状态组量化(PerSG)和排序与聚类(SnC)——通过以8位精度量化
、
和
,在最小化精度损失的同时解决了SsM中的这个问题。值得注意的是,
在SsM中仍然容易受到量化误差的影响,这与Chiang等人2025年的发现一致。作者的排序与聚类技术在该问题上优于裁剪(参考表6和11)。
6.2 W4A16的消融研究
作者研究了W4A16 Quamba2-8B情况下每个组件的影响(仅权重量化,即
{\overline{{W}}}X}
),并将结果展示在表9中。该表表明,Hadamard变换与每组分量化相结合(
)的精度高于GPTQ(Frantar等人,2023年)(
)。作者的分析表明,在输出投影的输入中使用Hadamard变换对于缩小SSM的仅权重量化性能差距至关重要。具体而言,Hadamard变换消除了半精度激活中的异常值,从而避免了输出投影中由大异常值放大的4位权重量化误差,使得
。通过结合所有方法(
),W4A16模型在LAMBADA数据集上缩小了半精度性能差距。
6.3 嵌入和输出头的量化
在表10中,作者对嵌入层、输出头以及W4A8模块进行了量化分析。由于所有层的权重均以4位表示,半精度嵌入层和输出头成为内存 Bottleneck ,阻碍了W4A8模型在边缘设备上的部署(参见图1 W4A8)。
因此,作者针对嵌入层、输出头以及W4A8模块进行了量化实验,并将结果展示在表10中。结果表明,较大模型对嵌入层和输出头的量化具有更强的鲁棒性,因为在LAMBADA数据集上的精度几乎保持不变。这一发现对于在内存有限的设备上部署大型模型特别有用。yinqing-Quamba2_2503为嵌入层和输出头提供了不同的位宽配置(即4位和8位),以满足在边缘设备上部署大型模型的需求。
unsetunset参考unsetunset
[1]. Quamba2: A Robust and Scalable Post-training Quantization. Framework for Selective State Space Models
点击上方卡片,关注「AI视界引擎」公众号