Task Calibration: Calibrating Large Language Models on Inference Tasks (2025)

Abstract: Large language models (LLMs) have exhibited impressive zero-shot performance on inference tasks. However, LLMs may suffer from spurious correlations between input texts and output labels, which limits LLMs’ ability to reason based purely on general language understanding. For example, in the natural language inference (NLI) task, LLMs may make predictions primarily based on premise or hypothesis, rather than both components. To address this problem that may lead to unexpected performance degradation, we propose task calibration (TC), a zero-shot and inference-only calibration method inspired by mutual information which recovers LLM performance through task reformulation. In NLI, TC encourages LLMs to reason based on both premise and hypothesis, while mitigating the models’ over-reliance on individual premise or hypothesis for inference. Experimental results show that TC achieves a substantial improvement on 13 different benchmarks in the zero-shot setup. We further validate the effectiveness of TC in few-shot setups and various natural language understanding tasks. Further analysis indicates that TC is also robust to prompt templates and has the potential to be integrated with other calibration methods. We publicly release our code to facilitate future research.

## introduction large language models (llms) @xcite @xcite @xcite have demonstrated strong generalization ability in a wide range of downstream tasks. in particular, prompt-based learning @xcite has been an effective paradigm, enabling zero-shot or few-shot learning @xcite . ideally, an llm with advanced language understanding capabilities could perform a new task such as natural language inference (nli) in a zeroshot setting without relying on annotated examples. however, research has shown that zero-shot capabilities of llms on inference tasks are currently constrained by the presence of spurious correlations that often lead to biased prediction @xcite . to mitigate spurious correlations, previous work @xcite @xcite @xcite has explored model calibration, which reweighs output probabilities based on various bias estimators. however, existing calibration methods fall short of addressing the bias that stems from llms' overreliance on certain parts of the context for prediction @xcite , which we call context preference bias. in nli, the parts of the context can be the premise and the hypothesis. in stance detection, the parts of the context can be the text and the topic. for convenience, we use nli as the main target for discussion in this paper, despite that this bias can be also observed in other tasks. figure 1 shows an example from qnli dataset @xcite , where the task is to determine whether a given context sentence contains the answer to a given question. we observe that the model prediction is incorrect because it relies excessively on the question itself when making the prediction. to address the content preference bias, we propose task calibration (tc), a zero-shot and inference-only calibration method. our work is inspired by mutual information @xcite . intuitively, for an nli task, proper use of mutual information can reveal how much more informative the combined presence of premise and hypothesis is concerning the label, compared to their individual presences. based on this insight, we reformulate llm inference by factoring out the probabilities of premise-only figure 1 : an example from qnli dataset @xcite . sentence-only, question-only and both indicate the inputs with only the sentence, question and using both components, respectively. while the initial model prediction is incorrect, potentially due to the influence of the question, we observe that task calibration finally leads to a correct prediction. and hypothesis-only inputs. tc requires no annotated data and is easy to implement, involving only two extra inference stages using premise-only and hypothesis-only inputs for each sample. as shown in figure 1 , although the model's initial answer is incorrect, it finally makes the correct prediction after task calibration, by using output probabilities derived from sentence-only, question-only, and combined inputs. experimental results demonstrate superior performance of tc over other calibration methods in the zero-shot setup, showcasing a noteworthy boost of two llms on 13 classification datasets. specifically, tc outperforms the best-performing baseline in 12 and 10 out of 13 datasets on the mistral-7b-instruct-v0.3 @xcite and phi-3-mini-4k-instruct (abdin et al., 2024) models, respectively. in addition, tc is robust to various prompt templates, demonstrating its effectiveness in few-shot setups and 4 different natural language understanding (nlu) benchmarks. finally, we find that the combination of tc and other calibration methods can yield better performance, which indicates their complementary strengths in fixing spurious correlations. to summarize, our contributions are as follows: calibration of language models. previous attempts to mitigate spurious correlations include training a debiased model with residual fitting @xcite or a debiased training set @xcite . however, these methods necessitate finetuning, and thus pose challenges for pursuing efficient llms. @xcite propose contextual calibration (cc), which first estimates the bias of language models with a content-free test input, and then counteracts the bias by calibrating the output distribution. @xcite find that different surface forms compete for probability mass. such competition can be greatly compensated by a scoring choice using domain conditional pointwise mutual information (dcpmi) that reweighs the model predictions. @xcite further identify the domain-label bias and propose a domaincontext calibration method (dc) that estimates the label bias using random in-domain words from the task corpus. @xcite propose prototypical calibration to learn a decision boundary with gaussian mixture models for zero-shot and few-shot classification. @xcite propose batch calibration (bc) to estimate the contextual bias for each class from a batch and obtain the calibrated probability by dividing the output probability over the contextual prior. in contrast, we tackle the problem from a different perspective of task reformulation, which mitigates bias while recovering model performance across challenging inference tasks. ## context preference bias formally, denote the input to a problem as @xmath0, x 2 , ..., x n , where x i represents the ith part of the input context such as a hypothesis in the nli problem. the context preference bias problem refers to a model's tendency to make a prediction based primarily on x i , rather than considering the entire x. in this work, n can be considered as 2, where x 1 and x 2 represent the premise and hypothesis for nli, the text and topic for stance detection, and different sentences for paraphrasing. without loss of generality, we consider nli as the main target in the following sections and use x p and x h to represent x 1 and x 2 , respectively. mckenna et al. (2023) identify the attestation bias for the nli task, which can be seen as a special case of context preference bias where llms falsely associate the hypothesis with entailment. we explore the context preference bias from a novel viewpoint, i.e., we examine whether llms can accurately predict not_entailment when the premise or hypothesis is absent from the input. specifically, we evaluate mistral-7b-instruct-v0.3 @xcite on binary nli tasks rte @xcite , scitail @xcite and qnli @xcite datasets where outputs include not_entailment or entailment. ideally, llms should be able to discern the absence of premise or hypothesis and make predictions on not_entailment. as shown in figure 2 , mistral-7b-instruct-v0.3 exhibits a tendency to associate premise-only or hypothesis-only inputs with labels other than not_entailment, as evidenced by the gap between the bars and the ideal value (i.e., 100%). it suggests the existence of spurious correlations (which we call context preference bias) that can distract llms from relying on both premise and hypothesis when making predictions. in addition, the performance of llms on premise-only and hypothesis-only inputs varies across datasets. for example, mistral-7b-instruct-v0.3 exhibits superior performance in the premise-only setting for scitail and performs better in the hypothesis-only setting for rte. building upon the observation, we further investigate the correlation between incorrect llm predictions (using both premise and hypothesis) and the labels derived from premise-only or hypothesisonly inputs. results are shown in figure 3 . we observe that llm predictions based solely on the premise or the hypothesis frequently align with incorrect predictions of using both components. for example, in the scitail dataset, over 90% of incorrect llm predictions align with the labels obtained from hypothesis-only inputs. it reveals that the llm excessively relies on the premise or hypothesis alone when making predictions. ## experimental setup datasets. we conduct experiments on 17 text classification datasets that cover a wide range of tasks. specifically, for standard inference task, we consider natural language inference: rte @xcite , wnli @xcite , sci-tail @xcite , cb (marneffe et al., 2019), mnli @xcite and qnli @xcite ; stance detection: perspectrum @xcite , ibm30k @xcite , ez-stance @xcite , iam @xcite and vast (allaway and mckeown, 2020); paraphrasing: paws @xcite and qqp. to indicate the effectiveness of tc on other tasks, we follow the experimental setting that adopts a textual entailment formulation in previous work @xcite and additionally consider sentiment classification: sst-2 @xcite ; offensive language identification: offenseval @xcite ; hate speech detection: hateval @xcite and @xcite . rte, wnli, cb, mnli, qnli and qqp datasets used for evaluation are drawn from the glue @xcite and superglue @xcite benchmarks. more details of these datasets can be found in baselines. we compare tc with the original lm and previous calibration methods, including cc @xcite , dcpmi @xcite , dc @xcite and bc @xcite . these methods are discussed in section 2 and their scoring functions are shown in model and implementation details. we conduct experiments mainly on two instruction-tuned models including mistral-7b-instruct-v0.3 @xcite and phi-3-mini-4k-instruct (3.8b) @xcite . for all experiments, unless stated otherwise, we perform the evaluation in the zeroshot setting. in the few-shot setting, we use @xmath1. the templates and label names used for all datasets can be found in