arXiv:2410.13229cs.LGcs.AI2024-10被引 33

提出针对状态空间模型的8位量化方法,显著降低推理延迟且保持高精度。

Quamba: A Post-Training Quantization Recipe for Selective State Space Models

  • 通过抑制输入激活最大值并用哈达玛变换处理输出,实现更精细量化
  • 2.8B参数Mamba模型在Orin Nano上生成延迟降低1.72倍,准确率仅降0.9%
  • 适用于云服务与边缘设备部署,对各类规模的SSM模型均有效

状态空间模型(SSMs)作为大语言模型的替代方案,以恒定内存复杂度实现最优准确率,支持比基于注意力机制的网络更长的上下文长度。尽管其在长序列建模中计算效率更高,但在请求密集型云服务和资源受限的边缘场景中进一步提升效率仍具挑战。模型量化是解决此问题的有效途径,但现有方法对SSMs不适用:因其选择性扫描机制中的特征图敏感且输出激活存在大量异常值。为此,我们提出一种静态8位每张量量化方法,通过抑制输入激活最大值提升量化精度,并利用哈达玛变换在无异常值空间中量化输出激活。该方法使2.8B参数的Mamba模型在Nvidia Orin Nano 8G上获得1.72倍的生成延迟降低,零样本任务平均准确率仅下降0.9%,同时支持硬件加速。实验验证了该方法在云与边缘平台部署各类尺寸SSM模型的有效性与实用性。

原文摘要 · Abstract (English)

State Space Models (SSMs) have emerged as an appealing alternative to Transformers for large language models, achieving state-of-the-art accuracy with constant memory complexity which allows for holding longer context lengths than attention-based networks. The superior computational efficiency of SSMs in long sequence modeling positions them favorably over Transformers in many scenarios. However, improving the efficiency of SSMs on request-intensive cloud-serving and resource-limited edge applications is still a formidable task. SSM quantization is a possible solution to this problem, making SSMs more suitable for wide deployment, while still maintaining their accuracy. Quantization is a common technique to reduce the model size and to utilize the low bit-width acceleration features on modern computing units, yet existing quantization techniques are poorly suited for SSMs. Most notably, SSMs have highly sensitive feature maps within the selective scan mechanism (i.e., linear recurrence) and massive outliers in the output activations which are not present in the output of token-mixing in the self-attention modules. To address this issue, we propose a static 8-bit per-tensor SSM quantization method which suppresses the maximum values of the input activations to the selective SSM for finer quantization precision and quantizes the output activations in an outlier-free space with Hadamard transform. Our 8-bit weight-activation quantized Mamba 2.8B SSM benefits from hardware acceleration and achieves a 1.72x lower generation latency on an Nvidia Orin Nano 8G, with only a 0.9% drop in average accuracy on zero-shot tasks. The experiments demonstrate the effectiveness and practical applicability of our approach for deploying SSM-based models of all sizes on both cloud and edge platforms.

量化状态空间模型推理优化边缘部署

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。