量化小规模状态空间模型,实现边缘AI低内存高效率部署
Quantizing Small-Scale State-Space Models for Edge AI
- 针对S4D模型,分析量化对状态矩阵和内部状态的影响
- QAT使8位精度下序列MNIST准确率从40%提升至96%
- 提出异构量化策略,内存减少6倍且不损失性能
状态空间模型(SSMs)因其高效建模长距离依赖的能力,在深度学习中备受关注,是边缘AI应用的有力候选。本文聚焦小规模SSM的量化,旨在降低内存与计算开销的同时保持任务性能。以S4D架构为例,首先研究了后训练量化(PTQ),发现状态矩阵A和内部状态x对量化特别敏感。进一步分析了不同量化技术对参数与激活的影响。为解决PTQ导致的性能下降,采用量化感知训练(QAT),在8位精度下将序列MNIST基准上的准确率从40%显著提升至96%。此外,验证了QAT在亚8位精度下的潜力,并评估了多种参数化方案以增强QAT稳定性。最后提出一种异构量化策略,为模型组件分配不同精度,使整体内存占用降低6倍,且性能不受影响。研究结果为资源受限环境中的量化SSM部署提供了可操作的指导。
原文摘要 · Abstract (English)
State-space models (SSMs) have recently gained attention in deep learning for their ability to efficiently model long-range dependencies, making them promising candidates for edge-AI applications. In this paper, we analyze the effects of quantization on small-scale SSMs with a focus on reducing memory and computational costs while maintaining task performance. Using the S4D architecture, we first investigate post-training quantization (PTQ) and show that the state matrix A and internal state x are particularly sensitive to quantization. Furthermore, we analyze the impact of different quantization techniques applied to the parameters and activations in the S4D architecture. To address the observed performance drop after Post-training Quantization (PTQ), we apply Quantization-aware Training (QAT), significantly improving performance from 40% (PTQ) to 96% on the sequential MNIST benchmark at 8-bit precision. We further demonstrate the potential of QAT in enabling sub-8-bit precisions and evaluate different parameterization schemes for QAT stability. Additionally, we propose a heterogeneous quantization strategy that assigns different precision levels to model components, reducing the overall memory footprint by a factor of 6x without sacrificing performance. Our results provide actionable insights for deploying quantized SSMs in resource-constrained environments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。