Source-linked AI summary
Towards Universal Calibration for In-Context Classification
Awik Dhar, Sunita Sarawagi, Arpit Agarwal
TL;DR
ICL’s effect on calibration is less studied than its ability to improve classification without fine-tuning, and existing calibration methods can require task- or model-specific training. CALIBFORMER uses prediction–label patterns in demonstrations to calibrate test confidence; it reduces average ECE on unseen tasks and models, including Gemini. Its evaluation of proprietary models, however, covers only Gemini-2.5-Flash.
Problem
ICL can impair calibration, while existing approaches may require task- or model-specific training that detracts from its no-fine-tuning convenience.
Source: §1 · p. 1Method
CALIBFORMER is a transformer calibrator that treats in-context examples as a latent validation set and uses prediction–label discrepancies to predict calibrated test confidence.
Source: §1 · p. 2Results
41–73% average ECE reduction relative to uncalibrated ICL; CALIBFORMER generalizes across held-out tasks and four unseen LLMs.
Source: §1 · p. 2Takeaways & Limitations
CALIBFORMER transfers to substantially different and black-box models, including Gemini-2.5-flash, where only token-level log-probabilities are available.
Source: §1 · p. 2Evaluation on proprietary models is limited to Gemini-2.5-Flash, constrained by inference costs and the limited log-probabilities exposed by APIs.
Source: §6 · p. 9
Abstract
from arXiv · showhide
In-Context Learning (ICL) has received a lot of interest recently due to its ability to improve classification accuracy for new tasks without having to fine-tune the model over these tasks. However, the impact of ICL on calibration has been far less studied. Recent work has shown that ICL can have a negative impact on calibration due to biased selection of in-context examples. The existing methods for improving ICL calibration require additional task and model-specific training, detracting from the no-fine-tuning convenience of ICL. In this work, we introduce CALIBFORMER, a universal ICL calibrator for classification tasks, that generalizes to new models and tasks without further fine-tuning on these tasks. Specifically, we design a transformer-based calibrator which uses the model predictions on in-context examples to predict a confidence score for the test input. We show that the calibrator trained on three models and ten tasks, is able to generalize/calibrate ICL predictions on several new tasks and four LLMs unseen during training. Even for frontier models like Gemini, CALIBFORMER is able to obtain over 50% reduction in calibration error.
1 Introduction
ICL adapts LLMs to classification tasks without parameter updates, but its demonstrations can leave confidence poorly aligned with correctness. CALIBFORMER uses those demonstrations to calibrate predictions and generalizes to unseen tasks and models.
1 Introduction
- ICL adapts LLMs to new classification tasks using labeled demonstrations without parameter updates, while calibrated probabilities support downstream decisions such as human review or external verification.Miscalibrated confidence can undermine reliable decision-making.Source: §1 · p. 1
- Accuracy often improves with more demonstrations while calibration error remains erratic or worsens, leaving confidence misaligned with correctness.The figure notes turbulent Wikitoxic ECE and worsening When2Call ECE despite stable accuracy.Source: §1 · p. 1, Figure 1 · p. 2
- Existing calibration approaches require task-specific estimation or retraining for each target LLM, and hidden-state methods cannot serve black-box models.Temperature scaling also treats correlated predictions at different context positions as independent.Source: §1 · pp. 1–2
- A causal transformer treats in-context examples as a latent validation set, using prediction–label discrepancies to predict calibrated confidence for the test input.Its task-agnostic features include top-class probabilities, running correctness, and ground-truth probabilities at each position.Source: §1 · p. 2
- 41–73% average ECE reduction relative to uncalibrated ICL; CALIBFORMER also outperforms strong calibration baselines and transfers to four unseen LLMs.The evaluation covers eight held-out tasks; the reported transfer includes Gemini-2.5-flash.Source: §1 · p. 2
2 Problem Formulation
The formulation concerns text-classification tasks presented through in-context demonstrations. It defines the model’s confidence and the goal of recalibrating its prediction using only the demonstrations and model outputs.
2 Problem Formulation
- A task τ consists of k labeled input–output pairs sampled from a latent task distribution, with labels drawn from a task-specific label space.The distribution’s conditional label probabilities are treated as unknown.Source: §2 · p. 2
- Given demonstrations and a test input, the LLM produces a label distribution; confidence is the probability assigned to its most likely label.The formulation denotes that predicted label by ŷ and its probability by the model confidence.Source: §2 · p. 2
- The goal is to transform the model’s predicted distribution into a better-calibrated one using only in-context examples and model outputs, without additional target-task labels.Calibration means that predictions assigned confidence p should be correct a fraction p of the time.Source: §2 · pp. 2–3
3 Universal Calibrator for ICL
CALIBFORMER uses task-agnostic signals from in-context predictions and labels to estimate confidence in a test prediction, aiming to transfer across tasks and LLMs. Its training combines diverse model–dataset pairs with data augmentation, while its coverage analysis suggests broader training diversity improves coverage of unseen combinations.
3 Universal Calibrator for ICL
- CALIBFORMER treats in-context examples as a latent validation set, using their prediction–label discrepancies to adjust confidence in the subsequent test prediction.It attends causally over task-agnostic features derived from model probabilities and known labels, with no further fine-tuning on unseen tasks or LLMs.Source: §3 · p. 3
3.4 Training Procedure
- Training uses diverse public classification datasets and open-weight LLMs, with held-out evaluation datasets and 20-shot prompts built using random or similarity-based demonstrations.The loss is binary cross-entropy over prediction correctness, applied only from position 6 onward so the calibrator relies on preceding in-context evidence.Source: §3 · p. 3, §3.4 · p. 4
3.1 Features
- Each position’s five features combine top-class probabilities with previous-example correctness, true-label probability, and squared loss.The features are bounded in [0, 1] and invariant to label-space size properties, supporting cross-task generalization.Source: §3.1 · pp. 3–4
3.2 CALIBFORMER Architecture
- A causal transformer maps the feature sequence to a calibrated confidence, with its test-position output attending to the full demonstration sequence.The architecture has 12 transformer layers, hidden dimension 64, 8 attention heads, and 0.61M parameters; a temperature is then found by binary search to match the predicted confidence.Source: §3.2 · p. 4
3.3 Characterizing Task Coverage
- Scaling training coverage from 2 LLMs × 5 datasets to 3 LLMs × 10 datasets substantially improves coverage of unseen combinations.The coverage analysis represents each LLM–dataset pair by shot-wise mean confidence and accuracy; Figure 2 shows the UMAP projections.Source: §3.3 · p. 4, Figure 2 · p. 5
3.4.1 Data Augmentation
- Two trajectory-level augmentations broaden calibration regimes: random logit-temperature scaling expands confidence patterns, while sampled-label replacement simulates near-calibrated and severely miscalibrated behavior.The label augmentation recomputes features 3–5 after replacing ground-truth labels with samples from a temperature-scaled predictive distribution.Source: §3.4.1 · p. 5
3.5 Extracting LLM’s Label Distributions
- For labels that span multiple tokens or share a prefix, a label-token trie supports probability extraction through forward passes along its paths.Source: §3.5 · p. 5
4 Empirical Evaluation
The evaluation tests CALIBFORMER on unseen models and tasks against inference-time calibration baselines, then examines noisy probabilities, shot counts, and design ablations. Across the main comparisons, it reports lower calibration error, with ablations identifying benefits from sequence modeling, diverse training data, and selected prediction features.
4.1 Baselines
- Baselines include semantic-prior correction, prompt-level temperature scaling, and averaging predictions across five context permutations.Two proposed recalibration baselines were excluded because they were dataset- or LLM-specific and did not generalize across the evaluated settings.Source: §4.1 · p. 6
4.2 LLMs; 4.3 Datasets
- The calibrator is trained on three open-weight LLMs and evaluated on four unseen LLMs, across held-out tasks in tool use, toxicity, intent, and knowledge or medical QA.The held-out models include Phi-4-mini-instruct, Gemma-4-E4B-IT, Qwen3.5-9B, and Gemini-2.5-Flash.Source: §4.2 · p. 6, §4.3 · p. 6
4.4 Experimental Setup; 4.5 Evaluation Metric
- Prompts use either cosine-similarity-selected or randomly sampled demonstrations; evaluation averages repeated-seed runs and reports smooth ECE, with Brier score in the appendix.Open-weight models use 10 seeds and 300 queries per setting, with specified lower counts for Massive Intent and Gemini-2.5-flash.Source: §4.4 · p. 6, §4.5 · pp. 6–7
4.6 Main Results
- 85% of cases, CALIBFORMER beats all baselines across three open-weight LLMs; on Gemini-2.5-flash it reduces uncalibrated ECE by over 50%.On Phi-4-mini-instruct, average ECE falls from 18.3% to 6.1% with similarity sampling and from 25.7% to 7.2% with random sampling.Source: Table 1 · p. 7, Table 2 · p. 7, §4.6 · p. 7
- Similarity sampling improves calibration for all methods, noisy-probability performance is robust relative to temperature scaling, and CALIBFORMER has the lowest average ECE across shot counts.Uncalibrated ECE improves only mildly with more shots, while recalibration methods improve; CALIBFORMER remains lowest across the full shot range.Source: §4.6 · pp. 7–8, Figure 3 · p. 9
4.7 Ablations
- Confidence and correctness alone outperform temperature-scaling baselines; three additional features help more with fewer shots and more classes, while halving training datasets increases ECE.The full feature set shows larger gains in the lower-shot and higher-class-count regimes.Source: §4.7 · p. 8, Table 4 · p. 8
- Unpermuted embeddings degrade generalization, temperature-output prediction increases ECE, and removing augmentation reduces average performance, chiefly on Qwen-3.5-9B.Permuting embeddings prevents memorization but performs slightly worse than the base model; Qwen-3.5-9B is described as reasonably calibrated out of the box.Source: §4.7 · pp. 8–9
- Sequence models outperform simpler calibrators, with the transformer slightly ahead of equal-resource LSTM and RNN models.The summary-statistic approach loses low-level sequence information; beta and isotonic calibration are described as sample-inefficient relative to temperature scaling.Source: Table 5 · p. 8, §4.7 · p. 9
5 Conclusion
CALIBFORMER is a lightweight calibrator for in-context classification that generalizes across unseen models and tasks without task-specific fine-tuning. It learns calibrated confidence scores from prediction–label discrepancies within the context window.
5 Conclusion
- Trained on 3 models and 10 tasks, CALIBFORMER significantly reduces ECE across unseen tasks and models.It uses a transformer-based architecture that treats in-context examples as a latent validation set.Source: §5 · p. 9
6 Limitations
The evaluation and method have several boundaries: proprietary-model coverage is limited, the approach targets classification, rare regressions occur on already well-calibrated models, and theoretical guarantees are absent.
6 Limitations
- Evaluation on proprietary models is limited to Gemini-2.5-Flash, constrained by inference cost and limited API log-probabilities.Repeated-sampling log-probability estimates are possible in principle but require many samples per partial prompt, further increasing cost.Source: §6 · p. 9
- Extending CALIBFORMER to open-ended generation is not straightforward because the label space is huge and its calibration target is ill-defined.Trie-based label-probability construction may be costly, and the correctness indicator would need replacement by a continuous quality signal.Source: §6 · p. 9
- In rare cases, CALIBFORMER can slightly worsen calibration error when the base model is already well calibrated.The authors note that temperature scaling suffers this issue to a greater extent.Source: §6 · p. 9
- The method lacks theoretical guarantees because existing analyses do not directly cover autoregressive demonstration predictions with correlated features.The authors leave this analysis to future work.Source: §6 · p. 9
A Dataset Details
The appendix collects dataset, calibrator, baseline, label-probability, and augmentation details. It describes a causal sequence model, compares related calibration approaches, and specifies how prediction and label perturbations are constructed.
A Dataset Details
- Table 6 consolidates training, validation, and test datasets, their class counts, and their sources.The supplied caption does not report the table’s entries.Source: Table 6 · p. 12
B Calibrator Overview
- CALIBFORMER is a causal transformer encoder with a linear embedding, 12 causal encoder layers, and a sigmoid output for prediction correctness.The layers use hidden dimension 64 and 8 attention heads; the sigmoid output serves as calibrated confidence.Source: B Calibrator Overview · p. 12
C Related Work
- Prior calibration approaches include prompt-bias correction, sequential prediction calibration, permutation averaging, and auxiliary models, with limitations in task/model transfer or inference cost.CALIBFORMER instead uses task-agnostic prediction-derived features and sequential structure to target generalization to unseen LLMs and tasks.Source: C Related Work · pp. 12–13
D Baseline Details
- The baseline details include in-context calibration, autoregressive and fixed-shot temperature scaling, and permutation averaging across 5 random demonstration orders.Temperature-scaling baselines fit per-sequence temperatures by minimizing negative log-likelihood over demonstration predictions.Source: D Baseline Details · p. 13
E Extracting Label Distribution with a Label Token Trie
- A label-token trie represents each label as a root-to-leaf token path and computes its probability from conditional probabilities at branching points.Restricting each softmax to valid next-token continuations yields probabilities that sum to 1 over the label set.Source: E Extracting Label Distribution with a… · p. 14
F.1 Data Augmentation
- The appendix introduces two data-augmentation strategies and defines mean prediction confidence across the original trajectory as a shared quantity.The supplied passages give details for prediction-confidence and label augmentation.Source: §F.1 · p. 14
F.1.1 Prediction Confidence Augmentation
- Prediction-confidence augmentation samples a temperature from 1 to Tmax and rescales logits at every position before recomputing features.For highly confident trajectories, a larger β gives more perturbation, while tanh saturation limits excessive temperatures.Source: §F.1.1 · p. 14
F.1.2 Label Augmentation
- Label augmentation simulates near-calibrated and severely miscalibrated regimes using different temperature ranges, then samples augmented labels from the scaled softmax.The near-calibrated regime keeps synthetic labels largely consistent with model confidence; the higher-temperature regime more often produces disagreement.Source: §F.1.2 · pp. 14–15
F.2 Training with embeddings
The embedding ablation permutes embeddings during training to discourage memorization, while the appendix also reports training-data generation and optimization details.
F.2 Training with embeddings
- Randomly permuting embeddings within each training sequence is intended to prevent sentence memorization and make the calibrator rely on in-context similarity.The authors report that training without permutation overfits and generalizes poorly; with permutation, calibration is still slightly worse than without embeddings.Source: §F.2 · p. 15
F.3 Additional Training Details
- Training data comprise 5,000 autoregressive prediction sequences per training dataset–LLM combination, each of length 21, with both demonstration-sampling strategies equally represented.Dataset dropout samples a fraction of combinations during training to prevent some from dominating and to help regularize generalization.Source: §F.3 · p. 15
F.4 Calibrator Training Hyperparameters
- A CALIBFORMER checkpoint takes less than 30 minutes to train on one Nvidia B200, and gradient clipping uses AutoClip at the 20th percentile.The supplied hyperparameter table caption identifies the training-hyperparameter table but does not state its entries.Source: §F.4 · p. 15, Table 7 · p. 16
G Evaluation Metrics
Calibration is evaluated with smooth ECE and Brier score. ECE compares empirical accuracy with mean predicted confidence across confidence bins, while the Brier score uses predicted probabilities and one-hot labels.
G Evaluation Metrics
- Smooth ECE estimates calibration by kernel-smoothing prediction–label pairs rather than using only the standard binned approach.The paper uses smooth ECE on evaluation datasets.Source: G Evaluation Metrics · p. 16
- ECE compares each confidence bin’s empirical accuracy with its mean predicted confidence.The bins are denoted Bm in the definition.Source: G Evaluation Metrics · p. 16
- The Brier score is defined over samples and classes using predicted probabilities and one-hot encoded true labels.N is the sample count, C the class count, p̂ij the predicted probability, and yij the true-label indicator.Source: G Evaluation Metrics · p. 16
H Additional Experimental Results
The appendix provides additional metric tables and calibration plots, plus examples of prompts used with chat templates.
H.1 Tables
- Tables report Brier scores for all LLMs, Phi-4-mini-instruct accuracy under both demonstration-sampling settings, and ECE and Brier-score variances for open-weight LLMs.The accuracy results are organized by similarity and random sampling; the variance tables cover open-weight models.Source: §H.1 · p. 16
H.2 Plots
- Figures 4–11 show Phi-4-mini-instruct calibration performance across methods and shots for all datasets.The figures include MedQA, Wikitoxic, Wildguard, When2call, MetaTool, Massive Intent, Amazon Counterfactual, and CommonsenseQA.Source: §H.2 · p. 16
I.1 Prompts
- The prompt appendix gives examples of one-shot Phi-4-mini-instruct prompts and two-shot Amazon Counterfactual prompts with chat templates applied.The base prompt specifies the label space for ICL.Source: §I.1 · p. 16
I.2 LLM inference
The inference appendix describes open-weight LLM serving, prompt-logprob extraction, and additional inference requirements for some model and label settings.
I.2 LLM inference
- The open-weight experiments use vLLM 0.19.0 with PyTorch 2.10 and CUDA 12.8, with a small sampler modification to mask non-label raw logits.The modification applies allowed-token filtering to raw_logprobs, enabling efficient inference with label tokens.Source: §I.2 · p. 16
- The appendix reports Brier-score tables for open-weight models and Gemini-2.5-flash, with Gemini values split by random and similarity-sampled contexts.It also includes Phi-4-mini-instruct accuracy tables for both sampling settings.Source: Table 8 · p. 17, Table 9 · p. 17, Table 10 · p. 17
- For black-box LLMs without prompt logprobs, calibrating a k-shot prompt requires k additional calls; shared first-token labels also require extra calls.The paper notes that most tasks tend to have simple label spaces.Source: §I.2 · p. 17
- Requesting prompt logprobs in vLLM disables prefix caching in this setting.The appendix attributes this to a lack of optimization demand for the setting.Source: §I.2 · p. 17
- Gemini experiments use gemini-2.5-flash-preview-09-2025, request the maximum allowed 20 logprobs, and generate one output token per call.Training and evaluation experiments ran on Nvidia B200 GPUs.Source: §I.2 · p. 17
I.3 Ablation Statistical Significance
The ablation tests find statistically significant calibration degradation for every tested ablation. The appendix also documents the testing procedure and supplementary results materials.
I.3 Ablation Statistical Significance
- Ablations are tested against CALIBFORMER using paired per-seed ECE and Brier differences across 30 seeds, with a paired t-test or Wilcoxon signed-rank test selected after a Shapiro–Wilk normality test.The Wilcoxon test is used when normality is not supported; this fallback occurs for ECE differences of Half Train Data and Temperature Output.Source: §I.3 · p. 18
- All ablations significantly degrade calibration error (p < 0.05); reducing features to confidence and correctness is the weakest ECE degradation (p = 0.03), while its Brier degradation remains significant (p < 0.001).The gaps widen in lower-shot and higher-class-count regimes; permuted embeddings and Temperature Output show the strongest degradations.Source: §I.3 · pp. 18–19
J LLM Usage Declaration
- The cited supplementary tables report calibration error (ECE) and Brier scores across calibration methods for Phi-4-mini-instruct, Qwen3.5-9B, and gemma-4-E4B-it, with standard deviations reported in parentheses.Source: Table 12 · p. 19, Table 13 · p. 20, Table 14 · p. 20
I.3 Ablation Statistical Significance
- The MLP and logistic-regression comparisons use seven-dimensional prediction summaries, without causal shifting of correctness and ground-truth-probability features.The summaries combine two query-position probabilities with five features averaged over in-context examples.Source: §I.3 · p. 19