Restoring Calibration for Aligned Large Language Models: A Calibration-Aware Fine-Tuning Approach

We thank reviewer hpMa for the comments and suggestions. Below we provide our answer to your questions.

Q1. The biggest problem is ... not propose any modifications to preference alignment.

A. Thank you for the comment. We would like to clarify that this is better to be considered as a realistic practical scenario instead of a problem.

To begin with, consider a practical deployment scenario, which is the primary focus of our work. A practitioner aiming to develop their own model typically starts by downloading an open-source LLM and adapting it to their specific needs. Importantly, for many such models, the intermediate checkpoints (e.g., SFT stages) are not publicly available. As a result, practitioners are often unable to modify or re-run the full preference alignment process. In such cases, post-alignment fine-tuning is the most accessible strategy. Our approach is designed for this realistic setting—providing calibration improvements without requiring access to earlier training stages. We have revised the paper to make this point clear and more explicit to readers.

In the scenario where the practitioner is able to go through the full training pipeline, this is also not an issue. We provide our answer below.

Q2. What is the overall post-training pipeline? Justify the long post-training pipeline.

A. CFT is a post-hoc fine-tuning step applied after RLHF, and we will clarify this in the revised manuscript. Following our previous response, in scenarios where developers will go through the full training pipeline, both replacing RLHF and post-hoc fine-tuning are valid options—what matters most is which method yields better performance.

Methods that aim to replace RLHF—such as PPO-C in [2]—are promising, but they are evaluated under different experimental settings, making a direct comparison difficult within a short timeframe. In our work, we demonstrate the effectiveness of our approach by outperforming the strong baseline of TS. We believe our method is robust and competitive with RLHF-replacement approaches.

Q3. It is an unfair comparison w/wo domain knowledge.

A. We believe that the comparison remains appropriate for several reasons.

First, regarding ECE, the concern about fairness seems to stem from intuition rooted in more standard metrics like accuracy, where additional domain knowledge can indeed lead to higher performance. However, ECE measures the discrepancy between confidence and accuracy, domain knowledge does not inherently bias ECE comparisons.

Second, we have included experiments on out-domain tasks to further mitigate any potential effects from domain-specific knowledge.

Finally, the ultimate goal of training language models is to obtain models that perform well in real-world settings. If incorporating additional knowledge leads to better-calibrated models, it should be considered a strength rather than a limitation.

Q4. It is unclear … hard to tell when people should use methods in Sec 5.1 and when to use methods in Sec 5.2.

A: Before addressing it directly, we would like to first clarify a potential misunderstanding regarding Section 5.

In Theorem 4.6, we prove that ECE ≤ TCE, which implies that minimizing TCE is sufficient for achieving low ECE. This reduces the goal of obtaining a well-calibrated model to finding one that is close to the target probabilistic model, formulated as the following constrained optimization problem: $\max \text{ACC}(\pi) \quad \text{s.t.} \quad \text{ECE}(\pi) = 0.$

In Section 5, we focus on how to approximate both accuracy and ECE in the context of LLMs—covered in Sections 5.1 and 5.2, respectively.

Now, returning to your question: Sections 5.1 and 5.2 do not present two separate methods, but rather provide a breakdown of our final algorithm.

Q5. I think the authors should not restrict to only four possible answers (only four-answer setting (A,B,C,D) is allowed), as that is not practical in real applications.

A. We chose to use the four-option format (A, B, C, D) for concreteness and readability, aiming to reduce the burden of mathematical notation and make the presentation more accessible to experimental researchers. This is explicitly stated in our paper (lines 113–114): “While we consider four options in our analysis for concreteness, the framework naturally extends to any number of alternatives.” We will further revise the text to make this generality clearer.

Q6. The original preference alignment ability.

In addition to the win rate results presented in our original submission, we have conducted three more experiments related to preference alignment: CFT vs. DPO, AlpacaEval, and Arena-Hard. Due to space constraints, we refer to our response to Q4 of QE3c and the end of our response to USq4 for the detailed results.

Q7. Related work, suggestions on explaining the two regimes at first and other writings.

A. We have included the 4 mentioned papers in the revised version of the paper and have revised the paper accordingly.

We thank reviewer USq4 for the insightful comments and questions. Below we provide our responses to the questions.

Q1. One limitation is that their evaluation focuses primarily on multiple-choice settings rather than free-form text generation, which would provide a more complete picture of calibration in real-world LLM deployments.

A. Thank you for the comment. We agree that calibration in free-form generation is important. We focus on multiple-choice settings as they provide a controlled and quantifiable evaluation of calibration. Extending our method to free-form generation is an exciting direction, and we will discuss it as future work.

Q2. the constant C in the lower bound lacks specific derivation,

A. Thank you for pointing this out. The constant C in the lower bound arises from solving a min-max optimization problem, which makes it difficult to obtain a closed-form expression. While we do not derive an explicit value for C, we will clarify its origin and theoretical role in the revised manuscript. Characterizing or approximating C more precisely is an interesting direction for future work.

Q3. The EM-algorithm for probability estimation is theoretically sound, though the proof for convergence is not provided.

A. Thank you very much for the question. When the optimization is performed over probability distributions (rather than neural network parameters), the convergence of our EM algorithm follows from standard EM theory, and we take this as given rather than presenting it as a contribution of our work. However, when optimizing over neural network parameters, convergence is no longer guaranteed due to the non-convexity and complexity of the underlying function space. We will include a discussion on the convergence properties of the algorithm in the revised manuscript to clarify these distinctions.

Q4. The experimental design is generally sound but has several limitations. While they use appropriate models (four different architectures with both DPO/RLHF alignment), metrics (ECE, accuracy, win rate), and testing conditions (in-domain and cross-domain), the analysis lacks statistical significance testing and confidence intervals. The ablation studies are minimal, only comparing LSFT2 vs LECE in the appendix. The regularization parameter λ is fixed at 1 without sensitivity analysis, and the win rate metric only captures binary preference alignment rather than nuanced quality dimensions.

A. Thank you very much for the helpful comments and suggestions. We have now conducted a more comprehensive ablation study on the hyperparameter λ to further support our analysis. The results of this study have been incorporated into the revised manuscript.

Q5. How might your calibration methods be adapted for cases where we only have black-box API access to LLMs without fine-tuning capabilities?

A. Thank you for the question. Our approach relies on fine-tuning the model, and is therefore not applicable to black-box APIs. Adapting our method to black-box models would go beyond a simple extension—it would effectively require developing an entirely new approach. For improving calibration in black-box settings, techniques such as prompt engineering or encouraging the model to explicitly express its confidence are more suitable. We will include this discussion in the revised version of the paper.

Q6. Have you considered how quantization or other efficiency techniques might impact calibration properties in aligned models?

A. Thank you very much—this is a great question. Previously, we did not explore the impact of quantization or other efficiency techniques on calibration. These methods can potentially alter the model's confidence estimates and thus affect calibration quality. We agree this is an important and practical direction, especially for deployment scenarios, and we will include it in our discussion of future work.

To all reviewers: Additional experiments of alignment ability. Due to limited time, only 2 models for Arena-hard.

We thank reviewer Ak1F for the insightful comments and questions. Below we provide our responses to the questions.

Q1. Why fine-tuning with a calibration objective superior to post-hoc calibration methods e.g. TS?

A. While temperature scaling (TS) is a strong baseline, our method's superiority stems from addressing fundamental limitations of post-hoc calibration in LLMs:

1. Our approach explicitly minimizes the discrepancy between accuracy and confidence—the definition of ECE—rather than relying on a single scaling parameter. This direct optimization improves calibration while maintaining or enhancing performance. This aligns with previous work on training classifiers with calibration objectives, such as the AvUC loss in Krishnan et al and the PAC-Bayes-based objective in Fujisawa & Futami, 2024, which demonstrated the superiority of optimization-based strategies over post-hoc methods.

2. Our approach generalizes better to unseen data and distribution shifts compared to TS, which risks overfitting to specific validation sets. Table 2 clearly demonstrates this advantage in out-domain scenarios—for example, with Olmo2-7B, our CFT method reduces out-domain cw-ECE to 0.0637 (vs. TS's 0.1196) while improving accuracy to 0.7085 (vs. DPO's 0.6635). Distribution shifts are common in language tasks, and TS is insufficient to handle such variations.

We will include these points in our revised manuscript.

Q2. Relationship between this paper and calibration-aware training in classification models?

A. The referenced works propose various calibration-aware training objectives for classification models, such as MMCE, AvUC, S-AvUC, and SB-ECE. Some of these are designed as accuracy–calibration trade-off objectives.

Our approach shares this general principle, as it can also be viewed as an accuracy–calibration trade-off method. The key difference is that our accuracy and calibration objectives are specifically designed for generative models and LLMs, rather than classification models.

For the accuracy objective, we use the SFT loss on next-token prediction. This loss not only promotes accuracy but also supports instruction following, knowledge grounding, and coherent text generation—all essential to LLMs.

For the ECE objective, we formulate calibration as a generative modeling problem and apply an EM algorithm to optimize it accordingly.

In summary, our method is tailored for LLMs, which represents a significant departure in both design and application.

Q3. Evaluation Metrics and ECE Bias and optimal bin size.

A. Thank you for the helpful suggestion. The two referenced papers on ECE bias and optimal bin size are indeed insightful.

In our work, we use the UWB strategy, and we will make this choice explicit in the revised manuscript.

According to the referenced work, the optimal bin size for the multiclass setting scales as $O(n^{1/3})$. However, to apply this practically, one needs the exact constant rather than just the asymptotic rate. Upon further examination, we found that the optimal bin size contains a Lipshitz constant of the model, in a rate of $(1 + L)^{2/3} n^{1/3}.$

In practice, estimating the Lipschitz constant for transformers is challenging, as it is known to be potentially very large and difficult to compute reliably. Nonetheless, we can still make use of the theoretical rate $n^{1/3}$ as in the two mentioned papers. In our experiments, the sample size is 3,000, which implies an optimal bin size on the order of O(14.4). The bin size of 10 used in our paper is reasonably close to this rate, suggesting that our choice is consistent with the theoretical guidance.

Q4. In which cases the proposed methods perform well and fail?

A. Our approach is specifically designed and performs well for LLMs which are originally well-calibrated but become poorly calibrated after preference alignment. In other scenarios—such as (1) traditional classification tasks, or (2) cases where the LLM is not well-calibrated to begin with—our method may not be as effective.

Q5. Certain types of LLM architectures, alignment techniques ... where this approach is less effective?

A. In our experiments, we evaluate the method across different architectures, alignment techniques, and datasets to demonstrate its generality. However, due to computational constraints, we did not experiment with large-scale models such as 70B. While the effectiveness may vary at that scale, we believe our method is conceptually scalable and can be applied to larger models with appropriate resources.

Q6. Computational costs? Compare to other techniques? ECE term differentiable?

A. On two A100 40GB GPUs, our approach takes approximately 1.5 hours to run for 5 epochs. The training time is comparable to label smoothing, but significantly slower than temperature scaling, as the latter does not require fine-tuning model weights. In addition, Our ECE term is differentiable.

We thank reviewer QE3c for the insightful comments and questions. Below we provide our responses to the questions.

Q1. Can we draw any senses about when the model falls into each case from a theoretical perspect?

A. This is a great question. Determining which regime a model falls into ultimately reduces to understanding whether its accuracy exceeds a certain threshold. This, in turn, depends on two key factors: 1. What is the accuracy threshold for a given neural network architecture? 2. Given such an architecture, which training algorithms lead the model to fall into each regime? Answering these questions requires a deeper theoretical analysis of the properties of transformers, which is currently an open and challenging direction. While we do not yet have a definitive theoretical answer, our experimental results provide some insights. We will add a discussion of this point in the final section and consider it an important direction for future work.

Q2. The theory contribution is unclear. If my understanding is correct, the theory mainly states that there are two possible scenarios and does not directly relate to other parts of the paper, such as algorithms. The insight drawn from the theory is vague.

A. Our theoretical results are directly connected to the algorithmic design and other components of the paper. Let us clarify. In Theorem 4.6, we prove that ECE ≤ TCE. In other words, obtaining a well-calibrated model reduces to finding one that is close to the target probabilistic model, i.e., the solution to the following constrained optimization problem:

In Section 5, we discuss how to approximate ACC and ECE in the context of LLMs—specifically in Sections 5.1 and 5.2, respectively. Having defined approximations for the ACC and ECE losses, we introduce an EM algorithm to learn the underlying probabilistic model. This choice is motivated by our theoretical formulation of calibration as a probabilistic inference problem, for which EM is a natural and widely used solution.

Q3. There is only one baseline (temp scaling); it would be better to compare with other baselines such as label smoothing, [1], etc.

A. Thank you for the question. First, we would like to emphasize that temperature scaling is a strong baseline for reducing ECE. Prior work [3,4] has shown that it is generally more effective than label smoothing and other techniques for improving calibration. Nonetheless, we conducted additional experiments and found that, similar to image classification tasks, label smoothing remains less effective than temperature scaling on LLMs as well. We will add the experiments to our revised paper.

Regarding the contextual calibration method introduced in [1], we have carefully reviewed the paper. Although the term "calibration" appears in its name, the method is primarily designed to improve task performance across various NLP tasks, rather than to reduce ECE specifically. As such, their objective differs from ours, and the approaches are not directly comparable.

[3] When Does Label Smoothing Help? Rafael Müller, Simon Kornblith, Geoffrey Hinton

[4] On Calibration of Modern Neural Networks Chuan Guo, Geoff Pleiss, Yu Sun, Kilian Q. Weinberger

Q4. A better comparison of the Win rate is between CFT and DPO-generated responses since response qualities may decrease but still outperform the one being compared in the current setting.

A: Thank you for the suggestion. We agree that comparing the win rate between CFT and DPO-generated responses provides more information. We have conducted these additional experiments on AlpacaEval dataset, and the results are now included in our revised paper and provided below for reference.

Experiments on Alpaca-Eval and Arena-hard are provided in the end of our response to Reviewer USq4.

Q5. Some relevant methods are not discussed: [1,2].

A. We have included both [1] and [2] in the revised version of the paper.