Token-based Decision Criteria Are Suboptimal in In-context Learning (2025)

Abstract: In-Context Learning (ICL) typically utilizes classification criteria from output probabilities of manually selected label tokens. However, we argue that such token-based classification criteria lead to suboptimal decision boundaries, despite delicate calibrations through translation and constrained rotation applied. To address this problem, we propose Hidden Calibration, which renounces token probabilities and uses the nearest centroid classifier on the LM’s last hidden states. In detail, we assign the label of the nearest centroid previously estimated from a calibration set to the test sample as the predicted label. Our experiments on 6 models and 10 classification datasets indicate that Hidden Calibration consistently outperforms current token-based baselines by about 20%~50%, achieving a strong state-of-the-art in ICL. Our further analysis demonstrates that Hidden Calibration finds better classification criteria with less inter-class overlap, and LMs provide linearly separable intra-class clusters with the help of demonstrations, which supports Hidden Calibration and gives new insights into the principle of ICL. Our official code implementation can be found at https://github.com/hc495/Hidden_Calibration.

## introduction in-context learning (icl) @xcite is a few-shot learning paradigm without model parameter updates on language models (lms). in detail, as shown in fig. 1-(a , b ), given a prompt consisting of demonstrations and a query, lms conduct causal language modeling operation from the prompt to assign probabilities to the label token candidates designed by hand, and icl chooses the one with the highest probability as the prediction. one well-known issue of icl is that the predicted probabilities are biased (under-calibrated), leading to a decrease in prediction performance @xcite @xcite + , e u -) to decode the h to prediction ŷ, then calibrations are used to adjust the predicted logits. c. our work uses the calibration dataset to calculate centroids ( h+ , h-) to decode the h. @xcite . to address this issue, previous work calibrates the predicted label token probabilities by performing affine transformations with estimated parameters to adjust these probabilities for more precise predictions. these previous works and also the vanilla icl are based on a potential assumption: the affine manifolds spanned by the decoding vectors in the lm head (un-embedding vectors) of manually selected label tokens are good subspaces of the hidden space to distinguish hidden vectors (i.e., the last hidden states to be fed to the lm head) onto various label appropriately, so that the label token probabilities decoded from these subspaces are accurate classification logits. however, although using the label un-embedding with task-related semantics (e.g. "positive" and "negative") seems intuitive, it should be noted that we have no reason to believe that these label un-embeddings have any explicit guarantee for decoding the last hidden states into accurate classification logits (token-based decision criterion), even if various delicate calibrations are used to move these boundaries inside the subspaces (see §3.1). also, some works have pointed out that randomly changing label spaces doesn't critically influence icl performance @xcite , which means the selected label subspaces are trivial and arbitrary, making a suspicion of: using manually selected label un-embeddings to decode the last hidden states, i.e., utilizing manually selected label probabilities as classification criteria may not be good icl practices. previous work has shown that using the output probabilities of the full vocabulary increases icl performance @xcite . this is a good start to avoid the manually selected classification criteria, but there is still doubt that output probability distributions are not informative enough for classification (see §3.2). therefore, we utilize the last hidden states instead, which are informative precursors of the token probabilities. concretely, we propose hidden calibration, training centroid classifiers on the last hidden states of icl prompts. as shown in fig. 3 , during the training, we build standard icl prompts similarly to fig. 1 -a from a supervised calibration set and input them into the lm to get the last hidden states of the last tokens of icl prompts. then, we calculate the centroids of the last hidden states w.r.t. the queries' label to get a centroid for each label, as an anchor for inference. during the inference, we input the test prompt, find the nearest centroid to the last hidden states of the test prompt, and assign the corresponding label of the centroid as the prediction. empirically, hidden calibration improves the icl performance by approximately more than 20% on 10 text classification datasets and 6 modern lms ( §4.1), with an equal computational cost with previous calibration methods. to the best of the author's knowledge, hidden calibration consistently outperforms the calibration baselines, achieving a strong state-of-the-art in icl. additional experiments indicate that hidden calibration effectively alleviates the stress of prompt engineering, performing robust accuracy under different prompt designs. moreover, our subsequent analysis indicates that hidden calibration does find better logits-mapping subspaces that effectively separate data points. in detail, we find that the distribution of classification logits calculated from hidden calibration have less inter-class overlapping than from label probabilities, while such overlapping is proportional to the lower bound of the classification error. this suggests hidden calibration finds subspaces with essentially better classification performance. furthermore, we investigate the principle of hidden calibration, that is, the reason why a simple centroid-based linear decision boundary can divide the icl hidden state properly. we find that lms provide linearly separable clusters in the hidden states w.r.t. query labels, while more demonstrations can promote such a process. our contributions can be summarized as: ## background this section reviews previous work on icl and denotes their mathematical descriptions as an introduction to the main motivation of this work. ## methodology based on the above background, in this section, we demonstrate the limitations of the above calibrations, and then propose hidden calibration to address such limitations fundamentally. ## experiments & main results in this section, we empirically prove the effectiveness of hidden calibration by classification performance on 6 models and 10 datasets. hidden calibration outperforms all the baselines and reaches a strong state-of-the-art of icl with high efficiency in calculation, data, and prompt engineering. ## related works in contrast, lightweight solutions focus on classification criteria-based method (calibration). such methods focus on re-calculating output label probabilities, keeping the main feed-forward calculation processes and model parameters un-modified. the original motivation for these works is to eliminate prior bias and unfaithful confidence in icl, by calibrating the output label probabilities @xcite @xcite @xcite @xcite . while, as described in the main text, some practices without the usage of labelspecific probabilities have also been proposed @xcite @xcite . also, a careful design of input prompts can help improve the icl performance. (1) demonstration selection. @xcite finds that selecting the demonstrations with lower perplexity improves the icl performance, similarly, @xcite generate the demonstrations from pre-trained lms, etc. (2) demonstration ordering. it is found that the ordering of demonstrations can significantly influence the performance @xcite @xcite , as also shown in our experiments in fig. 5 . specifically, @xcite detect the optimal demonstration ordering by some synthetic detecting sequences, while @xcite orders the demonstrations from easy to hard, following a curriculum learning form. ## limitations due to computability limitations, we are not able to compare the performance of hidden calibration with the baseline based on supervised finetuning. however, we believe that hidden calibration is not within the same methodology as the finetuning method, due to the significant difference in computational cost. so such a lack of comparison will not seriously hurt the soundness of this paper. we argue that human intuition in the label token choice is not reliable. however, we have not eliminated such human intuition completely from the icl loop: when we build prompts, we still choose the label token. how to automatically select the optimal label token in the prompt will be an important issue, leaving as future research directions for improving the performance of icl further. other label probability calibrations (e.g. batch calibration) can be combined with hidden calibration for further performance improvements, since the 0-point is not necessarily an exact classification boundary, as shown in fig. 8 . also, more complex prompts can be used. however, due to space constraints, we have not attempted this incremental approach, remaining it for an empirically possible practice. observation in §5.2 needs more theoretical and experimental analysis. as we can see, some models (gpt2-xl) do not benefit from demonstrations even through the lens of hidden state clustering or hidden calibration, which needs to be explained. an explanation of "why such clustering occurs or not", and "how to enhance the intra-class clustering by fine-tuning or prompt engineering" will be considerably beneficial for understanding icl.