Language Model Uncertainty Quantification with Attention Chain

We appreciate Reviewer ZMZs's feedback. We believe there might have been some misunderstandings, which we address below for clarity:

• Intuition behind Attention Weights: Attention weights intuitively represent the contributions of previous tokens toward predicting subsequent tokens, as discussed in Lines 49--55 and 143--145 of our manuscript. The efficacy of attention mechanisms has been validated since the very beginning of their proposal [1,2,3], both theoretically and empirically. Hence, we chose not to reiterate these theoretical foundations extensively within our paper.
• Scalability: We notice a potential divergence in the interpretation of "scalability". Commonly, scalability refers to a system's ability to maintain performance when facing increased input complexity or volume (e.g., longer sequences or varied input topics), rather than algorithmic simplicity. In this context, topic-specific fine-tuning is inherently unscalable, which we want to avoid.
• Weak Results: As clarified in Lines 44--46 and 104--113, fine-tuning methods necessitate task-specific labeled datasets, which are typically impractical to obtain in production environments or by end-users. Our study intentionally addresses UQ from a perspective distinct from fine-tuning. Nevertheless, we remain open to recommendations for alternative strong baselines within the scope of our work.

[1] Bahdanau, Dzmitry, Kyunghyun Cho, and Yoshua Bengio. "Neural machine translation by jointly learning to align and translate." arXiv preprint arXiv:1409.0473 (2014).
[2] Luong, Minh-Thang, Hieu Pham, and Christopher D. Manning. "Effective approaches to attention-based neural machine translation." arXiv preprint arXiv:1508.04025 (2015).
[3] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems 30 (2017).

We sincerely thank Reviewer 4Zzx for the supportive evaluation and insightful feedback. Below, we directly address the points raised:

• Ablation Study: Our initial goal was to establish a universal set of hyperparameters that can be broadly applicable, thus we limited detailed exploration of alternative parameter values. Recognizing the importance of your suggestion, we have conducted additional studies. Please refer to the Tables included in our response to Reviewer 5EYn above. In short, it is possible to see (much) better scores with hyper-parameters tuned for models and datasets, but the parameter set we provide in the paper is generally stable across different setups. We'll include the complete results and a more detailed discussion in the revised article.

• Consistency of Performance: We acknowledge the concern regarding performance consistency across different tasks, as highlighted in Table 1. This limitation was previously noted in the Limitation section of our paper (Lines 757--762). We are actively investigating enhancements to UQAC that will yield more stable and consistent improvements across diverse scenarios.

We appreciate your constructive feedback! It will significantly improve the clarity and robustness of our paper. Please let us know if you have further questions and concerns!

Hi Reviewer 4Zzx! Just a reminder that the discussion period for COLM papers has begun. Could you please take a look at the author response to review and let them know whether it addresses any of your outstanding questions?

Thanks, Your AC

We sincerely thank Reviewer 5EYn for the detailed review and constructive recommendations. Below, we address each point accordingly and describe the updates we will make to our manuscript:

• Typos: We appreciate the careful review and have corrected the indicated errors.

• Semantic Similarity Issues: We thank you for introducing these relevant works, which highlight potential improvements in semantic similarity evaluation. While we did not originally incorporate alternative similarity measures due to the scope of the initial development, we acknowledge the validity of your suggestion. We plan to explore these advanced semantic similarity methods in future iterations to further enhance the robustness of UQAC.

• Experiment Setup Details: To enhance reproducibility, we will add the detailed experimental settings in the updated manuscript, specifically:

  • Sampling temperature is set to 0 for all baselines except self-consistency, which uses a temperature of 0.5;
  • The parameter top_p is consistently set to 1; = Prompting schemes are briefly discussed in Lines 616--624. Specifically, on MATH and GSM8k, all models receive only the question as input, except for gemma-2, which includes the prompt "\n\nPut your final answer in \boxed{ }." For BBH, we use three provided in-context examples with answers enclosed in \boxed{ }. Further implementation details are available in our code repository. We'll include a detailed discussion on this topic in the revised manuscript

• Controlled Experiment on Reasoning Length: We appreciate this insightful suggestion. However, controlling reasoning sequence length without affecting the model’s inherent reasoning capabilities poses certain challenges. We have a short discussion on this topic in Lines 264--268 and appendix D.2 and remain open to suggestions on better experimental designs to address this variable.

• Further Analysis on Experimental Variations: We recognize the significant variation in performance across datasets and UQAC variants. We attribute differences across datasets primarily to inherent model calibration on the specific datasets. Notably, all three theoretically distinct methods—Self-Consistency, Verbalized Uncertainty, and UQAC—demonstrate notably lower performance on GSM8k in terms of AUROC. We discuss this phenomenon briefly in section D.1. Regarding variations between UQAC methods, the performance difference arises from approximations in marginalizing joint reasoning probabilities and answer probabilities, introducing complexities and noise. Ultimately, the marginal approximation ($\tilde{P}$) is more theoretically robust and consistently performs better in terms of stability across AUROC and ECE, as detailed in Lines 290--296.

• Hyper-Parameters: We conducted additional hyper-parameter studies and will include the results and analysis in the revised version. The results are in the Tables below. In short, It is possible to see (much) better scores with hyper-parameters tuned for models and datasets, but the parameter set we provide in the paper is generally stable across different setups.

Your feedback significantly improves our manuscript, and we thank you again for the thoughtful review. Please let us know about your thoughts!

Table 1. Hyper-parameter study of Llama 3.2 1B on GSM8k. The first column is organized as < n attention heads $K$ >-< attention weight $\theta$>-< similarity filtering target length $L'_{attn}$ >-< similarity filtering threshold >

Hi Reviewer 5EYn! Just a reminder that the discussion period for COLM papers has begun. Could you please take a look at the author response to review and let them know whether it addresses any of your outstanding questions?

Thanks, Your AC

We sincerely thank Reviewer tYEK for the positive assessment and constructive feedback. Below, we address your specific questions and suggestions in detail:

• Clarification on Baselines 1 & 2: Your understanding is correct. For baseline methods 1 and 2, we do not perform explicit marginalization over the entire vocabulary. Instead, we simplify by calculating the product of the probabilities for the actually sampled tokens, effectively sidestepping exhaustive marginalization. To further clarify the concern, we use \bm{x} as a sampled sequence (observation), and \mathbf{x} as a random variable, as stated in Table 3, category "Tokens & Sequences" and recommended by the conference.

• Code Accessibility: We apologize for the confusion regarding the anonymous GitHub repository. It appears this is an issue on the anonymous GitHub. All files have remained intact and accessible. We recommend downloading the full repository directly via the "Download Repository" button in the top-right corner. This may provide access to all the code.

• Transition to and from Related Works: We agree with your observation regarding the abrupt transitions. To enhance readability, we plan to reorganize our manuscript by relocating the Related Work section to follow the Experiments section, preceding the Conclusion. This restructuring should provide a smoother logical flow from Introduction to Method and facilitate easier comprehension.

• Figure 1 Visualization: Thank you for suggesting improvements to Figure 1. Instead of altering this figure, we will include additional examples in the appendix to better illustrate the quantified uncertainty measures explicitly. For the example provided, the uncertainty is indeed low, as the computed confidence is high ($\tilde{P}=0.92$).

• Dependence on CoT: Your interpretation is correct. UQAC explicitly requires the presence of a CoT reasoning sequence. When there is no reasoning chain, simpler methods such as baselines 2, 5, and even self consistency will become more applicable, lifting the necessity of applying UQAC.

• Redundancy in Lines 121--124: Thank you for pointing out the redundancy. Given our planned reorganization of the related work section, we believe keeping this brief repetition in Section 3 maintains narrative clarity for readers who might skip directly to specific sections.

• Justification for Probability-Based Uncertainty: We appreciate your suggestion regarding the justification for using probabilities as a measure of uncertainty. We will strengthen our manuscript by explicitly discussing prior empirical and theoretical justifications from existing literature to reinforce this common assumption.

• Hyper-Parameter Explanation: Indeed, we initially included a longer experiment setup section with more detailed explanations, but we compromised it to meet the length limitation. Nonetheless, we'll do our best to provide a clearer description or at least include a comprehensive explanation in the appendix.

• Qualitative Examples & Analysis: We acknowledge the value of providing qualitative examples. We plan to incorporate representative examples and accompanying analyses in the appendix of our revised manuscript, demonstrating concretely how UQAC improves calibration and provides intuitive uncertainty estimates.

Thank you again for your detailed comments and questions. We believe addressing these points significantly improves our paper's clarity and comprehensibility. Please let us know if your questions have been adequately answered or if you have any additional ideas or suggestions!

Hi Reviewer tYEK! Just a reminder that the discussion period for COLM papers has begun. Could you please take a look at the author response to review and let them know whether it addresses your concerns?

Thanks, Your AC