cs.LGApr 15, 2026

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

Authors: Jason Kong, Nilesh Prasad Pandey, Flavio Ponzina, Tajana Rosing

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

Figure 1: Quantization-sensitivity characteristics of Hymba. (a) Average ablation sensitivity Δ‾ ⁣PPL⁡\overline{\Delta}\!\operatorname{PPL} for each projection subtype, with mamba.x_proj dominating the profile. (b) Layer-wise cumulative sensitivity ΣΔ ⁣PPL⁡\Sigma\Delta\!\operatorname{PPL} , showing a sharp peak at block 31.
MetricOrder
Perplexity↓\downarrow
SQNR (dB)↑\uparrow
KL teacher→student{}_{\text{teacher}\to\text{student}}↓\downarrow
KL student→teacher{}_{\text{student}\to\text{teacher}}↓\downarrow
Δ\Delta Cross Entropy↓\downarrow
Table 1: Sort direction for each metric ( ↑\uparrow : ascending, ↓\downarrow : descending).
Figure 2: Quantization-sensitivity characteristics of Hymba.
MetricMamba-130MMamba-380MMamba-1.4BMamba2-130MHymbaZambaAvg.
SQNR (dB)0.69670.69670.73140.73140.79410.79410.48980.48980.71240.71240.84190.84190.71110.7111
KL teacher ⁣→ ⁣student{}_{\mathrm{teacher}\!\to\!\mathrm{student}}−0.1980-0.1980−0.2316-0.2316−0.2607-0.2607−0.0272-0.0272−0.1617-0.16170.11420.1142−0.1275-0.1275
KL student ⁣→ ⁣teacher{}_{\mathrm{student}\!\to\!\mathrm{teacher}}0.8419\mathbf{0.8419}0.7936\mathbf{0.7936}0.8327\mathbf{0.8327}0.8078\mathbf{0.8078}0.6646\mathbf{0.6646}0.8060\mathbf{0.8060}0.7911\mathbf{0.7911}
Δ\Delta CE−0.1400-0.1400−0.1662-0.1662−0.0679-0.06790.23640.2364−0.1260-0.1260−0.1234-0.1234−0.0645-0.0645
Table 2: Kendall’s τ\tau values for each metric across models.
MetricMamba-130MMamba-380MMamba-1.4BMamba2-130MHymbaZamba
SQNR (dB)4.97×10−244.97\times 10^{-24}1.5×10−511.5\times 10^{-51}1.91×10−601.91\times 10^{-60}6.869×10−76.869\times 10^{-7}5.87×10−575.87\times 10^{-57}4.57×10−524.57\times 10^{-52}
KL teacher ⁣→ ⁣student{}_{\mathrm{teacher}\!\to\!\mathrm{student}}4.06×10−34.06\times 10^{-3}1.729×10−61.729\times 10^{-6}7.30×10−87.30\times 10^{-8}7.827×10−17.827\times 10^{-1}3.07×10−43.07\times 10^{-4}3.95×10−23.95\times 10^{-2}
KL student ⁣→ ⁣teacher{}_{\mathrm{student}\!\to\!\mathrm{teacher}}2.51×10−342.51\times 10^{-34}2.29×10−602.29\times 10^{-60}2.84×10−662.84\times 10^{-66}2.637×10−162.637\times 10^{-16}8.31×10−508.31\times 10^{-50}7.1×10−487.1\times 10^{-48}
Δ\Delta CE4.21×10−24.21\times 10^{-2}5.97×10−45.97\times 10^{-4}1.61×10−11.61\times 10^{-1}1.656×10−21.656\times 10^{-2}4.92×10−34.92\times 10^{-3}2.61×10−22.61\times 10^{-2}
Table 3: p-values for each metric across models.
Figure 3: KL-Divergence vs. Perplexity correlation.
ModelConfigPPL ↓\downarrowSizeFPS ↑\uparrowLatency ↓\downarrow
Mamba-130M (CPU)FP1621.59493 MB22.644 ms
INT825.32126 MB35.528 ms
INT4*3e3584 MB39.725 ms
KL-MP p05 (ours)21.61136 MB35.928 ms
Mamba-1.4B (CPU)FP1611.225.2 GB2.6384 ms
INT811.251.3 GB5.1196 ms
Table 4: On-device results: best KL-MP point vs. baselines on Intel Lunar Lake. ↑\uparrow higher is better; ↓\downarrow lower is better. “–” * The 130M model exhibited substantially higher perplexity under uniform INT4 quantization. †FPS and latency are not reported, as INT8 execution is not supported for the certain relevant operations on Intel Lunar Lake
Figure 4: Perplexity vs. model size (CPU, Intel Lunar Lake). The x-axis is inverted so more compressed models appear to the right. KL-MP maintains near-FP16 accuracy across a wide compression range. The sharp cliff at p10 for Mamba-130M (diverged) and p10 for Mamba-1.4B ( PPL=60.6\mathrm{PPL}=60.6 ) reflects the quantization of a single high-sensitivity layer identified in our ablation ( Sec. 5.1 ); Uniform INT4 crosses this boundary by construction.
Figure 5: Inference latency vs. model size (CPU, Intel Lunar Lake). KL-MP p05 achieves the lowest latency for Mamba-1.4B (178 ms), undercutting both Uniform INT8 and INT4, demonstrating that sensitivity-guided precision allocation reduces not only memory but also compute overhead.
Figure 6: Throughput vs. model size (CPU, Intel Lunar Lake). KL-MP p05 matches or exceeds INT4 throughput for both models while preserving near-FP16 perplexity, occupying a region of the accuracy–efficiency space inaccessible to homogeneous quantization.
Figure 7: GPU profiling of Mamba2-130M on Intel Lunar Lake. (Left) Perplexity vs. throughput: KL-MP configurations p01 – p09 achieve up to 14.5×14.5\times throughput improvement over FP16 with negligible perplexity cost. (Right) Perplexity vs. model size: file sizes cluster near 45 MB due to FP16 export compression and INT8-only quantization; only p10 (red triangle) crosses the precision cliff, collapsing to Uniform INT8 perplexity.
More in this topic

Quantization Sensitivity Analysis

See all 20 papers

Explore similar work

CardsList
Sep 1, 2026

Triple-Bottom-Line Sustainability of Language Models for Edge AI: A Comparison Between SLMs and Quantized LLMs

Edge-AI model selection is commonly driven by one isolated metric - accuracy, latency, memory, energy, or safety, even though a deployable language model must balance all five. Our work focuses on answering the question whether na- tively trained small language models (SLMs) or large language models (LLMs) compressed through post-training quantization offer the more sustainable edge- deployment trade-off. We introduce a reproducible Holistic Sustainability Score (HSS) organized around the triple bottom line: an economic pillar for capability and systems efficiency, an environmental pillar for operational GPU energy and a social pillar for harmful-prompt robustness. Five BF16 SLMs and five LLMs under different quantization approaches - BF16, INT8, NF4 4-bit, GPTQ 4-bit, and GGUF Q4 produce 30 measured configurations. Capability is assessed on five zero-shot benchmarks; efficiency uses latency, throughput, peak VRAM and energy; and safety is approximated by attack success rate on five harmful prompts. Qwen3-30B-A3B/GGUF Q4 ranks first in the combined pool (93.38), followed by Mistral-Small-24B/GGUF Q4 (92.40), while Phi-4-mini/BF16 is the highest- ranked SLM in that pool (89.49). Thus, the hypothesis that native SLMs must be the most sustainable edge choice is not supported universally; optimized quantized LLMs can win overall, while SLMs remain competitive through lower resource demand. Quantization is a systems-level choice rather than a monotonic precision- efficiency trade-off and HSS remains relative to its comparison pool and proxy definitions.

Jainil Dharmil Shah

Elmore Family School of Electrical and Computer Engineering · Purdue University, West Lafayette, Indiana, USA

LLM QuantizationLLM EvaluationOn-Device Language Model InferenceSmall Language Models
AskPDF
Apr 10, 2026

EdgeRazor: A Lightweight Framework for Large Language Models via Mixed-Precision Quantization-Aware Distillation

Quantization has emerged as a mainstream approach for deploying Large Language Models (LLMs) on resource-constrained devices, yet compressing precision below 4-bit typically causes severe performance degradation or prohibitive retraining costs. In this paper, we propose EdgeRazor, a lightweight framework for LLMs via Mixed-Precision Quantization-Aware Distillation. It contains three modules: Structural Quantization with Mixed Precision for fine-grained control of bit-widths, Layer-Adaptive Feature Distillation that dynamically selects the most informative features for alignment, and Entropy-Aware KL Divergence for forward-reverse balance on both human-annotated and distilled datasets. Evaluations conducted on MobileLLM and Qwen families show that under weight-activation quantization, the 1.88-bit Qwen3-0.6B-EdgeRazor outperforms the state-of-the-art 2-bit baselines by 11.27 and surpasses the strongest 3-bit baselines by 4.38, while the quantized MobileLLM-350M-EdgeRazor requires a training budget 4-10×\times lower than the leading quantization-aware training method. In terms of efficiency, EdgeRazor achieves higher compression ratios at all bit-widths, and the 1.58-bit Qwen3-0.6B-EdgeRazor reduces storage from 1.11 GB to 0.19 GB while accelerating decoding by 15.16×\times over the 16-bit baseline. These results empirically validate the effectiveness and efficiency of EdgeRazor. The codes can be accessed from GitHub and Huggingface.

Shu-Hao Zhang, Le-Tong Huang, Xiang-Sheng Deng +5

National Key Laboratory for Novel Software Technology, Nanjing University, Nanjing 210063, China · School of Intelligent Science and Technology, Nanjing University, Suzhou 215163, China · Microsoft AI, Beijing 100080, China +1

Quantization-Aware DistillationLLM QuantizationLLM CompressionMixed-Precision QuantizationQuantization-Aware TrainingLLM Inference Acceleration+1
AskPDF
Mar 24, 2026

APreQEL: Adaptive Mixed Precision Quantization For Edge LLMs

Today, large language models have demonstrated their strengths in various tasks ranging from reasoning, code generation, and complex problem solving. However, this advancement comes with a high computational cost and memory requirements, making it challenging to deploy these models on edge devices to ensure real-time responses and data privacy. Quantization is one common approach to reducing memory use, but most methods apply it uniformly across all layers. This does not account for the fact that different layers may respond differently to reduced precision. Importantly, memory consumption and computational throughput are not necessarily aligned, further complicating deployment decisions. This paper proposes an adaptive mixed precision quantization mechanism that balances memory, latency, and accuracy in edge deployment under user-defined priorities. This is achieved by analyzing the layer-wise contribution and by inferring how different quantization types behave across the target hardware platform in order to assign the most suitable quantization type to each layer. This integration ensures that layer importance and the overall performance trade-offs are jointly respected in this design. Our work unlocks new configuration designs that uniform quantization cannot achieve, expanding the solution space to efficiently deploy the LLMs on resource-constrained devices.

Meriem Bouzouad, Yuan-Hao Chang, Jalil Boukhobza

Lab-STICC, CNRS UMR 6285 , ENSTA, Institut Polytechnique de Paris, Brest, France · National Taiwan University, Taipei, Taiwan

Hardware-Aware QuantizationLLM QuantizationNeural Network QuantizationMixed-Precision QuantizationOn-Device Language Model InferenceLLM Inference Acceleration
AskPDF