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

We sincerely thank the reviewer for their time reading our work, writing a thorough review and bringing up thoughtful comments.

We are encouraged that the reviewer finds our method intuitive, novel and well-motivated, that our experiments are deemed extensive, and that the paper is considered well-written and as not overclaiming the results.

Here are our responses to the weaknesses raised by the reviewer:

• We agree with the reviewer that it would be beneficial to more carefully analyze and represent the lm-eval-harness results. In the camera-ready version, we will add the results of each of the datasets separately and discuss the difference between them to those of the other datasets. Importantly, we note that according to a comment by a different reviewer, we have conducted an in-depth analysis of our model’s mistakes. The setup of the analysis is the following: We randomly sample 200 examples (out of all the datasets) on which the IDK-tuned model generates a wrong answer (without predicting the [IDK] token). We then categorizes it to one of these four categories: No Effect (both the original model and our model generate the same answer), Noise (the original model knows it, while after our training it doesn’t), White Noise (both the original model and ours don’t know it, though they generate different answers), and Abstain (when our model abstains from answering while generating text like “unknown” or “a mystery”). For this analysis we take three different models: Mistral-7B, Pythia-2.8B and Pythia-70M. The results are the following:

These results suggest that first, the bigger the model, the fewer changes our training approach causes in the model’s generations, and second, the bigger the model, the greater its ability to abstain from answering via words (which is generally equal to generating the new [IDK] token, though harder to evaluate automatically).

• We will definitely provide the code and checkpoints with the camera-ready version. Additionally, all the information needed to reproduce our model is demonstrated in the paper.

Here are our responses to the questions raised by the reviewer:

• “Multiple IDK tokens”: In principle, the single [IDK] token does cover all levels of uncertainty, as we measure a continuous amount of probability mass put on the token. However, we think multiple different [IDK] tokens are interesting, e.g., discriminating between different “categories” of uncertainty such as lack of knowledge, lack of context etc.
• “TriviaQA / PopQA reformatting”: Yes, we perform the reformatting to enable a straightforward evaluation.

We are very thankful to the reviewer for their time reading our work and writing a thorough review.

We are encouraged that the reviewer finds our paper to be well written, our approach to be novel, and our ablation experiments to be extensive and to properly demonstrate the behavior of our approach.

We will now address the reviewer's concerns:

Theory

• For the camera-ready version, we will modify all of the mathematical expressions in the paper for better coherence with previous common practices. Thank you for the suggestions!
• “The uncertainty factor is only bigger than $0$ if any other token gets assigned a higher probability than the $`[gold]`$ token”: 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. We plan to test a model that has been trained with a slight modification to our proposed training objective, reconsidering the true certainty it has on the gold token, even if this is the maximal probability token among the others. For example, we could take “1 - P(gold token)” as an additional term to be combined in the “Uncertainty Factor” definition, as well as to remove the separation we apply between cases where this factor is zero and the others, in terms of our objective. We will add some initial results of this alternative objective in the camera-ready version. We additionally think this is definitely an exciting follow-up to our work.

Evaluation

Thank you for this comment. We also believe that evaluating other uncertainty quantification methods might strengthen our work. We thus decided to add the results of the methods suggested by the reviewer to our paper. We will now provide the results of the “P(true)” and the “semantic entropy” methods using the Mistral-7B model, and will add these together with the results of the other models to our camera-ready version of the paper. For the “semantic entropy” method we use SOTA named entity recognition in order to extract only the “answer” itself from the model’s generation. It is important to note though that our method could potentially be applied directly during the normal pretraining of the model, by only modifying the pretraining objective to be ours, and thus the claim that our method requires additional training would not apply. Moreover, the “Semantic Entropy” method indeed does not require any additional training, though it requires significantly more inference calls and some post-processing of the generations (clustering etc.). The results are the following:

• LAMA Google-RE:

• LAMA T-Rex:

• LAMA SQuAD:

• TriviaQA:

• PopQA:

These results suggest that our approach still leads to the best precision and f1 scores compared to the new baselines too, though the gaps are smaller compared to the ones against the previous baselines.

Questions

• "Combining IDK-tuning with alignment”: While we do not explicitly study this in our work, we do believe it is an interesting direction for future work. Specifically, the line of work on Task Arithmetic [1] will be interesting: Can we extract and combine the “IDK” weight vector with an instruction-tuned model?
• "“The method is evaluated on a single model” / “generalization to bigger models”: We believe this question must be a misunderstanding, as we conduct very extensive experiments with eight different models, spanning encoder-only (BERT) and decoder-only architectures, as well as an explicit study of scaling behavior using the Pythia model suite and Mistral-7B.

In light of this, we kindly request you to reconsider and appropriately raise the score of your review if we have sufficiently addressed some or all of your concerns.

We first express our sincere gratitude to the reviewer for their time reading our work, writing a thorough review, and bringing up thoughtful comments.

We are encouraged that the reviewer finds our paper to be well written, our approach to be novel, and our experiments to be extensive and to properly demonstrate the robustness of our approach.

Here are our responses to the concerns raised by the reviewer:

• While our method does require finetuning, we also benefit from the vast amount of current work in the optimization and efficiency of LLM training. In any case, we argue that (moderate) finetuning does not constitute a technical weakness of our method.
• While we experiment with IDK-tuning for the reasons you mention, we believe integrating the [IDK]-flavor of uncertainty-aware training into pretraining from scratch presents an exciting direction for future work (e.g., via loss weight schedules). We do believe that with small adaptations of our loss function, our training method could be applied for pretraining from scratch. One of the main reasons we couldn’t check this effectively is the very extensive resources (time, compute and memory) these controlled experiments in this setting would have consumed (which would be impractical or impossible on typical academic budgets in a reasonable amount of time). Also, we believe this direction of work warrants a detailed study that will fill another whole paper.
• We agree that our method is not “free” (as it leads to somewhat lower recall), but argue that it presents a valuable tradeoff, given the importance of reducing hallucinations.

We first highly thank the reviewer for their time reading our work, writing a thorough review and bringing up thoughtful comments.

We are encouraged that the reviewer finds our method a creative and novel method to handle an existing and important problem current LLMs have - which is generating misinformation and hallucinations.

We would like to address the following weaknesses / limitations of our work raised in the review:

• “Using primarily The Pile”: We argue that using a general, English-centric dataset is well suited to demonstrate the efficacy of our method and not an inherent weakness of our method. Additionally, as you may assume, each of these training sessions is extremely resource-intensive. Exploring other languages and domains is a good direction for future work.
• “In-depth error analysis missing”: We do provide a very detailed analysis and ablation of precision-recall tradeoffs in Section 4.2. Based on the reviewer’s suggestion, we decided to conduct an in-depth error analysis of our model’s mistakes. We will provide the results here and will add them to the camera-ready version of the paper too. The setup of the analysis is the following: We randomly sample 200 examples (out of all the datasets) on which the IDK-tuned model generates a wrong answer (without predicting the [IDK] token). We then categorize these to one of four categories: No Effect (both the original model and our model generate the same answer), Noise (the original model knows it, while after our training it doesn’t), White Noise (both the original model and ours don’t know it, though they generate different answers), and Abstain (when our model abstains from answering while generating text like “unknown” or “a mystery”). For this analysis we take three different models: Mistral-7B, Pythia-2.8B and Pythia-70M. The results are the following: | model | No Effect | Noise | White Noise | Abstaining | |-------------|-----------|-------|-------------|------------| | Mistral-7B | 68.5 | 9 | 6.5 | 16 | | Pythia-2.8B | 59.5 | 13.5 | 12.5 | 14.5 | | Pythia-70M | 52 | 18.5 | 22 | 7.5 |

These results suggest that first, the bigger the model, the fewer changes our training approach causes in the model’s generations, and second, the bigger the model, the greater its ability to abstain from answering via words (which is generally equal to generating the new [IDK] token, though harder to evaluate automatically).

• “Bias from web-crawled pretraining data”: We agree and will use the extra page to discuss this in an extended Limitations section. However, while important to consider, we argue that pretraining data bias is (1) mostly inherent to the way LLMs are trained nowadays and (2) not a particular weakness of our method but an exciting direction for follow-up work (e.g., using synthetic data).
• "Potential for the model to contribute to the spread of misinformation": Indeed our extended Limitations section will also discuss the risk of misinformation. It is important to stress that we propose a single method, not a system design for safe deployment of LLMs. In practice, we anticipate our method to be coupled with other checks and balances, forming a safe system.

In light of this, we kindly request you to reconsider and appropriately raise the score of your review if we have sufficiently addressed some or all of your concerns.