A KL Lens on Quantization: Fast, Forward-Only Sensitivity for Mixed-Precision SSM-Transformer Models
KL divergence beats MSE for mixed-precision quantization of hybrid SSMs — and it runs on a single forward pass.
The Thesis
This paper argues that KL divergence is a theoretically and empirically superior sensitivity metric for deciding which layers to quantize in hybrid SSM-Transformer models, and that you don't need backprop or retraining to figure it out. The payoff: near-FP16 perplexity at INT4 model sizes on real Intel Lunar Lake hardware. For anyone shipping inference at the edge, this is a practical recipe to cut memory without torching accuracy.
Catalyst
Hybrid SSM-Transformer architectures (think Mamba-style state space models bolted onto attention layers) only matured into deployable form in the last 12-18 months, creating a gap in quantization tooling designed specifically for their mixed component types. Intel Lunar Lake's on-device AI silicon gave the authors a real edge target that didn't exist two years ago.
What's New
Prior quantization sensitivity methods relied on MSE or SQNR and required gradient computation or fine-tuning passes — expensive and often impossible under data-privacy constraints. This framework is forward-pass only, backprop-free, and the paper formally shows KL divergence better predicts per-layer sensitivity for language modeling objectives than those incumbents.
The Counter
KL divergence outperforming MSE/SQNR for quantization sensitivity is a reasonable theoretical claim, but the empirical delta matters more than the story — if the perplexity gap between KL-guided and MSE-guided mixed precision is 0.2 points, practitioners will just use whatever their existing toolchain already supports. The paper also validates on Intel Lunar Lake, which has negligible market share in actual edge AI deployments compared to Qualcomm and Apple silicon; shipping a method tuned to one vendor's profiling stack is not the same as solving edge inference broadly. More fundamentally, the hybrid SSM-Transformer architecture space is still in flux — Mamba 2, Jamba, Zamba and their descendants keep changing the component mix, meaning sensitivity profiles derived today could be stale within a product cycle. Finally, forward-pass-only calibration sounds like a feature but is also a limitation: you're trading compute for information, and in scenarios where even a small fine-tuning budget exists, gradient-based methods will likely still dominate on accuracy.
Longs
- INTC
- NVDA
- AMD
- MSFT
Shorts
- Cloud-only LLM inference providers whose moat depends on model size being a barrier to local deployment
- Quantization-as-a-service startups selling expensive gradient-based tooling
Enablers (Picks & Shovels)
- Intel (Lunar Lake NPU/CPU validation target)
- Hugging Face (model hub integration path)
- NVIDIA (GPU execution mode tested in paper)
Private Watchlist
- Mistral AI
- Together AI
- Neural Magic
The Paper
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.