Target-Aware Calibration Data Selection for Preserving Uncertainty in Quantized Language Models

AI summary

Overview

Research area: Model compression for large language models, specifically calibration-data selection for post-training quantization, with a focus on preserving the uncertainty behavior (confidence, margins, abstention/answerability decisions) of a full-precision model rather than only its top-1 accuracy.

Technical level: Intermediate. The paper includes formal risk definitions and two propositions, but the central idea — pick calibration examples that look like the decisions you care about — is explained in accessible terms.

One-sentence scope: The paper reframes calibration-data selection for quantization as a target-dependent uncertainty-preservation problem and introduces DPQ, a recipe family that mixes high-doubt examples with generic anchors, evaluated across 8 language models, 9 NLP benchmarks, and 22 comparison methods.

What This Paper Is About

Quantization lets large language models run on smaller hardware, but most compression evaluations only check whether the top-1 answer survives. In many deployments the model's scores matter too: confidence, the margin between the top two options, and whether the model decides to answer at all. The paper asks which calibration examples should be fed to a quantizer so that the quantized model keeps the same uncertainty behavior as the full-precision model, and shows that the best answer depends on which behavior a deployment needs to preserve.

Key Contributions

1. A reformulation of the problem. The authors formulate calibration-data selection for quantization as preserving full-precision uncertainty behavior, defining two distinct risks: a distributional risk (change in the option-probability distribution, measured with JSD) and a boundary risk (flips in the answer/no-answer or low-margin decision).

2. A theoretical argument for target dependence. They provide a mixture-mismatch argument (Proposition 2) showing that if two deployment targets weight boundary versus generic examples differently, no single boundary ratio minimizes mismatch for both, plus a boundary-fragility proposition (Proposition 1) explaining why small-margin examples are the ones that flip under quantization perturbation.

3. The DPQ method. Doubt-Preserving Quantization is a lightweight pre-quantization recipe family that scores candidates with the full-precision model, selects the top candidates by a doubt score, mixes them with generic anchors, and runs the unchanged quantizer. It alters only the calibration strings, not the quantizer kernel, bit width, reconstruction objective, or inference path.

4. A large target-dependent evaluation. Across 8 language models, 9 NLP benchmarks, and 22 comparison methods, the authors show empirically that the leading fixed recipe changes with the preservation target.

Main Findings

• Quantization frequently degrades calibration-sensitive metrics, not just accuracy. Table 1 reports the percentage of model–dataset settings worse than full precision. For BNB-NF4 this is 84.7 for accuracy, 66.7 for ECE, 84.7 for NLL, and 81.9 for Brier score; for GPTQ-WikiText it is 76.4, 58.3, 63.9, and 83.3 respectively. For DPQ-s128-r50 the figures are 79.2, 58.3, 65.3, and 80.6.

• SQuAD2 answerability-boundary preservation favors a high boundary ratio. DPQ-s128-r75 achieves the best aggregate rank of 6.20, FP agreement of 0.8495, boundary-accuracy deviation of 0.2125, answerability-rate deviation of 0.1000, and JSD of 0.0158. The corresponding GPTQ-WikiText numbers are rank 17.29, agreement 0.7446, boundary-accuracy deviation 0.4553, answerability-rate deviation 0.2271, and JSD 0.0387.

• Broad MCQA preservation favors milder or single-signal recipes. On the six answerable extra-MCQA datasets, DPQ-confidence-only ranks first at 8.96 (Top-5 count 7, agreement 0.8357, JSD 0.0284), followed by GPTQ-TaskRandom at 9.12, HighNLL-QA at 9.18, DPQ-entropy-only at 9.30, and Uncertainty-only at 9.31. DPQ-s128-r75 falls to rank 11.88 (agreement 0.8180, JSD 0.0324), while BNB-NF4 has the strongest agreement (0.8637) and JSD (0.0241) but a weaker aggregate rank of 10.53 — evidence that the two targets trade off rather than sharing one winner.

• Composition matters more than budget. Ablations show a non-monotonic ratio effect: r0 and r100 both trail r75, and scaling the calibration size to 64 or 256 examples does not match the 128-example mixed set. Boundary-only and boundary-random variants fall below r75, indicating that generic anchors stabilize the boundary-focused selection.

• Doubt is not the same as difficulty. Negative controls (HighNLL-QA and LowDoubt-QA) trail r75 by at least 0.09 on boundary-accuracy deviation on SQuAD2, yet are competitive on broad MCQA, matching the paper's distinction between hard examples and genuinely uncertain ones.

• The core SQuAD2 result survives removing a stress case. After dropping the Llama-3.2-1B model, DPQ-s128-r75 remains first with rank 6.26 and FP agreement rising to 0.9360.

• Post-hoc calibration optimizes a different objective. Adaptive temperature leaves accuracy (0.6833) and FP agreement (0.8245) unchanged but increases JSD from 0.0320 to 0.0445 and margin drift from 0.1533 to 0.2385. Flexible score-space calibrators (Vector, Matrix, Dirichlet, Isotonic) improve accuracy to roughly 0.699–0.703 while reducing FP agreement to 0.771–0.777.

• Partial transfer to AWQ. On 7B/8B-scale models, DPQ-r75 leads on AWQ mean rank (1.97), with 42.9% top-1 and 73.0% top-2 frequency, but per-dataset breakdowns show score-based metrics can still favor generic calibration on some datasets.

Methodology in Plain English

The authors start from how GPTQ-style quantization works: it estimates layer-wise activation statistics from a small calibration set, so the choice of calibration text determines which activation regions are represented accurately. Their insight is that calibration selection is therefore a way to steer which behaviors survive quantization.

They build a candidate pool of 512 examples from ARC-Challenge training-split questions (which supply answerable QA structure) and 512 SQuAD2 answerability examples (which supply balanced answerable and unanswerable boundary cases). The full-precision model scores every candidate once, under the same option-scoring protocol used at evaluation. From these scores they compute a doubt score, b(x) = 1 − (largest option probability − second-largest option probability), which is high for low-margin examples — exactly the fragile region their Proposition 1 identifies.

Given a calibration budget s and a mixture ratio r, DPQ takes the top ⌊s·r⌋ candidates by doubt score as boundary strings, each formed by concatenating the original prompt with the full-precision top option, and fills the remaining budget with generic anchors from WikiText-style text or RandomQA. The combined set is then passed to an otherwise unmodified GPTQ run. Alternatives replace the doubt score with confidence-only, entropy-only, answerability log-odds, or gold-label NLL to isolate which signal matters.

Evaluation covers three model families: Qwen2.5-0.5B/1.5B/3B/7B, Llama-3.2-1B/3B and Llama-3.1-8B, and Mistral-7B-v0.3. The benchmark suite has 9 datasets split into an "old-core" set (ARC-Challenge, SQuAD2 answerability, TruthfulQA) and an "extra-MCQA" set (ARC-Easy, BoolQ, PIQA, HellaSwag, OpenBookQA, CommonsenseQA). Ranking uses target-aligned metrics: for SQuAD2, boundary-accuracy deviation, answerability-rate deviation, FP agreement, and JSD; for broad MCQA, FP agreement, JSD, margin drift, confidence shift, accuracy deviation, and margin correlation, equally weighted as a neutral summary.

Why This Matters

Impact on research. Most compression work evaluates quantization by accuracy, perplexity, or benchmark score. This paper argues that a model can keep its top-1 answer while flipping its confidence from a 0.51/0.49 split to 0.95/0.05, which is invisible to accuracy checks but decisive for any downstream system that consumes scores. It adds two formal risks and a mismatch argument to a question — which calibration data to use — that is usually treated as an implementation detail rather than a design choice.

Real-world applications:

• Answerability and abstention systems. In SQuAD2-style QA, a small score shift determines whether the system answers an unanswerable question, making boundary preservation a safety-relevant property rather than a cosmetic one.
• Selective prediction and triage. Systems that defer to a human on low-confidence cases depend on calibrated margins; a quantized model that looks more confident than it is will defer less often than intended.
• Reranking, ensembling, and verification pipelines. These use option probabilities as inputs, so preserving the full-precision distribution matters more than preserving the argmax.
• Safety filtering and high-stakes assistants. The paper's deployment-risk discussion names abstention, triage, educational QA, medical or legal assistance, and safety filtering as settings where shifted confidence may cause a system to answer when it should defer.

Industry relevance. DPQ changes only calibration strings, adds no inference-time cost, and leaves the quantizer untouched, so it fits into existing post-training quantization pipelines as a drop-in selection step. The code is released at https://github.com/xi-xiaoran/DPQ. The finding that the best recipe depends on the target means teams need to pick a calibration recipe against the metric their product actually uses, not against benchmark accuracy alone.

Future Directions

• Move beyond option scoring. The paper operationalizes uncertainty through option probabilities, so it handles multiple-choice QA, selective prediction, abstention, and reranking well. Extending the same preservation view to verbalized confidence and long-form generation, where uncertainty appears in the generated text rather than in option scores, is named as an important direction.
• Test less aligned candidate pools. The current pool is deliberately target-oriented, containing answerable QA examples and answerability-style training cases. The authors call for testing less aligned or domain-shifted pools to see how broadly the selection principles transfer.
• Broaden the quantizer and bit-width sweep. The main experiments use GPTQ because it directly consumes calibration strings; AWQ and BNB-NF4 are included as family checks. A wider sweep is needed to refine deployment-specific recipes, since the paper reports that preservation signals partially transfer while score-based metrics remain quantizer-specific.
• Combine pre- and post-quantization interventions. The post-hoc analysis shows the two stages can optimize different objectives: temperature-style maps preserve the argmax but cannot repair flips, while flexible calibrators change decisions at the cost of moving away from full-precision behavior. How to combine target-aware calibration selection with post-hoc calibration is left open.

Target Audience

Researchers and engineers working on LLM quantization, model compression, and uncertainty quantification will find the target-dependent framing and risk definitions most directly useful. Practitioners deploying quantized models in pipelines that consume confidence scores — abstention, selective prediction, reranking, and safety filtering — benefit from the practical selection guidance and the ablation results showing composition matters more than budget. The paper assumes familiarity with post-training quantization and calibration metrics but builds its method from components that are straightforward to implement, so it is readable by graduate students and applied researchers entering the area.

Authors’ abstract

Quantization is widely used to deploy large language models, but its effect on uncertainty behavior, such as confidence, margins, and abstention, is rarely treated as a primary objective. We frame calibration-data selection for quantization as a target-dependent uncertainty-preservation problem. Different deployments emphasize different regions of the input distribution, yet prior work mainly optimizes accuracy-oriented compression metrics or adjusts scores after quantization. We formalize this goal with distributional and boundary preservation risks, and provide a simple mixture-mismatch argument explaining why no single calibration recipe should be expected to fit all targets. We introduce Doubt-Preserving Quantization (DPQ), a lightweight pre-quantization recipe family that uses full-precision predictions to construct target-aligned calibration mixtures of high-doubt examples and generic anchors. Across 8 language models, 9 NLP benchmarks, and 22 comparison methods, the leading fixed recipe changes with the preservation target: DPQ-r75 leads on SQuAD2 answerability-boundary preservation, while milder or single-signal variants, including DPQ-r50, confidence-only, and entropy-only, better preserve broad multiple-choice QA behavior. These results show that calibration data should be selected for the specific full-precision score behavior a deployment needs to preserve, rather than treated as a fixed quantization detail.