I Don't Know: Explicit Modeling of Uncertainty with an [IDK] Token

Summary

This paper proposes a training-based confidence calibration method, named IDK-tuning, for improving LLM factuality. Specifically, a special `[IDK]` ("I Don't Know") token is added to the model's vocabulary and an objective is introduced which shifts some probability mass of wrong predictions to the `[IDK]` token during continued pretraining, thereby encouraging the model to express uncertainty explicitly. Results of the IDK-tuned models are reported on commonly used factual benchmarks, showing the potential of this method for reducing hallucination while only suffering a slight sacrifice in terms of knowledge recall.

Strengths

* Confidence calibration and LM hallucination are both important topics and this paper connects them in an interesting way. * The proposed IDK-tuning method is intuitive and well-motivated. While there have been papers on supervised finetuning for calibration, most of them focus on aligning the model with human demonstration (which requires annotation) or synthetic data (which incur extra cost). The direct adaptation of the training objective that incorporates uncertainty seems quite novel to my best knowledge, and results seem promising with much higher precision and only slightly lower recall, which is suited for the current generally over-confident LLMs. * The authors perform extensive experiments on the scaling behavior and ablation for different components of the objective function. * The authors do not overclaim their contribution. Singularities of the experiment results (e.g. `NaN`s in loss and collapsed recall from the `pythia` models) are mentioned and analyzed, and interesting insights are drawn from them. * The writing is clear and the paper is easy to follow.

Weaknesses

* For multiple-choice QA tasks, the precision improvement as well as the absolute values appear lower than those for the factual sentence completion tasks (comparing Table 1 and 2). Given this, and that the multiple-choice QA tasks might be more applicable, it would be helpful if the authors could report more detailed results on each of the `lm-eval-harness` tasks, and investigate QA or other downstream tasks more carefully. * Code or tuned model checkpoints are not provided, although some details about the settings and resources are mentioned in the paper.

Questions

* Is it possible or beneficial to add multiple special tokens to express the different levels of uncertainty that are finer-grained, instead of only a single `[IDK]` token? * For the evaluation on `TriviaQA` and `PopQA`, why do you reformat them into sentence completion tasks? Is it because in this way the next token(s) can be directly modeled without being extracted from a complete answer?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors highlight several limitations, including the need for full pretraining on large corpus, which is costly, and that the method can slightly sacrifice recall or the overall performance on some downstream tasks such as text generation.

Summary

The paper proposes calibrating LLMs during a continued pertaining phase via an added [IDK] token to model uncertainty.

Strengths

- The paper is structured well and written coherently. - The introduction of the $\texttt{[IDK]}$ token to explicitly model uncertainty in LLMs is a novel approach - Through ablation studies, the behavior of the proposed method is investigated extensively.

Weaknesses

- **Theory**: - The approach is mathematically not very well grounded. Also, mathematical expressions such as $prob(\texttt{[tok]}, \hat{y})$ and $prob(argmax(\hat{y}), \hat{y})$ do not cohere with common practices [1,2,3,4]. The authors could consider something like $p(y_{t}=\texttt{[tok]} | y_{<t}, x)$ and $max_{i} \ p(y_{t}=i | y_{<t}, x)$. - The uncertainty factor is only bigger than $0$ if any other token gets assigned a higher probability than the $\texttt{[gold]}$ token. It does not account for the case where a model is uncertain about *any* token and thus predicts a (low) probability for all tokens. For instance, consider the tokens relating to ($\text{"Paris"}$, $\text{"Berlin"}$, $\text{"London"}$, $\text{"Rome"}$, $\text{"Vienna"}$). If the model predicts any of $p(y_{t} | \text{"The capital of France is"}) \in [(0.2, 0.2, 0.2, 0.2, 0.2), (0.3, 0.1, 0.2, 0.2, 0.2), ...]$, the uncertainty factor is $0$ no matter the hyperparameter $\Pi$, while it is clear that in all those cases the model is uncertain about the correct next token. The probability of the $\texttt{[IDK]}$ gets even decreased via the uncertainty regularization. - **Evaluation**: The authors do not compare against other uncertainty quantification methods, such as (length-normalized) predictive entropy [1], p(true) [2], or semantic entropy [3,4]. These methods do not require additional pertaining, and thus do not suffer from training instabilities, mode collapse, or high computational costs, but can directly be applied to "off-the-shelf" models. Additionally, there exist methods that consider fine-tuning models to express their lack of knowledge that have not been considered. --- [1] A. Malinin and M. Gales. Uncertainty estimation in autoregressive structured prediction. [2] S. Kadavath, T. Conerly, A. Askell, T. Henighan, D. Drain, E. Perez, N. Schiefer, Z. Hatfield-Dodds, N. DasSarma, E. Tran-Johnson, S. Johnston, S. El-Showk, A. Jones, N. Elhage, T. Hume, A. Chen, Y. Bai, S. Bowman, S. Fort, D. Ganguli, D. Hernandez, J. Jacobson, J. Kernion, S. Kravec, L. Lovitt, K. Ndousse, C. Olsson, S. Ringer, D. Amodei, T. Brown, J. Clark, N. Joseph, B. Mann, S. McCandlish, C. Olah, J. Kaplan. Language Models (Mostly) Know What They Know. [3] L. Kuhn, Y. Gal, and S. Farquhar. Semantic uncertainty: Linguistic invariances for uncertainty estimation in natural language generation. [4] L. Aichberger, K. Schweighofer, M. Ielanskyi, and S. Hochreiter. Semantically Diverse Language Generation for Uncertainty Estimation in Language Models.

Questions

- If an IDK-tuned model is to be aligned, how does the alignment hurt IDK-tuning? Also, if a model is IDK-tuned after alignment, how does the pertaining hurt alignment? - The method is evaluated on a single model. How do you guarantee that the results generalize to bigger models?

Rating

3

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors adequately addressed the limitations.

Summary

To allow LMs to express their uncertainty for generative tasks, the authors introduce a new special IDK token. The authors modify the cross-entropy training objective to assign part of the probability mass to the IDK token in cases where the model gets the prediction wrong. The token embedding is randomly initialized and then refined through additional fine-tuning of pretrained LLMs of various sizes and types (Pythia, Mistralv1, BERT). Through experiments on a range of completion, QA & MCQA datasets the authors show that IDK tuning positively affects precision at a slight cost of recall. The authors perform further ablation experiments solidifying their choices for the loss weight hyperparameter as well as the regularization term. The authors further show that IDK tuning does not significantly adversely affect other capabilities of the underlying LMs.

Strengths

- Well written and easy to follow - To the best of my knowledge, presents a novel approach to quantifying uncertainty in LMs - Exhaustive experimental evaluation & ablation over different hyperparameter choices demonstrating robustness of the proposed approach

Weaknesses

- The IDK-tuning setup requires sometimes prohibitive additonal fine-tuning of the base LM - It is unclear how the method would be applied when a model would be pretrained from scratch with the IDK token included - it is natural that LMs will be worse in predicting tokens as initial stages of training, so the pretrain - add IDK - tune paradigm seems as the only current option. The two previous points slightly limit the applicability of the metod. - While the reported F scores are generally higher compared to baselines and alternatives, the IDK-tuned models still suffer from tangibly lower recall.

Questions

None

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Summary

It introduces a novel method to address the issue of hallucinations in Large Language Models (LLMs). These models, despite their proficiency in capturing and generating human knowledge, can sometimes produce factually incorrect text. To combat this, the authors propose a calibration method that incorporates an [IDK] ("I don’t know") token into the model's vocabulary.

Strengths

The introduction of the [IDK] token is a creative solution to an existing problem in LLMs. It represents a novel way to handle uncertainty, which is not just a new definition or problem formulation but also a practical application within language models. The paper proposes a new pretraining objective that modifies the conventional cross-entropy loss to incorporate the [IDK] token. This is an original contribution to the field of natural language processing.

Weaknesses

The paper primarily uses The Pile for training, which may not be representative of all possible language use cases or domains. While the paper provides a good overview of the performance metrics, an in-depth error analysis could offer more insights into the types of errors the models are making and how the [IDK] token impacts these. As the model is trained on web-crawled data, there is a risk of learning and perpetuating societal biases present in that data.

Questions

See Weaknesses

Rating

4

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

While the paper notes the potential for bias in the training data, it could provide more details on how this might affect the model's predictions and decision-making. The paper could more explicitly discuss the potential for the model to contribute to the spread of misinformation, especially if it fails to correctly identify uncertain or incorrect information.

Theoretical and Performance Concerns

Thank you for the rebuttal. > *Review*: **The uncertainty factor is only bigger than 0 if any other token gets assigned a higher probability than the **[gold]** token. It does not account for the case where a model is uncertain about any token and thus predicts a (low) probability for all tokens.** > *Rebuttal*: We first thank the reviewer for this comment as well, which we strongly agree with. As each of the training runs is very resource-intensive, this idea has not been evaluated in a proper experimental setup. The main contribution of your work is the introduction of an objective function that shifts the probability mass to the **[IDK]** token for incorrect predictions. This is not just an "idea" from my perspective; it reveals a significant theoretical flaw in your work. Given that your method lacks a solid theoretical foundation and the empirical performance only shows marginal improvements over the current state-of-the-art uncertainty quantification methods, coupled with major drawbacks such as unknown behavior when combined with alignment, I must reject the paper at this point. Addressing the theoretical issues will presumably enhance the performance of your method. Also, insights into how your method can be effectively applied to instruction-tuned models will improve the work.