Batch Calibration: Rethinking Calibration for In-Context Learning and Prompt Engineering

Dear Reviewer tiQs,

Thank you for your prompt reply! We are glad that our response has alleviated your concerns.

Regarding the related work section, we thank the reviewer's suggestions. We will merge Sec 2 and Sec 5 in the next revision of the paper, as suggested.

Regarding BCL, there are two points we'd like to make:

• First, BCL, which features adjustable strength learned from labeled data, allows us to study the ground-truth optimality of the design boundary asserted by BC under zero shot. We can see that the performance of BC is consistently strong and, in many cases, indeed close to optimal -- see Fig. 5 where we validate different choices of $\gamma$ on labeled data that would not be otherwise available to a zero-shot method; we can see that $\gamma=1$, which corresponds to BC, is close to optimal; this provides empirical support to the design of BC described in Section 3.

• Second, while we acknowledge that the extent of improvement from BC to BCL is not always significant from Table 3, there are cases where BCL improves quite significantly (e.g., MRPC and CB). On the other hand, when BCL underperforms BC, the extent of performance decrease is mostly marginal. Thus, on balance, we felt it is still beneficial to include discussions of BCL as a potential improvement method that makes use of labeled data when they are indeed available (as contrasted to zero-shot).

Again, we will clarify the above points further and incorporate them into the next revision shortly.

We thank the reviewer for their constructive and valuable feedback and for recognizing the soundness, presentation, and novelty of this work to the research community. We appreciate the suggestions for including more $\gamma$ experiments and a section discussing border impacts and limitations. We have incorporated all the reviewer’s suggestions in the updated manuscript. We would be grateful if the reviewer could consider increasing their rating if they found our response has sufficiently addressed their concerns.

"Survival of the Most Influential Prompts: Efficient Black-Box Prompt Search via Clustering and Pruning" has been accepted by EMNLP-2023. Please cite it correctly. Unify the usage of abbreviated words. Tab. -> Table.

We thank the reviewer for pointing them out. We have corrected them in the updated manuscript.

Why is the related work in section 5? Introducing it in section 2 will make readers understand the background better.

We appreciate the reviewer’s comments on the structure of the presentation. Section 2 serves as a unified overview of the most relevant existing calibration methods while addressing two important research questions behind their design principles. Therefore, we found that having the related work with the current analysis section together may be chunky and could potentially dilute the attention of readers. Hence, we decided to move the related work to section 5 and present the main body of our work, section 3, to readers as early as possible for a better flow of content.

When introducing BCL, the author takes some examples to make readers understand the $\gamma$. I suggest authors provide more experiments on how different $\gamma$s affect the results.

We thank the reviewer for suggesting more experiments. We have now conducted additional experiments to illustrate the impact of the strength term on BCL and have included more plots in Figure 11, Appendix D, in the updated manuscript. To better understand the $\gamma$ term, We kindly refer the reviewer to Table 1 and Figure 2, which visualize the decision boundary for BC. The strength term is controlling the strength of the shift (BC) operation for the decision boundary, such that $p_0 = p_1 + \gamma \alpha$. When $\gamma =1$, the estimation of $\alpha$ by zero-shot BC serves as a strong but not optimal prior for shifting the decision boundary. Therefore, by leveraging labelled data to learn the strength term, we can further refine the decision boundary by learning $\gamma$.

As shown in Figure 3, the improvement from BC to BCL seems marginal.

We acknowledge that having labelled data brings marginal improvement to BC for certain tasks, which has also been reflected in Figure 5, Page 6, where the gain of performance between the BC (purple line) and BCL (green line) is observable but not significant. However, this empirical evidence exactly indicates the strength point of BC, such that BC at $\gamma=1$ already serves as a good prior for estimating the contextual bias from the target (test) statistics only. It further consolidates the effectiveness of BC even in the zero-shot setups. By incorporating an additional strength hyperparameter, we are enabled to estimate the source prior to source statistics in labelled data. This BCL design further refines the robustness of BC under covariate shift and shows a much more enhanced calibration performance in a range of tasks (SST-2, MRPC, QQP, CB) than BC.

Any border impact and limitation discussion?

We appreciate the reviewer’s suggestion for discussing the broader impacts of our approach and limitations. We have included a new limitation section in the Appendix in the updated manuscript. As a brief summary, we reiterate the limitations and the impact of BC below:

In terms of limitations, we recall that BC is a test-time method that relies on the target statistics from batched input. To mitigate any potential estimation errors from unlabelled data, we introduce the adjustable BCL extension to incorporate source statistics from labelled data. Though BC has shown remarkable sample efficiency in terms of batch sizes, it still requires a batch of inputs to estimate the contextual bias. We introduce a running estimation for BC from mini-batches, which subsequently stabilizes the prediction of LLMs when more mini-batches arrive.

Overall, BC is zero-shot, inference-only, and incurs negligible additional costs for driving toward robust LLM applications. BC mitigates unfair prediction of LLMs due to contextual bias while alleviating the sensitivity of LLMs with respect to prompt designs. We expect easy prompt engineering with BC for users building towards their own robust and responsible LLM applications.

In e.q.(3), how to obtain the contextual prior

In Equation 3, the contextual prior $\hat{\mathbf{p}}$ is obtained via Equation 2, which is equivalent by obtaining the estimated mean of log probability for each class. In the updated manuscript, we have updated the notation in Equation 2 to $\mathbf{p}(y|C)_{j}$ better to emphasize the link between Equations 2 & 3.

How would BC work on the task other than classification?

We thank the reviewer for discussing this appealing research direction that is worthy of being explored by the research community. We suspect that the contextual bias may also exist in short-form generation tasks. Motivated by [1], one possible solution for generative tasks is to calibrate over the logits of the first output token since the following tokens are likely deterministic based on the first token. We will endeavor to explore generative calibration. We thank the reviewer for this insight, and we include this idea in the future work section in Appendix in the updated manuscript.

[1] Guo, C., Pleiss, G., Sun, Y. and Weinberger, K.Q. On calibration of modern neural networks. ICML 2017.

It seems the calibration perform is not stable across different model architectures in 1-shot setting on some tasks. E.g., MNLI, 75.12/ 60.02/81.34 for PaLM-2-S/PaLM-2-M/PaLM-2-L. I understand that PaLM-2-S/L are based on 5 runs and PaLM-2-M result is from a single run. Could you help to understand, what causes the high perform variance, I may miss something here.

We acknowledge that different sizes of models preserve distinctive robustness features for different tasks, and it’s occasionally possible that a larger model underperforms a smaller model for specific tasks (this has also been reflected in, for example, the PaLM 2 technical report [2], Table 2) due to multiple possible factors such as pretraining strategies, type and size of input data used in training and instruction finetuning, and/or slight differences in templates, but in general, PaLM 2-M still shows a stronger 1-shot ICL performance than PaLM 2-S across the set of evaluation tasks, where BC also shows consistent performance improvement across all model scales.

[2] Anil, R., Dai, A.M., Firat, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., Chen, Z. and Chu, E. Palm 2 technical report. arXiv preprint arXiv:2305.10403.

We thank the reviewer for their insightful and constructive feedback. We appreciate the reviewer’s suggestion to present BCL plots in other tasks and detailed example performance in Figure 8. We have incorporated all suggestions made by the reviewer into the updated manuscript. We hope that our response sufficiently addresses the reviewer’s concerns and the reviewer could consider improving their score.

My main concern goes to the strength parameter in BCL; the current results show $\gamma$ seems to give strong performance (though not optimal) across tasks on CB and SST-2. Does this generalize to other tasks as well? It would be good to provide the optimal $\gamma$ for each task; maybe include it in Table 2? Besides, would it be more reasonable to sample $\gamma$ in [0, 5]?

We thank the reviewer for expressing interest in the strength parameter in BCL and the generalization of performance in terms of the strength term. We recall that when strength is at 1, Equation 5 becomes the zero-shot BC, and BC shows state-of-the-art performance across many tasks over ICL and calibration baselines in Table 2. To better understand the $\gamma$ term, We kindly refer the reviewer to Figure 2 that visualizes the decision boundary for BC. The strength term is controlling the strength of the shift (BC) operation for the decision boundary, such that $p_0 = p_1 + \gamma \alpha$. When $\gamma =1$, the estimation of contextual prior, $\alpha$, by zero-shot BC serves as a strong but not always optimal prior for shifting the decision boundary, which is due to the dependency on test-time statistics. Therefore, by leveraging labelled data to learn the strength term as a source prior, we can further refine the decision boundary while mitigating calibration errors.

We thank the reviewer for suggesting to sample $\gamma$ in [0, 5]. By having the $\gamma$ in the negative range, we can explore what the decision boundary looks like after shifting it in the opposite direction. We then retain both directions for a sufficient range to sample the optimal source prior. To provide a better intuition of $\gamma$ across other tasks, we have included more plots in Figure 11, Appendix E in the updated manuscript, where optimal $\gamma$ for the test set can be read from each task plot. BC ($\gamma=1$) demonstrates strong performance across tasks, where further slight gains can be achieved by BCL when we incorporate the source prior via labelled data. We welcome the reviewer for further discussion.

There are some statements that are not well presented or supported. In Sec. 4.3, the authors claimed BC retains the performance even when using emoji pairs as verbalizers. Is there an example or result related to this specific experiment?

We thank the reviewers for expressing their interest in individual test performance for each verbalizer choice in Figure 8. We kindly refer the reviewer to Table 10, which lists the detailed verbalizer choices and descriptions for each emoji pair. We agree that it will be interesting to interpret individual experimental results corresponding to each design choice. Therefore, we have updated our manuscript and included each test sample result before and after calibration in a new Table 7, Page 17. Empirically, we observe that BC can recover the performance of LLM for emoji pairs with opposed but context-relevant semantics. We hope this will address your remaining concerns.

Dear Reviewer jGRr,

Thank you for your engagement in the discussion! We are glad that our response has alleviated your concerns.

We would like to provide a further clarification that interprets the strength term $\gamma$ from our new Figure 11 following the suggestion of the reviewer:

The gap between the optimal strength term in the labelled data and the best test strength term indicates the distribution shift from the source to the target. This distribution shift exists due to the statistics of the dataset, not from any proposed algorithms, and the best test strength term is always an unknown prior for the test set. The gap can also be seen as, fundamentally, a manifestation of generalization gap that any learning algorithm will be subjected to rather than specific to us. Despite this, we still believe BCL yields empirical benefits, as discussed below.

Incorporating the strength term allows us to leverage an additional source prior when encountering a distribution shift from a uniform prior. In cases where, BCL improves quite significantly (e.g., MRPC and CB), we found the effectiveness of having the strength term to be reflected from the relative distance between the optimal strength term from labelled data (red) and the original ICL (blue), where BCL serves as a robust extension of BC with more significant performance enhancement over ICL across tasks. Empirically, we found BCL to often lead to large gains and even when it does not, the extent of performance change is minimal. Thus, on balance, we still believe BCL is an useful extension.

Finally, as we have responded to Reviewer tiQs, another reason of studying BCL, which features adjustable strength learned from labeled data, is it that allows us to study the ground-truth optimality of the design boundary asserted by BC under zero shot. We can see that the performance of BC is consistently strong and, in many cases, indeed close to optimal -- see Fig. 5 where we validate different choices of $\gamma$ on labeled data that would not be otherwise available to a zero-shot method; we can see that $\gamma=1$, which corresponds to BC, is close to optimal; this provides empirical support to the design of BC described in Section 3.

Therefore, we believe that BCL is a valuable extension of BC both empirically and conceptually for robustness, extending the applications of BC from zero-shot learning only to become more capable of few-shot learning. We appreciate the discussion from reviewer, and we are happy to provide further discussions.

We hope the reviewer could reconsider their rating if they believe their concerns have been clarified better with our latest response.

Thanks for the author response. I read the author response as well as others' reviews. My concerns are partially addressed. I am still not fully convinced by the effectiveness of learning the strength term as a source prior by looking at Figure 11, i.e., there is a gap between the learned optimal strength term and the best one at test time on some tasks. I'd like to keep my score for now.

We thank the reviewer for their insightful and constructive feedback, especially for valuable suggestions that further improves the presentation and clarity of terms in the paper. We have made changes in our updated manuscript accordingly. We hope that in light of the response that addresses all the reviewer’s concerns, the reviewer could consider improving their score.

While BC introduces minimal computational overhead, in highly resource-constrained environments, even small overheads might be significant.

We thank the reviewer for acknowledging the minimal computation overhead of our approach. While we agree in principle that small overheads can be significant in resource-constrained setup, it is true for any method that uses additional computations and/or storage and hence applies to all baselines unless we don’t do any calibration. In comparison to previous methods, batch calibration does not involve any additional forward (API) calls of LLMs that may incur observable inference latency for users, whereas prototypical calibration [1] can result in significant computation overhead with its hundreds of random initializations and iterations of EM. Furthermore, BC may also use an on-the-fly, running mean estimation in mini-batches, which only involves a post-hoc, linear operation on the log-probability of the model.

[1] Han, Z., Hao, Y., Dong, L., Sun, Y. and Wei, F. Prototypical calibration for few-shot learning of language models. ICLR 2023.

The term calibration can be confusing, maybe using bias is better. At least should add a footnote to the Introduction. Another widely used notion of calibration is from the perspective of uncertainty, explainability, and reliability.

We appreciate the reviewer’s suggestion for clarifying the notion of calibration. In our updated manuscript, we have included a clarification on the definition of calibration as a footnote in the introduction. We agree with the reviewer that the term “calibration” is usually aligned with the trustworthiness, reliability, and interpretability of neural networks [2]. In the context of our work, we follow the notion of calibration defined in [3] as a technique for mitigating the label shift and contextual bias in LLMs prediction.

[2] Guo, C., Pleiss, G., Sun, Y. and Weinberger, K.Q. On calibration of modern neural networks. ICML 2017.

[3] Zhao, Z., Wallace, E., Feng, S., Klein, D. and Singh, S. Calibrate before use: Improving few-shot performance of language models. ICML 2021.

Please include model sizes and details in the paper, as the performance of ICL and prompting is quite irreproducible and context-dependent.

We thank the reviewer for being aware of the nature of the context-dependency of prompting. We acknowledge the sensitivity of LLM prediction to the context and the prompt template. Therefore, we included detailed templates in Tables 8, 9, 10, & 11 for reproducibility and conducted all main experiments with 5 random seeds to reveal the variance for each study. Regarding the details of the model size, we included the PaLM 2-S, M, L variants, and we kindly refer the reviewer to [4] that lists the details for each variant, including experiments on PaLM 2 with comparable model sizes as PaLM, and PaLM 2-L shows competitive performance to GPT-4 [5]. We then follow the notion of model variants in [4] and hope this has addressed the reviewer’s concern.

[4] Anil, R., Dai, A.M., Firat, O., Johnson, M., Lepikhin, D., Passos, A., Shakeri, S., Taropa, E., Bailey, P., Chen, Z. and Chu, E. Palm 2 technical report. arXiv preprint arXiv:2305.10403.

[5] OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774.

Please adjust the margins of subfigures in fig. 2, it's overlapped.

We appreciate the reviewer's suggestions on the presentation. We have adjusted the margins and provided improved visualization in the updated manuscript.

The literature review of Test time tuning is a bit unnecessary as ICL itself is not a tuning method

In Appendix A, we included a discussion of test-time adaptation (TTA), where we explicitly focus on the previous literature using test-time batch statistics for BN without tuning. This line of research has shed light on the theoretical foundation for batch calibration, where contextual bias can also be framed as a form of covariate label shift [6, 7]. Batch calibration takes inspiration from test-time batch adaptation in the novel context of LLMs, and we would be grateful if researchers in the TTA community could be inspired to deliver advanced research ideas on top of batch calibration.

[6] Nado, Z., Padhy, S., Sculley, D., D'Amour, A., Lakshminarayanan, B. and Snoek, J., 2020. Evaluating prediction-time batch normalization for robustness under covariate shift. arXiv preprint arXiv:2006.10963.

[7] Sun, Q., Murphy, K., Ebrahimi, S. and D'Amour, A., 2022. Beyond invariance: Test-time label-shift adaptation for distributions with "spurious" correlations. arXiv preprint arXiv:2211.15646.

Dear Reviewer 99FV,

We thank you once again for your feedback! As the discussion period is ending, we wonder if you could take a look at our response if you haven't? We are eager to know whether our response has addressed your concerns: we'd be grateful if you could reconsider the rating if that's the case. Otherwise, we are always happy to answer further questions that you may have.

Best,

Authors

We thank the reviewer for their insightful review and the acknowledgment of the empirical effectiveness of our method. Please see our response below, which we believe has addressed all the concerns based on the contents we refer to in the paper. We hope that in light of the response, the reviewer could consider improving their score.

The proposed method is not thoroughly justified.

We’d like to refer the reviewer to Section 2, Pages 2-5, where we have conducted a comprehensive analysis of existing calibration methods through their decision boundaries. We identified the current deficiency of existing methods with individual case studies, and the design of each aspect of BC has been strongly justified with empirical evidence in Section 2.3 and motivated in Section 3. In particular, we addressed two important research questions for calibration methods with an interpretation of decision boundaries and model biases, as illustrated in Figures 2, 3, & 10, delivering novel guidelines for opting for a linear decision boundary and content-based calibration. We’d be grateful if the reviewer could point out what precisely, in their opinion, is lacking justification, and we are happy to provide further clarifications.

The writings could be improved in terms of the analysis of existing methods and the proposed method.

We thank the reviewer for their suggestion. We have revisited the analysis section and cleaned up any language problems. We’d be grateful if the reviewer could let us know of any further concerns, and we are happy to provide further clarifications.

The authors argued that linear decision boundaries produced by calibration methods can be more robust and generalizable across tasks. This argument is not well supported theoretically and empirically.

We stated in Section 2.3, as the alternative design choice to a linear decision boundary, the non-linear decision boundary by GMM tends to be susceptible to overfitting and vulnerable to instability, which, empirically, was further evidenced by the extensive results in Table 2 and discussed in Section 4.2 in the paper, where the non-linear decision boundary by PC exhibits the worst stability with the largest standard deviation even across the number of shots.

Theoretically, we have derived the non-linear boundary equation for GMM in Section 2.3, which was known to be unstable with outlier data points, and the robustness of GMM has been theoretically analyzed in previous work [1, 2]. Therefore, we opt for a linear decision boundary for BC due to its robustness, and we provide consistent and generalizable empirical evidence across tasks and modalities. We also conducted extensive robustness analysis to further validate the empirical robustness for the linear decision boundary design in Section 4.3, Page 8. We are open to any further discussion.

[1] Diakonikolas, I., Kane, D. M., & Stewart, A, 2017. Statistical query lower bounds for robust estimation of high-dimensional gaussians and gaussian mixtures. FOCS 2017. IEEE.

[2] Diakonikolas, I., Hopkins, S.B., Kane, D. and Karmalkar, S., 2020. Robustly learning any clusterable mixture of gaussians. arXiv preprint arXiv:2005.06417.

The advantage of Batch Calibration over existing methods is not thoroughly justified. Could you provide more intuitive descriptions and theoretical analysis for it?

We thank the reviewer for expressing their interest in a more intuitive description of BC. In contrast with existing methods that rely on content-free tokens and fail on multi-sentence classification tasks, BC serves as a content-based method using test-time statistics, thereby being more robust against distribution shifts and more generalizable across tasks. Intuitively, in Section 2.3, we discussed that the calibration can be framed as an unsupervised decision boundary learning problem. Figures 2 & 10 then provide a visualization of the decision boundary for BC and existing methods. To constitute a better decision boundary, we recall from Section 3 that the most critical component for calibration is to accurately estimate the contextual bias term $p(y|C)$ that shifts the original ICL decision boundary. This serves as a prior that can be approximately estimated by marginalizing the output likelihood over all samples within the batch in Equation 2. This forms the mathematical foundation of BC, and we welcome the reviewer for further discussion.

The derivations of Table 1 could be provided in more details in the main contents or supplementary. Currently, it is not easy to understand.

We appreciate the reviewer for pointing this out, and we have made updates in the caption of Table 1 to provide a supplementary explanation for each cell, including a reference to Section 2.3 that presents detailed derivations of decision boundaries for each method.