Improving Uncertainty Quantification in Large Language Models via Semantic Embeddings

Authors: Yashvir Singh Grewal, Edwin V. Bonilla, Thang D Bui
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: uncertainty estimation, large language models, natural language generation, variational inference

Abstract:
Accurately quantifying uncertainty in large language models (LLMs) is crucial for their reliable deployment, especially in high-stakes applications. Current state-of-the-art methods for measuring semantic uncertainty in LLMs rely on strict bidirectional entailment criteria between multiple generated responses and also depend on sequence likelihoods. While effective, these approaches often overestimate uncertainty due to their sensitivity to minor wording differences, additional correct information, and non-important words in the sequence. We propose a novel approach that leverages semantic embeddings to achieve smoother and more robust estimation of semantic uncertainty in LLMs. By capturing semantic similarities without depending on sequence likelihoods, our method inherently reduces any biases introduced by irrelevant words in the answers. Furthermore, we introduce an amortised version of our approach by explicitly modelling semantics as latent variables in a joint probabilistic model. This allows for uncertainty estimation in the embedding space with a single forward pass, significantly reducing computational overhead compared to existing multi-pass methods. Experiments across multiple question-answering datasets and frontier LLMs demonstrate that our embedding-based methods provide more accurate and nuanced uncertainty quantification than traditional approaches.

Decision: reject
Reviewer scores: 6, 5, 6

Reviews and discussion:
<thread>
creation date:
1730586340563

thread_type:
official_review

<thread_object>
creator: reviewer_tjgs

<content>
summary:
This paper identifies that bidirectional entailment used for semantic uncertainty is sensitive to minor variations in responses. 
To address this issue, the authors propose a refined method, semantic embedding uncertainty (SEU), which is based on the average pairwise cosine similarity of the generated responses’ embeddings. The SEU consistently outperforms other baselines on three short-answer datasets. However, SEU still requires multiple runs of LLMs to collect a set of responses. To estimate uncertainty within a single forward pass, Amortised SEU is presented, which introduces a latent variable model to impute latent semantic embeddings with regard to the embeddings of an input.  Experimental results show ASEU achieves better performance than another single forward-pass baseline, length normalized predictive entropy.

strengths:
1. the paper is well-written with a clear storyline
2. the Amortised SEU is able to estimate uncertainty within a single forward pass,  which is efficient and interesting

weaknesses:
1. Apart from the Amortised SEU, I find that the scientific contribution of SEU (average pairwise cosine similarity) is not sufficient. 
First, the proposed SEU is a simple adapted version of the original semantic uncertainty. Second,   The authors claim that bidirectional entailment cannot provide continuous scores, but it is doable to use the probability of the NLI model for this purpose. I would suggest adding another baseline that replaces cosine similarity in SEU with the NLI scores. 
2. The experimental setting should be improved. First, it is necessary to compare SEU to other embedding-based methods, e.g., INSIDE: LLMs' Internal States Retain the Power of Hallucination Detection (ICLR 2024). Second, it is useful to report multiple metrics for evaluating uncertainty methods, e.g., Brier and ECE. Last, even though the authors mention the problem of the ROUGE-L metric and only short-answer datasets, I would believe that it is easy to mitigate these issues, e.g., Bert score and evaluation of long-form questions like TruthfulQA
3. Amortised SEU required further fine-tuning while other baselines are unsupervised. The generality of ASEU is not clear

questions:
1. in line 185, "cosine similarity provides a continuous metric". In fact, the bidirectional entailment can also offer continuous scores by using the output probability (or logit). 
2. ROUGE-L cannot capture the semantic equivalence well and is insensitive to word orders. The evaluation step needs further improvement by using better metrics such as BERTScore and LLM-as-Judge (though stated in the conclusion)
3. It is important to report multiple metrics for uncertainty estimation, e.g., Brier and ECE.
4. a key baseline is missing: INSIDE: LLMs' Internal States Retain the Power of Hallucination Detection (ICLR 2024), which develops an EigenScore metric to measure the semantic consistency in the embedding space.  
5. In Table 2, it is hard to see the benefit of SEU since it leads to a worse FPR
</content>
</thread_object>

</thread>
<thread>
creation date:
1730688179901

thread_type:
official_review

<thread_object>
creator: reviewer_8m62

<content>
summary:
This paper introduce, the technical challenges include addressing "hallucinations" in LLMs, where models generate coherent but incorrect responses, and the difficulty of applying traditional uncertainty quantification methods to the open-ended nature of natural language generation. Our innovations introduce Semantic Embedding Uncertainty (SEU) to estimate semantic uncertainty more accurately by leveraging semantic embeddings, and an amortized version (ASEU) that reduces computational costs by estimating uncertainty in a single forward pass.

strengths:
1 The paper presents a novel approach to uncertainty quantification in large language models (LLMs) by introducing Semantic Embedding Uncertainty (SEU) and its amortized version (ASEU). These methods innovatively leverage semantic embeddings to address the limitations of traditional uncertainty estimation techniques, offering a more nuanced understanding of semantic uncertainty without relying on strict bidirectional entailment criteria.
2 The research demonstrates high quality through  empirical evaluation across multiple question-answering datasets and state-of-the-art LLMs. The paper provides a good comparison against existing methods, showcasing SEU's superior performance in accurately quantifying uncertainty. 
3 The paper is well-structured and clearly articulates the problem, the proposed solutions, and their significance.

weaknesses:
1. Generalization Limitations: The SEU method proposed in this paper has mainly been experimentally validated for question-answering tasks, and its effectiveness and generalization ability in other types of natural language processing tasks, such as text summarization or machine translation, have not been fully tested.
2. Computational Resource Consumption: Although the amortized SEU (ASEU) reduces computational overhead, during the model training phase, especially when it involves inference and optimization of latent variables, it may still require relatively high computational resources.
3. Dependence on Pre-trained Models: The SEU method relies on high-quality pre-trained language models to generate semantic embeddings. If the pre-trained models themselves have biases or inaccuracies, it could affect the accuracy and reliability of the SEU method.
4. Figure 1,2,3,  model size diversity analyses needed.  While the paper provides a comparison of uncertainty estimation methods across different models, it would be beneficial to include models of varying sizes to get a more comprehensive understanding of how the proposed methods scale. The current selection of models does not fully represent the spectrum of LLMs, particularly the very large models (e.g., 70 billion parameters) or smaller models (e.g., 3 billion parameters). Including a broader range of model sizes could provide insights into the scalability and generalizability of the SEU and ASEU methods.
5 . Close Model Variations analysing needed.  It would be insightful to see comparisons that include close variations of the same model architecture but with different sizes. This could help elucidate whether the performance of the uncertainty estimation methods is more dependent on the model architecture or its size. Such an analysis could reveal trends that are crucial for understanding the behavior of SEU and ASEU in different practical scenarios.

questions:
1. Generalization to Other NLP Tasks?  The SEU method has been validated primarily for question-answering tasks. Could the authors discuss the potential effectiveness and generalization of the SEU method to other natural language processing tasks such as text summarization or machine translation, where the context and requirements might differ significantly?

2. Computational Resources in Model Training? While ASEU is praised for reducing computational overhead during inference, what are the computational resource requirements during the model training phase, particularly when dealing with the inference and optimization of latent variables? How does this compare to traditional methods in terms of efficiency?

3. Reliability on Pre-trained Models? The SEU method's performance seems to heavily rely on the quality of the pre-trained language models for generating semantic embeddings. How sensitive is the SEU method to potential biases or inaccuracies in these pre-trained models, and what measures can be taken to mitigate such issues?

4. Diversity in Model Size Analysis? Figure 1, along with other figures, compares uncertainty estimation methods across a limited range of model sizes. Could the authors expand on the analysis to include a more diverse set of model sizes, especially very large models (e.g., 70 billion parameters) and smaller models (e.g., 3 billion parameters), to provide a comprehensive understanding of how the proposed methods scale and perform across different model sizes?

5. Analysis of Close Model Variations? It would be insightful to see an analysis that includes close variations of the same model architecture but with different sizes. How does the performance of the uncertainty estimation methods vary with model size while keeping the architecture constant? What trends can be observed that would help us understand the behavior of SEU and ASEU in practical applications with different model sizes?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730926318809

thread_type:
official_review

<thread_object>
creator: reviewer_vha8

<content>
summary:
This paper addresses the problem of quantifying uncertainty in large language models. Specifically, previous methods, such as Semantic Entropy (SE), use bidirectional entailment criteria to assess whether two responses share the same meaning. However, the authors argue that bidirectional entailment criteria are sensitive to linguistic variations. To address this, they propose SEU, which quantifies uncertainty using cosine similarities between the semantic embeddings of different responses. Additionally, since both SE and SEU require multiple forward passes, the authors introduce an amortized version of SEU that models the underlying semantics as latent variables, requiring only a single forward pass while still achieving strong results.

strengths:
1. **Important problem**: The paper focuses on an important and valuable problem: quantifying the uncertainty of large language models.

2. **Efficient method**: The authors propose an amortized version of SEU that models the underlying semantics as latent variables. This approach requires only a single forward pass but still demonstrates very strong performance.

3. **Experimental validation**: The authors conduct experiments on three QA datasets across three LLMs, demonstrating the effectiveness of the proposed SEU and ASEU methods.

weaknesses:
1. **Reliability of cosine similarity of embeddings as an absolute measure of semantic relatedness**: I have some concerns about using cosine similarity as an absolute measure to describe the similarity between two responses. As shown in [1], the embedding space is isotropic, meaning that some texts are very close together within a narrow cone, while others are far apart. Thus, cosine similarity may not function well as an absolute measure.

2. **Advantages of SEU compared to SE**: The authors claim that SE models are sensitive to linguistic variability and minor wording differences, whereas SEU models are more robust and focus on the underlying semantic content. However, the SE method uses NLI classifiers that take the combined input of both responses, allowing for more fine-grained interaction between them. In contrast, SEU encodes responses independently, lacking such fine-grained interaction. As such, I find the argument for SEU being superior to SE unconvincing. Additionally, the experiments are conducted on a single model for both SE and SEU, which may not reveal whether the observed advantages are due to model choice or differences in methodology.

3. **Bidirectional entailment criterion can only act as a binary measure**: The authors claim that bidirectional entailment criteria can only yield binary outcomes. However, since the criterion relies on an entailment classification, the output probability of the entailment class has the potential to express the degree of relatedness between different responses.

[1] Jun Gao, Di He, Xu Tan, Tao Qin, Liwei Wang, Tie-Yan Liu. Representation Degeneration Problem in Training Natural Language Generation Models. ICLR 2019.

questions:
Please see my comments on the Weakness section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732551746285

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely thank all reviewers for their thoughtful feedback and valuable suggestions, which have been instrumental in improving the quality of our work. We are encouraged that the reviewers appreciated the quality of writing and the efficiency benefits of our proposed ASEU method.

However, we acknowledge that some reviewers may have had questions regarding the proposed SEU method and its advantages. We have addressed these points comprehensively in our individual responses, along with the additional concerns and questions raised by each reviewer.

Furthermore, we have significantly enhanced our experimental evaluation following the reviewers' suggestions, including the addition of a new baseline and expanding our tests to include larger models up to 27B parameters. We believe these improvements help clarify the paper's main contributions while strengthening the overall quality of our experimental results.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper is interested in improving uncertainty quantification for large language models. The idea is to obtain better estimates by integrating an improved version of semantic entropy. This permits taking into account a more accurate notion of uncertainty compared to either (i) just trivially examining how different responses are or (ii) using bidirectional entailment. The authors introduce some improved mechanisms and then study the performance of these approaches, including an efficient version that is less computationally expensive. The main idea is to use cosine similarity between emebddings.

As strong points, the paper provides a very simple solution that directly offers improvements. This is always nice to see.

As weaknesses, my thoughts are the following. First, this paper studies a relatively narrow area; effectively, it is searching for mechanisms to improve a component of semantic entropy. To do so, the authors propose the most straightforward approach. This is not by itself bad, but it does make me think there is much more to do and learn in this direction.

As also noted by the reviewers, cosine similarity can be tricky to reliably use. This is not a problem because there are a ton of alternatives (or combinations) that could be used here, and I would like to see more and understand the pros and cons. In fact, I would love to see a combination of SE with the authors’ idea, since they should be complementary, each handling cases the other does poorly on. In addition, evaluation was a good start but could be far more extensive (and this should be relatively easy).

Ultimately I felt that the paper was on the right track, but there’s much more to add to make it a strong paper that passes the bar.