You Had One Job: Per-Task Quantization Using LLMs' Hidden Representations

AI summary

Overview

Research area: Efficient large language model inference — specifically post-training quantization (PTQ) combined with analysis of LLM hidden representations. The paper sits at the intersection of model compression, mechanistic interpretability, and systems-level deployment.

Technical level: Intermediate. Readers need familiarity with transformer layers, bit-widths, and basic quantization concepts (group-wise affine quantization, calibration sets), but the paper's central idea is conceptually simple: spend bits where the task needs them.

Scope: The paper introduces Task-Aware Quantization (TAQ), a training-free, weight-only mixed-precision PTQ framework that uses a small set of unlabeled task calibration prompts to assign higher precision to task-relevant transformer layers under a fixed bit budget, evaluated on four open-weight LLMs across knowledge, code, and reasoning benchmarks plus hardware measurements.

What This Paper Is About

Standard post-training quantization methods decide how many bits each layer gets using task-agnostic criteria such as weight reconstruction error, activation outliers, or global salience. But real deployments usually need only a narrow slice of an LLM's abilities (writing code, doing math, answering domain questions), so task-agnostic precision allocation can waste bits on layers that do not matter for the target task while over-compressing layers that do. The paper asks where precision should be spent given a frozen LLM, a small unlabeled calibration set from the target task, and a memory budget, and answers it with TAQ.

Key Contributions

1. Task-aware PTQ formulation. The authors reframe mixed-precision PTQ as a task-conditioned precision-allocation problem: choose a per-layer bit assignment from {4, 8, 16} (where 16 means FP16 retention) to minimize expected task loss subject to a total budget, approximated by scoring layers on a small unlabeled calibration set.

2. Three layer-importance scoring rules. TAQ-IS scores layers by combining activation information (eigenvalue-spectrum entropy of the layer covariance) with activation stability (negative variance), with weight α = β = 0.5. TAQ-KL injects a quantization-like uniform noise into one layer at a time and scores it by expected KL divergence between baseline and perturbed next-token distributions at temperature T = 1. TAQ-O is a label-informed oracle diagnostic that quantizes each layer individually to 4-bit and measures the drop in a task metric.

3. Accuracy–memory–latency gains with hardware validation. TAQ is compared against GPTQ, AWQ, and SliM-LLM on Qwen2.5-3B-Instruct, Qwen2.5-7B-Instruct, Qwen2-7B-Instruct, and Gemma-2-9B-it on MMLU-Pro, CodeMMLU, and TriviaQA, with throughput and latency measured on an NVIDIA A40 using real quantized kernels.

4. Robustness analysis. The paper studies calibration size, bit budget, scoring-signal ablations, structural control baselines (Last25, Uniform 8-bit), mixed-task calibration, cross-task allocation stability, and residual-stream error propagation.

Main Findings

• Task-agnostic baselines degrade sharply, TAQ does not. On CodeMMLU with Qwen2.5-3B, TAQ-KL reaches 46.83% EM at 2.25 GB, compared with GPTQ's 24.66% at 1.93 GB and AWQ's 20.61% at 2.50 GB. On MMLU-Pro with Qwen2.5-7B, TAQ-KL reaches 35.35% EM at 5.95 GB, which the paper describes as outperforming AWQ by nearly 12 points for a 0.76 GB premium over AWQ's 5.19 GB.

• TAQ beats a less-constrained mixed-precision baseline on most pairs. SliM-LLM has more allocation freedom and a higher realized effective bit budget — approximately 7.17, 6.13, 6.14, and 7.39 bits on Qwen2.5-3B, Qwen2.5-7B, Qwen2-7B, and Gemma-2-9B, versus roughly 4.87, 4.60, 4.61, and 5.05 bits for TAQ-IS/TAQ-KL. TAQ-IS/TAQ-KL average 5.00–5.05 bits per transformer layer under their 4/8-bit policy (only the top 25% of layers get 8-bit). Despite this, the best TAQ variant matches or exceeds SliM-LLM in 7 of the 8 MMLU-Pro/CodeMMLU model–task pairs while spending roughly 1–2 fewer effective bits per linear layer. On MMLU-Pro, TAQ wins on every backbone: 33.01 vs 30.86 (Qwen2.5-3B), 35.35 vs 31.01 (Qwen2.5-7B), 39.11 vs 38.48 (Qwen2-7B), 42.48 vs 39.89 (Gemma-2-9B).

• Label-free scores track the label-informed oracle. On CodeMMLU/Qwen2.5-3B, TAQ-KL nearly matches the label-informed TAQ-O while using less memory (46.83% vs 47.41%; 2.25 GB vs 3.20 GB). On MMLU-Pro/Qwen2.5-7B, TAQ-KL exceeds TAQ-O (35.35% vs 32.13% at lower memory). The authors state TAQ-O should be read as a marginal-sensitivity diagnostic, not a theoretical upper bound under joint quantization.

• Small calibration sets suffice, and more bits do not automatically help. Varying calibration size over {64, 128, 256, 512, 1024} changes EM by only about 1.3 points on Qwen2.5-7B/MMLU-Pro. Varying the 8-bit promotion fraction K over {10%, 25%, 50%} keeps performance in a narrow range: K = 25% is best on MMLU-Pro, K = 10% is slightly better on CodeMMLU, and raising the budget to 50% does not improve performance.

• Different scoring signals favor different tasks. With the model, calibration size, and K = 25% fixed, variance-only works best on MMLU-Pro, entropy-only works best on CodeMMLU, and the combined TAQ-IS score is most stable overall. Across scoring families, TAQ-KL leads on reasoning and TAQ-O leads on code.

• Task-aware layer choice beats simple structural choice. On MMLU-Pro, at least one TAQ variant outperforms the Last25 control on every backbone, for example 35.35 vs 31.35 on Qwen2.5-7B and 42.48 vs 40.97 on Gemma-2-9B. Against Uniform 8-bit, TAQ is not always higher in raw EM but is often competitive at substantially lower memory, and on Qwen2.5-7B and Gemma-2-9B the best TAQ variant also exceeds the available Uniform 8-bit result.

• Some quantized TAQ models exceed FP16 accuracy. The authors flag these cases cautiously and hypothesize that task-conditioned allocation preserves precision on task-relevant layers while compressing others, acting as an implicit task-conditioned denoiser or regularizer.

• Mixed-task calibration helps rather than hurts. Using a single TAQ-IS allocation computed on Qwen2.5-7B from 512 mixed prompts drawn from TriviaQA, CodeMMLU, and MMLU-Pro improved results over single-task allocation: CodeMMLU from 47.22 to 49.61 EM, MMLU-Pro from 33.35 to 34.28 EM, and TriviaQA from 12.55 to 17.09 EM. The authors hypothesize a broader calibration pool reduces overfitting.

• Layer selections are largely shared across tasks. Across the 5 × 3 model/dataset grid, 87.5% of layer assignments remain the same across calibration tasks, with cross-task Spearman correlation ρ = 0.757. Layers that do change are mostly near the first and last transformer blocks.

• Task-aligned calibration alone is not the source of the gain. On a broader AWQ-supported suite (Phi-4, Qwen3, Llama-3.1, Qwen2.5, Mistral), AWQ calibrated only on the target task (AWQ-OT) is unstable — it improves Llama-3.1 but degrades Phi-4 and Qwen2.5, and changes Mistral only slightly — while the best quantized results come from TAQ-IS or TAQ-O allocation.

• Hardware gains are real, not just accounting artifacts. On a single NVIDIA A40 with Qwen2.5-7B at batch size 1, TAQ-IS improves throughput from 27.63 to 37.29 tok/s relative to FP16, giving +35% throughput and −26% latency; the gain stays positive at batch size 8 (+11.6%). TAQ-IS also outperforms uniform 4-bit (28.71 tok/s), GPTQ-Int4 (4.40 tok/s), and AWQ-Int4 (16.12 tok/s). With real NF4/Int8 kernels, Qwen2.5-3B realizes 2.20 GB weight size and 2.31 GB peak memory, a 2.6× compression over FP16 — slightly less than the 2.25 GB accounting-based figure, due to backend-specific packing.

• Layer rankings survive upstream quantization mostly well. Comparing clean FP16 rankings to rankings recomputed after quantizing the first eight transformer blocks, Spearman correlations on TriviaQA are ρ = 0.997/0.999/0.990 on Qwen2.5-7B and ρ = 0.967/0.969/0.952 on Gemma-2-9B (TAQ-IS/TAQ-O/TAQ-KL). Qwen2.5-7B is also stable on CodeMMLU (ρ ≈ 0.75–0.80). The weakest case is Gemma-2-9B on CodeMMLU and MMLU-Pro (ρ ≈ 0.38–0.59), suggesting stronger task-dependent inter-layer interactions. Cosine similarity stabilizes around 0.93 after an initial drop, indicating residual-stream perturbations do not grow with depth.

• Compute cost is modest. All experiments ran on single-GPU NVIDIA A40 workers with 48 GB. A typical configuration took roughly 1.5–3 A40 GPU-hours for 3B models and 3–6 for 7B/9B models; hardware measurements about 0.5 A40 GPU-hours per configuration. Total reported compute was approximately 450 A40 GPU-hours, and approximately 750 A40 GPU-hours including preliminary, exploratory, and failed runs.

Methodology in Plain English

The approach rests on two observations the authors verify empirically. First, different tasks produce different layer-wise representation profiles inside an LLM — measuring spectral entropy of hidden states across depth shows the profiles vary by task and by layer. Second, quantization damages hidden representations unevenly across depth.

TAQ exploits this by measuring, for each layer, how much that layer seems to matter for the target task, using only a small set of unlabeled task prompts. TAQ-IS looks at the statistics of the token representations a layer produces: how spread out they are across directions (entropy of the covariance spectrum) and how steady they are (negative variance). TAQ-KL instead pokes each layer with small random noise roughly the size of a 4-bit quantization step and measures how much the model's output distribution shifts, using KL divergence. TAQ-O is a diagnostic that actually quantizes one layer at a time to 4-bit and measures the drop in task accuracy — it requires labels, so it is used only as a reference point, not as a deployable method. The first/last two layers are kept in FP16 for robustness in the oracle setting.

Given these scores, TAQ ranks layers and promotes the top portion to 8-bit while leaving the rest at 4-bit, with all layers quantized using the same group-wise affine weight-only operator with group size G = 128. No gradients are used anywhere: the calibration set is used only to compute scores and fit quantizer parameters. For TAQ-IS, a reservoir of r = 256 token vectors per layer is collected; for TAQ-O, n_sens = 16 held-out calibration examples measure sensitivity. Each benchmark used 512 calibration examples disjoint from 2,048 test examples.

Evaluation compares against GPTQ and AWQ (canonical 4-bit uniform baselines), SliM-LLM (a mixed-precision reference with more freedom and typically a higher realized bit budget), structural controls (Last25, Uniform 8-bit), and the FP16 model as a full-precision reference. Because baselines allocate bits differently, the paper reports Pareto comparisons across the accuracy–memory frontier rather than budget-matched sweeps.

Why This Matters

Impact on research. The paper shifts quantization from a model-centric compression problem to a task-conditioned allocation problem, and connects it to interpretability work showing that task-relevant computation is unevenly distributed across depth. It also provides a rare combination of accuracy benchmarks, memory accounting, and real-kernel throughput/latency on the same set of models, which is more than most PTQ papers report.

Real-world applications.

• Code assistants, where a model fine-tuned or deployed only for code completion can be compressed aggressively on layers that matter for other capabilities. CodeMMLU is the paper's proxy for this setting, and the baseline collapse (GPTQ 24.66% vs TAQ-KL 46.83% EM at a similar footprint on Qwen2.5-3B) is the most dramatic gap reported.
• Mathematical and analytical reasoning services, where the paper reports consistent MMLU-Pro gains over the stronger SliM-LLM baseline on every backbone.
• On-device or edge deployment, where the 2.6× compression over FP16 measured with real NF4/Int8 kernels on Qwen2.5-3B (2.20 GB weights, 2.31 GB peak memory) determines whether a model fits at all.
• Domain-specific question answering, where a single mixed calibration pool from several tasks produced one shared allocation that improved all three benchmarks in the paper's test.

Industry relevance. Serving costs scale with memory footprint, throughput, and latency. The paper reports +35% throughput and −26% latency versus FP16 at batch size 1 on an A40, and shows TAQ-IS beating GPTQ-Int4 and AWQ-Int4 in throughput. TAQ requires no gradients, no fine-tuning, and no labeled data, so it slots into existing PTQ pipelines as an allocation policy rather than a new training procedure — the authors explicitly note it can be combined with existing PTQ backends, and that all methods share the same quantization operator and differ only in layer ranking and allocation.

Future Directions

• Beyond layer-level allocation. The paper states that future work should extend TAQ past layer-level granularity, for example to channels or groups, and to model quantization error more directly rather than through local proxies.

• Interpretability-based error correction. The authors' linearity analysis suggests quantization error can appear as a persistent direction in the residual stream, motivating corrections such as activation steering or task-direction estimation.

• Understanding and mitigating the weak interaction cases. Gemma-2-9B showed markedly lower ranking stability after upstream quantization on CodeMMLU and MMLU-Pro (ρ ≈ 0.38–0.59), indicating task-dependent inter-layer interactions that clean FP16 scores do not capture.

• Distribution shift and joint allocation. The stated limitations are that TAQ depends on calibration prompts representative of the deployment distribution, and that the local scoring rules do not solve the full joint mixed-precision allocation problem, so higher-order interactions among quantized layers may be missed. Both remain open problems, and the authors also call for target-domain safety, privacy, fairness, and robustness evaluation before deployment in high-stakes settings.

Target Audience

This paper suits machine learning engineers and systems researchers who deploy LLMs under memory or latency constraints and want a training-free way to specialize compression to a known workload. It is also relevant to interpretability researchers interested in how task information is distributed across transformer depth, and to PTQ method developers looking for a task-conditioned alternative to global salience or curvature criteria. Readers without background in quantization will need to consult the background section or prior PTQ literature, since the scoring rules assume familiarity with hidden-state statistics, eigenvalue spectra, and KL divergence over output distributions.

Authors’ abstract

Many LLM applications require only narrow capabilities, yet standard post-training quantization (PTQ) methods allocate precision without considering the target task. This can waste bits on layers that are less relevant to the task signal while over-compressing layers that are critical for downstream behavior. We propose Task-Aware Quantization (TAQ), a training-free, weight-only mixed-precision PTQ framework that uses a small set of unlabeled task calibration prompts to allocate higher precision to task-relevant transformer layers under a fixed bit budget. TAQ estimates layer importance from hidden representations and output sensitivity, and we instantiate it with three scoring rules: TAQ-IS, based on activation information and stability; TAQ-KL, based on output-distribution sensitivity under a quantization-noise proxy; and TAQ-O, a label-informed oracle diagnostic for analyzing layer sensitivity. Across several benchmarks, TAQ outperforms task-agnostic baselines such in most settings, with especially strong gains in the accuracy--memory ratio. We further validate that these gains translate to real deployment behavior through hardware throughput and latency measurements, and analyze calibration robustness and residual-stream error propagation. Overall, TAQ turns mixed-precision PTQ from a model-centric compression step into a task-conditioned precision-allocation problem. A reference implementation is available at \href{https://anonymous.4open.science/r/TAQ-9217/README.md}{\includegraphics[height=1em]{imgs/github-mark.png}}.