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

2026-04-14

AI Summary

Deploying large language models on edge devices is constrained by limited computational and memory resources, and components within hybrid SSM-Transformer architectures exhibit markedly different sensitivities to quantization error. This work proposes a gradient-free, forward-pass sensitivity analysis framework that, for the first time, employs Kullback–Leibler (KL) divergence as a metric for quantization sensitivity—demonstrating superior effectiveness over conventional measures such as mean squared error (MSE) and signal-to-quantization-noise ratio (SQNR). The method requires neither backpropagation nor retraining, making it well-suited for data-scarce scenarios. Guided by KL-based sensitivity scores, a mixed-precision quantization strategy achieves perplexity close to FP16 while compressing model size to INT4 levels, significantly boosting CPU/GPU inference throughput on Intel Lunar Lake platforms.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Security and Privacy: Large-scale security measurements

Abstract

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.

Problem

• quantization
• mixed-precision
• SSM-Transformer
• edge deployment
• KL divergence

Innovation

• KL divergence
• mixed-precision quantization
• SSM-Transformer
• forward-only sensitivity
• edge deployment