Research map / Large language models / Language, safety and interpretability
I Don't Know: Explicit Modeling of Uncertainty with an [IDK] Token
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Adds a special [IDK] token and a continued-pretraining objective that shifts probability mass from wrong next-token predictions to [IDK], improving factual precision at a small recall cost.
Problem
LLMs hallucinate factually incorrect text and current calibration methods do not let them express uncertainty reliably in closed-book factual completion.
Approach
Introduce a new [IDK] vocabulary token and modify cross-entropy during continued pretraining so that, when the gold token is not the argmax, part of the target mass is moved to [IDK], scaled by an adaptive uncertainty factor (1 - p(gold)/max p) capped by a hyperparameter Pi=0.5, plus a false-positive regularizer when the model is correct.
Key results
- On Mistral-7B-v0.1, IDK-tuning raised TriviaQA precision from 52.0 to 72.5 (F1 52.0 to 60.6) and LAMA T-Rex precision from 71.2 to 88.5 with recall nearly unchanged (71.2 to 65.5).
- Average closed-book factual completion precision rose strongly for Mistral-7B, while recall dropped only modestly relative to the base model.
- Naively using the IDK loss without Pi and the FP regularizer caused training to collapse to predicting [IDK]; pythia-70m and pythia-160m diverged even with regularization.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply IDK-tuning on The Pile | Mistral-7B-v0.1 · LAMA Google-RE | F1: 50.0 -> 51.7 | Helped |
| Apply IDK-tuning on The Pile | Mistral-7B-v0.1 · LAMA T-Rex | F1: 75.0 -> 75.3 | Helped |
| Apply IDK-tuning on The Pile | Mistral-7B-v0.1 · LAMA SQuAD | F1: 54.4 -> 54.9 | Helped |
| Apply IDK-tuning on The Pile | Mistral-7B-v0.1 · TriviaQA | F1: 59.4 -> 60.6 | Helped |
| Apply IDK-tuning on The Pile | Mistral-7B-v0.1 · PopQA | F1: 33.7 -> 32.5 | Hurt |
| Apply IDK-tuning on The Pile | Mistral-7B-v0.1 · PopQA | Precision: 40.3 -> 78.1 | Helped |
| Apply IDK-tuning | Mistral-7B-v0.1 · lm-eval-harness | Precision: 45.0 -> 48.8 | Helped |
| Apply IDK-tuning | Mistral-7B-v0.1 · lm-eval-harness | Recall: 28.2 -> 20.8 | Hurt |
| Use adaptive rather than fixed λ (ablation) | IDK-tuning with Mistral-7B-v0.1 · closed-book factual sentence completion | IDK error rate | Helped |
| Add the L_FP-reg regularization term (ablation) | IDK-tuning with Mistral-7B-v0.1 · closed-book factual sentence completion | IDK error rate | Helped |
| Increase Π (ablation) | IDK-tuning with Mistral-7B-v0.1 · closed-book factual sentence completion | IDK recall | Helped |
| Increase Π (ablation) | IDK-tuning with Mistral-7B-v0.1 · closed-book factual sentence completion | IDK error rate | Hurt |
Surprising findings
- A naive IDK loss collapses training so the model predicts [IDK] almost always, requiring the Pi cap and FP regularizer.
- pythia-70m and pythia-160m collapsed in training even with regularization because the initial [IDK] probability underflowed to zero in float32.
Limitations
- Requires full continued pretraining on a large corpus, which is computationally expensive.
- Can slightly harm long-form text generation skills, and may not be effective for the smallest Pythia models whose recall collapses.
- Evaluated mainly on closed-book factual completion; multiple-choice tasks show smaller precision gains.
Experimental setup
Datasets: The Pile (uncopyrighted version) LAMA TriviaQA PopQA ARC HellaSwag MMLU TruthfulQA WinoGrande GSM8k Legal Plain English TLDR
Benchmarks: LAMA Google-RE LAMA T-Rex LAMA SQuAD TriviaQA PopQA ARC HellaSwag MMLU TruthfulQA WinoGrande GSM8k Legal Plain English
Metrics: precision recall F1 IDK recall IDK error rate RougeL
Baselines: Confidence Threshold P(True) Semantic Entropy The Pile continued training ablation
Builds on: cross-entropy language modeling P(True) semantic entropy
Closest papers by meaning
- Improving Uncertainty Estimation through Semantically Diverse Language Generation (ICLR 2025 · similarity 0.88)
- To Believe or Not to Believe Your LLM: Iterative Prompting for Estimating Epistemic Uncertainty (NeurIPS 2024 · similarity 0.88)
- Enhancing Uncertainty Modeling with Semantic Graph for Hallucination Detection (AAAI 2025 · similarity 0.87)
- In-Context Sharpness as Alerts: An Inner Representation Perspective for Hallucination Mitigation (ICML 2024 · similarity 0.87)
- LLMs Know More Than They Show: On the Intrinsic Representation of LLM Hallucinations (ICLR 2025 · similarity 0.87)
- Large Language Models Must Be Taught to Know What They Don’t Know (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
