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

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

    Sep 1, 2026Jainil Dharmil ShahLLM QuantizationLLM Evaluation

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

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

    Apr 10, 2026Shu-Hao Zhang, Le-Tong Huang, Xiang-Sheng Deng +5Quantization-Aware DistillationLLM Quantization

    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

  3. APreQEL: Adaptive Mixed Precision Quantization For Edge LLMs

    Mar 24, 2026Meriem Bouzouad, Yuan-Hao Chang, Jalil BoukhobzaHardware-Aware QuantizationLLM Quantization

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