💬 NLP

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

This paper introduces Doubt-Preserving Quantization (DPQ), a framework that selects calibration data based on specific deployment targets to better preserve uncertainty behaviors like confidence and abstention in quantized language models, demonstrating that optimal data selection varies depending on the desired preservation goal rather than following a fixed recipe.

Original authors: Zhen Yang, Sizai Hou, Kaiwen Zheng, Yaofang Liu, Liang He, Yixuan Chen, Kangning Cui

Published 2026-08-24

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhen Yang, Sizai Hou, Kaiwen Zheng, Yaofang Liu, Liang He, Yixuan Chen, Kangning Cui

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are the engines behind many modern artificial intelligence tools, capable of generating text, answering questions, and solving problems that once seemed beyond machines. To run these massive systems on everyday devices like smartphones or laptops, engineers often use a technique called quantization. This process shrinks the model's memory footprint by simplifying the numbers it uses to think, much like compressing a high-resolution photograph into a smaller file size so it loads faster. For years, the success of this compression has been judged almost entirely by whether the model still gets the right answer. If a model that once knew the capital of Australia was Canberra still says Canberra after being shrunk, it was considered a success. However, this narrow focus ignores a crucial layer of how these models work: their confidence. A model does not just output an answer; it also calculates how sure it is of that answer. This internal sense of certainty is vital for safety and reliability, especially when the model needs to decide whether to answer a question at all or admit it does not know.

A team of researchers from institutions including Yale University and the University of Oxford has discovered that the standard way of shrinking these models often breaks this internal sense of confidence, even when the final answers remain correct. They found that while a compressed model might still pick the right city, it might become wildly overconfident or strangely uncertain about its choice, a shift that can cause downstream systems to make dangerous errors. To fix this, the researchers developed a new method for selecting the specific data used to prepare the model for compression. Their work shows that there is no single "perfect" set of data to use for every situation. Instead, the data must be carefully chosen based on exactly what kind of behavior the user needs to preserve, whether that is the ability to spot unanswerable questions or the general reliability of answers across a wide range of topics.

The core of the problem lies in how these models handle uncertainty. In the real world, a helpful assistant should know when to stay silent. If a user asks a question that has no factual answer, the model should recognize the ambiguity and decline to guess, rather than confidently hallucinating a wrong fact. This ability to distinguish between a question it can answer and one it cannot is known as answerability. The researchers observed that standard compression techniques often scramble this delicate boundary. A model might see a question it cannot answer, and while the full-sized version would hesitate, the compressed version might suddenly become extremely confident in a wrong answer. This happens because the compression process introduces small errors that accumulate, pushing the model's internal scores across the threshold where it decides to speak up.

To solve this, the researchers treated the selection of training data not as a generic step, but as a targeted calibration task. They realized that different goals require different data. If the goal is to preserve the model's ability to detect unanswerable questions, the data used to prepare the model must be filled with examples that sit right on the edge of uncertainty—questions that are difficult to answer or where the model is naturally unsure. If the goal is to maintain general reliability across many different types of questions, the data needs to be more balanced, mixing those difficult edge cases with standard, straightforward examples. The researchers tested this idea across eight different language models and nine different benchmarks, comparing their approach against twenty-two other methods. They found that when they specifically selected data containing high levels of doubt and uncertainty, the compressed models were far better at preserving their ability to say "I don't know" when appropriate.

The study also revealed a surprising trade-off. The specific recipe that worked best for preserving the ability to spot unanswerable questions was not the same one that worked best for preserving general accuracy on standard multiple-choice tests. A method that heavily focused on difficult, uncertain examples excelled at protecting the model's decision boundaries, ensuring it didn't confidently answer questions it shouldn't. However, for broader tasks where the model is expected to answer most questions, a milder mix of data, or even data selected just for general confidence levels, performed better. This finding challenges the long-held assumption that there is a single, universal best way to prepare a model for compression. Instead, the researchers showed that the "best" data depends entirely on what the model will be used for in the real world.

The researchers introduced a new framework they call Doubt-Preserving Quantization to put this idea into practice. This method works by first asking the original, full-sized model to score a large pool of potential questions. It then identifies the specific questions where the model is most unsure or where the difference between the top two answers is very small. These "high-doubt" examples are then mixed with standard, generic examples to create a custom dataset for the compression process. By feeding this tailored mix into the compression algorithm, the resulting smaller model retains the original's nuanced sense of confidence. The study demonstrated that this approach significantly reduced the errors in confidence that usually plague compressed models, particularly in scenarios where knowing when to abstain is critical.

This work suggests a fundamental shift in how we think about making artificial intelligence smaller and faster. It is not enough to simply check if the model still gets the right answer; we must also ensure it still knows how sure it is of that answer. The researchers showed that by carefully curating the data used to shrink these models, we can preserve the very behaviors that make them safe and reliable. Whether a model needs to be a cautious medical assistant that refuses to guess on a diagnosis, or a general chatbot that answers a wide variety of questions, the data used to prepare it must be chosen with that specific purpose in mind. The era of one-size-fits-all compression is over; the future lies in targeted calibration that respects the specific uncertainties of the task at hand.

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

1. Problem Statement

Quantization is a standard technique for deploying Large Language Models (LLMs) under memory and compute constraints. However, existing evaluations primarily focus on preserving top-1 accuracy, perplexity, or benchmark scores, often neglecting the preservation of uncertainty behavior (e.g., confidence levels, decision margins, and abstention capabilities).

The paper identifies a critical issue: quantization can induce significant "uncertainty drift." As illustrated in the motivating example, a quantized model may retain the correct top-1 answer while drastically altering the probability distribution (e.g., shifting from a low-margin 51%/49% split to a high-confidence 95%/5% split). This shift compromises downstream tasks that rely on score behavior, such as:

Answerability detection: Determining whether a question is unanswerable.
Selective prediction: Deciding when to abstain from answering.
Reranking and ensembling: Utilizing confidence scores for model combination.

Current approaches to calibration data selection either optimize for reconstruction accuracy (ignoring uncertainty) or apply post-hoc calibration (adjusting scores after quantization), which often fails to repair decision flips or moves the model away from its full-precision (FP) behavior. The authors frame calibration-data selection as a target-dependent uncertainty-preservation problem, arguing that different deployment targets (e.g., broad multiple-choice QA vs. answerability boundary detection) require different calibration strategies.

2. Methodology: Doubt-Preserving Quantization (DPQ)

The authors propose Doubt-Preserving Quantization (DPQ), a lightweight pre-quantization framework that selects calibration data based on the specific uncertainty behavior a deployment needs to preserve.

Theoretical Foundation:

Preservation Risks: The paper defines two risks:
- Distributional Risk ( $R_{dist}$ ): Measures the divergence between quantized and FP option distributions (critical for broad MCQA).
- Boundary Risk ( $R_{bdry}$ ): Measures the frequency of decision flips, particularly for low-margin cases (critical for answerability/abstention).
Mixture-Mismatch Argument: The authors formalize that no single calibration mixture can optimally preserve all targets. They argue that if the target distribution requires a high ratio of boundary examples ( $\theta$ ), a calibration set with a different ratio ( $r$ ) will incur a mismatch proportional to $|\theta - r|$ .

DPQ Algorithm:
DPQ operates as a recipe family that modifies the calibration strings fed to standard quantizers (e.g., GPTQ) without altering the quantizer kernel or inference path.

Candidate Pool: A pool of examples is scored using the full-precision model to obtain option probabilities.
Doubt Scoring: A "doubt score" $b(x)$ is calculated as $1 - (p_{(1)} - p_{(2)})$ , where $p_{(1)}$ and $p_{(2)}$ are the top two option probabilities. High scores indicate low-margin (fragile) examples.
Mixture Construction: The calibration set $D$ $D$ is constructed as a mixture:
- Boundary Examples: The top $\lfloor sr \rfloor$ candidates with the highest doubt scores (or specific answerability boundaries for SQUAD2).
- Generic Anchors: The remaining $s - \lfloor sr \rfloor$ examples are drawn from generic sources (e.g., WikiText, RandomQA) to ensure stability.
- Ratio ( $r$ ): The mixing ratio $r \in [0, 1]$ is tunable based on the target. High $r$ (e.g., 0.75) targets boundary preservation; lower $r$ targets broad distributional preservation.
Quantization: The standard quantizer (e.g., GPTQ) runs on this constructed set $D$ .

3. Key Contributions

Problem Formulation: The paper formalizes calibration-data selection for quantization as preserving full-precision uncertainty behavior, distinguishing between distributional and boundary preservation risks.
Theoretical Insight: It provides a "mixture-mismatch" argument proving that a fixed calibration recipe cannot simultaneously optimize for targets with different boundary-to-generic ratios.
Method (DPQ): Introduction of a target-aware recipe family that uses FP predictions to select high-doubt or boundary-near examples, mixes them with generic anchors, and leaves the quantizer unchanged.
Empirical Evidence: Comprehensive evaluation across 8 language models, 9 NLP benchmarks, and 22 comparison methods demonstrating that the optimal recipe depends entirely on the preservation target.

4. Experimental Results

The evaluation covers 8 LMs (Qwen2.5, Llama-3.1/3.2, Mistral-7B) and compares DPQ variants against baselines like GPTQ (with WikiText, C4, RandomQA), BNB-NF4, and post-hoc calibration methods.

Key Findings:

Target Dependence:
- SQUAD2 (Answerability Boundary): The DPQ-r75 recipe (75% boundary examples, 25% anchors) significantly outperforms all baselines. It reduces boundary-accuracy deviation and JSD compared to standard GPTQ-WikiText. Pure boundary-only sets ( $r=1$ ) or generic sets ( $r=0$ ) perform worse, confirming the need for a mixed strategy.
- Broad MCQA (Answerable Tasks): For tasks like ARC-Challenge and TruthfulQA, where the goal is preserving broad option distributions, milder recipes perform better. Variants like DPQ-r50, confidence-only, and entropy-only outperform the high-boundary DPQ-r75.
Comparison to Baselines:
- Standard calibration data (WikiText, C4) often leads to significant drift in confidence and margins, even when accuracy is preserved.
- Post-hoc calibration (e.g., adaptive temperature) can improve label calibration but often increases distributional drift and reduces agreement with the FP model's original behavior.
- Negative controls (e.g., selecting high-NLL examples) show that "hard" examples are not synonymous with "doubtful" examples; high-NLL selection does not effectively preserve answerability boundaries.
Transferability: While the target-aware principle transfers to other quantizers (e.g., AWQ), the specific optimal ratio and metrics remain quantizer-specific.

5. Significance and Claims

The paper claims that calibration data should not be treated as a fixed quantization detail but rather selected specifically for the full-precision score behavior a deployment needs to preserve.

Practical Implication: In deployment, the choice of calibration recipe must align with the downstream task. If the system relies on answerability detection or abstention, a boundary-heavy mixture (like DPQ-r75) is essential. If the system relies on broad multiple-choice scoring, milder or single-signal recipes are superior.
Limitations: The authors acknowledge that their operationalization of uncertainty is limited to option scoring and answerability. Extending these preservation principles to verbalized confidence and long-form generation remains future work. They also note that their candidate pool is target-oriented, and future work should test the transferability of these principles to domain-shifted pools.

In conclusion, the paper argues that preserving the uncertainty of a model is as critical as preserving its accuracy for many real-world applications, and this requires a shift from generic calibration to target-aware calibration data selection.