Beyond Activation Alignment:The Alignment-Diversity Tradeoff in Task-Aware LLM Quantization

2026-07-01

AI Summary

Existing perplexity-based sensitivity analyses struggle to accurately identify critical layers in large language models for complex reasoning tasks, and calibration using data from a single task often degrades quantization performance. To address these limitations, this work proposes the TASA framework, which jointly optimizes the composition of calibration data and mixed-precision bit-width allocation. TASA uncovers a trade-off between alignment and diversity in calibration data, introduces a training-free gradient trace alignment criterion, and integrates both perplexity and reasoning-oriented sensitivity signals to guide quantization. Experiments demonstrate that 3.5-bit TASA significantly outperforms multiple 4-bit baselines on LLaMA-3-8B and Qwen2.5-7B, achieving over a 20-percentage-point improvement on GSM8K compared to the strongest W3 baseline.

Abstract

Mixed-precision quantization (MPQ) has become a key technique for deploying large language models under stringent memory and compute constraints. We first identify a phenomenon that we term the Perplexity Illusion: layers ranked as important by perplexity-based sensitivity show little rank correlation with those that are most influential for complex reasoning performance, with Kendall $τ\approx 0$ in our analysis. We further reveal an Alignment-Diversity Tradeoff: using only target-task calibration data can degrade post-quantization performance, whereas incorporating general-domain data stabilizes sensitivity estimation and improves robustness across tasks. Based on these observations, we propose TASA (Task-Aware Sensitivity Analysis), a two-level framework that jointly optimizes calibration-data composition and mixed-precision bit allocation. Specifically, TASA searches for a calibration-data mixture using a training-free gradient-trace alignment criterion, and then aggregates perplexity and reasoning-oriented sensitivity signals to guide both inter-layer and intra-layer bit allocation. Experiments on LLaMA-3-8B and Qwen2.5-7B reveal a precision inversion: appropriately allocated 3.5-bit models can match or surpass less task-aware 4-bit baselines. At an average precision of 3.5 bits, TASA matches or outperforms several competitive 4-bit uniform baselines in aggregate accuracy, and improves over the strongest W3 baseline on GSM8K by more than 20 absolute points on LLaMA-3-8B. These results show that calibration-data composition substantially affects task-sensitive quantization, a factor underexplored in prior work.

Problem

• mixed-precision quantization
• task-aware quantization
• sensitivity analysis
• calibration data
• large language models

Innovation

• mixed-precision quantization
• calibration-data composition
• task-aware sensitivity
• alignment-diversity tradeoff
• gradient-trace alignment