A KL Lens on Quantization: Fast, Forward-Only Sensitivity for Mixed-Precision SSM-Transformer Models

📄 arXiv: 2604.13440v1 📥 PDF

作者: Jason Kong, Nilesh Prasad Pandey, Flavio Ponzina, Tajana Rosing

分类: cs.LG, cs.AI

发布日期: 2026-04-15

🔗 代码/项目: GITHUB

💡 一句话要点

提出基于KL散度的前向敏感度分析方法，加速混合精度SSM-Transformer模型量化部署。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 量化 混合精度 边缘计算 大型语言模型 SSM-Transformer 敏感度分析 KL散度

📋 核心要点

  1. 边缘设备部署LLM面临计算和内存限制，现有量化方法对不同组件影响不均，需要精细管理。
  2. 提出一种仅依赖前向传递的敏感度分析框架，避免反向传播和重新训练，适用于数据受限场景。
  3. 实验表明，基于KL散度的敏感度分析能有效识别量化敏感组件，并指导混合精度量化，提升边缘设备性能。

📝 摘要（中文）

为了在边缘设备上部署大型语言模型(LLMs)，本文提出了一种轻量级的、无反向传播的、基于代理模型的敏感度分析框架，用于识别混合SSM-Transformer架构中对量化最敏感的组件。该方法仅依赖于前向传递的指标，避免了昂贵的梯度计算和重新训练，适用于因专有或隐私限制而无法访问领域内数据的情况。论文还提供了形式分析，表明Kullback-Leibler (KL)散度比均方误差(MSE)和信噪比(SQNR)等常用替代指标更能捕捉语言建模任务的量化敏感性。通过对SSM和混合架构的实验，消融研究证实了基于KL的排序与观察到的性能下降一致，并且优于其他指标。该框架能够以最小的精度损失在资源受限的边缘设备上实际部署先进的混合模型。最后，在Intel Lunar Lake硬件上的真实设备剖析验证了该方法，表明KL引导的混合精度在CPU和GPU执行模式下，以接近FP16的困惑度实现了与Uniform INT4相当的模型大小和吞吐量。

🔬 方法详解

问题定义：论文旨在解决在资源受限的边缘设备上部署混合SSM-Transformer模型时，如何高效地进行量化，以减小模型大小并加速推理的问题。现有的量化方法，特别是均匀量化，可能对模型不同组件产生不均衡的影响，导致精度显著下降。传统的敏感度分析方法依赖于反向传播和重新训练，计算成本高昂，且在数据受限场景下难以应用。

核心思路：论文的核心思路是利用前向传递的指标，特别是Kullback-Leibler (KL)散度，来评估模型不同组件对量化的敏感程度。KL散度能够更好地捕捉量化对语言建模任务的影响，从而指导混合精度量化策略，在保证精度的前提下，尽可能地减小模型大小和加速推理。这种方法避免了反向传播和重新训练，降低了计算成本，并适用于数据受限场景。

技术框架：该框架主要包含以下几个阶段：1) 前向传递：使用未量化的模型进行前向推理，获取模型各组件的输出。2) 量化模拟：对模型各组件进行量化模拟，例如将权重或激活值量化到INT4或INT8。3) 敏感度评估：使用KL散度等指标，比较量化前后模型输出的差异，评估各组件对量化的敏感程度。4) 混合精度量化：根据敏感度评估结果，对不同组件采用不同的量化精度，例如对敏感组件采用更高的精度，对不敏感组件采用更低的精度。5) 模型部署：将混合精度量化后的模型部署到边缘设备上。

关键创新：该论文最重要的技术创新点在于提出了一种基于KL散度的前向敏感度分析方法，用于指导混合精度量化。与传统的基于梯度或重训练的敏感度分析方法相比，该方法计算成本更低，且适用于数据受限场景。此外，论文还证明了KL散度比MSE和SQNR等常用指标更能捕捉语言建模任务的量化敏感性。

关键设计：论文的关键设计包括：1) 使用KL散度作为量化敏感度的度量指标，并提供了理论分析支持。2) 设计了一种轻量级的前向敏感度分析框架，避免了反向传播和重新训练。3) 针对混合SSM-Transformer模型，提出了具体的混合精度量化策略，例如对SSM层的状态矩阵采用更高的精度，对Transformer层的注意力机制采用更低的精度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，基于KL散度的敏感度分析方法能够有效识别量化敏感组件，并指导混合精度量化，在保证精度的前提下，显著减小模型大小和加速推理。在Intel Lunar Lake硬件上的真实设备剖析验证了该方法，表明KL引导的混合精度在CPU和GPU执行模式下，以接近FP16的困惑度实现了与Uniform INT4相当的模型大小和吞吐量。

🎯 应用场景

该研究成果可广泛应用于边缘设备上部署大型语言模型，例如智能手机、物联网设备和自动驾驶汽车等。通过高效的量化方法，可以在资源受限的设备上实现实时推理和本地智能，提升用户体验和设备功能。该方法还有助于保护用户隐私，因为模型可以在本地运行，而无需将数据上传到云端。

📄 摘要（原文）

Deploying Large Language Models (LLMs) on edge devices faces severe computational and memory constraints, limiting real-time processing and on-device intelligence. Hybrid architectures combining Structured State Space Models (SSMs) with transformer-based LLMs offer a balance of efficiency and performance. Aggressive quantization can drastically cut model size and speed up inference, but its uneven effects on different components require careful management. In this work, we propose a lightweight, backpropagation-free, surrogate-based sensitivity analysis framework to identify hybrid SSM-Transformer components most susceptible to quantization-induced degradation. Relying solely on forward-pass metrics, our method avoids expensive gradient computations and retraining, making it suitable for situations where access to in-domain data is limited due to proprietary restrictions or privacy constraints. We also provide a formal analysis showing that the Kullback-Leibler (KL) divergence metric better captures quantization sensitivity for Language modeling tasks than widely adopted alternatives such as mean squared error (MSE) and signal-to-quantization-noise ratio (SQNR). Through extensive experiments on SSM and hybrid architectures, our ablation studies confirm that KL-based rankings align with observed performance drops and outperform alternative metrics. This framework enables the practical deployment of advanced hybrid models on resource-constrained edge devices with minimal accuracy loss. We further validate our approach with real-world on-device profiling on Intel Lunar Lake hardware, demonstrating that KL-guided mixed-precision achieves near-FP16 perplexity with model sizes and throughput competitive with Uniform INT4 on both CPU and GPU execution modes. Code is available at https://github.com/jasonkongie/kl-ssm-quant.